diff --git a/.gitignore b/.gitignore index 900e5a53cbcf3bbb5e00389cca004c49f8600a66..c227f50d55299e873fe199f12234195db7cff3eb 100644 --- a/.gitignore +++ b/.gitignore @@ -4,15 +4,14 @@ node_modules /.bazelrc /.tf_configure.bazelrc /bazel-* -/third_party/py/numpy/numpy_include -/tools/bazel.rc +/bazel_pip /tools/python_bin_path.sh /tools/git/gen -/util/python/python_include -/util/python/python_lib /pip_test /_python_build *.pyc __pycache__ *.swp .vscode/ +cmake_build/ +.idea/** diff --git a/.mention-bot b/.mention-bot deleted file mode 100644 index 9e4858977f5da2992ccc4053dfbbda3f5f86ee90..0000000000000000000000000000000000000000 --- a/.mention-bot +++ /dev/null @@ -1,11 +0,0 @@ -{ - "maxReviewers": 2, - "numFilesToCheck": 10, - "userBlacklist": ["tensorflower-gardener"], - "requiredOrgs": ["tensorflow"], - "skipAlreadyAssignedPR": true, - "skipAlreadyMentionedPR": true, - "skipTitle": "Branch", - "delayed": true, - "delayedUntil": "10m" -} diff --git a/CODEOWNERS b/CODEOWNERS new file mode 100644 index 0000000000000000000000000000000000000000..0a12176aaa27fb2909d623e3bced896b04d96584 --- /dev/null +++ b/CODEOWNERS @@ -0,0 +1,53 @@ +# NOTE: Disabled temporarily because it's too noisy on pushes. +# Where component owners are known, add them here. + +#tensorflow/core/platform/windows/* @mrry +#tensorflow/java/* @asimshankar +#tensorflow/tensorboard/* @jart @dandelionmane +#tensorflow/tools/docs/* @markdaoust + +# contrib + +# NEED OWNER: tensorflow/contrib/avro/* +#tensorflow/contrib/batching/* @alextp @chrisolston +#tensorflow/contrib/bayesflow/* @ebrevdo @rsepassi @jvdillon +#tensorflow/contrib/cmake/* @mrry @benoitsteiner +#tensorflow/contrib/copy_graph/* @tucker @poxvoculi +#tensorflow/contrib/crf/* @kentonl +#tensorflow/contrib/data/* @mrry +#tensorflow/contrib/distributions/* @jvdillon @langmore @rsepassi +#tensorflow/contrib/factorization/* @agarwal-ashish @xavigonzalvo +#tensorflow/contrib/ffmpeg/* @fredbertsch +# NEED OWNER: tensorflow/contrib/framework/* +#tensorflow/contrib/graph_editor/* @purpledog +# NEED OWNER: tensorflow/contrib/grid_rnn/* +#tensorflow/contrib/hvx/* @satok16 +#tensorflow/contrib/imperative/* @keveman +#tensorflow/contrib/integrate/* @shoyer +#tensorflow/contrib/kernel_methods/* @petrosmol +#tensorflow/contrib/ios_examples/* @petewarden +#tensorflow/contrib/labeled_tensor/* @shoyer +#tensorflow/contrib/layers/* @fchollet @martinwicke +#tensorflow/contrib/learn/* @martinwicke @ispirmustafa @alextp +#tensorflow/contrib/linalg/* @langmore +#tensorflow/contrib/linear_optimizer/* @petrosmol @andreasst @katsiapis +#tensorflow/contrib/lookup/* @ysuematsu @andreasst +#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 +#tensorflow/contrib/pi_examples/* @maciekcc +#tensorflow/contrib/quantization/* @petewarden @cwhipkey @keveman +#tensorflow/contrib/rnn/* @ebrevdo +#tensorflow/contrib/saved_model/* @nfiedel @sukritiramesh +#tensorflow/contrib/seq2seq/* @lukaszkaiser +#tensorflow/contrib/session_bundle/* @nfiedel @sukritiramesh +#tensorflow/contrib/slim/* @sguada @thenbasilmanran +#tensorflow/contrib/stateless/* @girving +#tensorflow/contrib/tensor_forest/* @gilberthendry @thomascolthurst +#tensorflow/contrib/testing/* @dandelionmane +#tensorflow/contrib/timeseries/* @allenlavoie +#tensorflow/contrib/tpu/* @frankchn @saeta @jhseu +#tensorflow/contrib/training/* @joel-shor @ebrevdo +#tensorflow/contrib/util/* @sherrym diff --git a/CODE_OF_CONDUCT.md b/CODE_OF_CONDUCT.md new file mode 100644 index 0000000000000000000000000000000000000000..10fd595fec7f240c3fdc871e1f32cc83f2ffd46d --- /dev/null +++ b/CODE_OF_CONDUCT.md @@ -0,0 +1,70 @@ +# TensorFlow Code of Conduct + +In the interest of fostering an open and welcoming environment, we as contributors and maintainers pledge to making participation in our project and our community a harassment-free experience for everyone, regardless of age, body size, disability, ethnicity, gender identity and expression, level of experience, nationality, personal appearance, race, religion, or sexual identity and orientation. + + +## Our Standards + +Examples of behavior that contributes to creating a positive environment include: + +* Using welcoming and inclusive language +* Being respectful of differing viewpoints and experiences +* Gracefully accepting constructive criticism +* Focusing on what is best for the community +* Showing empathy towards other community members + +Examples of unacceptable behavior by participants include: + +* The use of sexualized language or imagery and unwelcome sexual attention or advances +* Trolling, insulting/derogatory comments, and personal or political attacks +* Public or private harassment +* Publishing others' private information, such as a physical or electronic address, without explicit permission +* Conduct which could reasonably be considered inappropriate for the forum in which it occurs. + +All TensorFlow forums and spaces are meant for professional interactions, and any behavior which could reasonably be considered inappropriate in a professional setting is unacceptable. + + +## Our Responsibilities + +Project maintainers are responsible for clarifying the standards of acceptable behavior and are expected to take appropriate and fair corrective action in response to any instances of unacceptable behavior. + +Project maintainers have the right and responsibility to remove, edit, or reject comments, commits, code, wiki edits, issues, and other contributions that are not aligned to this Code of Conduct, or to ban temporarily or permanently any contributor for other behaviors that they deem inappropriate, threatening, offensive, or harmful. + + +## Scope + +This Code of Conduct applies to all content on tensorflow.org, TensorFlow’s GitHub organization, or any other official TensorFlow web presence allowing for community interactions, as well as at all official TensorFlow events, whether offline or online. + +The Code of Conduct also applies within project spaces and in public spaces whenever an individual is representing TensorFlow or its community. Examples of representing a project or community include using an official project e-mail address, posting via an official social media account, or acting as an appointed or de facto representative at an online or offline event. + + +## Conflict Resolution + +Conflicts in an open source project can take many forms, from someone having a bad day and using harsh and hurtful language in the issue queue, to more serious instances such as sexist/racist statements or threats of violence, and everything in between. + +If the behaviour is threatening or harassing, or for other reasons requires immediate escalation, please see below. + +However, for the vast majority of issues, we aim to empower individuals to first resolve conflicts themselves, asking for help when needed, and only after that fails to escalate further. This approach gives people more control over the outcome of their dispute. + +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. + + +## Reporting Violations + +Violations of the Code of Conduct can be reported to TensorFlow’s Project Steward at conduct@tensorflow.org. The Project Steward will determine whether the Code of Conduct was violated, and will issue an appropriate sanction, possibly including a written warning or expulsion from the project, project sponsored spaces, or project forums. We ask that you make a good-faith effort to resolve your conflict via the conflict resolution policy before submitting a report. + +Violations of the Code of Conduct can occur in any setting, even those unrelated to the project. We will only consider complaints about conduct that has occurred within one year of the report. + + +## Enforcement + +If the Project Steward receives a report alleging a violation of the Code of Conduct, the Project Steward will notify the accused of the report, and provide them an opportunity to discuss the report before a sanction is issued. The Project Steward will do their utmost to keep the reporter anonymous. If the act is ongoing (such as someone engaging in harassment), or involves a threat to anyone's safety (e.g. threats of violence), the Project Steward may issue sanctions without notice. + + +## Attribution + +This Code of Conduct is adapted from the Contributor Covenant, version 1.4, available at http://contributor-covenant.org/version/1/4, and includes some aspects of the Geek Feminism Code of Conduct and the Drupal Code of Conduct. diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 5ae5c0fbbcd5b8da7e3f3f98e01f455e0c82e588..43abdaafbf45379430920cd027b26299cd62553b 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -27,3 +27,145 @@ contributions, often because we probably won't get to them right now. If you decide to start on an issue, leave a comment so that other people know that you're working on it. If you want to help out, but not alone, use the issue comment thread to coordinate. + +### Contribution guidelines and standards + +Before sending your pull request for +[review](https://github.com/tensorflow/tensorflow/pulls), +make sure your changes are consistent with the guidelines and follow the +TensorFlow coding style. + +#### General guidelines and philosophy for contribution + +* Include unit tests when you contribute new features, as they help to + a) prove that your code works correctly, b) guard against future breaking + changes to lower the maintenance cost. +* Bug fixes also generally require unit tests, because the presence of bugs + usually indicates insufficient test coverage. +* Keep API compatibility in mind when you change code in core TensorFlow, + e.g., code in [tensorflow/core](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/core) and [tensorflow/python](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/python). + TensorFlow has reached version 1 and hence cannot make + non-backward-compatible API changes without a major release. Reviewers of your + pull request will comment on any API compatibility issues. +* When you contribute a new feature to TensorFlow, the maintenance burden is (by + default) transferred to the TensorFlow team. This means that benefit of + contribution must be compared against the cost of maintaining the feature. +* Full new features (e.g., a new op implementing a cutting-edge algorithm) + typically will live in + [tensorflow/contrib](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/contrib) + to get some airtime before decision is made regarding whether they are to be + migrated to the core. + +#### License + +Include a license at the top of new files. + +* [C/C++ license example](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/framework/op.cc#L1) +* [Python license example](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/ops/nn.py#L1) +* [Java license example](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/java/src/main/java/org/tensorflow/Graph.java#L1) +* [Go license example](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/go/operation.go#L1) +* [Bash license example](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/tools/ci_build/ci_sanity.sh#L2) +* [HTML license example](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/tensorboard/dist/index.html#L2) +* [JavaScript/TypeScript license example](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/tensorboard/components/tf_backend/backend.ts#L1) + +Bazel BUILD files also need to include a license section, e.g., +[BUILD example](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/BUILD#L61). + +#### C++ coding style + +Changes to TensorFlow C++ code should conform to +[Google C++ Style Guide](https://google.github.io/styleguide/cppguide.html). + +Use `clang-tidy` to check your C/C++ changes. To install clang-tidy on ubuntu:16.04, do: + +```bash +apt-get install -y clang-tidy +``` + +You can check a C/C++ file by doing: + + +```bash +clang-format --style=google > /tmp/my_cc_file.cc +diff /tmp/my_cc_file.cc +``` + +#### Python coding style + +Changes to TensorFlow Python code should conform to +[Google Python Style Guide](https://google.github.io/styleguide/pyguide.html) + +Use `pylint` to check your Python changes. To install `pylint` and +retrieve TensorFlow's custom style definition: + +```bash +pip install pylint +wget -O /tmp/pylintrc https://raw.githubusercontent.com/tensorflow/tensorflow/master/tensorflow/tools/ci_build/pylintrc +``` + +To check a file with `pylint`: + +```bash +pylint --rcfile=/tmp/pylintrc myfile.py +``` + +#### Coding style for other languages + +* [Google Java Style Guide](https://google.github.io/styleguide/javaguide.html) +* [Google JavaScript Style Guide](https://google.github.io/styleguide/jsguide.html) +* [Google Shell Style Guide](https://google.github.io/styleguide/shell.xml) + +#### Running sanity check + +If you have Docker installed on your system, you can perform a sanity check on +your changes by running the command: + +```bash +tensorflow/tools/ci_build/ci_build.sh CPU tensorflow/tools/ci_build/ci_sanity.sh +``` + +This will catch most license, Python coding style and BUILD file issues that +may exist in your changes. + +#### Running unit tests + +There are two ways to run TensorFlow unit tests. + +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. + + 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 + + ```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" + ``` + + For example, to run all tests under tensorflow/python, do: + + ```bash + bazel test ${flags} //tensorflow/python/... + ``` + +2. Using [Docker](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. + diff --git a/ISSUE_TEMPLATE.md b/ISSUE_TEMPLATE.md index 50f67963bff9ab817915e56577169887c45da8eb..2bf2c754cf64ec3bac22a22fbafcebbd4dc54bf4 100644 --- a/ISSUE_TEMPLATE.md +++ b/ISSUE_TEMPLATE.md @@ -1,15 +1,38 @@ -NOTE: Issues that are not bugs or feature requests will be closed. Please ask usage questions on StackOverflow. +Please go to Stack Overflow for help and support: -### You must complete this information or else your issue will be closed -- *Have I written custom code (as opposed to using a stock example script provided in TensorFlow)?*: -- *TensorFlow installed from (source or binary)?*: -- *TensorFlow version*: -- *Bazel version (if compiling from source)*: -- *CUDA/cuDNN version*: -- *GPU Model and Memory*: -- *Exact command to reproduce*: +https://stackoverflow.com/questions/tagged/tensorflow -### Describe the problem clearly +If you open a GitHub issue, here is our policy: -### Source Code / 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. Try to reproducible test-case code the bare-minimum necessary to generate the problem +1. It must be a bug or a feature request. +2. The form below must be filled out. +3. It shouldn't be a TensorBoard issue. Those go [here](https://github.com/tensorflow/tensorboard/issues). + +**Here's why we have that policy**: TensorFlow developers respond to issues. We want to focus on work that benefits the whole community, e.g., fixing bugs and adding features. Support only helps individuals. GitHub also notifies thousands of people when issues are filed. We want them to see you communicating an interesting problem, rather than being redirected to Stack Overflow. + +------------------------ + +### System information +- **Have I written custom code (as opposed to using a stock example script provided in TensorFlow)**: +- **OS Platform and Distribution (e.g., Linux Ubuntu 16.04)**: +- **TensorFlow installed from (source or binary)**: +- **TensorFlow version (use command below)**: +- **Python version**: +- **Bazel version (if compiling from source)**: +- **CUDA/cuDNN version**: +- **GPU model and memory**: +- **Exact command to reproduce**: + +You can collect some of this information using our environment capture script: + +https://github.com/tensorflow/tensorflow/tree/master/tools/tf_env_collect.sh + +You can obtain the TensorFlow version with + +python -c "import tensorflow as tf; print(tf.GIT_VERSION, tf.VERSION)" + +### Describe the problem +Describe the problem clearly here. Be sure to convey here why it's a bug in TensorFlow or a feature request. + +### Source code / 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. Try to provide a reproducible test case that is the bare minimum necessary to generate the problem. diff --git a/README.md b/README.md index cd0bffde79646a77bd50b48cc2cd1495db2fe1e7..5a0739a603cb564aeb6223717545e3205235325f 100644 --- a/README.md +++ b/README.md @@ -9,37 +9,47 @@ | [![Build Status](https://ci.tensorflow.org/buildStatus/icon?job=tensorflow-master-cpu)](https://ci.tensorflow.org/job/tensorflow-master-cpu) | [![Build Status](https://ci.tensorflow.org/buildStatus/icon?job=tensorflow-master-linux-gpu)](https://ci.tensorflow.org/job/tensorflow-master-linux-gpu) | [![Build Status](https://ci.tensorflow.org/buildStatus/icon?job=tensorflow-master-mac)](https://ci.tensorflow.org/job/tensorflow-master-mac) | [![Build Status](https://ci.tensorflow.org/buildStatus/icon?job=tensorflow-master-win-cmake-py)](https://ci.tensorflow.org/job/tensorflow-master-win-cmake-py) | [![Build Status](https://ci.tensorflow.org/buildStatus/icon?job=tensorflow-master-android)](https://ci.tensorflow.org/job/tensorflow-master-android) | **TensorFlow** is an open source software library for numerical computation using -data flow graphs. Nodes in the graph represent mathematical operations, while +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 lets you deploy computation to one or more CPUs or GPUs in a desktop, server, or mobile device without rewriting code. TensorFlow also includes 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 +working on the Google Brain team within Google's Machine Intelligence Research organization for the purposes of conducting machine learning and deep neural networks research. The system is general enough to be applicable in a wide variety of other domains, as well. -**If you'd like to contribute to TensorFlow, be sure to review the [contribution -guidelines](CONTRIBUTING.md).** +**If you want to contribute to TensorFlow, be sure to review the [contribution +guidelines](CONTRIBUTING.md). This project adheres to TensorFlow's +[code of conduct](CODE_OF_CONDUCT.md). By participating, you are expected to +uphold this code.** **We use [GitHub issues](https://github.com/tensorflow/tensorflow/issues) for -tracking requests and bugs, but please see -[Community](tensorflow/docs_src/about/index.md#community) for general questions -and discussion.** +tracking requests and bugs. So please see +[TensorFlow Discuss](https://groups.google.com/a/tensorflow.org/forum/#!forum/discuss) for general questions +and discussion, and please direct specific questions to [Stack Overflow](https://stackoverflow.com/questions/tagged/tensorflow).** ## Installation -*See [Installing TensorFlow](https://www.tensorflow.org/install/) for instructions on how to install our release binaries or how to build from source.* +*See [Installing TensorFlow](https://www.tensorflow.org/get_started/os_setup.html) for instructions on how to install our release binaries or how to build from source.* People who are a little more adventurous can also try our nightly binaries: -* Linux CPU-only: [Python 2](https://ci.tensorflow.org/view/Nightly/job/nightly-matrix-cpu/TF_BUILD_IS_OPT=OPT,TF_BUILD_IS_PIP=PIP,TF_BUILD_PYTHON_VERSION=PYTHON2,label=cpu-slave/lastSuccessfulBuild/artifact/pip_test/whl/tensorflow-1.1.0rc1-cp27-none-linux_x86_64.whl) ([build history](https://ci.tensorflow.org/view/Nightly/job/nightly-matrix-cpu/TF_BUILD_IS_OPT=OPT,TF_BUILD_IS_PIP=PIP,TF_BUILD_PYTHON_VERSION=PYTHON2,label=cpu-slave)) / [Python 3.4](https://ci.tensorflow.org/view/Nightly/job/nightly-matrix-cpu/TF_BUILD_IS_OPT=OPT,TF_BUILD_IS_PIP=PIP,TF_BUILD_PYTHON_VERSION=PYTHON3,label=cpu-slave/lastSuccessfulBuild/artifact/pip_test/whl/tensorflow-1.1.0rc1-cp34-cp34m-linux_x86_64.whl) ([build history](https://ci.tensorflow.org/view/Nightly/job/nightly-matrix-cpu/TF_BUILD_IS_OPT=OPT,TF_BUILD_IS_PIP=PIP,TF_BUILD_PYTHON_VERSION=PYTHON3,label=cpu-slave/)) / [Python 3.5](https://ci.tensorflow.org/view/Nightly/job/nightly-python35-linux-cpu/lastSuccessfulBuild/artifact/pip_test/whl/tensorflow-1.1.0rc1-cp35-cp35m-linux_x86_64.whl) ([build history](https://ci.tensorflow.org/view/Nightly/job/nightly-python35-linux-cpu/)) -* Linux GPU: [Python 2](https://ci.tensorflow.org/view/Nightly/job/nightly-matrix-linux-gpu/TF_BUILD_IS_OPT=OPT,TF_BUILD_IS_PIP=PIP,TF_BUILD_PYTHON_VERSION=PYTHON2,label=gpu-linux/lastSuccessfulBuild/artifact/pip_test/whl/tensorflow_gpu-1.1.0rc1-cp27-none-linux_x86_64.whl) ([build history](https://ci.tensorflow.org/view/Nightly/job/nightly-matrix-linux-gpu/TF_BUILD_IS_OPT=OPT,TF_BUILD_IS_PIP=PIP,TF_BUILD_PYTHON_VERSION=PYTHON2,label=gpu-linux/)) / [Python 3.4](https://ci.tensorflow.org/view/Nightly/job/nightly-matrix-linux-gpu/TF_BUILD_IS_OPT=OPT,TF_BUILD_IS_PIP=PIP,TF_BUILD_PYTHON_VERSION=PYTHON3,label=gpu-linux/lastSuccessfulBuild/artifact/pip_test/whl/tensorflow_gpu-1.1.0rc1-cp34-cp34m-linux_x86_64.whl) ([build history](https://ci.tensorflow.org/view/Nightly/job/nightly-matrix-linux-gpu/TF_BUILD_IS_OPT=OPT,TF_BUILD_IS_PIP=PIP,TF_BUILD_PYTHON_VERSION=PYTHON3,label=gpu-linux/)) / [Python 3.5](https://ci.tensorflow.org/view/Nightly/job/nightly-matrix-linux-gpu/TF_BUILD_IS_OPT=OPT,TF_BUILD_IS_PIP=PIP,TF_BUILD_PYTHON_VERSION=PYTHON3.5,label=gpu-linux/lastSuccessfulBuild/artifact/pip_test/whl/tensorflow_gpu-1.1.0rc1-cp35-cp35m-linux_x86_64.whl) ([build history](https://ci.tensorflow.org/view/Nightly/job/nightly-matrix-linux-gpu/TF_BUILD_IS_OPT=OPT,TF_BUILD_IS_PIP=PIP,TF_BUILD_PYTHON_VERSION=PYTHON3.5,label=gpu-linux/)) -* Mac CPU-only: [Python 2](https://ci.tensorflow.org/view/Nightly/job/nightly-matrix-cpu/TF_BUILD_IS_OPT=OPT,TF_BUILD_IS_PIP=PIP,TF_BUILD_PYTHON_VERSION=PYTHON2,label=mac-slave/lastSuccessfulBuild/artifact/pip_test/whl/tensorflow-1.1.0rc1-py2-none-any.whl) ([build history](https://ci.tensorflow.org/view/Nightly/job/nightly-matrix-cpu/TF_BUILD_IS_OPT=OPT,TF_BUILD_IS_PIP=PIP,TF_BUILD_PYTHON_VERSION=PYTHON2,label=mac-slave/)) / [Python 3](https://ci.tensorflow.org/view/Nightly/job/nightly-matrix-cpu/TF_BUILD_IS_OPT=OPT,TF_BUILD_IS_PIP=PIP,TF_BUILD_PYTHON_VERSION=PYTHON3,label=mac-slave/lastSuccessfulBuild/artifact/pip_test/whl/tensorflow-1.1.0rc1-py3-none-any.whl) ([build history](https://ci.tensorflow.org/view/Nightly/job/nightly-matrix-cpu/TF_BUILD_IS_OPT=OPT,TF_BUILD_IS_PIP=PIP,TF_BUILD_PYTHON_VERSION=PYTHON3,label=mac-slave/)) -* Mac GPU: [Python 2](https://ci.tensorflow.org/view/Nightly/job/nightly-matrix-mac-gpu/TF_BUILD_IS_OPT=OPT,TF_BUILD_IS_PIP=PIP,TF_BUILD_PYTHON_VERSION=PYTHON2,label=gpu-mac/lastSuccessfulBuild/artifact/pip_test/whl/tensorflow_gpu-1.1.0rc1-py2-none-any.whl) ([build history](https://ci.tensorflow.org/view/Nightly/job/nightly-matrix-mac-gpu/TF_BUILD_IS_OPT=OPT,TF_BUILD_IS_PIP=PIP,TF_BUILD_PYTHON_VERSION=PYTHON2,label=gpu-mac/)) / [Python 3](https://ci.tensorflow.org/view/Nightly/job/nightly-matrix-mac-gpu/TF_BUILD_IS_OPT=OPT,TF_BUILD_IS_PIP=PIP,TF_BUILD_PYTHON_VERSION=PYTHON3,label=gpu-mac/lastSuccessfulBuild/artifact/pip_test/whl/tensorflow_gpu-1.1.0rc1-py3-none-any.whl) ([build history](https://ci.tensorflow.org/view/Nightly/job/nightly-matrix-mac-gpu/TF_BUILD_IS_OPT=OPT,TF_BUILD_IS_PIP=PIP,TF_BUILD_PYTHON_VERSION=PYTHON3,label=gpu-mac/)) -* Windows CPU-only: [Python 3.5 64-bit](https://ci.tensorflow.org/view/Nightly/job/nightly-win/DEVICE=cpu,OS=windows/lastSuccessfulBuild/artifact/cmake_build/tf_python/dist/tensorflow-1.1.0rc1-cp35-cp35m-win_amd64.whl) ([build history](https://ci.tensorflow.org/view/Nightly/job/nightly-win/DEVICE=cpu,OS=windows/)) -* Windows GPU: [Python 3.5 64-bit](https://ci.tensorflow.org/view/Nightly/job/nightly-win/DEVICE=gpu,OS=windows/lastSuccessfulBuild/artifact/cmake_build/tf_python/dist/tensorflow_gpu-1.1.0rc1-cp35-cp35m-win_amd64.whl) ([build history](https://ci.tensorflow.org/view/Nightly/job/nightly-win/DEVICE=gpu,OS=windows/)) +**Nightly pip packages** +* We are pleased to announce that TensorFlow now offers nightly pip packages +under the [tf-nightly](https://pypi.python.org/pypi/tf-nightly) project on pypi. +Simply run `pip install tf-nightly` in a clean environment to install the nightly +tensorflow build. We currently only support CPU-only packages on Linux and Mac. +GPU packages on all platforms and Windows CPU-only packages will arrive soon! + + +**Individual whl files** +* Linux CPU-only: [Python 2](https://ci.tensorflow.org/view/Nightly/job/nightly-matrix-cpu/TF_BUILD_IS_OPT=OPT,TF_BUILD_IS_PIP=PIP,TF_BUILD_PYTHON_VERSION=PYTHON2,label=cpu-slave/lastSuccessfulBuild/artifact/pip_test/whl/tensorflow-1.3.0-cp27-none-linux_x86_64.whl) ([build history](https://ci.tensorflow.org/view/Nightly/job/nightly-matrix-cpu/TF_BUILD_IS_OPT=OPT,TF_BUILD_IS_PIP=PIP,TF_BUILD_PYTHON_VERSION=PYTHON2,label=cpu-slave)) / [Python 3.4](https://ci.tensorflow.org/view/Nightly/job/nightly-matrix-cpu/TF_BUILD_IS_OPT=OPT,TF_BUILD_IS_PIP=PIP,TF_BUILD_PYTHON_VERSION=PYTHON3,label=cpu-slave/lastSuccessfulBuild/artifact/pip_test/whl/tensorflow-1.3.0-cp34-cp34m-linux_x86_64.whl) ([build history](https://ci.tensorflow.org/view/Nightly/job/nightly-matrix-cpu/TF_BUILD_IS_OPT=OPT,TF_BUILD_IS_PIP=PIP,TF_BUILD_PYTHON_VERSION=PYTHON3,label=cpu-slave/)) / [Python 3.5](https://ci.tensorflow.org/view/Nightly/job/nightly-python35-linux-cpu/lastSuccessfulBuild/artifact/pip_test/whl/tensorflow-1.3.0-cp35-cp35m-linux_x86_64.whl) ([build history](https://ci.tensorflow.org/view/Nightly/job/nightly-python35-linux-cpu/)) +* Linux GPU: [Python 2](https://ci.tensorflow.org/view/Nightly/job/nightly-matrix-linux-gpu/TF_BUILD_IS_OPT=OPT,TF_BUILD_IS_PIP=PIP,TF_BUILD_PYTHON_VERSION=PYTHON2,label=gpu-linux/lastSuccessfulBuild/artifact/pip_test/whl/tensorflow_gpu-1.3.0-cp27-none-linux_x86_64.whl) ([build history](https://ci.tensorflow.org/view/Nightly/job/nightly-matrix-linux-gpu/TF_BUILD_IS_OPT=OPT,TF_BUILD_IS_PIP=PIP,TF_BUILD_PYTHON_VERSION=PYTHON2,label=gpu-linux/)) / [Python 3.4](https://ci.tensorflow.org/view/Nightly/job/nightly-matrix-linux-gpu/TF_BUILD_IS_OPT=OPT,TF_BUILD_IS_PIP=PIP,TF_BUILD_PYTHON_VERSION=PYTHON3,label=gpu-linux/lastSuccessfulBuild/artifact/pip_test/whl/tensorflow_gpu-1.3.0-cp34-cp34m-linux_x86_64.whl) ([build history](https://ci.tensorflow.org/view/Nightly/job/nightly-matrix-linux-gpu/TF_BUILD_IS_OPT=OPT,TF_BUILD_IS_PIP=PIP,TF_BUILD_PYTHON_VERSION=PYTHON3,label=gpu-linux/)) / [Python 3.5](https://ci.tensorflow.org/view/Nightly/job/nightly-matrix-linux-gpu/TF_BUILD_IS_OPT=OPT,TF_BUILD_IS_PIP=PIP,TF_BUILD_PYTHON_VERSION=PYTHON3.5,label=gpu-linux/lastSuccessfulBuild/artifact/pip_test/whl/tensorflow_gpu-1.3.0-cp35-cp35m-linux_x86_64.whl) ([build history](https://ci.tensorflow.org/view/Nightly/job/nightly-matrix-linux-gpu/TF_BUILD_IS_OPT=OPT,TF_BUILD_IS_PIP=PIP,TF_BUILD_PYTHON_VERSION=PYTHON3.5,label=gpu-linux/)) +* Mac CPU-only: [Python 2](https://ci.tensorflow.org/view/Nightly/job/nightly-matrix-cpu/TF_BUILD_IS_OPT=OPT,TF_BUILD_IS_PIP=PIP,TF_BUILD_PYTHON_VERSION=PYTHON2,label=mac-slave/lastSuccessfulBuild/artifact/pip_test/whl/tensorflow-1.3.0-py2-none-any.whl) ([build history](https://ci.tensorflow.org/view/Nightly/job/nightly-matrix-cpu/TF_BUILD_IS_OPT=OPT,TF_BUILD_IS_PIP=PIP,TF_BUILD_PYTHON_VERSION=PYTHON2,label=mac-slave/)) / [Python 3](https://ci.tensorflow.org/view/Nightly/job/nightly-matrix-cpu/TF_BUILD_IS_OPT=OPT,TF_BUILD_IS_PIP=PIP,TF_BUILD_PYTHON_VERSION=PYTHON3,label=mac-slave/lastSuccessfulBuild/artifact/pip_test/whl/tensorflow-1.3.0-py3-none-any.whl) ([build history](https://ci.tensorflow.org/view/Nightly/job/nightly-matrix-cpu/TF_BUILD_IS_OPT=OPT,TF_BUILD_IS_PIP=PIP,TF_BUILD_PYTHON_VERSION=PYTHON3,label=mac-slave/)) +* Windows CPU-only: [Python 3.5 64-bit](https://ci.tensorflow.org/view/Nightly/job/nightly-win/M=windows,PY=35/lastSuccessfulBuild/artifact/cmake_build/tf_python/dist/tensorflow-1.3.0-cp35-cp35m-win_amd64.whl) ([build history](https://ci.tensorflow.org/view/Nightly/job/nightly-win/M=windows,PY=35/)) / [Python 3.6 64-bit](https://ci.tensorflow.org/view/Nightly/job/nightly-win/M=windows,PY=36/lastSuccessfulBuild/artifact/cmake_build/tf_python/dist/tensorflow-1.3.0-cp36-cp36m-win_amd64.whl) ([build history](https://ci.tensorflow.org/view/Nightly/job/nightly-win/M=windows,PY=36/)) +* Windows GPU: [Python 3.5 64-bit](https://ci.tensorflow.org/view/Nightly/job/nightly-win/M=windows-gpu,PY=35/lastSuccessfulBuild/artifact/cmake_build/tf_python/dist/tensorflow_gpu-1.3.0-cp35-cp35m-win_amd64.whl) ([build history](https://ci.tensorflow.org/view/Nightly/job/nightly-win/M=windows-gpu,PY=35/)) / [Python 3.6 64-bit](https://ci.tensorflow.org/view/Nightly/job/nightly-win/M=windows-gpu,PY=36/lastSuccessfulBuild/artifact/cmake_build/tf_python/dist/tensorflow_gpu-1.3.0-cp36-cp36m-win_amd64.whl) ([build history](https://ci.tensorflow.org/view/Nightly/job/nightly-win/M=windows-gpu,PY=36/)) * Android: [demo APK](https://ci.tensorflow.org/view/Nightly/job/nightly-android/lastSuccessfulBuild/artifact/out/tensorflow_demo.apk), [native libs](http://ci.tensorflow.org/view/Nightly/job/nightly-android/lastSuccessfulBuild/artifact/out/native/) ([build history](https://ci.tensorflow.org/view/Nightly/job/nightly-android/)) @@ -52,19 +62,24 @@ $ python >>> hello = tf.constant('Hello, TensorFlow!') >>> sess = tf.Session() >>> sess.run(hello) -Hello, TensorFlow! +'Hello, TensorFlow!' >>> a = tf.constant(10) >>> b = tf.constant(32) ->>> sess.run(a+b) +>>> sess.run(a + b) 42 ->>> +>>> sess.close() ``` ## For more information -* [TensorFlow website](http://tensorflow.org) -* [TensorFlow whitepaper](http://download.tensorflow.org/paper/whitepaper2015.pdf) +* [TensorFlow website](https://www.tensorflow.org) +* [TensorFlow White Papers](https://www.tensorflow.org/about/bib) * [TensorFlow Model Zoo](https://github.com/tensorflow/models) * [TensorFlow MOOC on Udacity](https://www.udacity.com/course/deep-learning--ud730) +* [TensorFlow course at Stanford](https://web.stanford.edu/class/cs20si) + +Learn more about the TensorFlow community at the [community page of tensorflow.org](https://www.tensorflow.org/community) for a few ways to participate. + +## License -The TensorFlow community has created amazing things with TensorFlow, please see the [resources section of tensorflow.org](https://www.tensorflow.org/about/#community) for an incomplete list. +[Apache License 2.0](LICENSE) diff --git a/RELEASE.md b/RELEASE.md index ebb5c28451dab19b14f577ef86252e786266fb4d..3d497dbaa965d2cf239cab8360109bf5804b6f6e 100644 --- a/RELEASE.md +++ b/RELEASE.md @@ -1,3 +1,304 @@ +# Release 1.4.0 + +## Major Features And Improvements + +## Bug Fixes and Other Changes +* `tf.nn.rnn_cell.DropoutWrapper` is now more careful about dropping out LSTM + states. Specifically, it no longer ever drops the `c` (memory) state of an + `LSTMStateTuple`. The new behavior leads to proper dropout behavior + for LSTMs and stacked LSTMs. This bug fix follows recommendations from + published literature, but is a behavioral change. State dropout behavior + may be customized via the new `dropout_state_filter_visitor` argument. +* Removed `tf.contrib.training.python_input`. The same behavior, in a more + flexible and reproducible package, is available via the new + `tf.contrib.data.Dataset.from_generator` method! + +# Release 1.3.0 + +See also [TensorBoard 0.1.4](https://github.com/tensorflow/tensorboard/releases/tag/0.1.4) release notes. + +## Major Features and Improvements +* Added canned estimators to Tensorflow library. List of added estimators: + * `DNNClassifier` + * `DNNRegressor` + * `LinearClassifier` + * `LinearRegressor` + * `DNNLinearCombinedClassifier` + * `DNNLinearCombinedRegressor`. +* All our prebuilt binaries have been built with cuDNN 6. We anticipate releasing TensorFlow 1.4 with cuDNN 7. +* `import tensorflow` now goes much faster. +* Adds a file cache to the GCS filesystem with configurable max staleness for file contents. This permits caching of file contents across close/open boundaries. +* Added an axis parameter to `tf.gather`. +* Added a `constant_values` keyword argument to `tf.pad`. +* Adds `Dataset.interleave` transformation. +* Add `ConcatenateDataset` to concatenate two datasets. +* Added Mobilenet support to TensorFlow for Poets training script. +* Adds a block cache to the GCS filesystem with configurable block size and count. +* SinhArcSinh bijector added. +* Added `Dataset.list_files` API. +* Introduces new operations and Python bindings for the Cloud TPU. +* Adding TensorFlow-iOS CocoaPod for symmetry with tensorflow-android. +* Introduces base implementations of ClusterResolvers. +* Unify memory representations of TensorShape and PartialTensorShape. As a consequence, tensors now have a maximum of 254 dimensions, not 255. +* Changed references to LIBXSMM to use version 1.8.1. +* TensorFlow Debugger (tfdbg): + * Display summaries of numeric tensor values with the `-s` flag to command `print_tensor` or `pt`. + * Display feed values with the `print_feed` or `pf` command and clickable links in the curses UI. + * Runtime profiler at the op level and the Python source line level with the `run -p` command. +* Initial release of the statistical distribution library `tf.distributions`. +* GPU kernels and speed improvements for unary `tf.where` and `tf.nn.top_k`. +* Monotonic Attention wrappers added to `tf.contrib.seq2seq`. +* Added `tf.contrib.signal`, a library for signal processing primitives. +* Added `tf.contrib.resampler`, containing CPU and GPU ops for differentiable resampling of images. + +## Breaking Changes to the API +* `tf.RewriterConfig` was removed from the Python API after being available in 1.2 release candidates (it was never in an actual release). Graph rewriting is still available, just not as `tf.RewriterConfig`. Instead add an explicit import. +* Breaking change to `tf.contrib.data.Dataset` APIs that expect a nested structure. Lists are now converted to `tf.Tensor` implicitly. You may need to change uses of lists to tuples in existing code. In addition, dicts are now supported as a nested structure. + +## Changes to contrib APIs +* 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 + +## Known Issues +* Tensorflow_gpu compilation fails with Bazel 0.5.3. + +## Bug Fixes and Other Changes +* Fixes `strides` and `begin` dtype mismatch when slicing using int64 Tensor index in python. +* Improved convolution padding documentation. +* Add a tag constant, gpu, to present graph with GPU support. +* `saved_model.utils` now support SparseTensors transparently. +* A more efficient implementation of non-max suppression. +* Add support for the shrinkage-type L2 to FtrlOptimizer in addition to the online L2 it already supports. +* Fix negative variance in moments calculation. +* Expand UniqueOp Benchmark Tests to cover more collision cases. +* Improves stability of GCS filesystem on Mac. +* Add time estimation to HloCostAnalysis. +* Fixed the bug in Estimator that params in constructor was not a deepcopy of the user provided one. This bugs inadvertently enabled user to mutate the params after the creation of Estimator, leading to potentially undefined behavior. +* Added None check for save_path in `saver.restore`. +* Register devices under their legacy names in device_mgr to ease the transition to clusterspec-propagated configurations. +* VectorExponential added to distributions. +* Add a bitwise module with bitwise_and, bitwise_or, bitwise_xor, and invert functions. +* Add fixed-grid ODE integration routines. +* Allow passing bounds to ScipyOptimizerInterface. +* Correctness fixes for fft_length parameter to `tf.spectral.rfft` & `tf.spectral.irfft`. +* Exported model signatures using the 'predict' method will no longer have their input and output keys silently ignored and rewritten to 'inputs' and 'outputs'. If a model was exported with different names before 1.2, and is now served with tensorflow/serving, it will accept requests using 'inputs' and 'outputs'. Starting at 1.2, such a model will accept the keys specified during export. Therefore, inference requests using 'inputs' and 'outputs' may start to fail. To fix this, either update any inference clients to send requests with the actual input and output keys used by the trainer code, or conversely, update the trainer code to name the input and output Tensors 'inputs' and 'outputs', respectively. Signatures using the 'classify' and 'regress' methods are not affected by this change; they will continue to standardize their input and output keys as before. +* Add in-memory caching to the Dataset API. +* Set default end_of_sequence variable in datasets iterators to false. +* [Performance] Increase performance of `tf.layers.conv2d` when setting use_bias=True by 2x by using nn.bias_add. +* Update iOS examples to use CocoaPods, and moved to tensorflow/examples/ios. +* Adds a family= attribute in `tf.summary` ops to allow controlling the tab name used in Tensorboard for organizing summaries. +* When GPU is configured, do not require --config=cuda, instead, automatically build for GPU if this is requested in the configure script. +* Fix incorrect sampling of small probabilities in CPU/GPU multinomial. +* Add a list_devices() API on sessions to list devices within a cluster. Additionally, this change augment the ListDevices master API to support specifying a session. +* Allow uses of over-parameterized separable convolution. +* TensorForest multi-regression bug fix. +* Framework now supports armv7, cocoapods.org now displays correct page. +* Script to create iOS framework for CocoaPods. +* Android releases of TensorFlow are now pushed to jcenter for easier integration into apps. See https://github.com/tensorflow/tensorflow/blob/master/tensorflow/contrib/android/README.md for more details. +* TensorFlow Debugger (tfdbg): + * Fixed a bug that prevented tfdbg from functioning with multi-GPU setups. + * Fixed a bug that prevented tfdbg from working with `tf.Session.make_callable`. + +## Thanks to our Contributors + +This release contains contributions from many people at Google, as well as: + +4F2E4A2E, Adriano Carmezim, Adrià Arrufat, Alan Yee, Alex Lattas, Alex Rothberg, +Alexandr Baranezky, Ali Siddiqui, Andreas Solleder, Andrei Costinescu, Andrew Hundt, +Androbin, Andy Kernahan, Anish Shah, Anthony Platanios, Arvinds-Ds, b1rd, Baptiste +Arnaud, Ben Mabey, Benedikt Linse, Beomsu Kim, Bo Wang, Boyuan Deng, Brett Koonce, +Bruno Rosa, Carl Thomé, Changming Sun, Chase Roberts, Chirag Bhatia, Chris Antaki, +Chris Hoyean Song, Chris Tava, Christos Nikolaou, Croath Liu, cxx, Czxck001, Daniel +Ylitalo, Danny Goodman, Darren Garvey, David Brailovsky, David Norman, DavidNorman, +davidpham87, ddurham2, Dhruv, DimanNe, Drew Hintz, Dustin Tran, Earthson Lu, ethiraj, +Fabian Winnen, Fei Sun, Freedom" Koan-Sin Tan, Fritz Obermeyer, Gao, Xiang, Gautam, +Guenther Schmuelling, Gyu-Ho Lee, Hauke Brammer, horance, Humanity123, J Alammar, +Jayeol Chun, Jeroen BéDorf, Jianfei Wang, jiefangxuanyan, Jing Jun Yin, Joan Puigcerver, +Joel Hestness, Johannes Mayer, John Lawson, Johnson145, Jon Malmaud, Jonathan Alvarez-Gutierrez, +Juang, Yi-Lin, Julian Viereck, Kaarthik Sivashanmugam, Karl Lessard, karl@kubx.ca, Kevin +Carbone, Kevin Van Der Burgt, Kongsea, ksellesk, lanhin, Lef Ioannidis, Liangliang He, +Louis Tiao, Luke Iwanski, LáSzló Csomor, magixsno, Mahmoud Abuzaina, Marcel Hlopko, Mark +Neumann, Maxwell Paul Brickner, mdfaijul, MichaëL Defferrard, Michał JastrzęBski, Michele +Colombo, Mike Brodie, Mosnoi Ion, mouradmourafiq, myPrecious, Nayana Thorat, +Neeraj Kashyap, Nelson Liu, Niranjan Hasabnis, Olivier Moindrot, orome, Pankaj Gupta, Paul +Van Eck, peeyush18, Peng Yu, Pierre, preciousdp11, qjivy, Raingo, raoqiyu, ribx, Richard S. +Imaoka, Rishabh Patel, Robert Walecki, Rockford Wei, Ryan Kung, Sahil Dua, Sandip Giri, Sayed +Hadi Hashemi, sgt101, Shitian Ni, Shuolongbj, Siim PõDer, Simon Perkins, sj6077, SOLARIS, +Spotlight0xff, Steffen Eberbach, Stephen Fox, superryanguo, Sven Mayer, Tapan Prakash, +Tiago Morais Morgado, Till Hoffmann, Tj Rana, Vadim Markovtsev, vhasanov, Wei Wu, +windead, Yan (Asta) Li, Yan Chen, Yann Henon, Yi Wang, Yong Tang, yorkie, Yuan (Terry) +Tang, Yuxin Wu, zhengjiajin, zhongzyd, 黄璞 + +We are also grateful to all who filed issues or helped resolve them, asked and +answered questions, and were part of inspiring discussions. + +# Release 1.2.1 + +## Bug Fixes and Other Changes +* Updating markdown version required to >= 2.6.8. +* Support tensors as dropout rates again, by removing the min(max(..)) + +# Release 1.2.0 + +## Major Features and Improvements +* Python 3.6 support on Windows. +* Added `tf.layers.conv3d_transpose` layer for spatio temporal deconvolution. +* Added `tf.Session.make_callable()`, which provides a lower overhead means of running a similar step multiple times. +* Added libverbs-based RDMA support to contrib (courtesy @junshi15 from Yahoo). +* Bring `tf.feature_column.*` into the API. Non-deprecated functionality from `tf.contrib.layers.*` is moved to `tf.feature_column.*` with cosmetic changes. +* `RNNCell` objects now subclass `tf.layers.Layer`. The strictness described + in the TensorFlow 1.1 release is gone: The first time an RNNCell is used, + it caches its scope. All future uses of the RNNCell will reuse variables from + that same scope. This is a breaking change from the behavior of RNNCells + in TensorFlow versions <= 1.0.1. TensorFlow 1.1 had checks in place to + ensure old code works correctly with the new semantics; this version + allows more flexible uses of RNNCell but can lead to subtle errors if + using code meant for TensorFlow <= 1.0.1. For example, writing: + `MultiRNNCell([lstm] * 5)` will now build a 5-layer LSTM stack where each + layer shares the **same** parameters. To get 5 layers each with their own + parameters, write: `MultiRNNCell([LSTMCell(...) for _ in range(5)])`. + If at all unsure, first test your code with TF 1.1; ensure it raises no + errors, and then upgrade to TF 1.2. +* RNNCells' variable names have been renamed for consistency with Keras layers. + Specifically, the previous variable names "weights" and "biases" have + been changed to "kernel" and "bias", respectively. + This may cause backward incompatibility with regard to your old + checkpoints containing such RNN cells, in which case you can use the tool + [checkpoint_convert script](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/contrib/rnn/python/tools/checkpoint_convert.py) + to convert the variable names in your old checkpoints. +* Many of the RNN functions and classes that were in the `tf.nn` namespace + before the 1.0 release and which were moved to `tf.contrib.rnn` have now + been moved back to the core namespace. This includes + `RNNCell`, `LSTMCell`, `GRUCell`, and a number of other cells. These + now reside in `tf.nn.rnn_cell` (with aliases in `tf.contrib.rnn` for backwards + compatibility). The original `tf.nn.rnn` function is now `tf.nn.static_rnn`, + and the bidirectional static and state saving static rnn functions are also + now back in the `tf.nn` namespace. + + Notable exceptions are the `EmbeddingWrapper`, `InputProjectionWrapper` and + `OutputProjectionWrapper`, which will slowly be moved to deprecation + in `tf.contrib.rnn`. These are inefficient wrappers that should often + be replaced by calling `embedding_lookup` or `layers.dense` as pre- or post- + processing of the rnn. For RNN decoding, this functionality has been replaced + with an alternative API in `tf.contrib.seq2seq`. +* Intel MKL Integration (https://software.intel.com/en-us/articles/tensorflow-optimizations-on-modern-intel-architecture). Intel developed a number of + optimized deep learning primitives: In addition to matrix multiplication and + convolution, these building blocks include: + Direct batched convolution + Pooling: maximum, minimum, average + Normalization: LRN, batch normalization + Activation: rectified linear unit (ReLU) + Data manipulation: multi-dimensional transposition (conversion), split, + concat, sum and scale. +* TensorForest Estimator now supports SavedModel export for serving. +* Support client-provided ClusterSpec's and propagate them to all workers to enable the creation of dynamic TensorFlow clusters. +* TensorFlow C library now available for Windows. +* We released a new open-source version of TensorBoard. +* [`SavedModel CLI`](https://www.tensorflow.org/versions/master/programmers_guide/saved_model_cli) tool available to inspect and execute MetaGraph in SavedModel +* Android releases of TensorFlow are now pushed to jcenter for easier + integration into apps. See + https://github.com/tensorflow/tensorflow/blob/master/tensorflow/contrib/android/README.md + for more details. + +## Deprecations + +* TensorFlow 1.2 may be the last time we build with cuDNN 5.1. Starting with + TensorFlow 1.3, we will try to build all our prebuilt binaries with cuDNN 6.0. + While we will try to keep our source code compatible with cuDNN 5.1, it will + be best effort. + +## Breaking Changes to the API +* `org.tensorflow.contrib.android.TensorFlowInferenceInterface` now throws exceptions where possible and has simplified method signatures. + +## Changes to contrib APIs +* Added `tf.contrib.util.create_example`. +* Added bilinear interpolation to `tf.contrib.image`. +* Add `tf.contrib.stateless` for random ops with custom seed control. +* MultivariateNormalFullCovariance added to contrib/distributions/ +* tensorflow/contrib/rnn undergoes RNN cell variable renaming for + consistency with Keras layers. Specifically, the previous variable names + "weights" and "biases" are changed to "kernel" and "bias", respectively. + This may cause backward incompatibility with regard to your old + checkpoints containing such RNN cells, in which case you can use the + [checkpoint_convert script](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/contrib/rnn/python/tools/checkpoint_convert.py) + to convert the variable names in your old checkpoints. +* Added `tf.contrib.kernel_methods` module with Ops and estimators for primal + (explicit) kernel methods in TensorFlow. + +## Bug Fixes and Other Changes +* In python, `Operation.get_attr` on type attributes returns the Python DType + version of the type to match expected get_attr documentation rather than the + protobuf enum. +* tensorflow/contrib/rnn undergoes RNN cell variable renaming for + consistency with Keras layers. Specifically, the previous variable names + "weights" and "biases" are changed to "kernel" and "bias", respectively. +* Changed MIN_SDK version to 8.0 when building iOS libraries. +* Fixed LIBXSMM integration. +* Make decode_jpeg/decode_png/decode_gif handle all formats, since users frequently try to decode an image as the wrong type. +* Improve implicit broadcasting lowering. +* Improving stability of GCS/Bigquery clients by a faster retrying of stale transmissions. +* Remove OpKernelConstruction::op_def() as part of minimizing proto dependencies. +* VectorLaplaceDiag distribution added. +* Android demo no longer requires libtensorflow_demo.so to run (libtensorflow_inference.so still required) +* Added `categorical_column_with_vocabulary_file`. +* Introduce ops for batching/unbatching tensors across Session::Run() calls. +* Add tf.log_sigmoid(x) = tf.log(tf.sigmoid(x)) = -tf.nn.softplus(-x). +* Changed hooks lists to immutable tuples, and now allow any iterable for the associated arguments. +* Introduce TFDecorator. +* Added an Mfcc op for speech feature generation. +* Improved DirectSession::Run() overhead and error checking. Feeding a value of the wrong type will now synchronously raise an INVALID_ARGUMENT error instead of asynchronously raising an INTERNAL error. Code that depends on the (undefined) behavior when feeding a tensor of the wrong type may need to be updated. +* Added unreduced NONE, and reduced MEAN options for losses. Removed "WEIGHTED_" prefix from other Reduction constants. +* assertAllClose now handles dicts. +* Added Gmock matcher for HloInstructions. +* Add var name to errors on variable restore. +* Added an AudioSpectrogram op for audio feature generation. +* Added `reduction` arg to losses. +* `tf.placeholder` can represent scalar shapes and partially known. +* Remove estimator_spec(mode) argument. +* Added an AudioSpectrogram op for audio feature generation. +* TensorBoard disables all runs by default if there are more than 40 runs. +* Removed old doc generator code. +* GCS file system integration now supports domain buckets, e.g gs://bucket.domain.com/path. +* Add `tf.summary.text` for outputting text to TensorBoard. +* The "run" command of tfdbg's command-line interface now supports filtering of tensors by node name, op type and tensor dtype. +* `tf.string_to_number` now supports int64 and float64 outputs. + +## Thanks to our Contributors + +This release contains contributions from many people at Google, as well as: + +4F2E4A2E, Aaron Schumacher, Abhi Agg, admcrae, Adriano Carmezim, Adrià Arrufat, +agramesh1, Akimitsu Seo, Alan Mosca, Alex Egg, Alex Rothberg, Alexander Heinecke, +Alexander Matyasko, Alexandr Baranezky, Alexandre Caulier, Ali Siddiqui, Anand Venkat, +Andrew Hundt, Androbin, Anmol Sharma, Arie, Arno Leist, Arron Cao, AuréLien Geron, Bairen Yi, +Beomsu Kim, Carl Thomé, cfperez, Changming Sun, Corey Wharton, critiqjo, Dalei Li, Daniel +Rasmussen, Daniel Trebbien, DaríO Hereñú, David Eng, David Norman, David Y. Zhang, Davy Song, ddurham2, +Deepak Subburam, Dmytro Kyrychuk, Dominic Rossi, Dominik SchlöSser, Dustin Tran, +Eduardo Pinho, Egil Martinsson, Elliot Saba, Eric Bigelow, Erik Smistad, Evan Klitzke, +Fabrizio Milo, Falcon Dai, Fei Gao, FloopCZ, Fung Lam, Gautam, GBLin5566, Greg Peatfield, +Gu Wang, Guenther Schmuelling, Hans Pabst, Harun Gunaydin, Huaizheng, Ido Shamay, Ikaro +Silva, Ilya Edrenkin, Immexxx, James Mishra, Jamie Cooke, Jay Young, Jayaram Bobba, +Jianfei Wang, jinghua2, Joey Meyer, John Maidens, Jonghoon Jin, Julian Villella, +Jun Kim, Jun Shi, Junwei Pan, jyegerlehner, Karan Desai, Karel Van De Plassche, +Kb Sriram, KhabarlakKonstantin, Koan-Sin Tan, krivard, Kwotsin, Leandro Gracia Gil, +Li Chen, Liangliang He, Louie Helm, lspvic, Luiz Henrique Soares, LáSzló Csomor, +Mark Wong, Mathew Wicks, Matthew Rahtz, Maxwell Paul Brickner, Michael Hofmann, Miguel +Flores Ruiz De Eguino, MikeTam1021, Mortada Mehyar, Mycosynth, Namnamseo, +Nate Harada, Neven Miculinic, Nghia Tran, Nick Lyu, Niranjan Hasabnis, Nishidha, Oleksii +Kuchaiev, Oyesh Mann Singh, Panmari, Patrick, Paul Van Eck, Piyush Chaudhary, Quim Llimona, +Raingo, Richard Davies, Ruben Vereecken, Sahit Chintalapudi, Sam Abrahams, Santiago Castro, +Scott Sievert, Sean O'Keefe, Sebastian Schlecht, Shane, Shubhankar Deshpande, Spencer Schaber, +Sunyeop Lee, t13m, td2014, Thomas H. P. Andersen, Toby Petty, Umang Mehta, +Vadim Markovtsev, Valentin Iovene, Vincent Zhao, Vit Stepanovs, Vivek Rane, Vu Pham, wannabesrevenge, +weipingpku, wuhaixutab, wydwww, Xiang Gao, Xiaolin Lin, xiaoyaozhuzi, Yaroslav Bulatov, Yi Liu, +Yoshihiro Sugi, Yuan (Terry) Tang, Yuming Wang, Yuxin Wu, Zader Zheng, Zhaojun Zhang, zhengjiajin, +ZhipengShen, Ziming Dong, zjj2wry + +We are also grateful to all who filed issues or helped resolve them, asked and +answered questions, and were part of inspiring discussions. + # Release 1.1.0 ## Major Features and Improvements @@ -16,6 +317,7 @@ * New navigation bar in Curses-based UI * NodeStepper (command `invoke_stepper`) now uses intermediate tensor dumps. It also uses `TensorHandles` as direct feeds during successive `cont` calls for improved performance and reduced memory consumption. * Initial release of installation guides for Java, C, and Go. +* Added Text Dashboard to TensorBoard. ## Deprecations @@ -71,6 +373,8 @@ * Command history now persists across runs. * Bug fix in graph validation related to `tf.while_loops`. * Java Maven fixes for bugs with Windows installation. +* Backport fixes and improvements from external keras. +* Keras config file handling fix. ## Thanks to our Contributors diff --git a/WORKSPACE b/WORKSPACE index cab8389a55ccfeddb9dc077c9b999edbe775f25d..a0fe67bf3189c1156c524aced5210e466e1d8f12 100644 --- a/WORKSPACE +++ b/WORKSPACE @@ -2,11 +2,11 @@ workspace(name = "org_tensorflow") http_archive( name = "io_bazel_rules_closure", - sha256 = "60fc6977908f999b23ca65698c2bb70213403824a84f7904310b6000d78be9ce", - strip_prefix = "rules_closure-5ca1dab6df9ad02050f7ba4e816407f88690cf7d", + sha256 = "25f5399f18d8bf9ce435f85c6bbf671ec4820bc4396b3022cc5dc4bc66303609", + strip_prefix = "rules_closure-0.4.2", urls = [ - "http://bazel-mirror.storage.googleapis.com/github.com/bazelbuild/rules_closure/archive/5ca1dab6df9ad02050f7ba4e816407f88690cf7d.tar.gz", # 2017-02-03 - "https://github.com/bazelbuild/rules_closure/archive/5ca1dab6df9ad02050f7ba4e816407f88690cf7d.tar.gz", + "http://mirror.bazel.build/github.com/bazelbuild/rules_closure/archive/0.4.2.tar.gz", # 2017-08-29 + "https://github.com/bazelbuild/rules_closure/archive/0.4.2.tar.gz", ], ) @@ -20,7 +20,7 @@ load("//tensorflow:workspace.bzl", "tf_workspace") #android_sdk_repository( # name = "androidsdk", # api_level = 23, -# # Ensure that you have the build_tools_version below installed in the +# # Ensure that you have the build_tools_version below installed in the # # SDK manager as it updates periodically. # build_tools_version = "25.0.2", # # Replace with path to Android SDK on your system @@ -31,7 +31,10 @@ load("//tensorflow:workspace.bzl", "tf_workspace") #android_ndk_repository( # name="androidndk", # path="", -# # This needs to be 14 or higher to compile TensorFlow. +# # This needs to be 14 or higher to compile TensorFlow. +# # Please specify API level to >= 21 to build for 64-bit +# # archtectures or the Android NDK will automatically select biggest +# # API level that it supports without notice. # # Note that the NDK version is not the API level. # api_level=14) @@ -39,485 +42,51 @@ load("//tensorflow:workspace.bzl", "tf_workspace") tf_workspace() new_http_archive( - name = "inception5h", - build_file = "models.BUILD", - url = "https://storage.googleapis.com/download.tensorflow.org/models/inception5h.zip", - sha256 = "d13569f6a98159de37e92e9c8ec4dae8f674fbf475f69fe6199b514f756d4364" -) - -new_http_archive( - name = "mobile_multibox", - build_file = "models.BUILD", - url = "https://storage.googleapis.com/download.tensorflow.org/models/mobile_multibox_v1a.zip", - sha256 = "859edcddf84dddb974c36c36cfc1f74555148e9c9213dedacf1d6b613ad52b96" -) - -new_http_archive( - name = "stylize", - build_file = "models.BUILD", - url = "https://storage.googleapis.com/download.tensorflow.org/models/stylize_v1.zip", - sha256 = "3d374a730aef330424a356a8d4f04d8a54277c425e274ecb7d9c83aa912c6bfa" -) - -# TENSORBOARD_BOWER_AUTOGENERATED_BELOW_THIS_LINE_DO_NOT_EDIT - -new_http_archive( - name = "d3", - build_file = "bower.BUILD", - url = "https://github.com/mbostock-bower/d3-bower/archive/v3.5.15.tar.gz", - strip_prefix = "d3-bower-3.5.15", -) - -new_http_archive( - name = "dagre", - build_file = "bower.BUILD", - url = "https://github.com/cpettitt/dagre/archive/v0.7.4.tar.gz", - strip_prefix = "dagre-0.7.4", -) - -new_http_archive( - name = "es6_promise", - build_file = "bower.BUILD", - url = "https://github.com/components/es6-promise/archive/v2.1.0.tar.gz", - strip_prefix = "es6-promise-2.1.0", -) - -new_http_archive( - name = "font_roboto", - build_file = "bower.BUILD", - url = "https://github.com/polymerelements/font-roboto/archive/v1.0.1.tar.gz", - strip_prefix = "font-roboto-1.0.1", -) - -new_http_archive( - name = "graphlib", - build_file = "bower.BUILD", - url = "https://github.com/cpettitt/graphlib/archive/v1.0.7.tar.gz", - strip_prefix = "graphlib-1.0.7", -) - -new_http_archive( - name = "iron_a11y_announcer", - build_file = "bower.BUILD", - url = "https://github.com/polymerelements/iron-a11y-announcer/archive/v1.0.5.tar.gz", - strip_prefix = "iron-a11y-announcer-1.0.5", -) - -new_http_archive( - name = "iron_a11y_keys_behavior", - build_file = "bower.BUILD", - url = "https://github.com/polymerelements/iron-a11y-keys-behavior/archive/v1.1.8.tar.gz", - strip_prefix = "iron-a11y-keys-behavior-1.1.8", -) - -new_http_archive( - name = "iron_ajax", - build_file = "bower.BUILD", - url = "https://github.com/polymerelements/iron-ajax/archive/v1.2.0.tar.gz", - strip_prefix = "iron-ajax-1.2.0", -) - -new_http_archive( - name = "iron_autogrow_textarea", - build_file = "bower.BUILD", - url = "https://github.com/polymerelements/iron-autogrow-textarea/archive/v1.0.12.tar.gz", - strip_prefix = "iron-autogrow-textarea-1.0.12", -) - -new_http_archive( - name = "iron_behaviors", - build_file = "bower.BUILD", - url = "https://github.com/polymerelements/iron-behaviors/archive/v1.0.17.tar.gz", - strip_prefix = "iron-behaviors-1.0.17", -) - -new_http_archive( - name = "iron_checked_element_behavior", - build_file = "bower.BUILD", - url = "https://github.com/polymerelements/iron-checked-element-behavior/archive/v1.0.4.tar.gz", - strip_prefix = "iron-checked-element-behavior-1.0.4", -) - -new_http_archive( - name = "iron_collapse", - build_file = "bower.BUILD", - url = "https://github.com/polymerelements/iron-collapse/archive/v1.0.8.tar.gz", - strip_prefix = "iron-collapse-1.0.8", -) - -new_http_archive( - name = "iron_dropdown", - build_file = "bower.BUILD", - url = "https://github.com/polymerelements/iron-dropdown/archive/v1.4.0.tar.gz", - strip_prefix = "iron-dropdown-1.4.0", -) - -new_http_archive( - name = "iron_fit_behavior", - build_file = "bower.BUILD", - url = "https://github.com/polymerelements/iron-fit-behavior/archive/v1.2.5.tar.gz", - strip_prefix = "iron-fit-behavior-1.2.5", -) - -new_http_archive( - name = "iron_flex_layout", - build_file = "bower.BUILD", - url = "https://github.com/polymerelements/iron-flex-layout/archive/v1.3.0.tar.gz", - strip_prefix = "iron-flex-layout-1.3.0", -) - -new_http_archive( - name = "iron_form_element_behavior", - build_file = "bower.BUILD", - url = "https://github.com/polymerelements/iron-form-element-behavior/archive/v1.0.6.tar.gz", - strip_prefix = "iron-form-element-behavior-1.0.6", -) - -new_http_archive( - name = "iron_icon", - build_file = "bower.BUILD", - url = "https://github.com/polymerelements/iron-icon/archive/v1.0.11.tar.gz", - strip_prefix = "iron-icon-1.0.11", -) - -new_http_archive( - name = "iron_icons", - build_file = "bower.BUILD", - url = "https://github.com/polymerelements/iron-icons/archive/v1.1.3.tar.gz", - strip_prefix = "iron-icons-1.1.3", -) - -new_http_archive( - name = "iron_iconset_svg", - build_file = "bower.BUILD", - url = "https://github.com/polymerelements/iron-iconset-svg/archive/v1.1.0.tar.gz", - strip_prefix = "iron-iconset-svg-1.1.0", -) - -new_http_archive( - name = "iron_input", - build_file = "bower.BUILD", - url = "https://github.com/polymerelements/iron-input/archive/1.0.10.tar.gz", - strip_prefix = "iron-input-1.0.10", -) - -new_http_archive( - name = "iron_list", - build_file = "bower.BUILD", - url = "https://github.com/polymerelements/iron-list/archive/v1.3.9.tar.gz", - strip_prefix = "iron-list-1.3.9", -) - -new_http_archive( - name = "iron_menu_behavior", - build_file = "bower.BUILD", - url = "https://github.com/polymerelements/iron-menu-behavior/archive/v1.1.10.tar.gz", - strip_prefix = "iron-menu-behavior-1.1.10", -) - -new_http_archive( - name = "iron_meta", - build_file = "bower.BUILD", - url = "https://github.com/polymerelements/iron-meta/archive/v1.1.1.tar.gz", - strip_prefix = "iron-meta-1.1.1", -) - -new_http_archive( - name = "iron_overlay_behavior", - build_file = "bower.BUILD", - url = "https://github.com/polymerelements/iron-overlay-behavior/archive/v1.10.1.tar.gz", - strip_prefix = "iron-overlay-behavior-1.10.1", -) - -new_http_archive( - name = "iron_range_behavior", - build_file = "bower.BUILD", - url = "https://github.com/polymerelements/iron-range-behavior/archive/v1.0.4.tar.gz", - strip_prefix = "iron-range-behavior-1.0.4", -) - -new_http_archive( - name = "iron_resizable_behavior", - build_file = "bower.BUILD", - url = "https://github.com/polymerelements/iron-resizable-behavior/archive/v1.0.3.tar.gz", - strip_prefix = "iron-resizable-behavior-1.0.3", -) - -new_http_archive( - name = "iron_scroll_target_behavior", - build_file = "bower.BUILD", - url = "https://github.com/polymerelements/iron-scroll-target-behavior/archive/v1.0.3.tar.gz", - strip_prefix = "iron-scroll-target-behavior-1.0.3", -) - -new_http_archive( - name = "iron_selector", - build_file = "bower.BUILD", - url = "https://github.com/polymerelements/iron-selector/archive/v1.5.2.tar.gz", - strip_prefix = "iron-selector-1.5.2", -) - -new_http_archive( - name = "iron_validatable_behavior", - build_file = "bower.BUILD", - url = "https://github.com/polymerelements/iron-validatable-behavior/archive/v1.1.1.tar.gz", - strip_prefix = "iron-validatable-behavior-1.1.1", -) - -new_http_archive( - name = "lodash", - build_file = "bower.BUILD", - url = "https://github.com/lodash/lodash/archive/3.8.0.tar.gz", - strip_prefix = "lodash-3.8.0", -) - -new_http_archive( - name = "neon_animation", - build_file = "bower.BUILD", - url = "https://github.com/polymerelements/neon-animation/archive/v1.2.2.tar.gz", - strip_prefix = "neon-animation-1.2.2", -) - -http_file( - name = "numericjs_numeric_min_js", - url = "https://cdnjs.cloudflare.com/ajax/libs/numeric/1.2.6/numeric.min.js", -) - -new_http_archive( - name = "paper_behaviors", - build_file = "bower.BUILD", - url = "https://github.com/polymerelements/paper-behaviors/archive/v1.0.12.tar.gz", - strip_prefix = "paper-behaviors-1.0.12", -) - -new_http_archive( - name = "paper_button", - build_file = "bower.BUILD", - url = "https://github.com/polymerelements/paper-button/archive/v1.0.11.tar.gz", - strip_prefix = "paper-button-1.0.11", -) - -new_http_archive( - name = "paper_checkbox", - build_file = "bower.BUILD", - url = "https://github.com/polymerelements/paper-checkbox/archive/v1.4.0.tar.gz", - strip_prefix = "paper-checkbox-1.4.0", -) - -new_http_archive( - name = "paper_dialog", - build_file = "bower.BUILD", - url = "https://github.com/polymerelements/paper-dialog/archive/v1.0.4.tar.gz", - strip_prefix = "paper-dialog-1.0.4", -) - -new_http_archive( - name = "paper_dialog_behavior", - build_file = "bower.BUILD", - url = "https://github.com/polymerelements/paper-dialog-behavior/archive/v1.2.5.tar.gz", - strip_prefix = "paper-dialog-behavior-1.2.5", -) - -new_http_archive( - name = "paper_dialog_scrollable", - build_file = "bower.BUILD", - url = "https://github.com/polymerelements/paper-dialog-scrollable/archive/1.1.5.tar.gz", - strip_prefix = "paper-dialog-scrollable-1.1.5", -) - -new_http_archive( - name = "paper_dropdown_menu", - build_file = "bower.BUILD", - url = "https://github.com/polymerelements/paper-dropdown-menu/archive/v1.4.0.tar.gz", - strip_prefix = "paper-dropdown-menu-1.4.0", -) - -new_http_archive( - name = "paper_header_panel", - build_file = "bower.BUILD", - url = "https://github.com/polymerelements/paper-header-panel/archive/v1.1.4.tar.gz", - strip_prefix = "paper-header-panel-1.1.4", -) - -new_http_archive( - name = "paper_icon_button", - build_file = "bower.BUILD", - url = "https://github.com/polymerelements/paper-icon-button/archive/v1.1.3.tar.gz", - strip_prefix = "paper-icon-button-1.1.3", -) - -new_http_archive( - name = "paper_input", - build_file = "bower.BUILD", - url = "https://github.com/polymerelements/paper-input/archive/v1.1.18.tar.gz", - strip_prefix = "paper-input-1.1.18", -) - -new_http_archive( - name = "paper_item", - build_file = "bower.BUILD", - url = "https://github.com/polymerelements/paper-item/archive/v1.1.4.tar.gz", - strip_prefix = "paper-item-1.1.4", -) - -new_http_archive( - name = "paper_listbox", - build_file = "bower.BUILD", - url = "https://github.com/polymerelements/paper-listbox/archive/v1.1.2.tar.gz", - strip_prefix = "paper-listbox-1.1.2", -) - -new_http_archive( - name = "paper_material", - build_file = "bower.BUILD", - url = "https://github.com/polymerelements/paper-material/archive/v1.0.6.tar.gz", - strip_prefix = "paper-material-1.0.6", -) - -new_http_archive( - name = "paper_menu", - build_file = "bower.BUILD", - url = "https://github.com/polymerelements/paper-menu/archive/v1.2.2.tar.gz", - strip_prefix = "paper-menu-1.2.2", -) - -new_http_archive( - name = "paper_menu_button", - build_file = "bower.BUILD", - url = "https://github.com/polymerelements/paper-menu-button/archive/v1.5.1.tar.gz", - strip_prefix = "paper-menu-button-1.5.1", -) - -new_http_archive( - name = "paper_progress", - build_file = "bower.BUILD", - url = "https://github.com/polymerelements/paper-progress/archive/v1.0.9.tar.gz", - strip_prefix = "paper-progress-1.0.9", -) - -new_http_archive( - name = "paper_radio_button", - build_file = "bower.BUILD", - url = "https://github.com/polymerelements/paper-radio-button/archive/v1.1.2.tar.gz", - strip_prefix = "paper-radio-button-1.1.2", -) - -new_http_archive( - name = "paper_radio_group", - build_file = "bower.BUILD", - url = "https://github.com/polymerelements/paper-radio-group/archive/v1.0.9.tar.gz", - strip_prefix = "paper-radio-group-1.0.9", -) - -new_http_archive( - name = "paper_ripple", - build_file = "bower.BUILD", - url = "https://github.com/polymerelements/paper-ripple/archive/v1.0.5.tar.gz", - strip_prefix = "paper-ripple-1.0.5", -) - -new_http_archive( - name = "paper_slider", - build_file = "bower.BUILD", - url = "https://github.com/polymerelements/paper-slider/archive/v1.0.10.tar.gz", - strip_prefix = "paper-slider-1.0.10", -) - -new_http_archive( - name = "paper_spinner", - build_file = "bower.BUILD", - url = "https://github.com/polymerelements/paper-spinner/archive/v1.1.1.tar.gz", - strip_prefix = "paper-spinner-1.1.1", -) - -new_http_archive( - name = "paper_styles", - build_file = "bower.BUILD", - url = "https://github.com/polymerelements/paper-styles/archive/v1.1.4.tar.gz", - strip_prefix = "paper-styles-1.1.4", -) - -new_http_archive( - name = "paper_tabs", - build_file = "bower.BUILD", - url = "https://github.com/polymerelements/paper-tabs/archive/v1.7.0.tar.gz", - strip_prefix = "paper-tabs-1.7.0", -) - -new_http_archive( - name = "paper_toast", - build_file = "bower.BUILD", - url = "https://github.com/polymerelements/paper-toast/archive/v1.3.0.tar.gz", - strip_prefix = "paper-toast-1.3.0", -) - -new_http_archive( - name = "paper_toggle_button", - build_file = "bower.BUILD", - url = "https://github.com/polymerelements/paper-toggle-button/archive/v1.2.0.tar.gz", - strip_prefix = "paper-toggle-button-1.2.0", -) - -new_http_archive( - name = "paper_toolbar", - build_file = "bower.BUILD", - url = "https://github.com/polymerelements/paper-toolbar/archive/v1.1.4.tar.gz", - strip_prefix = "paper-toolbar-1.1.4", -) - -new_http_archive( - name = "paper_tooltip", - build_file = "bower.BUILD", - url = "https://github.com/polymerelements/paper-tooltip/archive/v1.1.2.tar.gz", - strip_prefix = "paper-tooltip-1.1.2", -) - -new_http_archive( - name = "plottable", - build_file = "bower.BUILD", - url = "https://github.com/palantir/plottable/archive/v1.16.1.tar.gz", - strip_prefix = "plottable-1.16.1", + name = "inception5h", + build_file = "models.BUILD", + sha256 = "d13569f6a98159de37e92e9c8ec4dae8f674fbf475f69fe6199b514f756d4364", + urls = [ + "http://storage.googleapis.com/download.tensorflow.org/models/inception5h.zip", + "http://download.tensorflow.org/models/inception5h.zip", + ], ) new_http_archive( - name = "polymer", - build_file = "bower.BUILD", - url = "https://github.com/polymer/polymer/archive/v1.7.0.tar.gz", - strip_prefix = "polymer-1.7.0", + name = "mobile_ssd", + build_file = "models.BUILD", + sha256 = "bddd81ea5c80a97adfac1c9f770e6f55cbafd7cce4d3bbe15fbeb041e6b8f3e8", + urls = [ + "http://storage.googleapis.com/download.tensorflow.org/models/object_detection/ssd_mobilenet_v1_android_export.zip", + "http://download.tensorflow.org/models/object_detection/ssd_mobilenet_v1_android_export.zip", + ], ) new_http_archive( - name = "promise_polyfill", - build_file = "bower.BUILD", - url = "https://github.com/polymerlabs/promise-polyfill/archive/v1.0.0.tar.gz", - strip_prefix = "promise-polyfill-1.0.0", -) - -http_file( - name = "three_js_three_min_js", - url = "https://raw.githubusercontent.com/mrdoob/three.js/r77/build/three.min.js", -) - -http_file( - name = "three_js_orbitcontrols_js", - url = "https://raw.githubusercontent.com/mrdoob/three.js/r77/examples/js/controls/OrbitControls.js", + name = "mobile_multibox", + build_file = "models.BUILD", + sha256 = "859edcddf84dddb974c36c36cfc1f74555148e9c9213dedacf1d6b613ad52b96", + urls = [ + "http://storage.googleapis.com/download.tensorflow.org/models/mobile_multibox_v1a.zip", + "http://download.tensorflow.org/models/mobile_multibox_v1a.zip", + ], ) new_http_archive( - name = "web_animations_js", - build_file = "bower.BUILD", - url = "https://github.com/web-animations/web-animations-js/archive/2.2.1.tar.gz", - strip_prefix = "web-animations-js-2.2.1", + name = "stylize", + build_file = "models.BUILD", + sha256 = "3d374a730aef330424a356a8d4f04d8a54277c425e274ecb7d9c83aa912c6bfa", + urls = [ + "http://storage.googleapis.com/download.tensorflow.org/models/stylize_v1.zip", + "http://download.tensorflow.org/models/stylize_v1.zip", + ], ) new_http_archive( - name = "webcomponentsjs", - build_file = "bower.BUILD", - url = "https://github.com/webcomponents/webcomponentsjs/archive/v0.7.22.tar.gz", - strip_prefix = "webcomponentsjs-0.7.22", -) - -http_file( - name = "weblas_weblas_js", - url = "https://raw.githubusercontent.com/waylonflinn/weblas/v0.9.0/dist/weblas.js", + name = "speech_commands", + 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/arm_compiler.BUILD b/arm_compiler.BUILD new file mode 100644 index 0000000000000000000000000000000000000000..db2e9bbe1e1156d1da19edb68079c61bd6e1923b --- /dev/null +++ b/arm_compiler.BUILD @@ -0,0 +1,81 @@ +package(default_visibility = ["//visibility:public"]) + +filegroup( + name = "gcc", + srcs = [ + "bin/arm-linux-gnueabihf-gcc", + ], +) + +filegroup( + name = "ar", + srcs = [ + "bin/arm-linux-gnueabihf-ar", + ], +) + +filegroup( + name = "ld", + srcs = [ + "bin/arm-linux-gnueabihf-ld", + ], +) + +filegroup( + name = "nm", + srcs = [ + "bin/arm-linux-gnueabihf-nm", + ], +) + +filegroup( + name = "objcopy", + srcs = [ + "bin/arm-linux-gnueabihf-objcopy", + ], +) + +filegroup( + name = "objdump", + srcs = [ + "bin/arm-linux-gnueabihf-objdump", + ], +) + +filegroup( + name = "strip", + srcs = [ + "bin/arm-linux-gnueabihf-strip", + ], +) + +filegroup( + name = "as", + srcs = [ + "bin/arm-linux-gnueabihf-as", + ], +) + +filegroup( + name = "compiler_pieces", + srcs = glob([ + "arm-linux-gnueabihf/**", + "libexec/**", + "lib/gcc/arm-linux-gnueabihf/**", + "include/**", + ]), +) + +filegroup( + name = "compiler_components", + srcs = [ + ":ar", + ":as", + ":gcc", + ":ld", + ":nm", + ":objcopy", + ":objdump", + ":strip", + ], +) diff --git a/bower.BUILD b/bower.BUILD deleted file mode 100644 index eabd1d6450728aab37ebeca6366009d74c6984b6..0000000000000000000000000000000000000000 --- a/bower.BUILD +++ /dev/null @@ -1,645 +0,0 @@ -# AUTOGENERATED FILE by tensorboard_bower_dependency_sync.py - -package(default_visibility = ["//visibility:public"]) - -filegroup( - name = "d3", - srcs = [ - "d3.js", - "d3.min.js", - "package.js", - ], -) - -filegroup( - name = "dagre", - srcs = [ - "dist/dagre.core.js", - "dist/dagre.core.min.js", - ], -) - -filegroup( - name = "es6_promise", - srcs = [ - "promise.js", - "promise.min.js", - ], -) - -filegroup( - name = "font_roboto", - srcs = ["roboto.html"], -) - -filegroup( - name = "graphlib", - srcs = [ - "dist/graphlib.core.js", - "dist/graphlib.core.min.js", - ], -) - -filegroup( - name = "iron_a11y_announcer", - srcs = [ - "index.html", - "iron-a11y-announcer.html", - ], -) - -filegroup( - name = "iron_a11y_keys_behavior", - srcs = [ - "index.html", - "iron-a11y-keys-behavior.html", - ], -) - -filegroup( - name = "iron_ajax", - srcs = [ - "index.html", - "iron-ajax.html", - "iron-request.html", - ], -) - -filegroup( - name = "iron_autogrow_textarea", - srcs = [ - "index.html", - "iron-autogrow-textarea.html", - ], -) - -filegroup( - name = "iron_behaviors", - srcs = [ - "index.html", - "iron-button-state.html", - "iron-control-state.html", - ], -) - -filegroup( - name = "iron_checked_element_behavior", - srcs = [ - "index.html", - "iron-checked-element-behavior.html", - ], -) - -filegroup( - name = "iron_collapse", - srcs = [ - "index.html", - "iron-collapse.html", - ], -) - -filegroup( - name = "iron_dropdown", - srcs = [ - "index.html", - "iron-dropdown.html", - "iron-dropdown-scroll-manager.html", - ], -) - -filegroup( - name = "iron_fit_behavior", - srcs = [ - "index.html", - "iron-fit-behavior.html", - ], -) - -filegroup( - name = "iron_flex_layout", - srcs = [ - "classes/iron-flex-layout.html", - "classes/iron-shadow-flex-layout.html", - "index.html", - "iron-flex-layout.html", - "iron-flex-layout-classes.html", - ], -) - -filegroup( - name = "iron_form_element_behavior", - srcs = [ - "index.html", - "iron-form-element-behavior.html", - ], -) - -filegroup( - name = "iron_icon", - srcs = [ - "index.html", - "iron-icon.html", - ], -) - -filegroup( - name = "iron_icons", - srcs = [ - "av-icons.html", - "communication-icons.html", - "device-icons.html", - "editor-icons.html", - "hardware-icons.html", - "image-icons.html", - "index.html", - "iron-icons.html", - "maps-icons.html", - "notification-icons.html", - "places-icons.html", - "social-icons.html", - ], -) - -filegroup( - name = "iron_iconset_svg", - srcs = [ - "index.html", - "iron-iconset-svg.html", - ], -) - -filegroup( - name = "iron_input", - srcs = [ - "index.html", - "iron-input.html", - ], -) - -filegroup( - name = "iron_list", - srcs = [ - "index.html", - "iron-list.html", - "test/smoke/avg-worst-case.html", - "test/smoke/dummy-data.html", - "test/smoke/index.html", - "test/smoke/physical-count.html", - ], -) - -filegroup( - name = "iron_menu_behavior", - srcs = [ - "index.html", - "iron-menu-behavior.html", - "iron-menubar-behavior.html", - ], -) - -filegroup( - name = "iron_meta", - srcs = [ - "index.html", - "iron-meta.html", - ], -) - -filegroup( - name = "iron_overlay_behavior", - srcs = [ - "index.html", - "iron-focusables-helper.html", - "iron-overlay-backdrop.html", - "iron-overlay-behavior.html", - "iron-overlay-manager.html", - ], -) - -filegroup( - name = "iron_range_behavior", - srcs = [ - "index.html", - "iron-range-behavior.html", - ], -) - -filegroup( - name = "iron_resizable_behavior", - srcs = [ - "demo/src/x-app.html", - "index.html", - "iron-resizable-behavior.html", - ], -) - -filegroup( - name = "iron_scroll_target_behavior", - srcs = [ - "index.html", - "iron-scroll-target-behavior.html", - ], -) - -filegroup( - name = "iron_selector", - srcs = [ - "index.html", - "iron-multi-selectable.html", - "iron-selectable.html", - "iron-selection.html", - "iron-selector.html", - ], -) - -filegroup( - name = "iron_validatable_behavior", - srcs = [ - "index.html", - "iron-validatable-behavior.html", - ], -) - -filegroup( - name = "lodash", - srcs = [ - "lodash.js", - "lodash.min.js", - ], -) - -filegroup( - name = "neon_animation", - srcs = [ - "animations/cascaded-animation.html", - "animations/fade-in-animation.html", - "animations/fade-out-animation.html", - "animations/hero-animation.html", - "animations/opaque-animation.html", - "animations/reverse-ripple-animation.html", - "animations/ripple-animation.html", - "animations/scale-down-animation.html", - "animations/scale-up-animation.html", - "animations/slide-down-animation.html", - "animations/slide-from-bottom-animation.html", - "animations/slide-from-left-animation.html", - "animations/slide-from-right-animation.html", - "animations/slide-from-top-animation.html", - "animations/slide-left-animation.html", - "animations/slide-right-animation.html", - "animations/slide-up-animation.html", - "animations/transform-animation.html", - "demo/card/index.html", - "demo/card/x-card.html", - "demo/card/x-cards-list.html", - "demo/declarative/index.html", - "demo/doc/index.html", - "demo/doc/my-animatable.html", - "demo/doc/my-dialog.html", - "demo/dropdown/animated-dropdown.html", - "demo/dropdown/index.html", - "demo/grid/animated-grid.html", - "demo/grid/fullsize-page-with-card.html", - "demo/grid/index.html", - "demo/list/full-view.html", - "demo/list/index.html", - "demo/list/list-demo.html", - "demo/list/list-view.html", - "demo/load/animated-grid.html", - "demo/load/full-page.html", - "demo/load/index.html", - "demo/reprojection/animated-grid.html", - "demo/reprojection/fullsize-page-with-card.html", - "demo/reprojection/index.html", - "demo/reprojection/reprojected-pages.html", - "demo/tiles/circles-page.html", - "demo/tiles/index.html", - "demo/tiles/squares-page.html", - "index.html", - "neon-animatable.html", - "neon-animatable-behavior.html", - "neon-animated-pages.html", - "neon-animation.html", - "neon-animation-behavior.html", - "neon-animation-runner-behavior.html", - "neon-animations.html", - "neon-shared-element-animatable-behavior.html", - "neon-shared-element-animation-behavior.html", - "web-animations.html", - ], -) - -filegroup( - name = "paper_behaviors", - srcs = [ - "index.html", - "paper-button-behavior.html", - "paper-checked-element-behavior.html", - "paper-inky-focus-behavior.html", - "paper-ripple-behavior.html", - ], -) - -filegroup( - name = "paper_button", - srcs = [ - "index.html", - "paper-button.html", - ], -) - -filegroup( - name = "paper_checkbox", - srcs = [ - "index.html", - "paper-checkbox.html", - ], -) - -filegroup( - name = "paper_dialog", - srcs = [ - "index.html", - "paper-dialog.html", - ], -) - -filegroup( - name = "paper_dialog_behavior", - srcs = [ - "index.html", - "paper-dialog-behavior.html", - "paper-dialog-common.css", - "paper-dialog-shared-styles.html", - ], -) - -filegroup( - name = "paper_dialog_scrollable", - srcs = [ - "index.html", - "paper-dialog-scrollable.html", - ], -) - -filegroup( - name = "paper_dropdown_menu", - srcs = [ - "index.html", - "paper-dropdown-menu.html", - "paper-dropdown-menu-icons.html", - "paper-dropdown-menu-light.html", - "paper-dropdown-menu-shared-styles.html", - ], -) - -filegroup( - name = "paper_header_panel", - srcs = [ - "index.html", - "paper-header-panel.html", - ], -) - -filegroup( - name = "paper_icon_button", - srcs = [ - "index.html", - "paper-icon-button.html", - "paper-icon-button-light.html", - ], -) - -filegroup( - name = "paper_input", - srcs = [ - "all-imports.html", - "index.html", - "paper-input.html", - "paper-input-addon-behavior.html", - "paper-input-behavior.html", - "paper-input-char-counter.html", - "paper-input-container.html", - "paper-input-error.html", - "paper-textarea.html", - ], -) - -filegroup( - name = "paper_item", - srcs = [ - "all-imports.html", - "index.html", - "paper-icon-item.html", - "paper-item.html", - "paper-item-behavior.html", - "paper-item-body.html", - "paper-item-shared-styles.html", - ], -) - -filegroup( - name = "paper_listbox", - srcs = [ - "index.html", - "paper-listbox.html", - ], -) - -filegroup( - name = "paper_material", - srcs = [ - "index.html", - "paper-material.html", - "paper-material-shared-styles.html", - ], -) - -filegroup( - name = "paper_menu", - srcs = [ - "index.html", - "paper-menu.html", - "paper-menu-shared-styles.html", - "paper-submenu.html", - ], -) - -filegroup( - name = "paper_menu_button", - srcs = [ - "index.html", - "paper-menu-button.html", - "paper-menu-button-animations.html", - ], -) - -filegroup( - name = "paper_progress", - srcs = [ - "index.html", - "paper-progress.html", - ], -) - -filegroup( - name = "paper_radio_button", - srcs = [ - "index.html", - "paper-radio-button.html", - ], -) - -filegroup( - name = "paper_radio_group", - srcs = [ - "index.html", - "paper-radio-group.html", - ], -) - -filegroup( - name = "paper_ripple", - srcs = [ - "index.html", - "paper-ripple.html", - ], -) - -filegroup( - name = "paper_slider", - srcs = [ - "index.html", - "paper-slider.html", - ], -) - -filegroup( - name = "paper_spinner", - srcs = [ - "index.html", - "paper-spinner.html", - "paper-spinner-behavior.html", - "paper-spinner-lite.html", - "paper-spinner-styles.html", - ], -) - -filegroup( - name = "paper_styles", - srcs = [ - "classes/global.html", - "classes/shadow.html", - "classes/shadow-layout.html", - "classes/typography.html", - "color.html", - "default-theme.html", - "demo.css", - "demo-pages.html", - "index.html", - "paper-styles.html", - "paper-styles-classes.html", - "shadow.html", - "typography.html", - ], -) - -filegroup( - name = "paper_tabs", - srcs = [ - "index.html", - "paper-tab.html", - "paper-tabs.html", - "paper-tabs-icons.html", - ], -) - -filegroup( - name = "paper_toast", - srcs = [ - "index.html", - "paper-toast.html", - ], -) - -filegroup( - name = "paper_toggle_button", - srcs = [ - "index.html", - "paper-toggle-button.html", - ], -) - -filegroup( - name = "paper_toolbar", - srcs = [ - "index.html", - "paper-toolbar.html", - ], -) - -filegroup( - name = "paper_tooltip", - srcs = [ - "index.html", - "paper-tooltip.html", - ], -) - -filegroup( - name = "plottable", - srcs = [ - "plottable.css", - "plottable.js", - "plottable.min.js", - ], -) - -filegroup( - name = "polymer", - srcs = [ - "polymer.html", - "polymer-micro.html", - "polymer-mini.html", - ], -) - -filegroup( - name = "promise_polyfill", - srcs = [ - "Gruntfile.js", - "Promise.js", - "Promise.min.js", - "Promise-Statics.js", - "promise-polyfill.html", - "promise-polyfill-lite.html", - ], -) - -filegroup( - name = "web_animations_js", - srcs = [ - "web-animations.html", - "web-animations.min.js", - "web-animations-next.min.js", - "web-animations-next-lite.min.js", - ], -) - -filegroup( - name = "webcomponentsjs", - srcs = [ - "CustomElements.js", - "CustomElements.min.js", - "HTMLImports.js", - "HTMLImports.min.js", - "MutationObserver.js", - "MutationObserver.min.js", - "ShadowDOM.js", - "ShadowDOM.min.js", - "webcomponents.js", - "webcomponents.min.js", - "webcomponents-lite.js", - "webcomponents-lite.min.js", - ], -) diff --git a/configure b/configure index 6360641be2ca99c8c8cbe58c95fc2fd59f917744..9c21d2b03a27714f05094667691e74c16fa89f35 100755 --- a/configure +++ b/configure @@ -3,618 +3,12 @@ set -e set -o pipefail -# Find out the absolute path to where ./configure resides -pushd `dirname $0` > /dev/null -SOURCE_BASE_DIR=`pwd -P` -popd > /dev/null - -PLATFORM="$(uname -s | tr 'A-Z' 'a-z')" - -function is_linux() { - if [[ "${PLATFORM}" == "linux" ]]; then - true - else - false - fi -} - -function is_macos() { - if [[ "${PLATFORM}" == "darwin" ]]; then - true - else - false - fi -} - -function is_windows() { - # On windows, the shell script is actually running in msys - if [[ "${PLATFORM}" =~ msys_nt*|mingw*|cygwin*|uwin* ]]; then - true - else - false - fi -} - -function sed_hyphen_i() { - if is_macos; then - sed -i '' "$@" - else - sed -i "$@" - fi -} - -function write_to_bazelrc() { - echo "$1" >> .tf_configure.bazelrc -} - -function write_action_env_to_bazelrc() { - write_to_bazelrc "build --action_env $1=\"$2\"" -} - -# This file contains customized config settings. -rm -f .tf_configure.bazelrc -touch .tf_configure.bazelrc -touch .bazelrc -sed_hyphen_i "/tf_configure/d" .bazelrc -echo "import .tf_configure.bazelrc" >> .bazelrc - -# Delete any leftover BUILD files from the Makefile build, which would interfere -# with Bazel parsing. -MAKEFILE_DOWNLOAD_DIR=tensorflow/contrib/makefile/downloads -if [ -d "${MAKEFILE_DOWNLOAD_DIR}" ]; then - find ${MAKEFILE_DOWNLOAD_DIR} -type f -name '*BUILD' -delete -fi - -## Set up python-related environment settings -while true; do - fromuser="" - if [ -z "$PYTHON_BIN_PATH" ]; then - default_python_bin_path=$(which python || which python3 || true) - read -p "Please specify the location of python. [Default is $default_python_bin_path]: " PYTHON_BIN_PATH - fromuser="1" - if [ -z "$PYTHON_BIN_PATH" ]; then - PYTHON_BIN_PATH=$default_python_bin_path - fi - fi - if [ -e "$PYTHON_BIN_PATH" ]; then - break - fi - echo "Invalid python path. ${PYTHON_BIN_PATH} cannot be found" 1>&2 - if [ -z "$fromuser" ]; then - exit 1 - fi - PYTHON_BIN_PATH="" - # Retry -done - -## Set up MKL related environment settings -if false; then # Disable building with MKL for now - while [ "$TF_NEED_MKL" == "" ]; do - fromuser="" - read -p "Do you wish to build TensorFlow with MKL support? [y/N] " INPUT - fromuser="1" - case $INPUT in - [Yy]* ) echo "MKL support will be enabled for TensorFlow"; TF_NEED_MKL=1;; - [Nn]* ) echo "No MKL support will be enabled for TensorFlow"; TF_NEED_MKL=0;; - "" ) echo "No MKL support will be enabled for TensorFlow"; TF_NEED_MKL=0;; - * ) echo "Invalid selection: " $INPUT;; - esac - done - - OSNAME=`uname -s` - - if [ "$TF_NEED_MKL" == "1" ]; then # TF_NEED_MKL - DST=`dirname $0` - ARCHIVE_BASENAME=mklml_lnx_2017.0.2.20170209.tgz - GITHUB_RELEASE_TAG=v0.5 - MKLURL="https://github.com/01org/mkl-dnn/releases/download/$GITHUB_RELEASE_TAG/$ARCHIVE_BASENAME" - if ! [ -e "$DST/third_party/mkl/$ARCHIVE_BASENAME" ]; then - wget --no-check-certificate -P $DST/third_party/mkl/ $MKLURL - fi - tar -xzf $DST/third_party/mkl/$ARCHIVE_BASENAME -C $DST/third_party/mkl/ - extracted_dir_name="${ARCHIVE_BASENAME%.*}" - MKL_INSTALL_PATH=$DST/third_party/mkl/$extracted_dir_name - MKL_INSTALL_PATH=`${PYTHON_BIN_PATH} -c "import os; print(os.path.realpath(os.path.expanduser('${MKL_INSTALL_PATH}')))"` - - if [ "$OSNAME" == "Linux" ]; then - # Full MKL configuration - MKL_RT_LIB_PATH="lib/intel64/libmkl_rt.so" #${TF_MKL_EXT}#TODO version? - MKL_RT_OMP_LIB_PATH="../compiler/lib/intel64/libiomp5.so" #TODO VERSION? - - # MKL-ML configuration - MKL_ML_LIB_PATH="lib/libmklml_intel.so" #${TF_MKL_EXT}#TODO version? - MKL_ML_OMP_LIB_PATH="lib/libiomp5.so" #TODO VERSION? - elif [ "$OSNAME" == "Darwin" ]; then - echo "Darwin is unsupported yet"; - exit 1 - fi - - if [ -e "$MKL_INSTALL_PATH/${MKL_ML_LIB_PATH}" ]; then - ln -sf $MKL_INSTALL_PATH/${MKL_ML_LIB_PATH} third_party/mkl/ - ln -sf $MKL_INSTALL_PATH/${MKL_ML_OMP_LIB_PATH} third_party/mkl/ - ln -sf $MKL_INSTALL_PATH/include third_party/mkl/ - ln -sf $MKL_INSTALL_PATH/include third_party/eigen3/mkl_include - else - echo "ERROR: $MKL_INSTALL_PATH/${MKL_ML_LIB_PATH} does not exist"; - exit 1 - fi - - if [ -z "$fromuser" ]; then - exit 1 - fi - -cat > third_party/mkl/mkl.config <> tools/bazel.rc -for opt in $CC_OPT_FLAGS; do - echo "build:opt --cxxopt=$opt --copt=$opt" >> tools/bazel.rc -done - -# Run the gen_git_source to create links where bazel can track dependencies for -# git hash propagation -GEN_GIT_SOURCE=tensorflow/tools/git/gen_git_source.py -chmod a+x ${GEN_GIT_SOURCE} -"${PYTHON_BIN_PATH}" ${GEN_GIT_SOURCE} --configure "${SOURCE_BASE_DIR}" - -## Set up SYCL-related environment settings -while [ "$TF_NEED_OPENCL" == "" ]; do - read -p "Do you wish to build TensorFlow with OpenCL support? [y/N] " INPUT - case $INPUT in - [Yy]* ) echo "OpenCL support will be enabled for TensorFlow"; TF_NEED_OPENCL=1;; - [Nn]* ) echo "No OpenCL support will be enabled for TensorFlow"; TF_NEED_OPENCL=0;; - "" ) echo "No OpenCL support will be enabled for TensorFlow"; TF_NEED_OPENCL=0;; - * ) echo "Invalid selection: " $INPUT;; - esac -done - -## Set up Cuda-related environment settings - -while [ "$TF_NEED_CUDA" == "" ]; do - read -p "Do you wish to build TensorFlow with CUDA support? [y/N] " INPUT - case $INPUT in - [Yy]* ) echo "CUDA support will be enabled for TensorFlow"; TF_NEED_CUDA=1;; - [Nn]* ) echo "No CUDA support will be enabled for TensorFlow"; TF_NEED_CUDA=0;; - "" ) echo "No CUDA support will be enabled for TensorFlow"; TF_NEED_CUDA=0;; - * ) echo "Invalid selection: " $INPUT;; - esac -done - -export TF_NEED_CUDA -write_action_env_to_bazelrc "TF_NEED_CUDA" "$TF_NEED_CUDA" - -export TF_NEED_OPENCL - -if [ "$TF_NEED_CUDA" == "1" ]; then -while [[ "$TF_CUDA_CLANG" == "" ]]; do - read -p "Do you want to use clang as CUDA compiler? [y/N] " INPUT - case $INPUT in - [Yy]* ) echo "Clang will be used as CUDA compiler"; TF_CUDA_CLANG=1;; - [Nn]* ) echo "nvcc will be used as CUDA compiler"; TF_CUDA_CLANG=0;; - "" ) echo "nvcc will be used as CUDA compiler"; TF_CUDA_CLANG=0;; - * ) echo "Invalid selection: " $INPUT;; - esac -done - -export TF_CUDA_CLANG -write_action_env_to_bazelrc "TF_CUDA_CLANG" "$TF_CUDA_CLANG" - -# Set up which gcc nvcc should use as the host compiler -# No need to set this on Windows -while [[ "$TF_CUDA_CLANG" != "1" ]] && ! is_windows && true; do - fromuser="" - if [ -z "$GCC_HOST_COMPILER_PATH" ]; then - default_gcc_host_compiler_path=$(which gcc || true) - read -p "Please specify which gcc should be used by nvcc as the host compiler. [Default is $default_gcc_host_compiler_path]: " GCC_HOST_COMPILER_PATH - fromuser="1" - if [ -z "$GCC_HOST_COMPILER_PATH" ]; then - GCC_HOST_COMPILER_PATH="$default_gcc_host_compiler_path" - fi - fi - if [ -e "$GCC_HOST_COMPILER_PATH" ]; then - export GCC_HOST_COMPILER_PATH - write_action_env_to_bazelrc "GCC_HOST_COMPILER_PATH" "$GCC_HOST_COMPILER_PATH" - break - fi - echo "Invalid gcc path. ${GCC_HOST_COMPILER_PATH} cannot be found" 1>&2 - if [ -z "$fromuser" ]; then - exit 1 - fi - GCC_HOST_COMPILER_PATH="" - # Retry -done - -# Set up which clang we should use as the cuda / host compiler. -while [[ "$TF_CUDA_CLANG" == "1" ]] && true; do - fromuser="" - if [ -z "$CLANG_CUDA_COMPILER_PATH" ]; then - default_clang_host_compiler_path=$(which clang || true) - read -p "Please specify which clang should be used as device and host compiler. [Default is $default_clang_host_compiler_path]: " CLANG_CUDA_COMPILER_PATH - fromuser="1" - if [ -z "$CLANG_CUDA_COMPILER_PATH" ]; then - CLANG_CUDA_COMPILER_PATH="$default_clang_host_compiler_path" - fi - fi - if [ -e "$CLANG_CUDA_COMPILER_PATH" ]; then - export CLANG_CUDA_COMPILER_PATH - write_action_env_to_bazelrc "CLANG_CUDA_COMPILER_PATH" "$CLANG_CUDA_COMPILER_PATH" - break - fi - echo "Invalid clang path. ${CLANG_CUDA_COMPILER_PATH} cannot be found" 1>&2 - if [ -z "$fromuser" ]; then - exit 1 - fi - CLANG_CUDA_COMPILER_PATH="" - # Retry -done - -# Find out where the CUDA toolkit is installed -while true; do - # Configure the Cuda SDK version to use. - if [ -z "$TF_CUDA_VERSION" ]; then - read -p "Please specify the CUDA SDK version you want to use, e.g. 7.0. [Leave empty to use system default]: " TF_CUDA_VERSION - fi - - fromuser="" - if [ -z "$CUDA_TOOLKIT_PATH" ]; then - default_cuda_path=/usr/local/cuda - if is_windows; then - if [ -z "$CUDA_PATH" ]; then - default_cuda_path="C:/Program Files/NVIDIA GPU Computing Toolkit/CUDA/v8.0" - else - default_cuda_path="$(cygpath -m "$CUDA_PATH")" - fi - fi - read -p "Please specify the location where CUDA $TF_CUDA_VERSION toolkit is installed. Refer to README.md for more details. [Default is $default_cuda_path]: " CUDA_TOOLKIT_PATH - fromuser="1" - if [ -z "$CUDA_TOOLKIT_PATH" ]; then - CUDA_TOOLKIT_PATH="$default_cuda_path" - fi - fi - - if [[ -z "$TF_CUDA_VERSION" ]]; then - TF_CUDA_EXT="" - else - TF_CUDA_EXT=".$TF_CUDA_VERSION" - fi - - if is_windows; then - CUDA_RT_LIB_PATH="lib/x64/cudart.lib" - elif is_linux; then - CUDA_RT_LIB_PATH="lib64/libcudart.so${TF_CUDA_EXT}" - elif is_macos; then - CUDA_RT_LIB_PATH="lib/libcudart${TF_CUDA_EXT}.dylib" - fi - - if [ -e "${CUDA_TOOLKIT_PATH}/${CUDA_RT_LIB_PATH}" ]; then - export CUDA_TOOLKIT_PATH - write_action_env_to_bazelrc "CUDA_TOOLKIT_PATH" "$CUDA_TOOLKIT_PATH" - export TF_CUDA_VERSION - write_action_env_to_bazelrc "TF_CUDA_VERSION" "$TF_CUDA_VERSION" - break - fi - echo "Invalid path to CUDA $TF_CUDA_VERSION toolkit. ${CUDA_TOOLKIT_PATH}/${CUDA_RT_LIB_PATH} cannot be found" - - if [ -z "$fromuser" ]; then - exit 1 - fi - # Retry - TF_CUDA_VERSION="" - CUDA_TOOLKIT_PATH="" -done - -# Find out where the cuDNN library is installed -while true; do - # Configure the cuDNN version to use. - if [ -z "$TF_CUDNN_VERSION" ]; then - read -p "Please specify the cuDNN version you want to use. [Leave empty to use system default]: " TF_CUDNN_VERSION - fi - - fromuser="" - if [ -z "$CUDNN_INSTALL_PATH" ]; then - default_cudnn_path=${CUDA_TOOLKIT_PATH} - read -p "Please specify the location where cuDNN $TF_CUDNN_VERSION library is installed. Refer to README.md for more details. [Default is $default_cudnn_path]: " CUDNN_INSTALL_PATH - fromuser="1" - if [ -z "$CUDNN_INSTALL_PATH" ]; then - CUDNN_INSTALL_PATH=$default_cudnn_path - fi - # Result returned from "read" will be used unexpanded. That make "~" unuseable. - # Going through one more level of expansion to handle that. - CUDNN_INSTALL_PATH=`"${PYTHON_BIN_PATH}" -c "import os; print(os.path.realpath(os.path.expanduser('${CUDNN_INSTALL_PATH}')))"` - fi - - if [[ -z "$TF_CUDNN_VERSION" ]]; then - TF_CUDNN_EXT="" - else - TF_CUDNN_EXT=".$TF_CUDNN_VERSION" - fi - - if is_windows; then - CUDA_DNN_LIB_PATH="lib/x64/cudnn.lib" - CUDA_DNN_LIB_ALT_PATH="lib/x64/cudnn.lib" - elif is_linux; then - CUDA_DNN_LIB_PATH="lib64/libcudnn.so${TF_CUDNN_EXT}" - CUDA_DNN_LIB_ALT_PATH="libcudnn.so${TF_CUDNN_EXT}" - elif is_macos; then - CUDA_DNN_LIB_PATH="lib/libcudnn${TF_CUDNN_EXT}.dylib" - CUDA_DNN_LIB_ALT_PATH="libcudnn${TF_CUDNN_EXT}.dylib" - fi - - if [ -e "$CUDNN_INSTALL_PATH/${CUDA_DNN_LIB_ALT_PATH}" -o -e "$CUDNN_INSTALL_PATH/${CUDA_DNN_LIB_PATH}" ]; then - export TF_CUDNN_VERSION - write_action_env_to_bazelrc "TF_CUDNN_VERSION" "$TF_CUDNN_VERSION" - export CUDNN_INSTALL_PATH - write_action_env_to_bazelrc "CUDNN_INSTALL_PATH" "$CUDNN_INSTALL_PATH" - break - fi - - if is_linux; then - if ! type ldconfig > /dev/null 2>&1; then - LDCONFIG_BIN=/sbin/ldconfig - else - LDCONFIG_BIN=ldconfig - fi - CUDNN_PATH_FROM_LDCONFIG="$($LDCONFIG_BIN -p | sed -n 's/.*libcudnn.so .* => \(.*\)/\1/p')" - if [ -e "${CUDNN_PATH_FROM_LDCONFIG}${TF_CUDNN_EXT}" ]; then - export TF_CUDNN_VERSION - write_action_env_to_bazelrc "TF_CUDNN_VERSION" "$TF_CUDNN_VERSION" - export CUDNN_INSTALL_PATH="$(dirname ${CUDNN_PATH_FROM_LDCONFIG})" - write_action_env_to_bazelrc "CUDNN_INSTALL_PATH" "$CUDNN_INSTALL_PATH" - break - fi - fi - echo "Invalid path to cuDNN ${CUDNN_VERSION} toolkit. Neither of the following two files can be found:" - echo "${CUDNN_INSTALL_PATH}/${CUDA_DNN_LIB_PATH}" - echo "${CUDNN_INSTALL_PATH}/${CUDA_DNN_LIB_ALT_PATH}" - if is_linux; then - echo "${CUDNN_PATH_FROM_LDCONFIG}${TF_CUDNN_EXT}" - fi - - if [ -z "$fromuser" ]; then - exit 1 - fi - # Retry - TF_CUDNN_VERSION="" - CUDNN_INSTALL_PATH="" -done - -# Configure the compute capabilities that TensorFlow builds for. -# Since Cuda toolkit is not backward-compatible, this is not guaranteed to work. -while true; do - fromuser="" - default_cuda_compute_capabilities="3.5,5.2" - if [ -z "$TF_CUDA_COMPUTE_CAPABILITIES" ]; then -cat << EOF -Please specify a list of comma-separated Cuda compute capabilities you want to build with. -You can find the compute capability of your device at: https://developer.nvidia.com/cuda-gpus. -Please note that each additional compute capability significantly increases your build time and binary size. -EOF - read -p "[Default is: \"3.5,5.2\"]: " TF_CUDA_COMPUTE_CAPABILITIES - fromuser=1 - fi - if [ -z "$TF_CUDA_COMPUTE_CAPABILITIES" ]; then - TF_CUDA_COMPUTE_CAPABILITIES=$default_cuda_compute_capabilities - fi - # Check whether all capabilities from the input is valid - COMPUTE_CAPABILITIES=${TF_CUDA_COMPUTE_CAPABILITIES//,/ } - ALL_VALID=1 - for CAPABILITY in $COMPUTE_CAPABILITIES; do - if [[ ! "$CAPABILITY" =~ [0-9]+.[0-9]+ ]]; then - echo "Invalid compute capability: " $CAPABILITY - ALL_VALID=0 - break - fi - done - if [ "$ALL_VALID" == "0" ]; then - if [ -z "$fromuser" ]; then - exit 1 - fi - else - export TF_CUDA_COMPUTE_CAPABILITIES - write_action_env_to_bazelrc "TF_CUDA_COMPUTE_CAPABILITIES" "$TF_CUDA_COMPUTE_CAPABILITIES" - break - fi - TF_CUDA_COMPUTE_CAPABILITIES="" -done - -if is_windows; then - # The following three variables are needed for MSVC toolchain configuration in Bazel - export CUDA_PATH="$CUDA_TOOLKIT_PATH" - export CUDA_COMPUTE_CAPABILITIES="$TF_CUDA_COMPUTE_CAPABILITIES" - export NO_WHOLE_ARCHIVE_OPTION=1 - write_action_env_to_bazelrc "CUDA_PATH" "$CUDA_PATH" - write_action_env_to_bazelrc "CUDA_COMPUTE_CAPABILITIES" "$CUDA_COMPUTE_CAPABILITIES" - write_action_env_to_bazelrc "NO_WHOLE_ARCHIVE_OPTION" "1" -fi - -# end of if "$TF_NEED_CUDA" == "1" +if [ -z "$PYTHON_BIN_PATH" ]; then + PYTHON_BIN_PATH=$(which python || which python3 || true) fi -# OpenCL configuration - -if [ "$TF_NEED_OPENCL" == "1" ]; then - -# Determine which C++ compiler should be used as the host compiler -while true; do - fromuser="" - if [ -z "$HOST_CXX_COMPILER" ]; then - default_cxx_host_compiler=$(which clang++-3.6 || true) - read -p "Please specify which C++ compiler should be used as the host C++ compiler. [Default is $default_cxx_host_compiler]: " HOST_CXX_COMPILER - fromuser="1" - if [ -z "$HOST_CXX_COMPILER" ]; then - HOST_CXX_COMPILER=$default_cxx_host_compiler - fi - fi - if [ -e "$HOST_CXX_COMPILER" ]; then - export HOST_CXX_COMPILER - break - fi - echo "Invalid C++ compiler path. ${HOST_CXX_COMPILER} cannot be found" 1>&2 - if [ -z "$fromuser" ]; then - exit 1 - fi - HOST_CXX_COMPILER="" - # Retry -done +# Set all env variables +"$PYTHON_BIN_PATH" configure.py -# Determine which C compiler should be used as the host compiler -while true; do - fromuser="" - if [ -z "$HOST_C_COMPILER" ]; then - default_c_host_compiler=$(which clang-3.6 || true) - read -p "Please specify which C compiler should be used as the host C compiler. [Default is $default_c_host_compiler]: " HOST_C_COMPILER - fromuser="1" - if [ -z "$HOST_C_COMPILER" ]; then - HOST_C_COMPILER=$default_c_host_compiler - fi - fi - if [ -e "$HOST_C_COMPILER" ]; then - export HOST_C_COMPILER - break - fi - echo "Invalid C compiler path. ${HOST_C_COMPILER} cannot be found" 1>&2 - if [ -z "$fromuser" ]; then - exit 1 - fi - HOST_C_COMPILER="" - # Retry -done - -while true; do - # Configure the OPENCL version to use. - TF_OPENCL_VERSION="1.2" - - # Point to ComputeCpp root - if [ -z "$COMPUTECPP_TOOLKIT_PATH" ]; then - default_computecpp_toolkit_path=/usr/local/computecpp - read -p "Please specify the location where ComputeCpp for SYCL $TF_OPENCL_VERSION is installed. [Default is $default_computecpp_toolkit_path]: " COMPUTECPP_TOOLKIT_PATH - fromuser="1" - if [ -z "$COMPUTECPP_TOOLKIT_PATH" ]; then - COMPUTECPP_TOOLKIT_PATH=$default_computecpp_toolkit_path - fi - fi - - if is_linux; then - SYCL_RT_LIB_PATH="lib/libComputeCpp.so" - fi - - if [ -e "${COMPUTECPP_TOOLKIT_PATH}/${SYCL_RT_LIB_PATH}" ]; then - export COMPUTECPP_TOOLKIT_PATH - break - fi - echo "Invalid SYCL $TF_OPENCL_VERSION library path. ${COMPUTECPP_TOOLKIT_PATH}/${SYCL_RT_LIB_PATH} cannot be found" - - if [ -z "$fromuser" ]; then - exit 1 - fi - # Retry - TF_OPENCL_VERSION="" - COMPUTECPP_TOOLKIT_PATH="" -done - -# end of if "$TF_NEED_OPENCL" == "1" -fi - -# TODO(gunan): Remove once bazel correctly handles changes in remote repositories. -bazel clean echo "Configuration finished" + diff --git a/configure.py b/configure.py new file mode 100644 index 0000000000000000000000000000000000000000..186fdc9ddcef9569f6b9c6a3331f2792f045a427 --- /dev/null +++ b/configure.py @@ -0,0 +1,1021 @@ +# 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. +# ============================================================================== +"""configure script to get build parameters from user.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import errno +import os +import platform +import re +import subprocess +import sys + +try: + from shutil import which +except ImportError: + from distutils.spawn import find_executable as which + +_TF_BAZELRC = '.tf_configure.bazelrc' +_DEFAULT_CUDA_VERSION = '8.0' +_DEFAULT_CUDNN_VERSION = '6' +_DEFAULT_CUDA_COMPUTE_CAPABILITIES = '3.5,5.2' +_DEFAULT_CUDA_PATH = '/usr/local/cuda' +_DEFAULT_CUDA_PATH_LINUX = '/opt/cuda' +_DEFAULT_CUDA_PATH_WIN = ('C:/Program Files/NVIDIA GPU Computing ' + 'Toolkit/CUDA/v%s' % _DEFAULT_CUDA_VERSION) +_TF_OPENCL_VERSION = '1.2' +_DEFAULT_COMPUTECPP_TOOLKIT_PATH = '/usr/local/computecpp' + + +def is_windows(): + return platform.system() == 'Windows' + + +def is_linux(): + return platform.system() == 'Linux' + + +def is_macos(): + return platform.system() == 'Darwin' + + +def is_ppc64le(): + return platform.machine() == 'ppc64le' + + +def is_cygwin(): + return platform.system().startswith('CYGWIN_NT') + + +def get_input(question): + try: + try: + answer = raw_input(question) + except NameError: + answer = input(question) # pylint: disable=bad-builtin + except EOFError: + answer = '' + return answer + + +def symlink_force(target, link_name): + """Force symlink, equivalent of 'ln -sf'. + + Args: + target: items to link to. + link_name: name of the link. + """ + try: + os.symlink(target, link_name) + except OSError as e: + if e.errno == errno.EEXIST: + os.remove(link_name) + os.symlink(target, link_name) + else: + raise e + + +def sed_in_place(filename, old, new): + """Replace old string with new string in file. + + Args: + filename: string for filename. + old: string to replace. + new: new string to replace to. + """ + with open(filename, 'r') as f: + filedata = f.read() + newdata = filedata.replace(old, new) + with open(filename, 'w') as f: + f.write(newdata) + + +def remove_line_with(filename, token): + """Remove lines that contain token from file. + + Args: + filename: string for filename. + token: string token to check if to remove a line from file or not. + """ + with open(filename, 'r') as f: + filedata = f.read() + + with open(filename, 'w') as f: + for line in filedata.strip().split('\n'): + if token not in line: + f.write(line + '\n') + + +def write_to_bazelrc(line): + with open(_TF_BAZELRC, 'a') as f: + f.write(line + '\n') + + +def write_action_env_to_bazelrc(var_name, var): + write_to_bazelrc('build --action_env %s="%s"' % (var_name, str(var))) + + +def run_shell(cmd, allow_non_zero=False): + if allow_non_zero: + try: + output = subprocess.check_output(cmd) + except subprocess.CalledProcessError as e: + output = e.output + else: + output = subprocess.check_output(cmd) + return output.decode('UTF-8').strip() + + +def cygpath(path): + """Convert path from posix to windows.""" + return os.path.abspath(path).replace('\\', '/') + + +def get_python_path(environ_cp, python_bin_path): + """Get the python site package paths.""" + python_paths = [] + if environ_cp.get('PYTHONPATH'): + python_paths = environ_cp.get('PYTHONPATH').split(':') + try: + library_paths = run_shell( + [python_bin_path, '-c', + 'import site; print("\\n".join(site.getsitepackages()))']).split("\n") + except subprocess.CalledProcessError: + library_paths = [run_shell( + [python_bin_path, '-c', + 'from distutils.sysconfig import get_python_lib;' + 'print(get_python_lib())'])] + + all_paths = set(python_paths + library_paths) + + paths = [] + for path in all_paths: + if os.path.isdir(path): + paths.append(path) + return paths + + +def get_python_major_version(python_bin_path): + """Get the python major version.""" + return run_shell([python_bin_path, '-c', 'import sys; print(sys.version[0])']) + + +def setup_python(environ_cp, bazel_version): + """Setup python related env variables.""" + # Get PYTHON_BIN_PATH, default is the current running python. + default_python_bin_path = sys.executable + ask_python_bin_path = ('Please specify the location of python. [Default is ' + '%s]: ') % default_python_bin_path + while True: + python_bin_path = get_from_env_or_user_or_default( + environ_cp, 'PYTHON_BIN_PATH', ask_python_bin_path, + default_python_bin_path) + # Check if the path is valid + if os.path.isfile(python_bin_path) and os.access( + python_bin_path, os.X_OK): + break + elif not os.path.exists(python_bin_path): + print('Invalid python path: %s cannot be found.' % python_bin_path) + else: + print('%s is not executable. Is it the python binary?' % python_bin_path) + environ_cp['PYTHON_BIN_PATH'] = '' + + # Convert python path to Windows style before checking lib and version + if is_windows() or is_cygwin(): + python_bin_path = cygpath(python_bin_path) + + # Get PYTHON_LIB_PATH + python_lib_path = environ_cp.get('PYTHON_LIB_PATH') + if not python_lib_path: + python_lib_paths = get_python_path(environ_cp, python_bin_path) + if environ_cp.get('USE_DEFAULT_PYTHON_LIB_PATH') == '1': + python_lib_path = python_lib_paths[0] + else: + print('Found possible Python library paths:\n %s' % + '\n '.join(python_lib_paths)) + default_python_lib_path = python_lib_paths[0] + python_lib_path = get_input( + 'Please input the desired Python library path to use. ' + 'Default is [%s]\n' % python_lib_paths[0]) + if not python_lib_path: + python_lib_path = default_python_lib_path + environ_cp['PYTHON_LIB_PATH'] = python_lib_path + + python_major_version = get_python_major_version(python_bin_path) + + # Convert python path to Windows style before writing into bazel.rc + if is_windows() or is_cygwin(): + python_lib_path = cygpath(python_lib_path) + + # Set-up env variables used by python_configure.bzl + write_action_env_to_bazelrc('PYTHON_BIN_PATH', python_bin_path) + write_action_env_to_bazelrc('PYTHON_LIB_PATH', python_lib_path) + write_to_bazelrc('build --define PYTHON_BIN_PATH="%s"' % python_bin_path) + write_to_bazelrc('build --define PYTHON_LIB_PATH="%s"' % python_lib_path) + write_to_bazelrc('build --force_python=py%s' % python_major_version) + write_to_bazelrc('build --host_force_python=py%s' % python_major_version) + bazel_version_int = convert_version_to_int(bazel_version) + version_0_5_3_int = convert_version_to_int('0.5.3') + # If bazel_version_int is None, we are testing a release Bazel, then the + # version should be higher than 0.5.3 + # TODO(pcloudy): remove this after required min bazel version is higher + # than 0.5.3 + if not bazel_version_int or bazel_version_int >= version_0_5_3_int: + write_to_bazelrc('build --python_path=\"%s"' % python_bin_path) + else: + write_to_bazelrc('build --python%s_path=\"%s"' % (python_major_version, + python_bin_path)) + write_to_bazelrc('test --force_python=py%s' % python_major_version) + write_to_bazelrc('test --host_force_python=py%s' % python_major_version) + write_to_bazelrc('test --define PYTHON_BIN_PATH="%s"' % python_bin_path) + write_to_bazelrc('test --define PYTHON_LIB_PATH="%s"' % python_lib_path) + write_to_bazelrc('run --define PYTHON_BIN_PATH="%s"' % python_bin_path) + write_to_bazelrc('run --define PYTHON_LIB_PATH="%s"' % python_lib_path) + environ_cp['PYTHON_BIN_PATH'] = python_bin_path + + # Write tools/python_bin_path.sh + with open('tools/python_bin_path.sh', 'w') as f: + f.write('export PYTHON_BIN_PATH="%s"' % python_bin_path) + + +def reset_tf_configure_bazelrc(): + """Reset file that contains customized config settings.""" + open(_TF_BAZELRC, 'w').close() + + home = os.path.expanduser('~') + if not os.path.exists('.bazelrc'): + if os.path.exists(os.path.join(home, '.bazelrc')): + with open('.bazelrc', 'a') as f: + f.write('import %s/.bazelrc\n' % home) + else: + open('.bazelrc', 'w').close() + + remove_line_with('.bazelrc', 'tf_configure') + with open('.bazelrc', 'a') as f: + f.write('import %workspace%/.tf_configure.bazelrc\n') + + +def run_gen_git_source(environ_cp): + """Run the gen_git_source to create links. + + The links are for bazel to track dependencies for git hash propagation. + + Args: + environ_cp: copy of the os.environ. + """ + cmd = '"%s" tensorflow/tools/git/gen_git_source.py --configure %s' % ( + environ_cp.get('PYTHON_BIN_PATH'), os.getcwd()) + os.system(cmd) + + +def cleanup_makefile(): + """Delete any leftover BUILD files from the Makefile build. + + These files could interfere with Bazel parsing. + """ + makefile_download_dir = 'tensorflow/contrib/makefile/downloads' + if os.path.isdir(makefile_download_dir): + for root, _, filenames in os.walk(makefile_download_dir): + for f in filenames: + if f.endswith('BUILD'): + os.remove(os.path.join(root, f)) + + +def get_var(environ_cp, + var_name, + query_item, + enabled_by_default, + question=None, + yes_reply=None, + no_reply=None): + """Get boolean input from user. + + If var_name is not set in env, ask user to enable query_item or not. If the + response is empty, use the default. + + Args: + environ_cp: copy of the os.environ. + var_name: string for name of environment variable, e.g. "TF_NEED_HDFS". + query_item: string for feature related to the variable, e.g. "Hadoop File + System". + enabled_by_default: boolean for default behavior. + question: optional string for how to ask for user input. + yes_reply: optionanl string for reply when feature is enabled. + no_reply: optional string for reply when feature is disabled. + + Returns: + boolean value of the variable. + """ + if not question: + question = 'Do you wish to build TensorFlow with %s support?' % query_item + if not yes_reply: + yes_reply = '%s support will be enabled for TensorFlow.' % query_item + if not no_reply: + no_reply = 'No %s' % yes_reply + + yes_reply += '\n' + no_reply += '\n' + + if enabled_by_default: + question += ' [Y/n]: ' + else: + question += ' [y/N]: ' + + var = environ_cp.get(var_name) + while var is None: + user_input_origin = get_input(question) + user_input = user_input_origin.strip().lower() + if user_input == 'y': + print(yes_reply) + var = True + elif user_input == 'n': + print(no_reply) + var = False + elif not user_input: + if enabled_by_default: + print(yes_reply) + var = True + else: + print(no_reply) + var = False + else: + print('Invalid selection: %s' % user_input_origin) + return var + + +def set_build_var(environ_cp, var_name, query_item, option_name, + enabled_by_default): + """Set if query_item will be enabled for the build. + + Ask user if query_item will be enabled. Default is used if no input is given. + Set subprocess environment variable and write to .bazelrc if enabled. + + Args: + environ_cp: copy of the os.environ. + var_name: string for name of environment variable, e.g. "TF_NEED_HDFS". + query_item: string for feature related to the variable, e.g. "Hadoop File + System". + option_name: string for option to define in .bazelrc. + enabled_by_default: boolean for default behavior. + """ + + var = str(int(get_var(environ_cp, var_name, query_item, enabled_by_default))) + environ_cp[var_name] = var + if var == '1': + write_to_bazelrc('build --define %s=true' % option_name) + + +def set_action_env_var(environ_cp, + var_name, + query_item, + enabled_by_default, + question=None, + yes_reply=None, + no_reply=None): + """Set boolean action_env variable. + + Ask user if query_item will be enabled. Default is used if no input is given. + Set environment variable and write to .bazelrc. + + Args: + environ_cp: copy of the os.environ. + var_name: string for name of environment variable, e.g. "TF_NEED_HDFS". + query_item: string for feature related to the variable, e.g. "Hadoop File + System". + enabled_by_default: boolean for default behavior. + question: optional string for how to ask for user input. + yes_reply: optionanl string for reply when feature is enabled. + no_reply: optional string for reply when feature is disabled. + """ + var = int( + get_var(environ_cp, var_name, query_item, enabled_by_default, question, + yes_reply, no_reply)) + + write_action_env_to_bazelrc(var_name, var) + environ_cp[var_name] = str(var) + + +def convert_version_to_int(version): + """Convert a version number to a integer that can be used to compare. + + Version strings of the form X.YZ and X.Y.Z-xxxxx are supported. The + 'xxxxx' part, for instance 'homebrew' on OS/X, is ignored. + + Args: + version: a version to be converted + + Returns: + An integer if converted successfully, otherwise return None. + """ + version = version.split('-')[0] + version_segments = version.split('.') + for seg in version_segments: + if not seg.isdigit(): + return None + + version_str = ''.join(['%03d' % int(seg) for seg in version_segments]) + return int(version_str) + + +def check_bazel_version(min_version): + """Check installed bezel version is at least min_version. + + Args: + min_version: string for minimum bazel version. + + Returns: + The bazel version detected. + """ + if which('bazel') is None: + print('Cannot find bazel. Please install bazel.') + sys.exit(0) + curr_version = run_shell(['bazel', '--batch', 'version']) + + for line in curr_version.split('\n'): + if 'Build label: ' in line: + curr_version = line.split('Build label: ')[1] + break + + min_version_int = convert_version_to_int(min_version) + curr_version_int = convert_version_to_int(curr_version) + + # Check if current bazel version can be detected properly. + if not curr_version_int: + print('WARNING: current bazel installation is not a release version.') + print('Make sure you are running at least bazel %s' % min_version) + return curr_version + + print('You have bazel %s installed.' % curr_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) + return curr_version + + +def set_cc_opt_flags(environ_cp): + """Set up architecture-dependent optimization flags. + + Also append CC optimization flags to bazel.rc.. + + Args: + environ_cp: copy of the os.environ. + """ + if is_ppc64le(): + # gcc on ppc64le does not support -march, use mcpu instead + default_cc_opt_flags = '-mcpu=native' + else: + default_cc_opt_flags = '-march=native' + question = ('Please specify optimization flags to use during compilation when' + ' bazel option "--config=opt" is specified [Default is %s]: ' + ) % default_cc_opt_flags + cc_opt_flags = get_from_env_or_user_or_default(environ_cp, 'CC_OPT_FLAGS', + question, default_cc_opt_flags) + for opt in cc_opt_flags.split(): + write_to_bazelrc('build:opt --cxxopt=%s --copt=%s' % (opt, opt)) + + +def set_tf_cuda_clang(environ_cp): + """set TF_CUDA_CLANG action_env. + + Args: + environ_cp: copy of the os.environ. + """ + question = 'Do you want to use clang as CUDA compiler?' + yes_reply = 'Clang will be used as CUDA compiler.' + no_reply = 'nvcc will be used as CUDA compiler.' + set_action_env_var( + environ_cp, + 'TF_CUDA_CLANG', + None, + False, + question=question, + yes_reply=yes_reply, + no_reply=no_reply) + + +def get_from_env_or_user_or_default(environ_cp, var_name, ask_for_var, + var_default): + """Get var_name either from env, or user or default. + + If var_name has been set as environment variable, use the preset value, else + ask for user input. If no input is provided, the default is used. + + Args: + environ_cp: copy of the os.environ. + var_name: string for name of environment variable, e.g. "TF_NEED_HDFS". + ask_for_var: string for how to ask for user input. + var_default: default value string. + + Returns: + string value for var_name + """ + var = environ_cp.get(var_name) + if not var: + var = get_input(ask_for_var) + print('\n') + if not var: + var = var_default + return var + + +def set_clang_cuda_compiler_path(environ_cp): + """Set CLANG_CUDA_COMPILER_PATH.""" + default_clang_path = which('clang') or '' + ask_clang_path = ('Please specify which clang should be used as device and ' + 'host compiler. [Default is %s]: ') % default_clang_path + + while True: + clang_cuda_compiler_path = get_from_env_or_user_or_default( + environ_cp, 'CLANG_CUDA_COMPILER_PATH', ask_clang_path, + default_clang_path) + if os.path.exists(clang_cuda_compiler_path): + break + + # Reset and retry + print('Invalid clang path: %s cannot be found.' % clang_cuda_compiler_path) + environ_cp['CLANG_CUDA_COMPILER_PATH'] = '' + + # Set CLANG_CUDA_COMPILER_PATH + environ_cp['CLANG_CUDA_COMPILER_PATH'] = clang_cuda_compiler_path + write_action_env_to_bazelrc('CLANG_CUDA_COMPILER_PATH', + clang_cuda_compiler_path) + + +def set_gcc_host_compiler_path(environ_cp): + """Set GCC_HOST_COMPILER_PATH.""" + default_gcc_host_compiler_path = which('gcc') or '' + cuda_bin_symlink = '%s/bin/gcc' % environ_cp.get('CUDA_TOOLKIT_PATH') + + if os.path.islink(cuda_bin_symlink): + # os.readlink is only available in linux + default_gcc_host_compiler_path = os.path.realpath(cuda_bin_symlink) + + ask_gcc_path = ( + 'Please specify which gcc should be used by nvcc as the ' + 'host compiler. [Default is %s]: ') % default_gcc_host_compiler_path + while True: + gcc_host_compiler_path = get_from_env_or_user_or_default( + environ_cp, 'GCC_HOST_COMPILER_PATH', ask_gcc_path, + default_gcc_host_compiler_path) + + if os.path.exists(gcc_host_compiler_path): + break + + # Reset and retry + print('Invalid gcc path. %s cannot be found' % gcc_host_compiler_path) + environ_cp['GCC_HOST_COMPILER_PATH'] = '' + + # Set GCC_HOST_COMPILER_PATH + environ_cp['GCC_HOST_COMPILER_PATH'] = gcc_host_compiler_path + write_action_env_to_bazelrc('GCC_HOST_COMPILER_PATH', gcc_host_compiler_path) + + +def set_tf_cuda_version(environ_cp): + """Set CUDA_TOOLKIT_PATH and TF_CUDA_VERSION.""" + ask_cuda_version = ( + 'Please specify the CUDA SDK version you want to use, ' + 'e.g. 7.0. [Leave empty to default to CUDA %s]: ') % _DEFAULT_CUDA_VERSION + + while True: + # Configure the Cuda SDK version to use. + tf_cuda_version = get_from_env_or_user_or_default( + environ_cp, 'TF_CUDA_VERSION', ask_cuda_version, _DEFAULT_CUDA_VERSION) + + # Find out where the CUDA toolkit is installed + default_cuda_path = _DEFAULT_CUDA_PATH + if is_windows() or is_cygwin(): + default_cuda_path = cygpath( + environ_cp.get('CUDA_PATH', _DEFAULT_CUDA_PATH_WIN)) + elif is_linux(): + # If the default doesn't exist, try an alternative default. + if (not os.path.exists(default_cuda_path) + ) and os.path.exists(_DEFAULT_CUDA_PATH_LINUX): + default_cuda_path = _DEFAULT_CUDA_PATH_LINUX + ask_cuda_path = ('Please specify the location where CUDA %s toolkit is' + ' installed. Refer to README.md for more details. ' + '[Default is %s]: ') % (tf_cuda_version, default_cuda_path) + cuda_toolkit_path = get_from_env_or_user_or_default( + environ_cp, 'CUDA_TOOLKIT_PATH', ask_cuda_path, default_cuda_path) + + if is_windows(): + cuda_rt_lib_path = 'lib/x64/cudart.lib' + elif is_linux(): + cuda_rt_lib_path = 'lib64/libcudart.so.%s' % tf_cuda_version + elif is_macos(): + cuda_rt_lib_path = 'lib/libcudart.%s.dylib' % tf_cuda_version + + cuda_toolkit_path_full = os.path.join(cuda_toolkit_path, cuda_rt_lib_path) + if os.path.exists(cuda_toolkit_path_full): + break + + # Reset and retry + print('Invalid path to CUDA %s toolkit. %s cannot be found' % + (tf_cuda_version, cuda_toolkit_path_full)) + environ_cp['TF_CUDA_VERSION'] = '' + environ_cp['CUDA_TOOLKIT_PATH'] = '' + + # Set CUDA_TOOLKIT_PATH and TF_CUDA_VERSION + environ_cp['CUDA_TOOLKIT_PATH'] = cuda_toolkit_path + write_action_env_to_bazelrc('CUDA_TOOLKIT_PATH', cuda_toolkit_path) + environ_cp['TF_CUDA_VERSION'] = tf_cuda_version + write_action_env_to_bazelrc('TF_CUDA_VERSION', tf_cuda_version) + + +def set_tf_cunn_version(environ_cp): + """Set CUDNN_INSTALL_PATH and TF_CUDNN_VERSION.""" + ask_cudnn_version = ( + 'Please specify the cuDNN version you want to use. ' + '[Leave empty to default to cuDNN %s.0]: ') % _DEFAULT_CUDNN_VERSION + + while True: + tf_cudnn_version = get_from_env_or_user_or_default( + environ_cp, 'TF_CUDNN_VERSION', ask_cudnn_version, + _DEFAULT_CUDNN_VERSION) + + default_cudnn_path = environ_cp.get('CUDA_TOOLKIT_PATH') + ask_cudnn_path = (r'Please specify the location where cuDNN %s library is ' + 'installed. Refer to README.md for more details. [Default' + ' is %s]:') % (tf_cudnn_version, default_cudnn_path) + cudnn_install_path = get_from_env_or_user_or_default( + environ_cp, 'CUDNN_INSTALL_PATH', ask_cudnn_path, default_cudnn_path) + + # Result returned from "read" will be used unexpanded. That make "~" + # unusable. Going through one more level of expansion to handle that. + cudnn_install_path = os.path.realpath( + os.path.expanduser(cudnn_install_path)) + if is_windows() or is_cygwin(): + cudnn_install_path = cygpath(cudnn_install_path) + + if is_windows(): + cuda_dnn_lib_path = 'lib/x64/cudnn.lib' + cuda_dnn_lib_alt_path = 'lib/x64/cudnn.lib' + elif is_linux(): + cuda_dnn_lib_path = 'lib64/libcudnn.so.%s' % tf_cudnn_version + cuda_dnn_lib_alt_path = 'libcudnn.so.%s' % tf_cudnn_version + elif is_macos(): + cuda_dnn_lib_path = 'lib/libcudnn.%s.dylib' % tf_cudnn_version + cuda_dnn_lib_alt_path = 'libcudnn.%s.dylib' % tf_cudnn_version + + cuda_dnn_lib_path_full = os.path.join(cudnn_install_path, cuda_dnn_lib_path) + cuda_dnn_lib_alt_path_full = os.path.join(cudnn_install_path, + cuda_dnn_lib_alt_path) + if os.path.exists(cuda_dnn_lib_path_full) or os.path.exists( + cuda_dnn_lib_alt_path_full): + break + + # Try another alternative for Linux + if is_linux(): + ldconfig_bin = which('ldconfig') or '/sbin/ldconfig' + cudnn_path_from_ldconfig = run_shell([ldconfig_bin, '-p']) + cudnn_path_from_ldconfig = re.search('.*libcudnn.so .* => (.*)', + cudnn_path_from_ldconfig).group(1) + if os.path.exists('%s.%s' % (cudnn_path_from_ldconfig, tf_cudnn_version)): + cudnn_install_path = os.path.dirname(cudnn_path_from_ldconfig) + break + + # Reset and Retry + print( + 'Invalid path to cuDNN %s toolkit. None of the following files can be ' + 'found:' % tf_cudnn_version) + print(cuda_dnn_lib_path_full) + print(cuda_dnn_lib_alt_path_full) + if is_linux(): + print('%s.%s' % (cudnn_path_from_ldconfig, tf_cudnn_version)) + + environ_cp['TF_CUDNN_VERSION'] = '' + + # Set CUDNN_INSTALL_PATH and TF_CUDNN_VERSION + environ_cp['CUDNN_INSTALL_PATH'] = cudnn_install_path + write_action_env_to_bazelrc('CUDNN_INSTALL_PATH', cudnn_install_path) + environ_cp['TF_CUDNN_VERSION'] = tf_cudnn_version + write_action_env_to_bazelrc('TF_CUDNN_VERSION', tf_cudnn_version) + + +def get_native_cuda_compute_capabilities(environ_cp): + """Get native cuda compute capabilities. + + Args: + environ_cp: copy of the os.environ. + Returns: + string of native cuda compute capabilities, separated by comma. + """ + device_query_bin = os.path.join( + environ_cp.get('CUDA_TOOLKIT_PATH'), 'extras/demo_suite/deviceQuery') + if os.path.isfile(device_query_bin) and os.access(device_query_bin, os.X_OK): + try: + output = run_shell(device_query_bin).split('\n') + pattern = re.compile('[0-9]*\\.[0-9]*') + output = [pattern.search(x) for x in output if 'Capability' in x] + output = ','.join(x.group() for x in output if x is not None) + except subprocess.CalledProcessError: + output = '' + else: + output = '' + return output + + +def set_tf_cuda_compute_capabilities(environ_cp): + """Set TF_CUDA_COMPUTE_CAPABILITIES.""" + while True: + native_cuda_compute_capabilities = get_native_cuda_compute_capabilities( + environ_cp) + if not native_cuda_compute_capabilities: + default_cuda_compute_capabilities = _DEFAULT_CUDA_COMPUTE_CAPABILITIES + else: + default_cuda_compute_capabilities = native_cuda_compute_capabilities + + ask_cuda_compute_capabilities = ( + 'Please specify a list of comma-separated ' + 'Cuda compute capabilities you want to ' + 'build with.\nYou can find the compute ' + 'capability of your device at: ' + 'https://developer.nvidia.com/cuda-gpus.\nPlease' + ' note that each additional compute ' + 'capability significantly increases your ' + 'build time and binary size. [Default is: %s]' % + default_cuda_compute_capabilities) + tf_cuda_compute_capabilities = get_from_env_or_user_or_default( + environ_cp, 'TF_CUDA_COMPUTE_CAPABILITIES', + ask_cuda_compute_capabilities, default_cuda_compute_capabilities) + # Check whether all capabilities from the input is valid + all_valid = True + for compute_capability in tf_cuda_compute_capabilities.split(','): + m = re.match('[0-9]+.[0-9]+', compute_capability) + if not m: + print('Invalid compute capability: ' % compute_capability) + all_valid = False + else: + ver = int(m.group(0).split('.')[0]) + if ver < 3: + print('Only compute capabilities 3.0 or higher are supported.') + all_valid = False + + if all_valid: + break + + # Reset and Retry + environ_cp['TF_CUDA_COMPUTE_CAPABILITIES'] = '' + + # Set TF_CUDA_COMPUTE_CAPABILITIES + environ_cp['TF_CUDA_COMPUTE_CAPABILITIES'] = tf_cuda_compute_capabilities + write_action_env_to_bazelrc('TF_CUDA_COMPUTE_CAPABILITIES', + tf_cuda_compute_capabilities) + + +def set_other_cuda_vars(environ_cp): + """Set other CUDA related variables.""" + if is_windows(): + # The following three variables are needed for MSVC toolchain configuration + # in Bazel + environ_cp['CUDA_PATH'] = environ_cp.get('CUDA_TOOLKIT_PATH') + environ_cp['CUDA_COMPUTE_CAPABILITIES'] = environ_cp.get( + 'TF_CUDA_COMPUTE_CAPABILITIES') + environ_cp['NO_WHOLE_ARCHIVE_OPTION'] = 1 + write_action_env_to_bazelrc('CUDA_PATH', environ_cp.get('CUDA_PATH')) + write_action_env_to_bazelrc('CUDA_COMPUTE_CAPABILITIE', + environ_cp.get('CUDA_COMPUTE_CAPABILITIE')) + write_action_env_to_bazelrc('NO_WHOLE_ARCHIVE_OPTION', + environ_cp.get('NO_WHOLE_ARCHIVE_OPTION')) + write_to_bazelrc('build --config=win-cuda') + write_to_bazelrc('test --config=win-cuda') + else: + # If CUDA is enabled, always use GPU during build and test. + if environ_cp.get('TF_CUDA_CLANG') == '1': + write_to_bazelrc('build --config=cuda_clang') + write_to_bazelrc('test --config=cuda_clang') + else: + write_to_bazelrc('build --config=cuda') + write_to_bazelrc('test --config=cuda') + + +def set_host_cxx_compiler(environ_cp): + """Set HOST_CXX_COMPILER.""" + default_cxx_host_compiler = which('g++') or '' + ask_cxx_host_compiler = ( + 'Please specify which C++ compiler should be used as' + ' the host C++ compiler. [Default is %s]: ') % default_cxx_host_compiler + + while True: + host_cxx_compiler = get_from_env_or_user_or_default( + environ_cp, 'HOST_CXX_COMPILER', ask_cxx_host_compiler, + default_cxx_host_compiler) + if os.path.exists(host_cxx_compiler): + break + + # Reset and retry + print('Invalid C++ compiler path. %s cannot be found' % host_cxx_compiler) + environ_cp['HOST_CXX_COMPILER'] = '' + + # Set HOST_CXX_COMPILER + environ_cp['HOST_CXX_COMPILER'] = host_cxx_compiler + write_action_env_to_bazelrc('HOST_CXX_COMPILER', host_cxx_compiler) + + +def set_host_c_compiler(environ_cp): + """Set HOST_C_COMPILER.""" + default_c_host_compiler = which('gcc') or '' + ask_c_host_compiler = ( + 'Please specify which C compiler should be used as the' + ' host C compiler. [Default is %s]: ') % default_c_host_compiler + + while True: + host_c_compiler = get_from_env_or_user_or_default( + environ_cp, 'HOST_C_COMPILER', ask_c_host_compiler, + default_c_host_compiler) + if os.path.exists(host_c_compiler): + break + + # Reset and retry + print('Invalid C compiler path. %s cannot be found' % host_c_compiler) + environ_cp['HOST_C_COMPILER'] = '' + + # Set HOST_C_COMPILER + environ_cp['HOST_C_COMPILER'] = host_c_compiler + write_action_env_to_bazelrc('HOST_C_COMPILER', host_c_compiler) + + +def set_computecpp_toolkit_path(environ_cp): + """Set COMPUTECPP_TOOLKIT_PATH.""" + ask_computecpp_toolkit_path = ('Please specify the location where ComputeCpp ' + 'for SYCL %s is installed. [Default is %s]: ' + ) % (_TF_OPENCL_VERSION, + _DEFAULT_COMPUTECPP_TOOLKIT_PATH) + + while True: + computecpp_toolkit_path = get_from_env_or_user_or_default( + environ_cp, 'COMPUTECPP_TOOLKIT_PATH', ask_computecpp_toolkit_path, + _DEFAULT_COMPUTECPP_TOOLKIT_PATH) + if is_linux(): + sycl_rt_lib_path = 'lib/libComputeCpp.so' + else: + sycl_rt_lib_path = '' + + sycl_rt_lib_path_full = os.path.join(computecpp_toolkit_path, + sycl_rt_lib_path) + if os.path.exists(sycl_rt_lib_path_full): + break + + print('Invalid SYCL %s library path. %s cannot be found' % + (_TF_OPENCL_VERSION, sycl_rt_lib_path_full)) + environ_cp['COMPUTECPP_TOOLKIT_PATH'] = '' + + # Set COMPUTECPP_TOOLKIT_PATH + environ_cp['COMPUTECPP_TOOLKIT_PATH'] = computecpp_toolkit_path + write_action_env_to_bazelrc('COMPUTECPP_TOOLKIT_PATH', + computecpp_toolkit_path) + + +def set_mpi_home(environ_cp): + """Set MPI_HOME.""" + default_mpi_home = which('mpirun') or which('mpiexec') or '' + default_mpi_home = os.path.dirname(os.path.dirname(default_mpi_home)) + + ask_mpi_home = ('Please specify the MPI toolkit folder. [Default is %s]: ' + ) % default_mpi_home + while True: + mpi_home = get_from_env_or_user_or_default(environ_cp, 'MPI_HOME', + ask_mpi_home, default_mpi_home) + + if os.path.exists(os.path.join(mpi_home, 'include')) and os.path.exists( + os.path.join(mpi_home, 'lib')): + break + + 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')))) + environ_cp['MPI_HOME'] = '' + + # Set MPI_HOME + environ_cp['MPI_HOME'] = str(mpi_home) + + +def set_other_mpi_vars(environ_cp): + """Set other MPI related variables.""" + # Link the MPI header files + mpi_home = environ_cp.get('MPI_HOME') + symlink_force('%s/include/mpi.h' % mpi_home, 'third_party/mpi/mpi.h') + + # Determine if we use OpenMPI or MVAPICH, these require different header files + # to be included here to make bazel dependency checker happy + if os.path.exists(os.path.join(mpi_home, 'include/mpi_portable_platform.h')): + symlink_force( + os.path.join(mpi_home, 'include/mpi_portable_platform.h'), + 'third_party/mpi/mpi_portable_platform.h') + # TODO(gunan): avoid editing files in configure + sed_in_place('third_party/mpi/mpi.bzl', 'MPI_LIB_IS_OPENMPI=False', + 'MPI_LIB_IS_OPENMPI=True') + else: + # MVAPICH / MPICH + symlink_force( + os.path.join(mpi_home, 'include/mpio.h'), 'third_party/mpi/mpio.h') + symlink_force( + os.path.join(mpi_home, 'include/mpicxx.h'), 'third_party/mpi/mpicxx.h') + # TODO(gunan): avoid editing files in configure + sed_in_place('third_party/mpi/mpi.bzl', 'MPI_LIB_IS_OPENMPI=True', + 'MPI_LIB_IS_OPENMPI=False') + + 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') + else: + raise ValueError('Cannot find the MPI library file in %s/lib' % mpi_home) + + +def set_mkl(): + write_to_bazelrc('build:mkl --define using_mkl=true') + write_to_bazelrc('build:mkl -c opt') + write_to_bazelrc('build:mkl --copt="-DEIGEN_USE_VML"') + print( + 'Add "--config=mkl" to your bazel command to build with MKL ' + 'support.\nPlease note that MKL on MacOS or windows is still not ' + 'supported.\nIf you would like to use a local MKL instead of ' + 'downloading, please set the environment variable \"TF_MKL_ROOT\" every ' + 'time before build.') + + +def main(): + # Make a copy of os.environ to be clear when functions and getting and setting + # environment variables. + environ_cp = dict(os.environ) + + bazel_version = check_bazel_version('0.4.5') + + reset_tf_configure_bazelrc() + cleanup_makefile() + setup_python(environ_cp, bazel_version) + run_gen_git_source(environ_cp) + + if is_windows(): + environ_cp['TF_NEED_GCP'] = '0' + environ_cp['TF_NEED_HDFS'] = '0' + environ_cp['TF_NEED_JEMALLOC'] = '0' + environ_cp['TF_NEED_OPENCL'] = '0' + environ_cp['TF_CUDA_CLANG'] = '0' + + if is_macos(): + environ_cp['TF_NEED_JEMALLOC'] = '0' + + set_build_var(environ_cp, 'TF_NEED_JEMALLOC', 'jemalloc as malloc', + 'with_jemalloc', True) + set_build_var(environ_cp, 'TF_NEED_GCP', 'Google Cloud Platform', + 'with_gcp_support', False) + set_build_var(environ_cp, 'TF_NEED_HDFS', 'Hadoop File System', + 'with_hdfs_support', False) + set_build_var(environ_cp, 'TF_ENABLE_XLA', 'XLA JIT', 'with_xla_support', + False) + set_build_var(environ_cp, 'TF_NEED_GDR', 'GDR', 'with_gdr_support', + False) + set_build_var(environ_cp, 'TF_NEED_VERBS', 'VERBS', 'with_verbs_support', + False) + + set_action_env_var(environ_cp, 'TF_NEED_OPENCL', 'OpenCL', False) + if environ_cp.get('TF_NEED_OPENCL') == '1': + set_host_cxx_compiler(environ_cp) + set_host_c_compiler(environ_cp) + set_computecpp_toolkit_path(environ_cp) + + set_action_env_var(environ_cp, 'TF_NEED_CUDA', 'CUDA', False) + if environ_cp.get('TF_NEED_CUDA') == '1': + set_tf_cuda_version(environ_cp) + set_tf_cunn_version(environ_cp) + set_tf_cuda_compute_capabilities(environ_cp) + + set_tf_cuda_clang(environ_cp) + if environ_cp.get('TF_CUDA_CLANG') == '1': + # Set up which clang we should use as the cuda / host compiler. + set_clang_cuda_compiler_path(environ_cp) + else: + # Set up which gcc nvcc should use as the host compiler + # No need to set this on Windows + if not is_windows(): + set_gcc_host_compiler_path(environ_cp) + set_other_cuda_vars(environ_cp) + + set_build_var(environ_cp, 'TF_NEED_MPI', 'MPI', 'with_mpi_support', False) + if environ_cp.get('TF_NEED_MPI') == '1': + set_mpi_home(environ_cp) + set_other_mpi_vars(environ_cp) + + set_cc_opt_flags(environ_cp) + set_mkl() + + +if __name__ == '__main__': + main() diff --git a/tensorflow/BUILD b/tensorflow/BUILD index e437987112b36cb9d2de1d4150693dcd8fd3f305..5b6a18b6a693264c6409bd8dd1953799f6de9d48 100644 --- a/tensorflow/BUILD +++ b/tensorflow/BUILD @@ -9,14 +9,15 @@ licenses(["notice"]) # Apache 2.0 exports_files([ "LICENSE", "ACKNOWLEDGMENTS", + # The leakr files are used by //third_party/cloud_tpu. + "leakr_badwords.dic", + "leakr_badfiles.dic", ]) # Config setting for determining if we are building for Android. config_setting( name = "android", - values = { - "crosstool_top": "//external:android/crosstool", - }, + values = {"crosstool_top": "//external:android/crosstool"}, visibility = ["//visibility:public"], ) @@ -38,6 +39,15 @@ config_setting( visibility = ["//visibility:public"], ) +config_setting( + name = "android_armeabi", + values = { + "crosstool_top": "//external:android/crosstool", + "cpu": "armeabi", + }, + visibility = ["//visibility:public"], +) + config_setting( name = "android_arm", values = { @@ -56,6 +66,24 @@ config_setting( visibility = ["//visibility:public"], ) +config_setting( + name = "android_mips", + values = { + "crosstool_top": "//external:android/crosstool", + "cpu": "mips", + }, + visibility = ["//visibility:public"], +) + +config_setting( + name = "android_mips64", + values = { + "crosstool_top": "//external:android/crosstool", + "cpu": "mips64", + }, + visibility = ["//visibility:public"], +) + config_setting( name = "darwin", values = {"cpu": "darwin"}, @@ -64,6 +92,12 @@ config_setting( config_setting( name = "windows", + values = {"cpu": "x64_windows"}, + visibility = ["//visibility:public"], +) + +config_setting( + name = "windows_msvc", values = {"cpu": "x64_windows_msvc"}, visibility = ["//visibility:public"], ) @@ -76,9 +110,7 @@ config_setting( config_setting( name = "ios", - values = { - "crosstool_top": "//tools/osx/crosstool:crosstool", - }, + values = {"crosstool_top": "//tools/osx/crosstool:crosstool"}, visibility = ["//visibility:public"], ) @@ -88,6 +120,12 @@ config_setting( visibility = ["//visibility:public"], ) +config_setting( + name = "linux_ppc64le", + values = {"cpu": "ppc"}, + visibility = ["//visibility:public"], +) + config_setting( name = "debug", values = { @@ -112,7 +150,7 @@ config_setting( # TODO(jhseu): Enable on other platforms other than Linux. config_setting( - name = "with_jemalloc", + name = "with_jemalloc_linux_x86_64", values = { "cpu": "k8", "define": "with_jemalloc=true", @@ -120,6 +158,15 @@ config_setting( visibility = ["//visibility:public"], ) +config_setting( + name = "with_jemalloc_linux_ppc64le", + values = { + "cpu": "ppc", + "define": "with_jemalloc=true", + }, + visibility = ["//visibility:public"], +) + config_setting( name = "with_gcp_support", values = {"define": "with_gcp_support=true"}, @@ -138,9 +185,30 @@ config_setting( visibility = ["//visibility:public"], ) +config_setting( + name = "with_gdr_support", + values = {"define": "with_gdr_support=true"}, + visibility = ["//visibility:public"], +) + +config_setting( + name = "with_verbs_support", + values = {"define": "with_verbs_support=true"}, + visibility = ["//visibility:public"], +) + +config_setting( + name = "with_mpi_support", + values = {"define": "with_mpi_support=true"}, + visibility = ["//visibility:public"], +) + package_group( name = "internal", - packages = ["//tensorflow/..."], + packages = [ + "//learning/protonn/llgtm/...", + "//tensorflow/...", + ], ) filegroup( @@ -178,14 +246,16 @@ filegroup( "//tensorflow/compiler/jit/kernels:all_files", "//tensorflow/compiler/jit/legacy_flags:all_files", "//tensorflow/compiler/jit/ops:all_files", + "//tensorflow/compiler/plugin/executor:all_files", "//tensorflow/compiler/tests:all_files", "//tensorflow/compiler/tf2xla:all_files", + "//tensorflow/compiler/tf2xla/cc:all_files", "//tensorflow/compiler/tf2xla/kernels:all_files", + "//tensorflow/compiler/tf2xla/ops:all_files", "//tensorflow/compiler/xla:all_files", "//tensorflow/compiler/xla/client:all_files", "//tensorflow/compiler/xla/client/lib:all_files", "//tensorflow/compiler/xla/legacy_flags:all_files", - "//tensorflow/compiler/xla/port:all_files", "//tensorflow/compiler/xla/service:all_files", "//tensorflow/compiler/xla/service/cpu:all_files", "//tensorflow/compiler/xla/service/gpu:all_files", @@ -196,28 +266,41 @@ filegroup( "//tensorflow/contrib:all_files", "//tensorflow/contrib/android:all_files", "//tensorflow/contrib/batching:all_files", + "//tensorflow/contrib/batching/kernels:all_files", "//tensorflow/contrib/batching/test_util:all_files", "//tensorflow/contrib/batching/util:all_files", "//tensorflow/contrib/bayesflow:all_files", "//tensorflow/contrib/boosted_trees:all_files", + "//tensorflow/contrib/boosted_trees/estimator_batch:all_files", "//tensorflow/contrib/boosted_trees/lib:all_files", "//tensorflow/contrib/boosted_trees/proto:all_files", "//tensorflow/contrib/boosted_trees/resources:all_files", "//tensorflow/contrib/cloud:all_files", "//tensorflow/contrib/cloud/kernels:all_files", + "//tensorflow/contrib/cluster_resolver:all_files", "//tensorflow/contrib/compiler:all_files", "//tensorflow/contrib/copy_graph:all_files", "//tensorflow/contrib/crf:all_files", "//tensorflow/contrib/cudnn_rnn:all_files", + "//tensorflow/contrib/data:all_files", + "//tensorflow/contrib/data/python/framework:all_files", + "//tensorflow/contrib/data/python/kernel_tests:all_files", + "//tensorflow/contrib/data/python/ops:all_files", + "//tensorflow/contrib/data/python/util:all_files", + "//tensorflow/contrib/decision_trees/proto:all_files", "//tensorflow/contrib/distributions:all_files", + "//tensorflow/contrib/eager/python:all_files", "//tensorflow/contrib/factorization:all_files", "//tensorflow/contrib/factorization/kernels:all_files", "//tensorflow/contrib/ffmpeg:all_files", "//tensorflow/contrib/ffmpeg/default:all_files", "//tensorflow/contrib/framework:all_files", + "//tensorflow/contrib/fused_conv:all_files", + "//tensorflow/contrib/gan:all_files", "//tensorflow/contrib/graph_editor:all_files", "//tensorflow/contrib/grid_rnn:all_files", "//tensorflow/contrib/hooks:all_files", + "//tensorflow/contrib/hvx/hvx_ops_support_checker:all_files", "//tensorflow/contrib/image:all_files", "//tensorflow/contrib/imperative:all_files", "//tensorflow/contrib/input_pipeline:all_files", @@ -234,29 +317,50 @@ filegroup( "//tensorflow/contrib/linear_optimizer:all_files", "//tensorflow/contrib/lookup:all_files", "//tensorflow/contrib/losses:all_files", + "//tensorflow/contrib/meta_graph_transform:all_files", "//tensorflow/contrib/metrics:all_files", "//tensorflow/contrib/ndlstm:all_files", + "//tensorflow/contrib/nearest_neighbor:all_files", "//tensorflow/contrib/nn:all_files", "//tensorflow/contrib/opt:all_files", + "//tensorflow/contrib/predictor:all_files", + "//tensorflow/contrib/receptive_field:all_files", + "//tensorflow/contrib/reduce_slice_ops:all_files", + "//tensorflow/contrib/remote_fused_graph/pylib:all_files", + "//tensorflow/contrib/resampler:all_files", "//tensorflow/contrib/rnn:all_files", "//tensorflow/contrib/saved_model:all_files", + "//tensorflow/contrib/saved_model/cc/saved_model:all_files", "//tensorflow/contrib/seq2seq:all_files", "//tensorflow/contrib/session_bundle:all_files", "//tensorflow/contrib/session_bundle/example:all_files", + "//tensorflow/contrib/signal:all_files", "//tensorflow/contrib/slim:all_files", "//tensorflow/contrib/slim/python/slim/data:all_files", "//tensorflow/contrib/slim/python/slim/nets:all_files", "//tensorflow/contrib/solvers:all_files", "//tensorflow/contrib/sparsemax:all_files", "//tensorflow/contrib/specs:all_files", + "//tensorflow/contrib/staging:all_files", "//tensorflow/contrib/stat_summarizer:all_files", + "//tensorflow/contrib/stateless:all_files", + "//tensorflow/contrib/summary:all_files", "//tensorflow/contrib/tensor_forest:all_files", "//tensorflow/contrib/tensor_forest/hybrid:all_files", + "//tensorflow/contrib/tensor_forest/kernels/v4:all_files", + "//tensorflow/contrib/tensor_forest/proto:all_files", "//tensorflow/contrib/tensorboard:all_files", "//tensorflow/contrib/testing:all_files", - "//tensorflow/contrib/tfprof/python/tools/tfprof:all_files", + "//tensorflow/contrib/text:all_files", + "//tensorflow/contrib/tfprof:all_files", + "//tensorflow/contrib/timeseries:all_files", + "//tensorflow/contrib/timeseries/examples:all_files", + "//tensorflow/contrib/timeseries/python/timeseries:all_files", + "//tensorflow/contrib/timeseries/python/timeseries/state_space_models:all_files", + "//tensorflow/contrib/tpu:all_files", "//tensorflow/contrib/training:all_files", "//tensorflow/contrib/util:all_files", + "//tensorflow/contrib/verbs:all_files", "//tensorflow/contrib/xla_tf_graph:all_files", "//tensorflow/core:all_files", "//tensorflow/core/debug:all_files", @@ -267,51 +371,52 @@ filegroup( "//tensorflow/core/grappler/costs:all_files", "//tensorflow/core/grappler/inputs:all_files", "//tensorflow/core/grappler/optimizers:all_files", + "//tensorflow/core/grappler/utils:all_files", "//tensorflow/core/kernels:all_files", + "//tensorflow/core/kernels/fuzzing:all_files", "//tensorflow/core/kernels/hexagon:all_files", + "//tensorflow/core/kernels/neon:all_files", "//tensorflow/core/ops/compat:all_files", "//tensorflow/core/platform/cloud:all_files", "//tensorflow/core/platform/default/build_config:all_files", "//tensorflow/core/platform/hadoop:all_files", + "//tensorflow/core/profiler:all_files", + "//tensorflow/core/profiler/internal:all_files", + "//tensorflow/core/profiler/internal/advisor:all_files", "//tensorflow/core/util/ctc:all_files", "//tensorflow/core/util/tensor_bundle:all_files", "//tensorflow/examples/android:all_files", + "//tensorflow/examples/benchmark:all_files", + "//tensorflow/examples/get_started/regression:all_files", "//tensorflow/examples/how_tos/reading_data:all_files", "//tensorflow/examples/image_retraining:all_files", "//tensorflow/examples/label_image:all_files", "//tensorflow/examples/learn:all_files", "//tensorflow/examples/saved_model:all_files", + "//tensorflow/examples/speech_commands:all_files", "//tensorflow/examples/tutorials/estimators:all_files", "//tensorflow/examples/tutorials/mnist:all_files", "//tensorflow/examples/tutorials/word2vec:all_files", + "//tensorflow/examples/wav_to_spectrogram:all_files", "//tensorflow/go:all_files", "//tensorflow/java:all_files", "//tensorflow/java/src/main/java/org/tensorflow/examples:all_files", "//tensorflow/java/src/main/native:all_files", "//tensorflow/python:all_files", "//tensorflow/python/debug:all_files", + "//tensorflow/python/eager:all_files", "//tensorflow/python/estimator:all_files", + "//tensorflow/python/feature_column:all_files", "//tensorflow/python/kernel_tests:all_files", + "//tensorflow/python/kernel_tests/distributions:all_files", + "//tensorflow/python/ops/distributions:all_files", + "//tensorflow/python/profiler:all_files", + "//tensorflow/python/profiler/internal:all_files", "//tensorflow/python/saved_model:all_files", "//tensorflow/python/tools:all_files", - "//tensorflow/tensorboard:all_files", - "//tensorflow/tensorboard/app:all_files", - "//tensorflow/tensorboard/backend:all_files", - "//tensorflow/tensorboard/backend/event_processing:all_files", - "//tensorflow/tensorboard/components:all_files", - "//tensorflow/tensorboard/components/tf_text_dashboard:all_files", - "//tensorflow/tensorboard/components/vz_data_summary:all_files", - "//tensorflow/tensorboard/components/vz_line_chart:all_files", - "//tensorflow/tensorboard/components/vz_line_chart/demo:all_files", - "//tensorflow/tensorboard/components/vz_projector:all_files", - "//tensorflow/tensorboard/components/vz_sorting:all_files", - "//tensorflow/tensorboard/components/vz_sorting/test:all_files", - "//tensorflow/tensorboard/lib:all_files", - "//tensorflow/tensorboard/plugins:all_files", - "//tensorflow/tensorboard/plugins/debugger:all_files", - "//tensorflow/tensorboard/plugins/projector:all_files", - "//tensorflow/tensorboard/plugins/text:all_files", - "//tensorflow/tensorboard/scripts:all_files", + "//tensorflow/tools/api/golden:all_files", + "//tensorflow/tools/api/lib:all_files", + "//tensorflow/tools/api/tests:all_files", "//tensorflow/tools/common:all_files", "//tensorflow/tools/compatibility:all_files", "//tensorflow/tools/dist_test/server:all_files", @@ -319,11 +424,10 @@ filegroup( "//tensorflow/tools/docker/notebooks:all_files", "//tensorflow/tools/docs:all_files", "//tensorflow/tools/git:all_files", + "//tensorflow/tools/mlpbtxt:all_files", "//tensorflow/tools/proto_text:all_files", "//tensorflow/tools/quantization:all_files", "//tensorflow/tools/test:all_files", - "//tensorflow/tools/tfprof:all_files", - "//tensorflow/tools/tfprof/internal:all_files", "//tensorflow/user_ops:all_files", "//third_party/hadoop:all_files", "//third_party/sycl:all_files", @@ -346,14 +450,36 @@ filegroup( ), ) +filegroup( + name = "docs_src", + data = glob(["docs_src/**/*.md"]), +) + # ------------------------------------------- # New rules should be added above this target. # ------------------------------------------- cc_binary( name = "libtensorflow.so", + linkopts = select({ + "//tensorflow:darwin": [ + "-Wl,-exported_symbols_list", # This line must be directly followed by the exported_symbols.lds file + "//tensorflow/c:exported_symbols.lds", + ], + "//tensorflow:windows": [], + "//tensorflow:windows_msvc": [], + "//conditions:default": [ + "-z defs", + "-s", + "-Wl,--version-script", # This line must be directly followed by the version_script.lds file + "//tensorflow/c:version_script.lds", + ], + }), linkshared = 1, deps = [ "//tensorflow/c:c_api", + "//tensorflow/c:exported_symbols.lds", + "//tensorflow/c:version_script.lds", + "//tensorflow/c/eager:c_api", "//tensorflow/core:tensorflow", ], ) @@ -363,6 +489,7 @@ cc_binary( linkshared = 1, deps = [ "//tensorflow/c:c_api", + "//tensorflow/c/eager:c_api", "//tensorflow/cc:cc_ops", "//tensorflow/cc:client_session", "//tensorflow/cc:scope", diff --git a/tensorflow/__init__.py b/tensorflow/__init__.py index 0bca6f8fb8051925908db5e86f30d97d534e60f4..083634bd7964b0c12e10a1f3c71be5eab597a6c4 100644 --- a/tensorflow/__init__.py +++ b/tensorflow/__init__.py @@ -24,19 +24,9 @@ from __future__ import print_function from tensorflow.python import * # pylint: enable=wildcard-import -# Lazily import the `tf.contrib` module. This avoids loading all of the -# dependencies of `tf.contrib` at `import tensorflow` time. -class _LazyContribLoader(object): - - def __getattr__(self, item): - global contrib - # Replace the lazy loader with the imported module itself. - import importlib # pylint: disable=g-import-not-at-top - contrib = importlib.import_module('tensorflow.contrib') - return getattr(contrib, item) - - -contrib = _LazyContribLoader() +from tensorflow.python.util.lazy_loader import LazyLoader +contrib = LazyLoader('contrib', globals(), 'tensorflow.contrib') +del LazyLoader del absolute_import del division diff --git a/tensorflow/c/BUILD b/tensorflow/c/BUILD index 0019dfeeb13f5e591d44dd37d73a93ce64a92d95..1822e235eba3f9919f2d3e19c628fc7160dd1977 100644 --- a/tensorflow/c/BUILD +++ b/tensorflow/c/BUILD @@ -26,25 +26,64 @@ filegroup( visibility = ["//tensorflow:__subpackages__"], ) +tf_cuda_library( + name = "c_api_internal", + srcs = ["c_api.h"], + hdrs = ["c_api_internal.h"], + visibility = ["//tensorflow/c:__subpackages__"], + deps = select({ + "//tensorflow:android": [ + "//tensorflow/core:android_tensorflow_lib_lite", + ], + "//conditions:default": [ + "//tensorflow/core:core_cpu", + "//tensorflow/core:framework", + "//tensorflow/core:lib", + ], + }), +) + tf_cuda_library( name = "c_api", - srcs = ["c_api.cc"], - hdrs = ["c_api.h"], + srcs = [ + "c_api.cc", + "c_api_function.cc", + ], + hdrs = [ + "c_api.h", + ], copts = tf_copts(), visibility = ["//visibility:public"], deps = select({ "//tensorflow:android": [ + ":c_api_internal", "//tensorflow/core:android_tensorflow_lib_lite", ], "//conditions:default": [ + ":c_api_internal", "//tensorflow/cc/saved_model:loader", + "//tensorflow/cc:gradients", + "//tensorflow/cc:ops", + "//tensorflow/cc:grad_ops", + "//tensorflow/cc:scope_internal", + "//tensorflow/cc:while_loop", "//tensorflow/core:core_cpu", + "//tensorflow/core:core_cpu_internal", "//tensorflow/core:framework", + "//tensorflow/core:protos_all_cc", "//tensorflow/core:lib", ], }), ) +exports_files( + [ + "version_script.lds", + "exported_symbols.lds", + ], + visibility = ["//visibility:public"], +) + tf_cuda_library( name = "tf_status_helper", srcs = ["tf_status_helper.cc"], @@ -52,6 +91,7 @@ tf_cuda_library( visibility = ["//visibility:public"], deps = [ ":c_api", + ":c_api_internal", "//tensorflow/core:lib", ], ) @@ -72,6 +112,19 @@ tf_cuda_library( # ----------------------------------------------------------------------------- # Tests +tf_cuda_library( + name = "c_test_util", + testonly = 1, + srcs = ["c_test_util.cc"], + hdrs = ["c_test_util.h"], + deps = [ + ":c_api", + "//tensorflow/core:lib", + "//tensorflow/core:protos_all_cc", + "//tensorflow/core:test", + ], +) + tf_cc_test( name = "c_api_test", size = "small", @@ -89,21 +142,51 @@ tf_cc_test( # linkstatic = tf_kernel_tests_linkstatic(), deps = [ ":c_api", + ":c_test_util", + "//tensorflow/cc:cc_ops", + "//tensorflow/cc:grad_ops", "//tensorflow/cc/saved_model:signature_constants", "//tensorflow/cc/saved_model:tag_constants", "//tensorflow/core:core_cpu_internal", "//tensorflow/core:direct_session", "//tensorflow/core:framework", + "//tensorflow/core:framework_internal", "//tensorflow/core:lib", "//tensorflow/core:proto_text", "//tensorflow/core:protos_all_cc", "//tensorflow/core:test", "//tensorflow/core:test_main", - "//tensorflow/core:testlib", "//tensorflow/core/kernels:array", "//tensorflow/core/kernels:control_flow_ops", "//tensorflow/core/kernels:math", - "//third_party/eigen3", + ], +) + +tf_cc_test( + name = "c_api_function_test", + size = "small", + srcs = ["c_api_function_test.cc"], + deps = [ + ":c_api", + ":c_test_util", + "//tensorflow/core:lib", + "//tensorflow/core:lib_internal", + "//tensorflow/core:protos_all_cc", + "//tensorflow/core:test", + "//tensorflow/core:test_main", + ], +) + +tf_cc_test( + name = "while_loop_test", + size = "small", + srcs = ["while_loop_test.cc"], + deps = [ + ":c_api", + ":c_test_util", + "//tensorflow/core:lib", + "//tensorflow/core:test", + "//tensorflow/core:test_main", ], ) @@ -112,6 +195,20 @@ tf_custom_op_library( srcs = ["test_op.cc"], ) +# ----------------------------------------------------------------------------- +# Python API target + +tf_cuda_library( + name = "python_api", + srcs = ["python_api.cc"], + hdrs = ["python_api.h"], + visibility = ["//tensorflow/python:__pkg__"], + deps = [ + ":c_api", + ":c_api_internal", + ], +) + # ----------------------------------------------------------------------------- # Google-internal targets. diff --git a/tensorflow/c/c_api.cc b/tensorflow/c/c_api.cc index d4bcc01b6b89329ad8149e2e98ac2df5d1c15882..c454c94249bbc59a491eb00675d2c188b6bf9d1b 100644 --- a/tensorflow/c/c_api.cc +++ b/tensorflow/c/c_api.cc @@ -21,15 +21,25 @@ limitations under the License. #include #ifndef __ANDROID__ +#include "tensorflow/cc/framework/gradients.h" +#include "tensorflow/cc/framework/ops.h" +#include "tensorflow/cc/framework/scope_internal.h" +#include "tensorflow/cc/ops/while_loop.h" #include "tensorflow/cc/saved_model/loader.h" #endif +#include "tensorflow/c/c_api_internal.h" +#include "tensorflow/core/common_runtime/device_mgr.h" #include "tensorflow/core/common_runtime/shape_refiner.h" +#include "tensorflow/core/framework/allocation_description.pb.h" #include "tensorflow/core/framework/log_memory.h" #include "tensorflow/core/framework/node_def_util.h" #include "tensorflow/core/framework/op_kernel.h" #include "tensorflow/core/framework/partial_tensor_shape.h" #include "tensorflow/core/framework/tensor.h" #include "tensorflow/core/framework/tensor_shape.h" +#include "tensorflow/core/framework/tensor_shape.pb.h" +#include "tensorflow/core/framework/types.h" +#include "tensorflow/core/framework/versions.pb.h" #include "tensorflow/core/graph/graph.h" #include "tensorflow/core/graph/graph_constructor.h" #include "tensorflow/core/graph/node_builder.h" @@ -49,37 +59,34 @@ limitations under the License. // The implementation below is at the top level instead of the // brain namespace because we are defining 'extern "C"' functions. -using tensorflow::error::Code; -using tensorflow::errors::InvalidArgument; -using tensorflow::gtl::ArraySlice; -using tensorflow::strings::StrCat; using tensorflow::AllocationDescription; using tensorflow::DataType; -using tensorflow::Env; using tensorflow::Graph; using tensorflow::GraphDef; -using tensorflow::mutex; -using tensorflow::mutex_lock; using tensorflow::NameRangeMap; using tensorflow::NameRangesForNode; using tensorflow::NewSession; using tensorflow::Node; -using tensorflow::NodeDef; using tensorflow::NodeBuilder; +using tensorflow::NodeDef; using tensorflow::OpDef; using tensorflow::OpRegistry; using tensorflow::PartialTensorShape; -using tensorflow::Reset; using tensorflow::RunMetadata; using tensorflow::RunOptions; using tensorflow::Session; -using tensorflow::SessionOptions; using tensorflow::Status; using tensorflow::Tensor; using tensorflow::TensorBuffer; using tensorflow::TensorId; using tensorflow::TensorShape; using tensorflow::TensorShapeProto; +using tensorflow::error::Code; +using tensorflow::errors::FailedPrecondition; +using tensorflow::errors::InvalidArgument; +using tensorflow::gtl::ArraySlice; +using tensorflow::mutex_lock; +using tensorflow::strings::StrCat; extern "C" { @@ -93,9 +100,6 @@ size_t TF_DataTypeSize(TF_DataType dt) { } // -------------------------------------------------------------------------- -struct TF_Status { - Status status; -}; TF_Status* TF_NewStatus() { return new TF_Status; } @@ -143,7 +147,7 @@ class TF_ManagedBuffer : public TensorBuffer { void* allocate_tensor(const char* operation, size_t len) { void* data = tensorflow::cpu_allocator()->AllocateRaw(EIGEN_MAX_ALIGN_BYTES, len); - if (tensorflow::LogMemory::IsEnabled()) { + if (tensorflow::LogMemory::IsEnabled() && data != nullptr) { tensorflow::LogMemory::RecordRawAllocation( operation, tensorflow::LogMemory::EXTERNAL_TENSOR_ALLOCATION_STEP_ID, len, data, tensorflow::cpu_allocator()); @@ -152,7 +156,7 @@ void* allocate_tensor(const char* operation, size_t len) { } void deallocate_buffer(void* data, size_t len, void* arg) { - if (tensorflow::LogMemory::IsEnabled()) { + if (tensorflow::LogMemory::IsEnabled() && data != nullptr) { tensorflow::LogMemory::RecordRawDeallocation( "TensorFlow C Api", tensorflow::LogMemory::EXTERNAL_TENSOR_ALLOCATION_STEP_ID, data, @@ -161,29 +165,9 @@ void deallocate_buffer(void* data, size_t len, void* arg) { tensorflow::cpu_allocator()->DeallocateRaw(data); } -Status MessageToBuffer(const tensorflow::protobuf::Message& in, - TF_Buffer* out) { - if (out->data != nullptr) { - return InvalidArgument("Passing non-empty TF_Buffer is invalid."); - } - const auto proto_size = in.ByteSize(); - void* buf = tensorflow::port::Malloc(proto_size); - in.SerializeToArray(buf, proto_size); - out->data = buf; - out->length = proto_size; - out->data_deallocator = [](void* data, size_t length) { - tensorflow::port::Free(data); - }; - return Status::OK(); -} - } // namespace -struct TF_Tensor { - TF_DataType dtype; - TensorShape shape; - TensorBuffer* buffer; -}; +TF_Tensor::~TF_Tensor() { buffer->Unref(); } TF_Tensor* TF_AllocateTensor(TF_DataType dtype, const int64_t* dims, int num_dims, size_t len) { @@ -203,9 +187,16 @@ TF_Tensor* TF_NewTensor(TF_DataType dtype, const int64_t* dims, int num_dims, TF_ManagedBuffer* buf = new TF_ManagedBuffer; buf->len_ = len; - if (reinterpret_cast(data) % EIGEN_MAX_ALIGN_BYTES != 0) { - // Copy the data into a buffer that satisfies Eigen's alignment - // requirements. + if (dtype != TF_STRING && dtype != TF_RESOURCE && + tensorflow::DataTypeCanUseMemcpy(static_cast(dtype)) && + reinterpret_cast(data) % EIGEN_MAX_ALIGN_BYTES != 0) { + // TF_STRING and TF_RESOURCE tensors have a different representation in + // TF_Tensor than they do in tensorflow::Tensor. So a copy here is a waste + // (any alignement requirements will be taken care of by TF_TensorToTensor + // and TF_TensorFromTensor). + // + // Other types have the same represntation, 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; @@ -220,11 +211,20 @@ TF_Tensor* TF_NewTensor(TF_DataType dtype, const int64_t* dims, int num_dims, return new TF_Tensor{dtype, TensorShape(dimvec), buf}; } -void TF_DeleteTensor(TF_Tensor* t) { - t->buffer->Unref(); - delete t; +TF_Tensor* TF_TensorMaybeMove(TF_Tensor* tensor) { + // It is safe to move the Tensor if and only if we own the unique reference to + // it. In that case, we might as well not delete and reallocate, but a future + // implementation might need to do so. + TensorBuffer* buf = tensor->buffer; + if (buf->RefCountIsOne() && buf->root_buffer()->RefCountIsOne() && + buf->OwnsMemory()) { + return tensor; + } + return nullptr; } +void TF_DeleteTensor(TF_Tensor* t) { delete t; } + TF_DataType TF_TensorType(const TF_Tensor* t) { return t->dtype; } int TF_NumDims(const TF_Tensor* t) { return t->shape.dims(); } int64_t TF_Dim(const TF_Tensor* t, int dim_index) { @@ -252,24 +252,27 @@ size_t TF_StringEncode(const char* src, size_t src_len, char* dst, return sz; } -size_t TF_StringDecode(const char* src, size_t src_len, const char** dst, - size_t* dst_len, TF_Status* status) { +static Status TF_StringDecode_Impl(const char* src, size_t src_len, + const char** dst, size_t* dst_len) { tensorflow::uint64 len64 = 0; const char* p = tensorflow::core::GetVarint64Ptr(src, src + src_len, &len64); if (p == nullptr) { - status->status = - InvalidArgument("invalid string encoding or truncated src buffer"); - return 0; + return InvalidArgument("invalid string encoding or truncated src buffer"); } if (len64 > std::numeric_limits::max()) { - status->status = - InvalidArgument("encoded string is ", len64, - "-bytes, which is too large for this architecture"); - return 0; + return InvalidArgument("encoded string is ", len64, + "-bytes, which is too large for this architecture"); } *dst = p; *dst_len = static_cast(len64); - return static_cast(p - src) + *dst_len; + return Status::OK(); +} + +size_t TF_StringDecode(const char* src, size_t src_len, const char** dst, + size_t* dst_len, TF_Status* status) { + status->status = TF_StringDecode_Impl(src, src_len, dst, dst_len); + if (!status->status.ok()) return 0; + return static_cast(*dst - src) + *dst_len; } size_t TF_StringEncodedSize(size_t len) { @@ -277,9 +280,6 @@ size_t TF_StringEncodedSize(size_t len) { } // -------------------------------------------------------------------------- -struct TF_SessionOptions { - SessionOptions options; -}; TF_SessionOptions* TF_NewSessionOptions() { return new TF_SessionOptions; } void TF_DeleteSessionOptions(TF_SessionOptions* opt) { delete opt; } @@ -320,9 +320,6 @@ void TF_DeleteBuffer(TF_Buffer* buffer) { TF_Buffer TF_GetBuffer(TF_Buffer* buffer) { return *buffer; } // -------------------------------------------------------------------------- -struct TF_DeprecatedSession { - Session* session; -}; TF_DeprecatedSession* TF_NewDeprecatedSession(const TF_SessionOptions* opt, TF_Status* status) { @@ -332,7 +329,7 @@ TF_DeprecatedSession* TF_NewDeprecatedSession(const TF_SessionOptions* opt, return new TF_DeprecatedSession({session}); } else { DCHECK_EQ(nullptr, session); - return NULL; + return nullptr; } } @@ -391,16 +388,36 @@ void TF_Reset(const TF_SessionOptions* opt, const char** containers, namespace tensorflow { -// Non-static for testing. -bool TF_Tensor_DecodeStrings(TF_Tensor* src, Tensor* dst, TF_Status* status) { +Status TF_TensorToTensor(const TF_Tensor* src, Tensor* dst) { + if (src->dtype == TF_RESOURCE) { + if (src->shape.dims() != 0) { + return InvalidArgument( + "Malformed TF_RESOURCE tensor: expected a scalar, got a tensor with " + "shape ", + src->shape.DebugString()); + } + *dst = Tensor(DT_RESOURCE, src->shape); + if (!dst->scalar()().ParseFromString( + string(static_cast(TF_TensorData(src)), + TF_TensorByteSize(src)))) { + return InvalidArgument( + "Malformed TF_RESOUCE tensor: unable to parse resource handle"); + } + return Status::OK(); + } + if (src->dtype != TF_STRING) { + *dst = TensorCApi::MakeTensor(src->dtype, src->shape, src->buffer); + return Status::OK(); + } + // TF_STRING tensors require copying since Tensor class expects a sequence of + // string objects. const tensorflow::int64 num_elements = src->shape.num_elements(); const char* input = reinterpret_cast(TF_TensorData(src)); const size_t src_size = TF_TensorByteSize(src); if (static_cast(src_size / sizeof(tensorflow::uint64)) < num_elements) { - status->status = InvalidArgument( + return InvalidArgument( "Malformed TF_STRING tensor; too short to hold number of elements"); - return false; } const char* data_start = input + sizeof(tensorflow::uint64) * num_elements; const char* limit = input + src_size; @@ -411,24 +428,73 @@ bool TF_Tensor_DecodeStrings(TF_Tensor* src, Tensor* dst, TF_Status* status) { tensorflow::uint64 offset = reinterpret_cast(input)[i]; if (static_cast(offset) >= (limit - data_start)) { - status->status = InvalidArgument("Malformed TF_STRING tensor; element ", - i, " out of range"); - return false; + return InvalidArgument("Malformed TF_STRING tensor; element ", i, + " out of range"); } size_t len; const char* p; const char* srcp = data_start + offset; - TF_StringDecode(srcp, limit - srcp, &p, &len, status); - if (!status->status.ok()) { - return false; - } + Status status = TF_StringDecode_Impl(srcp, limit - srcp, &p, &len); + if (!status.ok()) return status; dstarray(i).assign(p, len); } - return true; + return Status::OK(); +} + +// Create an empty tensor of type 'dtype'. 'shape' can be arbitrary, but has to +// result in a zero-sized tensor. +static TF_Tensor* EmptyTensor(TF_DataType dtype, const TensorShape& shape) { + static char empty; + tensorflow::int64 nelems = 1; + std::vector dims; + for (int i = 0; i < shape.dims(); ++i) { + dims.push_back(shape.dim_size(i)); + nelems *= shape.dim_size(i); + } + 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); } // Non-static for testing. -TF_Tensor* TF_Tensor_EncodeStrings(const Tensor& src) { +TF_Tensor* TF_TensorFromTensor(const tensorflow::Tensor& src, + TF_Status* status) { + if (!src.IsInitialized()) { + status->status = FailedPrecondition( + "attempt to use a tensor with an uninitialized value"); + return nullptr; + } + if (src.NumElements() == 0) { + return EmptyTensor(static_cast(src.dtype()), src.shape()); + } + if (src.dtype() == DT_RESOURCE) { + if (src.shape().dims() != 0) { + status->status = InvalidArgument( + "Unexpected non-scalar DT_RESOURCE tensor seen (shape: ", + src.shape().DebugString(), + "). Please file a bug at " + "https://github.com/tensorflow/tensorflow/issues/new, " + "ideally with a " + "short code snippet that reproduces this error."); + return nullptr; + } + const string str = src.scalar()().SerializeAsString(); + TF_Tensor* t = TF_AllocateTensor(TF_RESOURCE, {}, 0, str.size()); + std::memcpy(TF_TensorData(t), str.c_str(), str.size()); + return t; + } + if (src.dtype() != DT_STRING) { + TensorBuffer* buf = TensorCApi::Buffer(src); + buf->Ref(); + return new TF_Tensor{static_cast(src.dtype()), src.shape(), + buf}; + } + // DT_STRING tensors require a copying since TF_Tensor.buffer expects a flatly + // encoded sequence of strings. + // Compute bytes needed for encoding. size_t size = 0; const auto& srcarray = src.flat(); @@ -444,18 +510,26 @@ TF_Tensor* TF_Tensor_EncodeStrings(const Tensor& src) { char* dst = data_start; // Where next string is encoded. size_t dst_len = size - static_cast(data_start - base); tensorflow::uint64* offsets = reinterpret_cast(base); - TF_Status status; for (int i = 0; i < srcarray.size(); ++i) { *offsets = (dst - data_start); offsets++; const tensorflow::string& s = srcarray(i); - size_t consumed = - TF_StringEncode(s.data(), s.size(), dst, dst_len, &status); - CHECK(status.status.ok()); + size_t consumed = TF_StringEncode(s.data(), s.size(), dst, dst_len, status); + if (!status->status.ok()) { + status->status = InvalidArgument( + "invalid string tensor encoding (string #", i, " of ", + srcarray.size(), "): ", status->status.error_message()); + return nullptr; + } dst += consumed; dst_len -= consumed; } - CHECK_EQ(dst, base + size); + if (dst != base + size) { + status->status = InvalidArgument( + "invalid string tensor encoding (decoded ", (dst - base), + " bytes, but the tensor is encoded in ", size, " bytes"); + return nullptr; + } auto dims = src.shape().dim_sizes(); std::vector dimvec(dims.size()); @@ -469,31 +543,25 @@ TF_Tensor* TF_Tensor_EncodeStrings(const Tensor& src) { dimvec.size(), base, size, DeleteArray, base); } -class TensorCApi { - public: - static TensorBuffer* Buffer(const Tensor& tensor) { return tensor.buf_; } - static Tensor MakeTensor(TF_DataType type, const TensorShape& shape, - TensorBuffer* buf) { - return Tensor(static_cast(type), shape, buf); +Status MessageToBuffer(const tensorflow::protobuf::Message& in, + TF_Buffer* out) { + if (out->data != nullptr) { + return InvalidArgument("Passing non-empty TF_Buffer is invalid."); } -}; - -// Create an empty tensor of type 'dtype'. 'shape' can be arbitrary, but has to -// result in a zero-sized tensor. -static TF_Tensor* EmptyTensor(TF_DataType dtype, const TensorShape& shape) { - static char empty; - tensorflow::int64 nelems = 1; - std::vector dims; - for (int i = 0; i < shape.dims(); ++i) { - dims.push_back(shape.dim_size(i)); - nelems *= shape.dim_size(i); + const size_t proto_size = in.ByteSizeLong(); + void* buf = tensorflow::port::Malloc(proto_size); + if (buf == nullptr) { + return tensorflow::errors::ResourceExhausted( + "Failed to allocate memory to serialize message of type '", + in.GetTypeName(), "' and size ", proto_size); } - 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); + in.SerializeToArray(buf, proto_size); + out->data = buf; + out->length = proto_size; + out->data_deallocator = [](void* data, size_t length) { + tensorflow::port::Free(data); + }; + return Status::OK(); } // Helpers for loading a TensorFlow plugin (a .so file). @@ -506,7 +574,7 @@ static void TF_Run_Setup(int noutputs, TF_Tensor** c_outputs, TF_Status* status) { status->status = Status::OK(); for (int i = 0; i < noutputs; ++i) { - c_outputs[i] = NULL; + c_outputs[i] = nullptr; } } @@ -516,16 +584,8 @@ static bool TF_Run_Inputs( TF_Status* status) { const int ninputs = input_pairs->size(); for (int i = 0; i < ninputs; ++i) { - TF_Tensor* src = c_inputs[i]; - if (c_inputs[i]->dtype != TF_STRING) { - (*input_pairs)[i].second = tensorflow::TensorCApi::MakeTensor( - src->dtype, src->shape, src->buffer); - } else if (!tensorflow::TF_Tensor_DecodeStrings( - src, &(*input_pairs)[i].second, status)) { - // TF_STRING tensors require copying since Tensor class expects - // a sequence of string objects. - return false; - } + status->status = TF_TensorToTensor(c_inputs[i], &(*input_pairs)[i].second); + if (!status->status.ok()) return false; } return true; } @@ -579,19 +639,12 @@ static void TF_Run_Helper( for (int i = 0; i < noutputs; ++i) { const Tensor& src = outputs[i]; if (!src.IsInitialized() || src.NumElements() == 0) { - c_outputs[i] = tensorflow::EmptyTensor( - static_cast(src.dtype()), src.shape()); + c_outputs[i] = + EmptyTensor(static_cast(src.dtype()), src.shape()); continue; } - if (src.dtype() != tensorflow::DT_STRING) { - // Share the underlying buffer. - TensorBuffer* buf = tensorflow::TensorCApi::Buffer(src); - buf->Ref(); - c_outputs[i] = new TF_Tensor{static_cast(src.dtype()), - src.shape(), buf}; - } else { - c_outputs[i] = tensorflow::TF_Tensor_EncodeStrings(src); - } + c_outputs[i] = TF_TensorFromTensor(src, status); + if (!status->status.ok()) return; } } @@ -631,7 +684,7 @@ void TF_PRunSetup(TF_DeprecatedSession* s, // Target nodes const char** c_target_oper_names, int ntargets, const char** handle, TF_Status* status) { - status->status = Status::OK(); + *handle = nullptr; std::vector input_names(ninputs); std::vector output_names(noutputs); @@ -646,15 +699,12 @@ void TF_PRunSetup(TF_DeprecatedSession* s, target_oper_names[i] = c_target_oper_names[i]; } tensorflow::string new_handle; - Status result; - result = s->session->PRunSetup(input_names, output_names, target_oper_names, - &new_handle); - if (result.ok()) { + status->status = s->session->PRunSetup(input_names, output_names, + target_oper_names, &new_handle); + if (status->status.ok()) { char* buf = new char[new_handle.size() + 1]; memcpy(buf, new_handle.c_str(), new_handle.size() + 1); *handle = buf; - } else { - status->status = result; } } @@ -685,11 +735,6 @@ void TF_PRun(TF_DeprecatedSession* s, const char* handle, c_outputs, target_oper_names, nullptr, status); } -struct TF_Library { - void* lib_handle; - TF_Buffer op_list; -}; - TF_Library* TF_LoadLibrary(const char* library_filename, TF_Status* status) { TF_Library* lib_handle = new TF_Library; status->status = tensorflow::LoadLibrary( @@ -721,71 +766,54 @@ TF_Buffer* TF_GetAllOpList() { return ret; } -} // end extern "C" - // -------------------------------------------------------------------------- -// New Graph and Session API +// ListDevices & SessionListDevices API -// Structures ----------------------------------------------------------------- +void TF_DeleteDeviceList(TF_DeviceList* s) { delete s; } -extern "C" { +TF_DeviceList* TF_SessionListDevices(TF_Session* session, TF_Status* status) { + TF_DeviceList* response = new TF_DeviceList; + status->status = session->session->ListDevices(&response->response); + return response; +} -struct TF_Graph { - TF_Graph() - : graph(OpRegistry::Global()), - refiner(graph.versions().producer(), graph.op_registry()), - num_sessions(0), - delete_requested(false), - parent(nullptr), - parent_inputs(nullptr) {} - mutex mu; - Graph graph GUARDED_BY(mu); - - // Runs shape inference. - tensorflow::ShapeRefiner refiner GUARDED_BY(mu); - - // Maps from name of an operation to the Node* in 'graph'. - std::unordered_map name_map GUARDED_BY(mu); - - // TF_Graph may only / must be deleted when - // num_sessions == 0 && delete_requested == true - - // num_sessions incremented by TF_NewSession, and decremented by - // TF_DeleteSession. - int num_sessions GUARDED_BY(mu); - bool delete_requested GUARDED_BY(mu); // set true by TF_DeleteGraph - - // Used to link graphs contained in TF_WhileParams to the parent graph that - // will eventually contain the full while loop. - TF_Graph* parent; - TF_Output* parent_inputs; -}; +TF_DeviceList* TF_DeprecatedSessionListDevices(TF_DeprecatedSession* session, + TF_Status* status) { + TF_DeviceList* response = new TF_DeviceList; + status->status = session->session->ListDevices(&response->response); + return response; +} -struct TF_OperationDescription { - TF_OperationDescription(TF_Graph* g, const char* op_type, - const char* node_name) - : node_builder(node_name, op_type, g->graph.op_registry()), graph(g) {} +int TF_DeviceListCount(const TF_DeviceList* list) { + return list->response.size(); +} - NodeBuilder node_builder; - TF_Graph* graph; - std::vector colocation_constraints; -}; +#define TF_DEVICELIST_METHOD(return_type, method_name, accessor, err_val) \ + return_type method_name(const TF_DeviceList* list, const int index, \ + TF_Status* status) { \ + if (list == nullptr) { \ + status->status = InvalidArgument("list is null!"); \ + return err_val; \ + } \ + if (index < 0 || index >= list->response.size()) { \ + status->status = InvalidArgument("index out of bounds"); \ + return err_val; \ + } \ + return list->response[index].accessor; \ + } -struct TF_Operation { - Node node; -}; +TF_DEVICELIST_METHOD(const char*, TF_DeviceListName, name().c_str(), nullptr); +TF_DEVICELIST_METHOD(const char*, TF_DeviceListType, device_type().c_str(), + nullptr); +TF_DEVICELIST_METHOD(int64_t, TF_DeviceListMemoryBytes, memory_limit(), -1); -struct TF_Session { - TF_Session(Session* s, TF_Graph* g) - : session(s), graph(g), last_num_graph_nodes(0) {} - Session* session; - TF_Graph* graph; - mutex mu; - int last_num_graph_nodes; -}; +#undef TF_DEVICELIST_METHOD } // end extern "C" +// -------------------------------------------------------------------------- +// New Graph and Session API + // Helper functions ----------------------------------------------------------- namespace { @@ -801,8 +829,7 @@ tensorflow::string OutputName(const TF_Output& output) { const tensorflow::AttrValue* GetAttrValue(TF_Operation* oper, const char* attr_name, TF_Status* status) { - const tensorflow::AttrValue* attr = - tensorflow::AttrSlice(oper->node.def()).Find(attr_name); + const tensorflow::AttrValue* attr = oper->node.attrs().Find(attr_name); if (attr == nullptr) { status->status = InvalidArgument("Operation has no attr named '", attr_name, "'."); @@ -810,6 +837,30 @@ const tensorflow::AttrValue* GetAttrValue(TF_Operation* oper, return attr; } +TensorId ToTensorId(const TF_Output& output) { + return TensorId(output.oper->node.name(), output.index); +} + +#ifndef __ANDROID__ +std::vector OutputsFromTFOutputs(TF_Output* tf_outputs, + int n) { + std::vector outputs(n); + for (int i = 0; i < n; ++i) { + outputs[i] = + tensorflow::Output(&tf_outputs[i].oper->node, tf_outputs[i].index); + } + return outputs; +} + +void TFOutputsFromOutputs(const std::vector& outputs, + TF_Output* tf_outputs) { + for (int i = 0; i < outputs.size(); i++) { + tf_outputs[i].oper = ToOperation(outputs[i].node()); + tf_outputs[i].index = outputs[i].index(); + } +} +#endif // __ANDROID__ + } // namespace // Shape functions ----------------------------------------------------------- @@ -830,6 +881,7 @@ void TF_GraphSetTensorShape(TF_Graph* graph, TF_Output output, } std::vector dim_vec; + dim_vec.reserve(num_dims); for (int i = 0; i < num_dims; ++i) { dim_vec.push_back(ic->MakeDim(dims[i])); } @@ -944,7 +996,7 @@ void TF_AddControlInput(TF_OperationDescription* desc, TF_Operation* input) { } void TF_ColocateWith(TF_OperationDescription* desc, TF_Operation* op) { - desc->colocation_constraints.emplace_back( + desc->colocation_constraints.emplace( StrCat(tensorflow::kColocationGroupPrefix, op->node.name())); } @@ -957,12 +1009,20 @@ void TF_SetAttrString(TF_OperationDescription* desc, const char* attr_name, void TF_SetAttrStringList(TF_OperationDescription* desc, const char* attr_name, const void* const* values, const size_t* lengths, int num_values) { - std::vector v; - v.reserve(num_values); - for (int i = 0; i < num_values; ++i) { - v.emplace_back(static_cast(values[i]), lengths[i]); + if (strcmp(attr_name, tensorflow::kColocationAttrName) == 0) { + desc->colocation_constraints.clear(); + for (int i = 0; i < num_values; ++i) { + desc->colocation_constraints.emplace(static_cast(values[i]), + lengths[i]); + } + } else { + std::vector v; + v.reserve(num_values); + for (int i = 0; i < num_values; ++i) { + v.emplace_back(static_cast(values[i]), lengths[i]); + } + desc->node_builder.Attr(attr_name, v); } - desc->node_builder.Attr(attr_name, v); } void TF_SetAttrInt(TF_OperationDescription* desc, const char* attr_name, @@ -1096,20 +1156,9 @@ void TF_SetAttrTensorShapeProtoList(TF_OperationDescription* desc, void TF_SetAttrTensor(TF_OperationDescription* desc, const char* attr_name, TF_Tensor* value, TF_Status* status) { - status->status = Status::OK(); Tensor t; - bool ok = true; - - if (value->dtype != TF_STRING) { - t = tensorflow::TensorCApi::MakeTensor(value->dtype, value->shape, - value->buffer); - } else { - // TF_STRING tensors require copying since Tensor class expects - // a sequence of string objects. - ok = tensorflow::TF_Tensor_DecodeStrings(value, &t, status); - } - - if (ok) desc->node_builder.Attr(attr_name, t); + status->status = TF_TensorToTensor(value, &t); + if (status->status.ok()) desc->node_builder.Attr(attr_name, t); } void TF_SetAttrTensorList(TF_OperationDescription* desc, const char* attr_name, @@ -1118,33 +1167,42 @@ void TF_SetAttrTensorList(TF_OperationDescription* desc, const char* attr_name, status->status = Status::OK(); std::vector t; t.reserve(num_values); - bool ok = true; - for (int i = 0; i < num_values && ok; ++i) { - if (values[i]->dtype != TF_STRING) { - t.emplace_back(tensorflow::TensorCApi::MakeTensor( - values[i]->dtype, values[i]->shape, values[i]->buffer)); - } else { - t.emplace_back(::tensorflow::DT_STRING); - // TF_STRING tensors require copying since Tensor class expects - // a sequence of string objects. - ok = tensorflow::TF_Tensor_DecodeStrings(values[i], &t.back(), status); - } + for (int i = 0; i < num_values && status->status.ok(); ++i) { + Tensor v; + status->status = TF_TensorToTensor(values[i], &v); + t.emplace_back(v); } - if (ok) desc->node_builder.Attr(attr_name, t); + if (status->status.ok()) desc->node_builder.Attr(attr_name, t); } void TF_SetAttrValueProto(TF_OperationDescription* desc, const char* attr_name, const void* proto, size_t proto_len, TF_Status* status) { tensorflow::AttrValue attr_value; - if (attr_value.ParseFromArray(proto, proto_len)) { - desc->node_builder.Attr(attr_name, attr_value); - status->status = Status::OK(); - } else { + if (!attr_value.ParseFromArray(proto, proto_len)) { status->status = InvalidArgument("Unparseable AttrValue proto"); + return; + } + + if (strcmp(attr_name, tensorflow::kColocationAttrName) == 0) { + if (attr_value.value_case() != tensorflow::AttrValue::kList && + attr_value.value_case() != tensorflow::AttrValue::VALUE_NOT_SET) { + status->status = + InvalidArgument("Expected \"list\" field for \"", + tensorflow::kColocationAttrName, "\" attribute"); + return; + } + desc->colocation_constraints.clear(); + for (const tensorflow::string& location : attr_value.list().s()) { + desc->colocation_constraints.insert(location); + } + } else { + desc->node_builder.Attr(attr_name, attr_value); } + + status->status = Status::OK(); } static TF_Operation* TF_FinishOperationLocked(TF_OperationDescription* desc, @@ -1156,22 +1214,24 @@ static TF_Operation* TF_FinishOperationLocked(TF_OperationDescription* desc, status->status = InvalidArgument("Duplicate node name in graph: '", desc->node_builder.node_name(), "'"); } else { - std::sort(desc->colocation_constraints.begin(), - desc->colocation_constraints.end()); - desc->node_builder.Attr(tensorflow::kColocationAttrName, - desc->colocation_constraints); + if (!desc->colocation_constraints.empty()) { + desc->node_builder.Attr( + tensorflow::kColocationAttrName, + std::vector(desc->colocation_constraints.begin(), + desc->colocation_constraints.end())); + } status->status = desc->node_builder.Finalize(&desc->graph->graph, &ret); if (status->status.ok()) { // Run shape inference function for newly added node. - // - // TODO(b/28152992): Enable returning the result of this - // code-path once we have converted all python shape functions - // to call their C++ versions. - desc->graph->refiner.AddNode(ret).IgnoreError(); - + status->status = desc->graph->refiner.AddNode(ret); + } + if (status->status.ok()) { // Add the node to the name-to-node mapping. desc->graph->name_map[ret->name()] = ret; + } else if (ret != nullptr) { + desc->graph->graph.RemoveNode(ret); + ret = nullptr; } } @@ -1198,7 +1258,7 @@ const char* TF_OperationOpType(TF_Operation* oper) { } const char* TF_OperationDevice(TF_Operation* oper) { - return oper->node.def().device().c_str(); + return oper->node.requested_device().c_str(); } int TF_OperationNumOutputs(TF_Operation* oper) { @@ -1213,8 +1273,8 @@ TF_DataType TF_OperationOutputType(TF_Output oper_out) { int TF_OperationOutputListLength(TF_Operation* oper, const char* arg_name, TF_Status* status) { NameRangeMap name_ranges; - status->status = NameRangesForNode(oper->node.def(), oper->node.op_def(), - nullptr, &name_ranges); + status->status = + NameRangesForNode(oper->node, oper->node.op_def(), nullptr, &name_ranges); if (!status->status.ok()) return -1; auto iter = name_ranges.find(arg_name); if (iter == name_ranges.end()) { @@ -1235,8 +1295,8 @@ TF_DataType TF_OperationInputType(TF_Input oper_in) { int TF_OperationInputListLength(TF_Operation* oper, const char* arg_name, TF_Status* status) { NameRangeMap name_ranges; - status->status = NameRangesForNode(oper->node.def(), oper->node.op_def(), - &name_ranges, nullptr); + status->status = + NameRangesForNode(oper->node, oper->node.op_def(), &name_ranges, nullptr); if (!status->status.ok()) return -1; auto iter = name_ranges.find(arg_name); if (iter == name_ranges.end()) { @@ -1474,26 +1534,27 @@ void TF_OperationGetAttrStringList(TF_Operation* oper, const char* attr_name, } } -#define DEFINE_GETATTR(func, c_type, cpp_type, list_field) \ - void func(TF_Operation* oper, const char* attr_name, c_type* value, \ - TF_Status* status) { \ - cpp_type v; \ - status->status = tensorflow::GetNodeAttr(oper->node.def(), attr_name, &v); \ - *value = static_cast(v); \ - } \ - void func##List(TF_Operation* oper, const char* attr_name, c_type* values, \ - int max_values, TF_Status* status) { \ - const auto* attr = GetAttrValue(oper, attr_name, status); \ - if (!status->status.ok()) return; \ - if (attr->value_case() != tensorflow::AttrValue::kList) { \ - status->status = \ - InvalidArgument("Value for '", attr_name, "' is not a list."); \ - return; \ - } \ - const auto len = std::min(max_values, attr->list().list_field##_size()); \ - for (int i = 0; i < len; ++i) { \ - values[i] = static_cast(attr->list().list_field(i)); \ - } \ +#define DEFINE_GETATTR(func, c_type, cpp_type, list_field) \ + void func(TF_Operation* oper, const char* attr_name, c_type* value, \ + TF_Status* status) { \ + cpp_type v; \ + status->status = \ + tensorflow::GetNodeAttr(oper->node.attrs(), attr_name, &v); \ + *value = static_cast(v); \ + } \ + void func##List(TF_Operation* oper, const char* attr_name, c_type* values, \ + int max_values, TF_Status* status) { \ + const auto* attr = GetAttrValue(oper, attr_name, status); \ + if (!status->status.ok()) return; \ + if (attr->value_case() != tensorflow::AttrValue::kList) { \ + status->status = \ + InvalidArgument("Value for '", attr_name, "' is not a list."); \ + return; \ + } \ + const auto len = std::min(max_values, attr->list().list_field##_size()); \ + for (int i = 0; i < len; ++i) { \ + values[i] = static_cast(attr->list().list_field(i)); \ + } \ } DEFINE_GETATTR(TF_OperationGetAttrInt, int64_t, tensorflow::int64, i); DEFINE_GETATTR(TF_OperationGetAttrFloat, float, float, f); @@ -1504,7 +1565,8 @@ DEFINE_GETATTR(TF_OperationGetAttrType, TF_DataType, DataType, type); void TF_OperationGetAttrShape(TF_Operation* oper, const char* attr_name, int64_t* value, int num_dims, TF_Status* status) { PartialTensorShape shape; - status->status = tensorflow::GetNodeAttr(oper->node.def(), attr_name, &shape); + status->status = + tensorflow::GetNodeAttr(oper->node.attrs(), attr_name, &shape); if (!status->status.ok()) return; auto len = std::min(shape.dims(), num_dims); for (int i = 0; i < len; ++i) { @@ -1518,7 +1580,7 @@ void TF_OperationGetAttrShapeList(TF_Operation* oper, const char* attr_name, int storage_size, TF_Status* status) { std::vector shapes; status->status = - tensorflow::GetNodeAttr(oper->node.def(), attr_name, &shapes); + tensorflow::GetNodeAttr(oper->node.attrs(), attr_name, &shapes); if (!status->status.ok()) return; auto len = std::min(static_cast(shapes.size()), max_values); int64_t* p = storage; @@ -1585,25 +1647,20 @@ void TF_OperationGetAttrTensor(TF_Operation* oper, const char* attr_name, TF_Tensor** value, TF_Status* status) { *value = nullptr; Tensor t; - status->status = tensorflow::GetNodeAttr(oper->node.def(), attr_name, &t); + status->status = tensorflow::GetNodeAttr(oper->node.attrs(), attr_name, &t); if (!status->status.ok()) return; - *value = new TF_Tensor{static_cast(t.dtype()), t.shape(), - tensorflow::TensorCApi::Buffer(t)}; - (*value)->buffer->Ref(); + *value = TF_TensorFromTensor(t, status); } void TF_OperationGetAttrTensorList(TF_Operation* oper, const char* attr_name, TF_Tensor** values, int max_values, TF_Status* status) { std::vector ts; - status->status = tensorflow::GetNodeAttr(oper->node.def(), attr_name, &ts); + status->status = tensorflow::GetNodeAttr(oper->node.attrs(), attr_name, &ts); if (!status->status.ok()) return; const auto len = std::min(max_values, static_cast(ts.size())); for (int i = 0; i < len; ++i) { - const Tensor& t = ts[i]; - values[i] = new TF_Tensor{static_cast(t.dtype()), t.shape(), - tensorflow::TensorCApi::Buffer(t)}; - values[i]->buffer->Ref(); + values[i] = TF_TensorFromTensor(ts[i], status); } } @@ -1622,6 +1679,14 @@ void TF_OperationToNodeDef(TF_Operation* oper, TF_Buffer* output_node_def, // TF_Graph functions --------------------------------------------------------- +TF_Graph::TF_Graph() + : graph(tensorflow::OpRegistry::Global()), + refiner(graph.versions().producer(), graph.op_registry()), + num_sessions(0), + delete_requested(false), + parent(nullptr), + parent_inputs(nullptr) {} + TF_Graph* TF_NewGraph() { return new TF_Graph; } void TF_DeleteGraph(TF_Graph* g) { @@ -1675,10 +1740,6 @@ void TF_GraphToGraphDef(TF_Graph* graph, TF_Buffer* output_graph_def, status->status = MessageToBuffer(def, output_graph_def); } -struct TF_ImportGraphDefOptions { - tensorflow::ImportGraphDefOptions opts; -}; - TF_ImportGraphDefOptions* TF_NewImportGraphDefOptions() { return new TF_ImportGraphDefOptions; } @@ -1690,14 +1751,6 @@ void TF_ImportGraphDefOptionsSetPrefix(TF_ImportGraphDefOptions* opts, opts->opts.prefix = prefix; } -namespace { - -TensorId ToTensorId(const TF_Output& output) { - return TensorId(output.oper->node.name(), output.index); -} - -} // namespace - void TF_ImportGraphDefOptionsAddInputMapping(TF_ImportGraphDefOptions* opts, const char* src_name, int src_index, TF_Output dst) { @@ -1781,6 +1834,11 @@ void TF_GraphImportGraphDef(TF_Graph* graph, const TF_Buffer* graph_def, // While loop functions ------------------------------------------------------- namespace { + +#ifndef __ANDROID__ + +// Creates a placeholder representing an input to the cond or body graph. +// TODO(skyewm): remove these from final graph bool CreateInput(const TF_Output& parent_input, TF_Graph* g, const char* name, TF_Output* input, TF_Status* status) { TF_OperationDescription* desc = TF_NewOperation(g, "Placeholder", name); @@ -1792,128 +1850,50 @@ bool CreateInput(const TF_Output& parent_input, TF_Graph* g, const char* name, return true; } -bool CreateEnter(TF_Graph* g, const char* node_name, const char* frame_name, - const TF_Output& input, TF_Output* enter, TF_Status* status) - EXCLUSIVE_LOCKS_REQUIRED(g->mu) { - TF_OperationDescription* desc = TF_NewOperationLocked(g, "Enter", node_name); - TF_AddInput(desc, input); - TF_SetAttrString(desc, "frame_name", frame_name, strlen(frame_name)); - TF_Operation* oper = TF_FinishOperationLocked(desc, status); - if (!status->status.ok()) return false; - *enter = {oper, 0}; - return true; -} - -bool CreateMerge(TF_Graph* g, const char* name, const TF_Output& input, - const char* backedge_name, int backedge_index, - TF_Output* merge, TF_Status* status) - EXCLUSIVE_LOCKS_REQUIRED(g->mu) { - TF_OperationDescription* desc = TF_NewOperationLocked(g, "Merge", name); - - // The merge nodes accept the while loop's back edges as an input. Use the - // underlying NodeBuilder API directly to create an input to the - // not-yet-created back edge. - std::vector input_list; - input_list.push_back(NodeBuilder::NodeOut(&input.oper->node, input.index)); - // All merge inputs must have same type - DataType type = input.oper->node.output_type(input.index); - input_list.push_back( - NodeBuilder::NodeOut(backedge_name, backedge_index, type)); - - desc->node_builder.Input(input_list); - - TF_Operation* oper = TF_FinishOperationLocked(desc, status); - if (!status->status.ok()) return false; - *merge = {oper, 0}; - return true; -} - -bool CreateSwitch(TF_Graph* g, const char* name, const TF_Output& input, - const TF_Output& predicate, TF_Output* switch_true, - TF_Output* switch_false, TF_Status* status) - EXCLUSIVE_LOCKS_REQUIRED(g->mu) { - TF_OperationDescription* desc = TF_NewOperationLocked(g, "Switch", name); - TF_AddInput(desc, input); - TF_AddInput(desc, predicate); - TF_Operation* oper = TF_FinishOperationLocked(desc, status); - if (!status->status.ok()) return false; - *switch_false = {oper, 0}; - *switch_true = {oper, 1}; - return true; -} - -bool CreateNext(TF_Graph* g, const char* name, const TF_Output& input, - TF_Output* next, TF_Status* status) - EXCLUSIVE_LOCKS_REQUIRED(g->mu) { - TF_OperationDescription* desc = - TF_NewOperationLocked(g, "NextIteration", name); - TF_AddInput(desc, input); - TF_Operation* oper = TF_FinishOperationLocked(desc, status); - if (!status->status.ok()) return false; - *next = {oper, 0}; - return true; -} - -bool CreateExit(TF_Graph* g, const char* name, const TF_Output& input, - TF_Output* exit, TF_Status* status) - EXCLUSIVE_LOCKS_REQUIRED(g->mu) { - TF_OperationDescription* desc = TF_NewOperationLocked(g, "Exit", name); - TF_AddInput(desc, input); - TF_Operation* oper = TF_FinishOperationLocked(desc, status); - if (!status->status.ok()) return false; - *exit = {oper, 0}; - return true; -} - -class ScopedImportGraphDefOptions { - public: - ScopedImportGraphDefOptions() { opts_ = TF_NewImportGraphDefOptions(); } - ~ScopedImportGraphDefOptions() { TF_DeleteImportGraphDefOptions(opts_); } - - TF_ImportGraphDefOptions* get() const { return opts_; } - - private: - TF_ImportGraphDefOptions* opts_; - - TF_DISALLOW_COPY_AND_ASSIGN(ScopedImportGraphDefOptions); -}; - // Copies `src_graph` into `dst_graph`. Any node in `src_graph` with input -// `src_inputs[i]` will have that input replaced with `dst_inputs[i]`. -// `prefix` will be prepended to copied node names. `return_nodes` are nodes -// in `src_graph`, and the new corresponding nodes in `dst_graph` will be -// returned. `return_nodes` should be preallocated to size `nreturn_nodes`. -bool CopyGraph(TF_Graph* src_graph, TF_Graph* dst_graph, - const TF_Output* src_inputs, - const std::vector& dst_inputs, const char* prefix, - const TF_Output* nodes_to_return, int nreturn_nodes, - TF_Output* return_nodes, TF_Status* s) - EXCLUSIVE_LOCKS_REQUIRED(dst_graph->mu) { +// `src_inputs[i]` will have that input replaced with `dst_inputs[i]`. `prefix` +// will be prepended to copied node names. `control_deps` are nodes in +// `dst_graph` that the copied `src_graph` nodes will have control dependencies +// on. `return_nodes` are nodes in `src_graph`, and the new corresponding nodes +// in `dst_graph` will be returned. `return_nodes` must be non-null. +Status CopyGraph(Graph* src_graph, Graph* dst_graph, + tensorflow::ShapeRefiner* dst_refiner, + const TF_Output* src_inputs, + const std::vector& dst_inputs, + const tensorflow::string& prefix, + const std::vector& control_deps, + const TF_Output* nodes_to_return, int nreturn_nodes, + std::vector* return_nodes) { + DCHECK(return_nodes != nullptr); GraphDef gdef; - src_graph->graph.ToGraphDef(&gdef); + src_graph->ToGraphDef(&gdef); - ScopedImportGraphDefOptions opts; - TF_ImportGraphDefOptionsSetPrefix(opts.get(), prefix); + tensorflow::ImportGraphDefOptions opts; + opts.prefix = prefix; for (int i = 0; i < dst_inputs.size(); ++i) { - TensorId src = ToTensorId(src_inputs[i]); - TF_ImportGraphDefOptionsAddInputMapping(opts.get(), src.first.data(), - src.second, dst_inputs[i]); + opts.input_map[ToTensorId(src_inputs[i])] = + TensorId(dst_inputs[i].node()->name(), dst_inputs[i].index()); + } + opts.skip_mapped_nodes = true; + + for (const tensorflow::Operation& op : control_deps) { + opts.control_dependencies.push_back(op.node()->name()); } - // We use the pivot node to control constants in `src_graph` - TF_Operation* pivot = dst_inputs[0].oper; - TF_ImportGraphDefOptionsAddControlDependency(opts.get(), pivot); for (int i = 0; i < nreturn_nodes; ++i) { - TF_ImportGraphDefOptionsAddReturnOutput( - opts.get(), nodes_to_return[i].oper->node.name().c_str(), - nodes_to_return[i].index); + opts.return_tensors.push_back(ToTensorId(nodes_to_return[i])); } - GraphImportGraphDefLocked(dst_graph, gdef, opts.get(), return_nodes, - nreturn_nodes, s); - if (TF_GetCode(s) != TF_OK) return false; - return true; + // TOOD(skyewm): change to OutputTensor + std::vector> return_tensors; + TF_RETURN_IF_ERROR( + ImportGraphDef(opts, gdef, dst_graph, dst_refiner, &return_tensors)); + + for (const auto& pair : return_tensors) { + return_nodes->emplace_back(pair.first, pair.second); + } + return Status::OK(); } bool ValidateConstWhileParams(const TF_WhileParams& params, TF_Status* s) { @@ -1949,6 +1929,8 @@ bool ValidateInputWhileParams(const TF_WhileParams& params, TF_Status* s) { return true; } +#endif // __ANDROID__ + void FreeWhileResources(const TF_WhileParams* params) { TF_DeleteGraph(params->cond_graph); TF_DeleteGraph(params->body_graph); @@ -1966,6 +1948,13 @@ TF_WhileParams EmptyWhileParams() { TF_WhileParams TF_NewWhile(TF_Graph* g, TF_Output* inputs, int ninputs, TF_Status* status) { +#ifdef __ANDROID__ + status->status = tensorflow::errors::Unimplemented( + "Creating while loops is not supported in Android. File a bug at " + "https://github.com/tensorflow/tensorflow/issues if this feature is " + "important to you"); + return EmptyWhileParams(); +#else if (ninputs == 0) { status->status = InvalidArgument("TF_NewWhile() must be passed at least one input"); @@ -2006,8 +1995,10 @@ TF_WhileParams TF_NewWhile(TF_Graph* g, TF_Output* inputs, int ninputs, return EmptyWhileParams(); } return params; +#endif // __ANDROID__ } +#ifndef __ANDROID__ namespace { // TODO(skyewm): make nodes in while loop unfetchable like in Python version @@ -2017,92 +2008,143 @@ void TF_FinishWhileHelper(const TF_WhileParams* params, TF_Status* status, TF_Graph* parent = params->cond_graph->parent; TF_Output* parent_inputs = params->cond_graph->parent_inputs; - int n = params->ninputs; + int num_loop_vars = params->ninputs; mutex_lock l(parent->mu); - // Create Enter nodes - std::vector enter_nodes(n); - for (int i = 0; i < n; ++i) { - if (!CreateEnter(parent, StrCat(params->name, "/enter", i).c_str(), - params->name, parent_inputs[i], &enter_nodes[i], status)) { - return; - } - } - - // Create Merge nodes - std::vector merge_nodes(n); - for (int i = 0; i < n; ++i) { - if (!CreateMerge(parent, StrCat(params->name, "/merge", i).c_str(), - enter_nodes[i], StrCat(params->name, "/next", i).c_str(), - 0, &merge_nodes[i], status)) { - return; - } - } - - // Copy cond_graph to parent and replace input placeholders with merge node - // outputs, and get handle to new cond output - tensorflow::string cond_prefix = StrCat(params->name, "/cond"); - TF_Output cond_output; - if (!CopyGraph(params->cond_graph, parent, params->cond_inputs, merge_nodes, - cond_prefix.c_str(), ¶ms->cond_output, 1, &cond_output, - status)) { - return; - } - - // Create Switch nodes - std::vector switch_trues(n); - std::vector switch_falses(n); - for (int i = 0; i < n; ++i) { - if (!CreateSwitch(parent, StrCat(params->name, "/switch", i).c_str(), - merge_nodes[i], cond_output, &switch_trues[i], - &switch_falses[i], status)) { - return; - } - } - - // Copy body_graph to parent, replace input placeholders with switch node - // true outputs, and get handles to new body outputs - tensorflow::string body_prefix = StrCat(params->name, "/body"); - std::vector body_outputs(n); - if (!CopyGraph(params->body_graph, parent, params->body_inputs, switch_trues, - body_prefix.c_str(), params->body_outputs, n, - body_outputs.data(), status)) { - return; - } - - // Create Next nodes - std::vector next_nodes(n); - for (int i = 0; i < n; ++i) { - if (!CreateNext(parent, StrCat(params->name, "/next", i).c_str(), - body_outputs[i], &next_nodes[i], status)) { - return; - } - } - - // Create Exit nodes (which are the outputs of the while loop) - for (int i = 0; i < n; ++i) { - if (!CreateExit(parent, StrCat(params->name, "/exit", i).c_str(), - switch_falses[i], &outputs[i], status)) { - return; - } + // 'cond_fn' copies the cond graph into the parent graph. + tensorflow::ops::CondGraphBuilderFn cond_fn = + [params, parent](const tensorflow::Scope& scope, + const std::vector& inputs, + tensorflow::Output* output) { + DCHECK_EQ(scope.graph(), &parent->graph); + std::vector cond_output; + TF_RETURN_IF_ERROR(CopyGraph( + ¶ms->cond_graph->graph, &parent->graph, &parent->refiner, + params->cond_inputs, inputs, scope.impl()->name(), + scope.impl()->control_deps(), ¶ms->cond_output, + /* nreturn_nodes */ 1, &cond_output)); + *output = cond_output[0]; + return Status::OK(); + }; + + // 'body_fn' copies the body graph into the parent graph. + tensorflow::ops::BodyGraphBuilderFn body_fn = + [params, parent, num_loop_vars]( + const tensorflow::Scope& scope, + const std::vector& inputs, + std::vector* outputs) { + DCHECK_EQ(scope.graph(), &parent->graph); + TF_RETURN_IF_ERROR( + CopyGraph(¶ms->body_graph->graph, &parent->graph, + &parent->refiner, params->body_inputs, inputs, + scope.impl()->name(), scope.impl()->control_deps(), + params->body_outputs, num_loop_vars, outputs)); + return Status::OK(); + }; + + // Create the while loop using an internal scope. + tensorflow::Scope scope = + NewInternalScope(&parent->graph, &status->status, &parent->refiner) + .NewSubScope(params->name); + + const int first_new_node_id = parent->graph.num_node_ids(); + + tensorflow::OutputList loop_outputs; + status->status = tensorflow::ops::BuildWhileLoop( + scope, OutputsFromTFOutputs(parent_inputs, num_loop_vars), cond_fn, + body_fn, params->name, &loop_outputs); + + // Update name_map with newly-created ops. + // TODO(skyewm): right now BuildWhileLoop() may alter the graph if it returns + // a bad status. Once we fix this, we may want to return early instead of + // executing the following code. + for (int i = first_new_node_id; i < parent->graph.num_node_ids(); ++i) { + Node* new_node = parent->graph.FindNodeId(i); + if (new_node == nullptr) continue; + parent->name_map[new_node->name()] = new_node; + } + + // Populate 'outputs'. + DCHECK_LE(loop_outputs.size(), num_loop_vars); + for (int i = 0; i < loop_outputs.size(); ++i) { + outputs[i] = {ToOperation(loop_outputs[i].node()), loop_outputs[i].index()}; } } } // namespace +#endif // __ANDROID__ void TF_FinishWhile(const TF_WhileParams* params, TF_Status* status, TF_Output* outputs) { +#ifdef __ANDROID__ + status->status = tensorflow::errors::Unimplemented( + "Creating while loops is not supported in Android. File a bug at " + "https://github.com/tensorflow/tensorflow/issues if this feature is " + "important to you"); +#else // If it appears the caller created or modified `params`, don't free resources if (!ValidateConstWhileParams(*params, status)) return; TF_FinishWhileHelper(params, status, outputs); FreeWhileResources(params); +#endif // __ANDROID__ } void TF_AbortWhile(const TF_WhileParams* params) { FreeWhileResources(params); } +void TF_AddGradients(TF_Graph* g, TF_Output* y, int ny, TF_Output* x, int nx, + TF_Output* dx, TF_Status* status, TF_Output* dy) { +#ifdef __ANDROID__ + status->status = tensorflow::errors::Unimplemented( + "Adding gradients is not supported in Android. File a bug at " + "https://github.com/tensorflow/tensorflow/issues if this feature is " + "important to you"); +#else + std::vector y_arg = OutputsFromTFOutputs(y, ny); + std::vector x_arg = OutputsFromTFOutputs(x, nx); + std::vector dy_arg; + + { + // We need to hold on to the lock while we have a scope that uses TF_Graph. + mutex_lock graph_lock(g->mu); + + const int first_new_node_id = g->graph.num_node_ids(); + + tensorflow::Scope scope = + NewInternalScope(&g->graph, &status->status, &g->refiner) + .NewSubScope("gradients"); + + if (dx != nullptr) { + std::vector dx_arg = OutputsFromTFOutputs(dx, ny); + status->status = + AddSymbolicGradients(scope, y_arg, x_arg, dx_arg, &dy_arg); + } else { + status->status = AddSymbolicGradients(scope, y_arg, x_arg, &dy_arg); + } + + // Update g->name_map with the name_map from the scope, which will contain + // the new gradient ops. + for (int i = first_new_node_id; i < g->graph.num_node_ids(); ++i) { + Node* n = g->graph.FindNodeId(i); + if (n == nullptr) continue; + g->name_map[n->name()] = n; + } + } + + // Unpack the results from grad_outputs_arg. + TFOutputsFromOutputs(dy_arg, dy); +#endif // __ANDROID__ +} + // TF_Session functions ---------------------------------------------- +TF_Session::TF_Session(tensorflow::Session* s, TF_Graph* g) + : session(s), graph(g), last_num_graph_nodes(0), device_mgr(nullptr) { + if (s->LocalDeviceManager(&device_mgr).ok()) { + devices = device_mgr->ListDevices(); + } +} + TF_Session* TF_NewSession(TF_Graph* graph, const TF_SessionOptions* opt, TF_Status* status) { Session* session; @@ -2115,7 +2157,7 @@ TF_Session* TF_NewSession(TF_Graph* graph, const TF_SessionOptions* opt, return new TF_Session(session, graph); } else { DCHECK_EQ(nullptr, session); - return NULL; + return nullptr; } } @@ -2133,7 +2175,6 @@ TF_Session* TF_LoadSessionFromSavedModel( return nullptr; #else mutex_lock l(graph->mu); - if (!graph->name_map.empty()) { status->status = InvalidArgument("Graph is non-empty."); return nullptr; @@ -2211,7 +2252,7 @@ static bool ExtendSessionGraphHelper(TF_Session* session, TF_Status* status) { const auto num_nodes = graph.num_node_ids(); if (session->last_num_graph_nodes < num_nodes) { GraphDef graph_def; - graph_def.mutable_versions()->CopyFrom(graph.versions()); + *graph_def.mutable_versions() = graph.versions(); // Fill graph_def with nodes with ids in the range // [session->last_num_graph_nodes, num_nodes), that is the nodes // added since the last TF_SessionRun() call. @@ -2222,8 +2263,8 @@ static bool ExtendSessionGraphHelper(TF_Session* session, TF_Status* status) { *node_def = node->def(); } } + *graph_def.mutable_library() = graph.flib_def().ToProto(); session->graph->mu.unlock(); - // TODO(josh11b): Also send the function library if needed. status->status = session->session->Extend(graph_def); if (!status->status.ok()) { // Contract is we always delete input_values[i]. @@ -2283,6 +2324,8 @@ void TF_SessionPRunSetup(TF_Session* session, const TF_Output* inputs, int ninputs, const TF_Output* outputs, int noutputs, const TF_Operation* const* target_opers, int ntargets, const char** handle, TF_Status* status) { + *handle = nullptr; + if (!ExtendSessionGraphHelper(session, status)) { return; } diff --git a/tensorflow/c/c_api.h b/tensorflow/c/c_api.h index f837b68d76c34ba836720df820daaae5bc29c93c..ee110d88cea50614515b3b3f42af1db1aaee9012 100644 --- a/tensorflow/c/c_api.h +++ b/tensorflow/c/c_api.h @@ -64,6 +64,25 @@ limitations under the License. // and the API just provides high level controls over the number of // devices of each type. +// Macro to control visibility of exported symbols in the shared library (.so, +// .dylib, .dll). +// This duplicates the TF_EXPORT macro definition in +// tensorflow/core/platform/macros.h in order to keep this .h file independent +// of any other includes.$a +#ifdef SWIG +#define TF_CAPI_EXPORT +#else +#if defined(COMPILER_MSVC) +#ifdef TF_COMPILE_LIBRARY +#define TF_CAPI_EXPORT __declspec(dllexport) +#else +#define TF_CAPI_EXPORT __declspec(dllimport) +#endif // TF_COMPILE_LIBRARY +#else +#define TF_CAPI_EXPORT __attribute__((visibility("default"))) +#endif // COMPILER_MSVC +#endif // SWIG + #ifdef __cplusplus extern "C" { #endif @@ -71,12 +90,12 @@ extern "C" { // -------------------------------------------------------------------------- // TF_Version returns a string describing version information of the // TensorFlow library. TensorFlow using semantic versioning. -extern const char* TF_Version(); +TF_CAPI_EXPORT extern const char* TF_Version(); // -------------------------------------------------------------------------- // TF_DataType holds the type for a scalar value. E.g., one slot in a tensor. // The enum values here are identical to corresponding values in types.proto. -typedef enum { +typedef enum TF_DataType { TF_FLOAT = 1, TF_DOUBLE = 2, TF_INT32 = 3, // Int32 tensors are always in 'host' memory. @@ -98,17 +117,18 @@ typedef enum { TF_COMPLEX128 = 18, // Double-precision complex TF_HALF = 19, TF_RESOURCE = 20, + TF_VARIANT = 21, } TF_DataType; // TF_DataTypeSize returns the sizeof() for the underlying type corresponding // to the given TF_DataType enum value. Returns 0 for variable length types // (eg. TF_STRING) or on failure. -extern size_t TF_DataTypeSize(TF_DataType dt); +TF_CAPI_EXPORT extern size_t TF_DataTypeSize(TF_DataType dt); // -------------------------------------------------------------------------- // TF_Code holds an error code. The enum values here are identical to // corresponding values in error_codes.proto. -typedef enum { +typedef enum TF_Code { TF_OK = 0, TF_CANCELLED = 1, TF_UNKNOWN = 2, @@ -134,23 +154,24 @@ typedef enum { typedef struct TF_Status TF_Status; // Return a new status object. -extern TF_Status* TF_NewStatus(); +TF_CAPI_EXPORT extern TF_Status* TF_NewStatus(); // Delete a previously created status object. -extern void TF_DeleteStatus(TF_Status*); +TF_CAPI_EXPORT extern void TF_DeleteStatus(TF_Status*); // Record in *s. Any previous information is lost. // A common use is to clear a status: TF_SetStatus(s, TF_OK, ""); -extern void TF_SetStatus(TF_Status* s, TF_Code code, const char* msg); +TF_CAPI_EXPORT extern void TF_SetStatus(TF_Status* s, TF_Code code, + const char* msg); // Return the code record in *s. -extern TF_Code TF_GetCode(const TF_Status* s); +TF_CAPI_EXPORT extern TF_Code TF_GetCode(const TF_Status* s); // Return a pointer to the (null-terminated) error message in *s. The // return value points to memory that is only usable until the next // mutation to *s. Always returns an empty string if TF_GetCode(s) is // TF_OK. -extern const char* TF_Message(const TF_Status* s); +TF_CAPI_EXPORT extern const char* TF_Message(const TF_Status* s); // -------------------------------------------------------------------------- // TF_Buffer holds a pointer to a block of data and its associated length. @@ -168,14 +189,15 @@ typedef struct TF_Buffer { // Makes a copy of the input and sets an appropriate deallocator. Useful for // passing in read-only, input protobufs. -extern TF_Buffer* TF_NewBufferFromString(const void* proto, size_t proto_len); +TF_CAPI_EXPORT extern TF_Buffer* TF_NewBufferFromString(const void* proto, + size_t proto_len); // Useful for passing *out* a protobuf. -extern TF_Buffer* TF_NewBuffer(); +TF_CAPI_EXPORT extern TF_Buffer* TF_NewBuffer(); -extern void TF_DeleteBuffer(TF_Buffer*); +TF_CAPI_EXPORT extern void TF_DeleteBuffer(TF_Buffer*); -extern TF_Buffer TF_GetBuffer(TF_Buffer* buffer); +TF_CAPI_EXPORT extern TF_Buffer TF_GetBuffer(TF_Buffer* buffer); // -------------------------------------------------------------------------- // TF_Tensor holds a multi-dimensional array of elements of a single data type. @@ -202,11 +224,10 @@ typedef struct TF_Tensor TF_Tensor; // (*deallocator)(data, len, deallocator_arg) // Clients must provide a custom deallocator function so they can pass in // memory managed by something like numpy. -extern TF_Tensor* TF_NewTensor(TF_DataType, const int64_t* dims, int num_dims, - void* data, size_t len, - void (*deallocator)(void* data, size_t len, - void* arg), - void* deallocator_arg); +TF_CAPI_EXPORT extern TF_Tensor* TF_NewTensor( + TF_DataType, const int64_t* dims, int num_dims, void* data, size_t len, + void (*deallocator)(void* data, size_t len, void* arg), + void* deallocator_arg); // Allocate and return a new Tensor. // @@ -217,27 +238,32 @@ extern TF_Tensor* TF_NewTensor(TF_DataType, const int64_t* dims, int num_dims, // // The caller must set the Tensor values by writing them to the pointer returned // by TF_TensorData with length TF_TensorByteSize. -extern TF_Tensor* TF_AllocateTensor(TF_DataType, const int64_t* dims, - int num_dims, size_t len); +TF_CAPI_EXPORT extern TF_Tensor* TF_AllocateTensor(TF_DataType, + const int64_t* dims, + int num_dims, size_t len); + +// Deletes `tensor` and returns a new TF_Tensor with the same content if +// possible. Returns nullptr and leaves `tensor` untouched if not. +TF_CAPI_EXPORT extern TF_Tensor* TF_TensorMaybeMove(TF_Tensor* tensor); // Destroy a tensor. -extern void TF_DeleteTensor(TF_Tensor*); +TF_CAPI_EXPORT extern void TF_DeleteTensor(TF_Tensor*); // Return the type of a tensor element. -extern TF_DataType TF_TensorType(const TF_Tensor*); +TF_CAPI_EXPORT extern TF_DataType TF_TensorType(const TF_Tensor*); // Return the number of dimensions that the tensor has. -extern int TF_NumDims(const TF_Tensor*); +TF_CAPI_EXPORT extern int TF_NumDims(const TF_Tensor*); // Return the length of the tensor in the "dim_index" dimension. // REQUIRES: 0 <= dim_index < TF_NumDims(tensor) -extern int64_t TF_Dim(const TF_Tensor* tensor, int dim_index); +TF_CAPI_EXPORT extern int64_t TF_Dim(const TF_Tensor* tensor, int dim_index); // Return the size of the underlying data in bytes. -extern size_t TF_TensorByteSize(const TF_Tensor*); +TF_CAPI_EXPORT extern size_t TF_TensorByteSize(const TF_Tensor*); // Return a pointer to the underlying data buffer. -extern void* TF_TensorData(const TF_Tensor*); +TF_CAPI_EXPORT extern void* TF_TensorData(const TF_Tensor*); // -------------------------------------------------------------------------- // Encode the string `src` (`src_len` bytes long) into `dst` in the format @@ -247,8 +273,9 @@ extern void* TF_TensorData(const TF_Tensor*); // // On success returns the size in bytes of the encoded string. // Returns an error into `status` otherwise. -extern size_t TF_StringEncode(const char* src, size_t src_len, char* dst, - size_t dst_len, TF_Status* status); +TF_CAPI_EXPORT extern size_t TF_StringEncode(const char* src, size_t src_len, + char* dst, size_t dst_len, + TF_Status* status); // Decode a string encoded using TF_StringEncode. // @@ -258,19 +285,20 @@ extern size_t TF_StringEncode(const char* src, size_t src_len, char* dst, // `*dst` and `*dst_len` are undefined and an error is set in `status`. // // Does not read memory more than `src_len` bytes beyond `src`. -extern size_t TF_StringDecode(const char* src, size_t src_len, const char** dst, - size_t* dst_len, TF_Status* status); +TF_CAPI_EXPORT extern size_t TF_StringDecode(const char* src, size_t src_len, + const char** dst, size_t* dst_len, + TF_Status* status); // Return the size in bytes required to encode a string `len` bytes long into a // TF_STRING tensor. -extern size_t TF_StringEncodedSize(size_t len); +TF_CAPI_EXPORT extern size_t TF_StringEncodedSize(size_t len); // -------------------------------------------------------------------------- // TF_SessionOptions holds options that can be passed during session creation. typedef struct TF_SessionOptions TF_SessionOptions; // Return a new options object. -extern TF_SessionOptions* TF_NewSessionOptions(); +TF_CAPI_EXPORT extern TF_SessionOptions* TF_NewSessionOptions(); // Set the target in TF_SessionOptions.options. // target can be empty, a single entry, or a comma separated list of entries. @@ -278,17 +306,19 @@ extern TF_SessionOptions* TF_NewSessionOptions(); // "local" // ip:port // host:port -extern void TF_SetTarget(TF_SessionOptions* options, const char* target); +TF_CAPI_EXPORT extern void TF_SetTarget(TF_SessionOptions* options, + const char* target); // Set the config in TF_SessionOptions.options. // config should be a serialized tensorflow.ConfigProto proto. // If config was not parsed successfully as a ConfigProto, record the // error information in *status. -extern void TF_SetConfig(TF_SessionOptions* options, const void* proto, - size_t proto_len, TF_Status* status); +TF_CAPI_EXPORT extern void TF_SetConfig(TF_SessionOptions* options, + const void* proto, size_t proto_len, + TF_Status* status); // Destroy an options object. -extern void TF_DeleteSessionOptions(TF_SessionOptions*); +TF_CAPI_EXPORT extern void TF_DeleteSessionOptions(TF_SessionOptions*); // TODO(jeff,sanjay): // - export functions to set Config fields @@ -301,11 +331,11 @@ extern void TF_DeleteSessionOptions(TF_SessionOptions*); typedef struct TF_Graph TF_Graph; // Return a new graph object. -extern TF_Graph* TF_NewGraph(); +TF_CAPI_EXPORT extern TF_Graph* TF_NewGraph(); // Destroy an options object. Graph will be deleted once no more // TFSession's are referencing it. -extern void TF_DeleteGraph(TF_Graph*); +TF_CAPI_EXPORT extern void TF_DeleteGraph(TF_Graph*); // Operation being built. The underlying graph must outlive this. typedef struct TF_OperationDescription TF_OperationDescription; @@ -327,12 +357,20 @@ typedef struct TF_Output { int index; // The index of the output within oper. } TF_Output; +// TF_Function is a grouping of operations with defined inputs and outputs. +// Once created and added to graphs, functions can be invoked by creating an +// operation whose operation type matches the function name. +typedef struct TF_Function TF_Function; + +// Function definition options. TODO(iga): Define and implement +typedef struct TF_FunctionOptions TF_FunctionOptions; + // Sets the shape of the Tensor referenced by `output` in `graph` to // the shape described by `dims` and `num_dims`. // -// If the number of dimensions is unknown, `num_dims` must be -// set to -1 and dims can be null. If a dimension is unknown, -// the corresponding entry in the `dims` array must be -1. +// If the number of dimensions is unknown, `num_dims` must be set to +// -1 and `dims` can be null. If a dimension is unknown, the +// corresponding entry in the `dims` array must be -1. // // This does not overwrite the existing shape associated with `output`, // but merges the input shape with the existing shape. For example, @@ -343,9 +381,11 @@ typedef struct TF_Output { // * `output` is not in `graph`. // * An invalid shape is being set (e.g., the shape being set // is incompatible with the existing shape). -extern void TF_GraphSetTensorShape(TF_Graph* graph, TF_Output output, - const int64_t* dims, const int num_dims, - TF_Status* status); +TF_CAPI_EXPORT extern void TF_GraphSetTensorShape(TF_Graph* graph, + TF_Output output, + const int64_t* dims, + const int num_dims, + TF_Status* status); // Returns the number of dimensions of the Tensor referenced by `output` // in `graph`. @@ -354,8 +394,9 @@ extern void TF_GraphSetTensorShape(TF_Graph* graph, TF_Output output, // // Returns an error into `status` if: // * `output` is not in `graph`. -extern int TF_GraphGetTensorNumDims(TF_Graph* graph, TF_Output output, - TF_Status* status); +TF_CAPI_EXPORT extern int TF_GraphGetTensorNumDims(TF_Graph* graph, + TF_Output output, + TF_Status* status); // Returns the shape of the Tensor referenced by `output` in `graph` // into `dims`. `dims` must be an array large enough to hold `num_dims` @@ -369,20 +410,21 @@ extern int TF_GraphGetTensorNumDims(TF_Graph* graph, TF_Output output, // Returns an error into `status` if: // * `output` is not in `graph`. // * `num_dims` does not match the actual number of dimensions. -extern void TF_GraphGetTensorShape(TF_Graph* graph, TF_Output output, - int64_t* dims, int num_dims, - TF_Status* status); +TF_CAPI_EXPORT extern void TF_GraphGetTensorShape(TF_Graph* graph, + TF_Output output, + int64_t* dims, int num_dims, + TF_Status* status); // Operation will only be added to *graph when TF_FinishOperation() is // called (assuming TF_FinishOperation() does not return an error). // *graph must not be deleted until after TF_FinishOperation() is // called. -extern TF_OperationDescription* TF_NewOperation(TF_Graph* graph, - const char* op_type, - const char* oper_name); +TF_CAPI_EXPORT extern TF_OperationDescription* TF_NewOperation( + TF_Graph* graph, const char* op_type, const char* oper_name); // Specify the device for `desc`. Defaults to empty, meaning unconstrained. -extern void TF_SetDevice(TF_OperationDescription* desc, const char* device); +TF_CAPI_EXPORT extern void TF_SetDevice(TF_OperationDescription* desc, + const char* device); // The calls to TF_AddInput and TF_AddInputList must match (in number, // order, and type) the op declaration. For example, the "Concat" op @@ -405,101 +447,115 @@ extern void TF_SetDevice(TF_OperationDescription* desc, const char* device); // TF_AddInputList(desc, values_inputs, 5); // For inputs that take a single tensor. -extern void TF_AddInput(TF_OperationDescription* desc, TF_Output input); +TF_CAPI_EXPORT extern void TF_AddInput(TF_OperationDescription* desc, + TF_Output input); // For inputs that take a list of tensors. // inputs must point to TF_Output[num_inputs]. -extern void TF_AddInputList(TF_OperationDescription* desc, - const TF_Output* inputs, int num_inputs); +TF_CAPI_EXPORT extern void TF_AddInputList(TF_OperationDescription* desc, + const TF_Output* inputs, + int num_inputs); // Call once per control input to `desc`. -extern void TF_AddControlInput(TF_OperationDescription* desc, - TF_Operation* input); +TF_CAPI_EXPORT extern void TF_AddControlInput(TF_OperationDescription* desc, + TF_Operation* input); // Request that `desc` be co-located on the device where `op` // is placed. // // Use of this is discouraged since the implementation of device placement is // subject to change. Primarily intended for internal libraries -extern void TF_ColocateWith(TF_OperationDescription* desc, TF_Operation* op); +TF_CAPI_EXPORT extern void TF_ColocateWith(TF_OperationDescription* desc, + TF_Operation* op); // Call some TF_SetAttr*() function for every attr that is not // inferred from an input and doesn't have a default value you wish to // keep. // `value` must point to a string of length `length` bytes. -extern void TF_SetAttrString(TF_OperationDescription* desc, - const char* attr_name, const void* value, - size_t length); +TF_CAPI_EXPORT extern void TF_SetAttrString(TF_OperationDescription* desc, + const char* attr_name, + const void* value, size_t length); // `values` and `lengths` each must have lengths `num_values`. // `values[i]` must point to a string of length `lengths[i]` bytes. -extern void TF_SetAttrStringList(TF_OperationDescription* desc, - const char* attr_name, - const void* const* values, - const size_t* lengths, int num_values); -extern void TF_SetAttrInt(TF_OperationDescription* desc, const char* attr_name, - int64_t value); -extern void TF_SetAttrIntList(TF_OperationDescription* desc, - const char* attr_name, const int64_t* values, - int num_values); -extern void TF_SetAttrFloat(TF_OperationDescription* desc, - const char* attr_name, float value); -extern void TF_SetAttrFloatList(TF_OperationDescription* desc, - const char* attr_name, const float* values, - int num_values); -extern void TF_SetAttrBool(TF_OperationDescription* desc, const char* attr_name, - unsigned char value); -extern void TF_SetAttrBoolList(TF_OperationDescription* desc, - const char* attr_name, - const unsigned char* values, int num_values); -extern void TF_SetAttrType(TF_OperationDescription* desc, const char* attr_name, - TF_DataType value); -extern void TF_SetAttrTypeList(TF_OperationDescription* desc, - const char* attr_name, const TF_DataType* values, - int num_values); +TF_CAPI_EXPORT extern void TF_SetAttrStringList(TF_OperationDescription* desc, + const char* attr_name, + const void* const* values, + const size_t* lengths, + int num_values); +TF_CAPI_EXPORT extern void TF_SetAttrInt(TF_OperationDescription* desc, + const char* attr_name, int64_t value); +TF_CAPI_EXPORT extern void TF_SetAttrIntList(TF_OperationDescription* desc, + const char* attr_name, + const int64_t* values, + int num_values); +TF_CAPI_EXPORT extern void TF_SetAttrFloat(TF_OperationDescription* desc, + const char* attr_name, float value); +TF_CAPI_EXPORT extern void TF_SetAttrFloatList(TF_OperationDescription* desc, + const char* attr_name, + const float* values, + int num_values); +TF_CAPI_EXPORT extern void TF_SetAttrBool(TF_OperationDescription* desc, + const char* attr_name, + unsigned char value); +TF_CAPI_EXPORT extern void TF_SetAttrBoolList(TF_OperationDescription* desc, + const char* attr_name, + const unsigned char* values, + int num_values); +TF_CAPI_EXPORT extern void TF_SetAttrType(TF_OperationDescription* desc, + const char* attr_name, + TF_DataType value); +TF_CAPI_EXPORT extern void TF_SetAttrTypeList(TF_OperationDescription* desc, + const char* attr_name, + const TF_DataType* values, + int num_values); // Set `num_dims` to -1 to represent "unknown rank". Otherwise, // `dims` points to an array of length `num_dims`. `dims[i]` must be // >= -1, with -1 meaning "unknown dimension". -extern void TF_SetAttrShape(TF_OperationDescription* desc, - const char* attr_name, const int64_t* dims, - int num_dims); +TF_CAPI_EXPORT extern void TF_SetAttrShape(TF_OperationDescription* desc, + const char* attr_name, + const int64_t* dims, int num_dims); // `dims` and `num_dims` must point to arrays of length `num_shapes`. // Set `num_dims[i]` to -1 to represent "unknown rank". Otherwise, // `dims[i]` points to an array of length `num_dims[i]`. `dims[i][j]` // must be >= -1, with -1 meaning "unknown dimension". -extern void TF_SetAttrShapeList(TF_OperationDescription* desc, - const char* attr_name, - const int64_t* const* dims, const int* num_dims, - int num_shapes); +TF_CAPI_EXPORT extern void TF_SetAttrShapeList(TF_OperationDescription* desc, + const char* attr_name, + const int64_t* const* dims, + const int* num_dims, + int num_shapes); // `proto` must point to an array of `proto_len` bytes representing a // binary-serialized TensorShapeProto. -extern void TF_SetAttrTensorShapeProto(TF_OperationDescription* desc, - const char* attr_name, const void* proto, - size_t proto_len, TF_Status* status); +TF_CAPI_EXPORT extern void TF_SetAttrTensorShapeProto( + TF_OperationDescription* desc, const char* attr_name, const void* proto, + size_t proto_len, TF_Status* status); // `protos` and `proto_lens` must point to arrays of length `num_shapes`. // `protos[i]` must point to an array of `proto_lens[i]` bytes // representing a binary-serialized TensorShapeProto. -extern void TF_SetAttrTensorShapeProtoList(TF_OperationDescription* desc, - const char* attr_name, - const void* const* protos, - const size_t* proto_lens, - int num_shapes, TF_Status* status); - -extern void TF_SetAttrTensor(TF_OperationDescription* desc, - const char* attr_name, TF_Tensor* value, - TF_Status* status); -extern void TF_SetAttrTensorList(TF_OperationDescription* desc, - const char* attr_name, - TF_Tensor* const* values, int num_values, - TF_Status* status); +TF_CAPI_EXPORT extern void TF_SetAttrTensorShapeProtoList( + TF_OperationDescription* desc, const char* attr_name, + const void* const* protos, const size_t* proto_lens, int num_shapes, + TF_Status* status); + +TF_CAPI_EXPORT extern void TF_SetAttrTensor(TF_OperationDescription* desc, + const char* attr_name, + TF_Tensor* value, + TF_Status* status); +TF_CAPI_EXPORT extern void TF_SetAttrTensorList(TF_OperationDescription* desc, + const char* attr_name, + TF_Tensor* const* values, + int num_values, + TF_Status* status); // `proto` should point to a sequence of bytes of length `proto_len` // representing a binary serialization of an AttrValue protocol // buffer. -extern void TF_SetAttrValueProto(TF_OperationDescription* desc, - const char* attr_name, const void* proto, - size_t proto_len, TF_Status* status); +TF_CAPI_EXPORT extern void TF_SetAttrValueProto(TF_OperationDescription* desc, + const char* attr_name, + const void* proto, + size_t proto_len, + TF_Status* status); // If this function succeeds: // * *status is set to an OK value, @@ -511,37 +567,38 @@ extern void TF_SetAttrValueProto(TF_OperationDescription* desc, // * the graph is not modified, // * a null value is returned. // In either case, it deletes `desc`. -extern TF_Operation* TF_FinishOperation(TF_OperationDescription* desc, - TF_Status* status); +TF_CAPI_EXPORT extern TF_Operation* TF_FinishOperation( + TF_OperationDescription* desc, TF_Status* status); // TF_Operation functions. Operations are immutable once created, so // these are all query functions. -extern const char* TF_OperationName(TF_Operation* oper); -extern const char* TF_OperationOpType(TF_Operation* oper); -extern const char* TF_OperationDevice(TF_Operation* oper); +TF_CAPI_EXPORT extern const char* TF_OperationName(TF_Operation* oper); +TF_CAPI_EXPORT extern const char* TF_OperationOpType(TF_Operation* oper); +TF_CAPI_EXPORT extern const char* TF_OperationDevice(TF_Operation* oper); -extern int TF_OperationNumOutputs(TF_Operation* oper); -extern TF_DataType TF_OperationOutputType(TF_Output oper_out); -extern int TF_OperationOutputListLength(TF_Operation* oper, - const char* arg_name, - TF_Status* status); +TF_CAPI_EXPORT extern int TF_OperationNumOutputs(TF_Operation* oper); +TF_CAPI_EXPORT extern TF_DataType TF_OperationOutputType(TF_Output oper_out); +TF_CAPI_EXPORT extern int TF_OperationOutputListLength(TF_Operation* oper, + const char* arg_name, + TF_Status* status); -extern int TF_OperationNumInputs(TF_Operation* oper); -extern TF_DataType TF_OperationInputType(TF_Input oper_in); -extern int TF_OperationInputListLength(TF_Operation* oper, const char* arg_name, - TF_Status* status); +TF_CAPI_EXPORT extern int TF_OperationNumInputs(TF_Operation* oper); +TF_CAPI_EXPORT extern TF_DataType TF_OperationInputType(TF_Input oper_in); +TF_CAPI_EXPORT extern int TF_OperationInputListLength(TF_Operation* oper, + const char* arg_name, + TF_Status* status); // In this code: // TF_Output producer = TF_OperationInput(consumer); // There is an edge from producer.oper's output (given by // producer.index) to consumer.oper's input (given by consumer.index). -extern TF_Output TF_OperationInput(TF_Input oper_in); +TF_CAPI_EXPORT extern TF_Output TF_OperationInput(TF_Input oper_in); // Get the number of current consumers of a specific output of an // operation. Note that this number can change when new operations // are added to the graph. -extern int TF_OperationOutputNumConsumers(TF_Output oper_out); +TF_CAPI_EXPORT extern int TF_OperationOutputNumConsumers(TF_Output oper_out); // Get list of all current consumers of a specific output of an // operation. `consumers` must point to an array of length at least @@ -550,24 +607,24 @@ extern int TF_OperationOutputNumConsumers(TF_Output oper_out); // modification of the graph can increase the number of consumers of // an operation. Returns the number of output consumers (should match // TF_OperationOutputNumConsumers(oper_out)). -extern int TF_OperationOutputConsumers(TF_Output oper_out, TF_Input* consumers, - int max_consumers); +TF_CAPI_EXPORT extern int TF_OperationOutputConsumers(TF_Output oper_out, + TF_Input* consumers, + int max_consumers); // Get the number of control inputs to an operation. -extern int TF_OperationNumControlInputs(TF_Operation* oper); +TF_CAPI_EXPORT extern int TF_OperationNumControlInputs(TF_Operation* oper); // Get list of all control inputs to an operation. `control_inputs` must // point to an array of length `max_control_inputs` (ideally set to // TF_OperationNumControlInputs(oper)). Returns the number of control // inputs (should match TF_OperationNumControlInputs(oper)). -extern int TF_OperationGetControlInputs(TF_Operation* oper, - TF_Operation** control_inputs, - int max_control_inputs); +TF_CAPI_EXPORT extern int TF_OperationGetControlInputs( + TF_Operation* oper, TF_Operation** control_inputs, int max_control_inputs); // Get the number of operations that have `*oper` as a control input. // Note that this number can change when new operations are added to // the graph. -extern int TF_OperationNumControlOutputs(TF_Operation* oper); +TF_CAPI_EXPORT extern int TF_OperationNumControlOutputs(TF_Operation* oper); // Get the list of operations that have `*oper` as a control input. // `control_outputs` must point to an array of length at least @@ -576,12 +633,12 @@ extern int TF_OperationNumControlOutputs(TF_Operation* oper); // modification of the graph can increase the number of control // outputs. Returns the number of control outputs (should match // TF_OperationNumControlOutputs(oper)). -extern int TF_OperationGetControlOutputs(TF_Operation* oper, - TF_Operation** control_outputs, - int max_control_outputs); +TF_CAPI_EXPORT extern int TF_OperationGetControlOutputs( + TF_Operation* oper, TF_Operation** control_outputs, + int max_control_outputs); // TF_AttrType describes the type of the value of an attribute on an operation. -typedef enum { +typedef enum TF_AttrType { TF_ATTR_STRING = 0, TF_ATTR_INT = 1, TF_ATTR_FLOAT = 2, @@ -625,17 +682,18 @@ typedef struct TF_AttrMetadata { } TF_AttrMetadata; // Returns metadata about the value of the attribute `attr_name` of `oper`. -extern TF_AttrMetadata TF_OperationGetAttrMetadata(TF_Operation* oper, - const char* attr_name, - TF_Status* status); +TF_CAPI_EXPORT extern TF_AttrMetadata TF_OperationGetAttrMetadata( + TF_Operation* oper, const char* attr_name, TF_Status* status); // Fills in `value` with the value of the attribute `attr_name`. `value` must // point to an array of length at least `max_length` (ideally set to // TF_AttrMetadata.total_size from TF_OperationGetAttrMetadata(oper, // attr_name)). -extern void TF_OperationGetAttrString(TF_Operation* oper, const char* attr_name, - void* value, size_t max_length, - TF_Status* status); +TF_CAPI_EXPORT extern void TF_OperationGetAttrString(TF_Operation* oper, + const char* attr_name, + void* value, + size_t max_length, + TF_Status* status); // Get the list of strings in the value of the attribute `attr_name`. Fills in // `values` and `lengths`, each of which must point to an array of length at @@ -648,64 +706,78 @@ extern void TF_OperationGetAttrString(TF_Operation* oper, const char* attr_name, // attr_name). // // Fails if storage_size is too small to hold the requested number of strings. -extern void TF_OperationGetAttrStringList(TF_Operation* oper, - const char* attr_name, void** values, - size_t* lengths, int max_values, - void* storage, size_t storage_size, - TF_Status* status); +TF_CAPI_EXPORT extern void TF_OperationGetAttrStringList( + TF_Operation* oper, const char* attr_name, void** values, size_t* lengths, + int max_values, void* storage, size_t storage_size, TF_Status* status); -extern void TF_OperationGetAttrInt(TF_Operation* oper, const char* attr_name, - int64_t* value, TF_Status* status); +TF_CAPI_EXPORT extern void TF_OperationGetAttrInt(TF_Operation* oper, + const char* attr_name, + int64_t* value, + TF_Status* status); // Fills in `values` with the value of the attribute `attr_name` of `oper`. // `values` must point to an array of length at least `max_values` (ideally set // TF_AttrMetadata.list_size from TF_OperationGetAttrMetadata(oper, // attr_name)). -extern void TF_OperationGetAttrIntList(TF_Operation* oper, - const char* attr_name, int64_t* values, - int max_values, TF_Status* status); +TF_CAPI_EXPORT extern void TF_OperationGetAttrIntList(TF_Operation* oper, + const char* attr_name, + int64_t* values, + int max_values, + TF_Status* status); -extern void TF_OperationGetAttrFloat(TF_Operation* oper, const char* attr_name, - float* value, TF_Status* status); +TF_CAPI_EXPORT extern void TF_OperationGetAttrFloat(TF_Operation* oper, + const char* attr_name, + float* value, + TF_Status* status); // Fills in `values` with the value of the attribute `attr_name` of `oper`. // `values` must point to an array of length at least `max_values` (ideally set // to TF_AttrMetadata.list_size from TF_OperationGetAttrMetadata(oper, // attr_name)). -extern void TF_OperationGetAttrFloatList(TF_Operation* oper, - const char* attr_name, float* values, - int max_values, TF_Status* status); +TF_CAPI_EXPORT extern void TF_OperationGetAttrFloatList(TF_Operation* oper, + const char* attr_name, + float* values, + int max_values, + TF_Status* status); -extern void TF_OperationGetAttrBool(TF_Operation* oper, const char* attr_name, - unsigned char* value, TF_Status* status); +TF_CAPI_EXPORT extern void TF_OperationGetAttrBool(TF_Operation* oper, + const char* attr_name, + unsigned char* value, + TF_Status* status); // Fills in `values` with the value of the attribute `attr_name` of `oper`. // `values` must point to an array of length at least `max_values` (ideally set // to TF_AttrMetadata.list_size from TF_OperationGetAttrMetadata(oper, // attr_name)). -extern void TF_OperationGetAttrBoolList(TF_Operation* oper, - const char* attr_name, - unsigned char* values, int max_values, - TF_Status* status); +TF_CAPI_EXPORT extern void TF_OperationGetAttrBoolList(TF_Operation* oper, + const char* attr_name, + unsigned char* values, + int max_values, + TF_Status* status); -extern void TF_OperationGetAttrType(TF_Operation* oper, const char* attr_name, - TF_DataType* value, TF_Status* status); +TF_CAPI_EXPORT extern void TF_OperationGetAttrType(TF_Operation* oper, + const char* attr_name, + TF_DataType* value, + TF_Status* status); // Fills in `values` with the value of the attribute `attr_name` of `oper`. // `values` must point to an array of length at least `max_values` (ideally set // to TF_AttrMetadata.list_size from TF_OperationGetAttrMetadata(oper, // attr_name)). -extern void TF_OperationGetAttrTypeList(TF_Operation* oper, - const char* attr_name, - TF_DataType* values, int max_values, - TF_Status* status); +TF_CAPI_EXPORT extern void TF_OperationGetAttrTypeList(TF_Operation* oper, + const char* attr_name, + TF_DataType* values, + int max_values, + TF_Status* status); // Fills in `value` with the value of the attribute `attr_name` of `oper`. // `values` must point to an array of length at least `num_dims` (ideally set to // TF_Attr_Meta.size from TF_OperationGetAttrMetadata(oper, attr_name)). -extern void TF_OperationGetAttrShape(TF_Operation* oper, const char* attr_name, - int64_t* value, int num_dims, - TF_Status* status); +TF_CAPI_EXPORT extern void TF_OperationGetAttrShape(TF_Operation* oper, + const char* attr_name, + int64_t* value, + int num_dims, + TF_Status* status); // Fills in `dims` with the list of shapes in the attribute `attr_name` of // `oper` and `num_dims` with the corresponding number of dimensions. On return, @@ -720,35 +792,32 @@ extern void TF_OperationGetAttrShape(TF_Operation* oper, const char* attr_name, // attr_name). // // Fails if storage_size is insufficient to hold the requested shapes. -extern void TF_OperationGetAttrShapeList(TF_Operation* oper, - const char* attr_name, int64_t** dims, - int* num_dims, int num_shapes, - int64_t* storage, int storage_size, - TF_Status* status); +TF_CAPI_EXPORT extern void TF_OperationGetAttrShapeList( + TF_Operation* oper, const char* attr_name, int64_t** dims, int* num_dims, + int num_shapes, int64_t* storage, int storage_size, TF_Status* status); // Sets `value` to the binary-serialized TensorShapeProto of the value of // `attr_name` attribute of `oper`'. -extern void TF_OperationGetAttrTensorShapeProto(TF_Operation* oper, - const char* attr_name, - TF_Buffer* value, - TF_Status* status); +TF_CAPI_EXPORT extern void TF_OperationGetAttrTensorShapeProto( + TF_Operation* oper, const char* attr_name, TF_Buffer* value, + TF_Status* status); // Fills in `values` with binary-serialized TensorShapeProto values of the // attribute `attr_name` of `oper`. `values` must point to an array of length at // least `num_values` (ideally set to TF_AttrMetadata.list_size from // TF_OperationGetAttrMetadata(oper, attr_name)). -extern void TF_OperationGetAttrTensorShapeProtoList(TF_Operation* oper, - const char* attr_name, - TF_Buffer** values, - int max_values, - TF_Status* status); +TF_CAPI_EXPORT extern void TF_OperationGetAttrTensorShapeProtoList( + TF_Operation* oper, const char* attr_name, TF_Buffer** values, + int max_values, TF_Status* status); // Gets the TF_Tensor valued attribute of `attr_name` of `oper`. // // Allocates a new TF_Tensor which the caller is expected to take // ownership of (and can deallocate using TF_DeleteTensor). -extern void TF_OperationGetAttrTensor(TF_Operation* oper, const char* attr_name, - TF_Tensor** value, TF_Status* status); +TF_CAPI_EXPORT extern void TF_OperationGetAttrTensor(TF_Operation* oper, + const char* attr_name, + TF_Tensor** value, + TF_Status* status); // Fills in `values` with the TF_Tensor values of the attribute `attr_name` of // `oper`. `values` must point to an array of TF_Tensor* of length at least @@ -757,22 +826,22 @@ extern void TF_OperationGetAttrTensor(TF_Operation* oper, const char* attr_name, // // The caller takes ownership of all the non-null TF_Tensor* entries in `values` // (which can be deleted using TF_DeleteTensor(values[i])). -extern void TF_OperationGetAttrTensorList(TF_Operation* oper, - const char* attr_name, - TF_Tensor** values, int max_values, - TF_Status* status); +TF_CAPI_EXPORT extern void TF_OperationGetAttrTensorList(TF_Operation* oper, + const char* attr_name, + TF_Tensor** values, + int max_values, + TF_Status* status); // Sets `output_attr_value` to the binary-serialized AttrValue proto // representation of the value of the `attr_name` attr of `oper`. -extern void TF_OperationGetAttrValueProto(TF_Operation* oper, - const char* attr_name, - TF_Buffer* output_attr_value, - TF_Status* status); +TF_CAPI_EXPORT extern void TF_OperationGetAttrValueProto( + TF_Operation* oper, const char* attr_name, TF_Buffer* output_attr_value, + TF_Status* status); // Returns the operation in the graph with `oper_name`. Returns nullptr if // no operation found. -extern TF_Operation* TF_GraphOperationByName(TF_Graph* graph, - const char* oper_name); +TF_CAPI_EXPORT extern TF_Operation* TF_GraphOperationByName( + TF_Graph* graph, const char* oper_name); // Iterate through the operations of a graph. To use: // size_t pos = 0; @@ -780,7 +849,8 @@ extern TF_Operation* TF_GraphOperationByName(TF_Graph* graph, // while ((oper = TF_GraphNextOperation(graph, &pos)) != nullptr) { // DoSomethingWithOperation(oper); // } -extern TF_Operation* TF_GraphNextOperation(TF_Graph* graph, size_t* pos); +TF_CAPI_EXPORT extern TF_Operation* TF_GraphNextOperation(TF_Graph* graph, + size_t* pos); // Write out a serialized representation of `graph` (as a GraphDef protocol // message) to `output_graph_def` (allocated by TF_NewBuffer()). @@ -788,25 +858,27 @@ extern TF_Operation* TF_GraphNextOperation(TF_Graph* graph, size_t* pos); // is called. // // May fail on very large graphs in the future. -extern void TF_GraphToGraphDef(TF_Graph* graph, TF_Buffer* output_graph_def, - TF_Status* status); +TF_CAPI_EXPORT extern void TF_GraphToGraphDef(TF_Graph* graph, + TF_Buffer* output_graph_def, + TF_Status* status); // TF_ImportGraphDefOptions holds options that can be passed to // TF_GraphImportGraphDef. typedef struct TF_ImportGraphDefOptions TF_ImportGraphDefOptions; -extern TF_ImportGraphDefOptions* TF_NewImportGraphDefOptions(); -extern void TF_DeleteImportGraphDefOptions(TF_ImportGraphDefOptions* opts); +TF_CAPI_EXPORT extern TF_ImportGraphDefOptions* TF_NewImportGraphDefOptions(); +TF_CAPI_EXPORT extern void TF_DeleteImportGraphDefOptions( + TF_ImportGraphDefOptions* opts); // Set the prefix to be prepended to the names of nodes in `graph_def` that will // be imported into `graph`. -extern void TF_ImportGraphDefOptionsSetPrefix(TF_ImportGraphDefOptions* opts, - const char* prefix); +TF_CAPI_EXPORT extern void TF_ImportGraphDefOptionsSetPrefix( + TF_ImportGraphDefOptions* opts, const char* prefix); // Set any imported nodes with input `src_name:src_index` to have that input // replaced with `dst`. `src_name` refers to a node in the graph to be imported, // `dst` references a node already existing in the graph being imported into. -extern void TF_ImportGraphDefOptionsAddInputMapping( +TF_CAPI_EXPORT extern void TF_ImportGraphDefOptionsAddInputMapping( TF_ImportGraphDefOptions* opts, const char* src_name, int src_index, TF_Output dst); @@ -814,23 +886,23 @@ extern void TF_ImportGraphDefOptionsAddInputMapping( // replaced with `dst`. `src_name` refers to a node in the graph to be imported, // `dst` references an operation already existing in the graph being imported // into. -extern void TF_GraphImportGraphDefOptionsRemapControlDependency( +TF_CAPI_EXPORT extern void TF_ImportGraphDefOptionsRemapControlDependency( TF_ImportGraphDefOptions* opts, const char* src_name, TF_Operation* dst); // Cause the imported graph to have a control dependency on `oper`. `oper` // should exist in the graph being imported into. -extern void TF_ImportGraphDefOptionsAddControlDependency( +TF_CAPI_EXPORT extern void TF_ImportGraphDefOptionsAddControlDependency( TF_ImportGraphDefOptions* opts, TF_Operation* oper); // Add an output in `graph_def` to be returned via the `return_outputs` output // parameter of TF_GraphImportGraphDef(). If the output is remapped via an input // mapping, the corresponding existing tensor in `graph` will be returned. -extern void TF_ImportGraphDefOptionsAddReturnOutput( +TF_CAPI_EXPORT extern void TF_ImportGraphDefOptionsAddReturnOutput( TF_ImportGraphDefOptions* opts, const char* oper_name, int index); // Returns the number of return outputs added via // TF_ImportGraphDefOptionsAddReturnOutput(). -extern int TF_ImportGraphDefOptionsNumReturnOutputs( +TF_CAPI_EXPORT extern int TF_ImportGraphDefOptionsNumReturnOutputs( const TF_ImportGraphDefOptions* opts); // Import the graph serialized in `graph_def` into `graph`. @@ -839,22 +911,31 @@ extern int TF_ImportGraphDefOptionsNumReturnOutputs( // result of TF_ImportGraphDefOptionsNumReturnOutputs()). If // `num_return_outputs` is non-zero, `return_outputs` must be of length // `num_return_outputs`. Otherwise it can be null. -extern void TF_GraphImportGraphDefWithReturnOutputs( +TF_CAPI_EXPORT extern void TF_GraphImportGraphDefWithReturnOutputs( TF_Graph* graph, const TF_Buffer* graph_def, const TF_ImportGraphDefOptions* options, TF_Output* return_outputs, int num_return_outputs, TF_Status* status); // Import the graph serialized in `graph_def` into `graph`. // Convenience function for when no return outputs have been added. -extern void TF_GraphImportGraphDef(TF_Graph* graph, const TF_Buffer* graph_def, - const TF_ImportGraphDefOptions* options, - TF_Status* status); +TF_CAPI_EXPORT extern void TF_GraphImportGraphDef( + TF_Graph* graph, const TF_Buffer* graph_def, + const TF_ImportGraphDefOptions* options, TF_Status* status); + +// Add `function` to graph `g`. Once `function` is added to `g`, +// it can be called by creating an operation using the function's name. +// +// If successful, status is set to OK and function is added to g +// Otherwise, status is set to the encountered error and g is unmodified +TF_CAPI_EXPORT extern void TF_GraphAddFunction(TF_Graph* g, + const TF_Function* function, + TF_Status* status); // Note: The following function may fail on very large protos in the future. -extern void TF_OperationToNodeDef(TF_Operation* oper, - TF_Buffer* output_node_def, - TF_Status* status); +TF_CAPI_EXPORT extern void TF_OperationToNodeDef(TF_Operation* oper, + TF_Buffer* output_node_def, + TF_Status* status); typedef struct TF_WhileParams { // The number of inputs to the while loop, i.e. the number of loop variables. @@ -894,12 +975,13 @@ typedef struct TF_WhileParams { // TF_FinishWhile() or TF_AbortWhile(). // // Missing functionality (TODO): -// - Gradients (not yet implmented for any ops) +// - Gradients // - Reference-type inputs // - Directly referencing external tensors from the cond/body graphs (this is // possible in the Python API) -TF_WhileParams TF_NewWhile(TF_Graph* g, TF_Output* inputs, int ninputs, - TF_Status* status); +TF_CAPI_EXPORT extern TF_WhileParams TF_NewWhile(TF_Graph* g, TF_Output* inputs, + int ninputs, + TF_Status* status); // Builds the while loop specified by `params` and returns the output tensors of // the while loop in `outputs`. `outputs` should be allocated to size @@ -909,15 +991,131 @@ TF_WhileParams TF_NewWhile(TF_Graph* g, TF_Output* inputs, int ninputs, // // Either this or TF_AbortWhile() must be called after a successful // TF_NewWhile() call. -void TF_FinishWhile(const TF_WhileParams* params, TF_Status* status, - TF_Output* outputs); +TF_CAPI_EXPORT extern void TF_FinishWhile(const TF_WhileParams* params, + TF_Status* status, + TF_Output* outputs); // Frees `params`s resources without building a while loop. `params` is no // longer valid after this returns. Either this or TF_FinishWhile() must be // called after a successful TF_NewWhile() call. -void TF_AbortWhile(const TF_WhileParams* params); +TF_CAPI_EXPORT extern void TF_AbortWhile(const TF_WhileParams* params); + +// Adds operations to compute the partial derivatives of sum of `y`s w.r.t `x`s, +// i.e., d(y_1 + y_2 + ...)/dx_1, d(y_1 + y_2 + ...)/dx_2... +// `dx` are used as initial gradients (which represent the symbolic partial +// derivatives of some loss function `L` w.r.t. `y`). +// `dx` must be nullptr or have size `ny`. +// If `dx` is nullptr, the implementation will use dx of `OnesLike` for all +// shapes in `y`. +// The partial derivatives are returned in `dy`. `dy` should be allocated to +// size `nx`. +// +// WARNING: This function does not yet support all the gradients that python +// supports. See +// https://www.tensorflow.org/code/tensorflow/cc/gradients/README.md +// for instructions on how to add C++ more gradients. +TF_CAPI_EXPORT void TF_AddGradients(TF_Graph* g, TF_Output* y, int ny, + TF_Output* x, int nx, TF_Output* dx, + TF_Status* status, TF_Output* dy); + +// Create a TF_Function from a TF_Graph +// +// Params: +// fn_body - the graph whose operations (or subset of whose operations) will be +// converted to TF_Function. +// fn_name - the name of the new TF_Function. Should match the operation +// name (OpDef.name) regexp [A-Z][A-Za-z0-9_.\\-/]* and be distinct +// from other operation names (at least those registered in graphs +// where this function will be used). +// TODO(iga): Allow null in here and have C API come up with +// a unique name with high probability (similarly to +// _create_hash_str in function.py) +// num_opers - `num_opers` contains the number of elements in the `opers` array +// or a special value of -1 meaning that no array is given. +// The distinction between an empty array of operations and no +// array of operations is necessary to distinguish the case of +// creating a function with no body (e.g. identity or permutation) +// and the case of creating a function whose body contains all +// the nodes in the graph (except for the automatic skipping, see +// below). +// opers - Array of operations to become the body of the function or null. +// - If no array is given (`num_opers` = -1), all the +// operations in `fn_body` will become part of the function +// except operations referenced in `inputs`. These operations +// must have a single output (these operations are typically +// placeholders created for the sole purpose of representing +// an input. We can relax this constraint if there are +// compelling use cases). +// - If an array is given (`num_opers` >= 0), all operations +// in it will become part of the function. In particular, no +// automatic skipping of dummy input operations is performed. +// ninputs - number of elements in `inputs` array +// inputs - array of TF_Outputs that specify the inputs to the function. +// If `ninputs` is zero (the function takes no inputs), `inputs` +// can be null. The names used for function inputs are normalized +// names of the operations (usually placeholders) pointed to by +// `inputs`. These operation names should start with a letter. +// Normalization will convert all letters to lowercase and +// non-alphanumeric characters to '_' to make resulting names match +// the "[a-z][a-z0-9_]*" pattern for operation argument names. +// `inputs` cannot contain the same tensor twice. +// noutputs - number of elements in `outputs` array +// outputs - array of TF_Outputs that specify the outputs of the function. +// If `noutputs` is zero (the function returns no outputs), `outputs` +// can be null. `outputs` can contain the same tensor more than once. +// output_names - The names of the function's outputs. `output_names` array +// must either have the same length as `outputs` +// (i.e. `noutputs`) or be null. In the former case, +// the names should match the regular expression for ArgDef +// names - "[a-z][a-z0-9_]*". In the latter case, +// names for outputs will be generated automatically. +// opts - various options for the function, e.g. XLA's inlining control. +// status - Set to OK on success and an appropriate error on failure. +// +// Note that when the same TF_Output is listed as both an input and an output, +// the corresponding function's output will equal to this input, +// instead of the original node's output. +// +// Callers must also satisfy the following constraints: +// - `inputs` cannot refer to TF_Outputs within a control flow context. For +// example, one cannot use the output of "switch" node as input. +// - No TF_Output of a function (inside any of `inputs`, `outputs`, `fn_body`) +// is allowed to have a reference type. Reference types are not exposed +// through C API and are being deprecated. +// - Every node in the function's body must have all of its inputs (including +// control inputs). In other words, for every node in the body, each input +// must be either listed in `inputs` or must come from another node in +// the body. In particular, it is an error to have a control edge going from +// a node outside of the body into a node in the body. This applies to control +// edges going from nodes referenced in `inputs` to nodes in the body when +// the former nodes are not in the body (automatically skipped or not +// included in explicitly specified body). +// +// Returns: +// On successful, a newly created TF_Function instance. It must be deleted by +// calling TF_DeleteFunction. +// +// On failure, null. +// +// TODO(iga): Add input_names argument and get output_names working (they are +// currently ignored) +TF_CAPI_EXPORT extern TF_Function* TF_GraphToFunction( + const TF_Graph* fn_body, const char* fn_name, int num_opers, + const TF_Operation* const* opers, int ninputs, const TF_Output* inputs, + int noutputs, const TF_Output* outputs, const char* const* output_names, + const TF_FunctionOptions* opts, TF_Status* status); + +// Write out a serialized representation of `func` (as a FunctionDef protocol +// message) to `output_func_def` (allocated by TF_NewBuffer()). +// `output_func_def`'s underlying buffer will be freed when TF_DeleteBuffer() +// is called. +// +// May fail on very large graphs in the future. +TF_CAPI_EXPORT extern void TF_FunctionToFunctionDef(TF_Function* func, + TF_Buffer* output_func_def, + TF_Status* status); -// TODO(andydavis): Function to add gradients to a graph. +TF_CAPI_EXPORT extern void TF_DeleteFunction(TF_Function*); // TODO(josh11b): Register OpDef, available to all operations added // to this graph. @@ -936,8 +1134,9 @@ typedef struct TF_Session TF_Session; // *graph must be a valid graph (not deleted or nullptr). This function will // prevent the graph from being deleted until TF_DeleteSession() is called. // Does not take ownership of opts. -extern TF_Session* TF_NewSession(TF_Graph* graph, const TF_SessionOptions* opts, - TF_Status* status); +TF_CAPI_EXPORT extern TF_Session* TF_NewSession(TF_Graph* graph, + const TF_SessionOptions* opts, + TF_Status* status); // This function creates a new TF_Session (which is created on success) using // `session_options`, and then initializes state (restoring tensors and other @@ -953,7 +1152,7 @@ extern TF_Session* TF_NewSession(TF_Graph* graph, const TF_SessionOptions* opts, // // If successful, populates `graph` with the contents of the Graph and // `meta_graph_def` with the MetaGraphDef of the loaded model. -TF_Session* TF_LoadSessionFromSavedModel( +TF_CAPI_EXPORT extern TF_Session* TF_LoadSessionFromSavedModel( const TF_SessionOptions* session_options, const TF_Buffer* run_options, const char* export_dir, const char* const* tags, int tags_len, TF_Graph* graph, TF_Buffer* meta_graph_def, TF_Status* status); @@ -962,7 +1161,7 @@ TF_Session* TF_LoadSessionFromSavedModel( // // Contacts any other processes associated with the session, if applicable. // May not be called after TF_DeleteSession(). -extern void TF_CloseSession(TF_Session*, TF_Status* status); +TF_CAPI_EXPORT extern void TF_CloseSession(TF_Session*, TF_Status* status); // Destroy a session object. // @@ -970,7 +1169,7 @@ extern void TF_CloseSession(TF_Session*, TF_Status* status); // local resources associated with the session. The session may not be used // during or after this call (and the session drops its reference to the // corresponding graph). -extern void TF_DeleteSession(TF_Session*, TF_Status* status); +TF_CAPI_EXPORT extern void TF_DeleteSession(TF_Session*, TF_Status* status); // Run the graph associated with the session starting with the supplied inputs // (inputs[0,ninputs-1] with corresponding values in input_values[0,ninputs-1]). @@ -996,21 +1195,20 @@ extern void TF_DeleteSession(TF_Session*, TF_Status* status); // to the caller, which must eventually call TF_DeleteTensor on them. // // On failure, output_values[] contains NULLs. -extern void TF_SessionRun(TF_Session* session, - // RunOptions - const TF_Buffer* run_options, - // Input tensors - const TF_Output* inputs, - TF_Tensor* const* input_values, int ninputs, - // Output tensors - const TF_Output* outputs, TF_Tensor** output_values, - int noutputs, - // Target operations - const TF_Operation* const* target_opers, int ntargets, - // RunMetadata - TF_Buffer* run_metadata, - // Output status - TF_Status*); +TF_CAPI_EXPORT extern void TF_SessionRun( + TF_Session* session, + // RunOptions + const TF_Buffer* run_options, + // Input tensors + const TF_Output* inputs, TF_Tensor* const* input_values, int ninputs, + // Output tensors + const TF_Output* outputs, TF_Tensor** output_values, int noutputs, + // Target operations + const TF_Operation* const* target_opers, int ntargets, + // RunMetadata + TF_Buffer* run_metadata, + // Output status + TF_Status*); // Set up the graph with the intended feeds (inputs) and fetches (outputs) for a // sequence of partial run calls. @@ -1020,40 +1218,36 @@ extern void TF_SessionRun(TF_Session* session, // needed. // // On failure, out_status contains a tensorflow::Status with an error -// message. -// NOTE: This is EXPERIMENTAL and subject to change. -extern void TF_SessionPRunSetup(TF_Session*, - // Input names - const TF_Output* inputs, int ninputs, - // Output names - const TF_Output* outputs, int noutputs, - // Target operations - const TF_Operation* const* target_opers, - int ntargets, - // Output handle - const char** handle, - // Output status - TF_Status*); +// message. *handle is set to nullptr. +TF_CAPI_EXPORT extern void TF_SessionPRunSetup( + TF_Session*, + // Input names + const TF_Output* inputs, int ninputs, + // Output names + const TF_Output* outputs, int noutputs, + // Target operations + const TF_Operation* const* target_opers, int ntargets, + // Output handle + const char** handle, + // Output status + TF_Status*); // Continue to run the graph with additional feeds and fetches. The // execution state is uniquely identified by the handle. -// NOTE: This is EXPERIMENTAL and subject to change. -extern void TF_SessionPRun(TF_Session*, const char* handle, - // Input tensors - const TF_Output* inputs, - TF_Tensor* const* input_values, int ninputs, - // Output tensors - const TF_Output* outputs, TF_Tensor** output_values, - int noutputs, - // Target operations - const TF_Operation* const* target_opers, - int ntargets, - // Output status - TF_Status*); +TF_CAPI_EXPORT extern void TF_SessionPRun( + TF_Session*, const char* handle, + // Input tensors + const TF_Output* inputs, TF_Tensor* const* input_values, int ninputs, + // Output tensors + const TF_Output* outputs, TF_Tensor** output_values, int noutputs, + // Target operations + const TF_Operation* const* target_opers, int ntargets, + // Output status + TF_Status*); // Deletes a handle allocated by TF_SessionPRunSetup. // Once called, no more calls to TF_SessionPRun should be made. -extern void TF_DeletePRunHandle(const char* handle); +TF_CAPI_EXPORT extern void TF_DeletePRunHandle(const char* handle); // -------------------------------------------------------------------------- // The deprecated session API. Please switch to the above instead of @@ -1062,39 +1256,96 @@ extern void TF_DeletePRunHandle(const char* handle); typedef struct TF_DeprecatedSession TF_DeprecatedSession; -extern TF_DeprecatedSession* TF_NewDeprecatedSession(const TF_SessionOptions*, +TF_CAPI_EXPORT extern TF_DeprecatedSession* TF_NewDeprecatedSession( + const TF_SessionOptions*, TF_Status* status); +TF_CAPI_EXPORT extern void TF_CloseDeprecatedSession(TF_DeprecatedSession*, TF_Status* status); -extern void TF_CloseDeprecatedSession(TF_DeprecatedSession*, TF_Status* status); -extern void TF_DeleteDeprecatedSession(TF_DeprecatedSession*, - TF_Status* status); -extern void TF_Reset(const TF_SessionOptions* opt, const char** containers, - int ncontainers, TF_Status* status); +TF_CAPI_EXPORT extern void TF_DeleteDeprecatedSession(TF_DeprecatedSession*, + TF_Status* status); +TF_CAPI_EXPORT extern void TF_Reset(const TF_SessionOptions* opt, + const char** containers, int ncontainers, + TF_Status* status); // Treat the bytes proto[0,proto_len-1] as a serialized GraphDef and // add the nodes in that GraphDef to the graph for the session. // // Prefer use of TF_Session and TF_GraphImportGraphDef over this. -extern void TF_ExtendGraph(TF_DeprecatedSession*, const void* proto, - size_t proto_len, TF_Status*); +TF_CAPI_EXPORT extern void TF_ExtendGraph(TF_DeprecatedSession*, + const void* proto, size_t proto_len, + TF_Status*); // See TF_SessionRun() above. -extern void TF_Run(TF_DeprecatedSession*, const TF_Buffer* run_options, - const char** input_names, TF_Tensor** inputs, int ninputs, - const char** output_names, TF_Tensor** outputs, int noutputs, - const char** target_oper_names, int ntargets, - TF_Buffer* run_metadata, TF_Status*); +TF_CAPI_EXPORT extern void TF_Run(TF_DeprecatedSession*, + const TF_Buffer* run_options, + const char** input_names, TF_Tensor** inputs, + int ninputs, const char** output_names, + TF_Tensor** outputs, int noutputs, + const char** target_oper_names, int ntargets, + TF_Buffer* run_metadata, TF_Status*); // See TF_SessionPRunSetup() above. -extern void TF_PRunSetup(TF_DeprecatedSession*, const char** input_names, - int ninputs, const char** output_names, int noutputs, - const char** target_oper_names, int ntargets, - const char** handle, TF_Status*); +TF_CAPI_EXPORT extern void TF_PRunSetup(TF_DeprecatedSession*, + const char** input_names, int ninputs, + const char** output_names, int noutputs, + const char** target_oper_names, + int ntargets, const char** handle, + TF_Status*); // See TF_SessionPRun above. -extern void TF_PRun(TF_DeprecatedSession*, const char* handle, - const char** input_names, TF_Tensor** inputs, int ninputs, - const char** output_names, TF_Tensor** outputs, - int noutputs, const char** target_oper_names, int ntargets, - TF_Status*); +TF_CAPI_EXPORT extern void TF_PRun(TF_DeprecatedSession*, const char* handle, + const char** input_names, TF_Tensor** inputs, + int ninputs, const char** output_names, + TF_Tensor** outputs, int noutputs, + const char** target_oper_names, int ntargets, + TF_Status*); + +typedef struct TF_DeviceList TF_DeviceList; + +// Lists all devices in a TF_Session. +// +// Caller takes ownership of the returned TF_DeviceList* which must eventually +// be freed with a call to TF_DeleteDeviceList. +TF_CAPI_EXPORT extern TF_DeviceList* TF_SessionListDevices(TF_Session* session, + TF_Status* status); + +// Lists all devices in a TF_Session. +// +// Caller takes ownership of the returned TF_DeviceList* which must eventually +// be freed with a call to TF_DeleteDeviceList. +TF_CAPI_EXPORT extern TF_DeviceList* TF_DeprecatedSessionListDevices( + TF_DeprecatedSession* session, TF_Status* status); + +// Deallocates the device list. +TF_CAPI_EXPORT extern void TF_DeleteDeviceList(TF_DeviceList* list); + +// Counts the number of elements in the device list. +TF_CAPI_EXPORT extern int TF_DeviceListCount(const TF_DeviceList* list); + +// Retrieves the full name of the device (e.g. /job:worker/replica:0/...) +// The return value will be a pointer to a null terminated string. The caller +// must not modify or delete the string. It will be deallocated upon a call to +// TF_DeleteDeviceList. +// +// If index is out of bounds, an error code will be set in the status object, +// and a null pointer will be returned. +TF_CAPI_EXPORT extern const char* TF_DeviceListName(const TF_DeviceList* list, + int index, TF_Status*); + +// Retrieves the type of the device at the given index. +// +// The caller must not modify or delete the string. It will be deallocated upon +// a call to TF_DeleteDeviceList. +// +// If index is out of bounds, an error code will be set in the status object, +// and a null pointer will be returned. +TF_CAPI_EXPORT extern const char* TF_DeviceListType(const TF_DeviceList* list, + int index, TF_Status*); + +// Retrieve the amount of memory associated with a given device. +// +// If index is out of bounds, an error code will be set in the status object, +// and -1 will be returned. +TF_CAPI_EXPORT extern int64_t TF_DeviceListMemoryBytes( + const TF_DeviceList* list, int index, TF_Status*); // -------------------------------------------------------------------------- // Load plugins containing custom ops and kernels @@ -1113,19 +1364,19 @@ typedef struct TF_Library TF_Library; // The caller owns the library handle. // // On failure, place an error status in status and return NULL. -extern TF_Library* TF_LoadLibrary(const char* library_filename, - TF_Status* status); +TF_CAPI_EXPORT extern TF_Library* TF_LoadLibrary(const char* library_filename, + TF_Status* status); // Get the OpList of OpDefs defined in the library pointed by lib_handle. // // Returns a TF_Buffer. The memory pointed to by the result is owned by // lib_handle. The data in the buffer will be the serialized OpList proto for // ops defined in the library. -extern TF_Buffer TF_GetOpList(TF_Library* lib_handle); +TF_CAPI_EXPORT extern TF_Buffer TF_GetOpList(TF_Library* lib_handle); // Frees the memory associated with the library handle. // Does NOT unload the library. -extern void TF_DeleteLibraryHandle(TF_Library* lib_handle); +TF_CAPI_EXPORT extern void TF_DeleteLibraryHandle(TF_Library* lib_handle); // Get the OpList of all OpDefs defined in this address space. // Returns a TF_Buffer, ownership of which is transferred to the caller @@ -1133,7 +1384,7 @@ 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. -extern TF_Buffer* TF_GetAllOpList(); +TF_CAPI_EXPORT extern TF_Buffer* TF_GetAllOpList(); #ifdef __cplusplus } /* end extern "C" */ diff --git a/tensorflow/c/c_api_function.cc b/tensorflow/c/c_api_function.cc new file mode 100644 index 0000000000000000000000000000000000000000..b4c6397d0b4d34b4745f0f5115426b166354f570 --- /dev/null +++ b/tensorflow/c/c_api_function.cc @@ -0,0 +1,496 @@ +/* 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/c_api_internal.h" + +#include +#include +#include + +#include "tensorflow/core/framework/attr_value_util.h" +#include "tensorflow/core/framework/function.pb.h" +#include "tensorflow/core/framework/node_def.pb.h" +#include "tensorflow/core/framework/node_def_util.h" +#include "tensorflow/core/framework/types.h" +#include "tensorflow/core/graph/graph.h" +#include "tensorflow/core/lib/strings/strcat.h" + +namespace tensorflow { +namespace { + +// Class that maintains a one-to-one original node name -> new node name +// mapping. We normalize the names used as input and output arguments to match +// regexp "[a-z][a-z0-9_]*" specified in definition of ArgDef.name. +// Once we rename them, we risk creating a name collision with the other +// node names, so if necessary we add a suffix to make +// names unique. If we have an input named "A" and a node in the function +// body named "a", they will be renamed to "a" and "a_0". +class NodeNameMapping { + public: + NodeNameMapping() = default; + + // Normalize the input/output name and make it unique. + string GetIOName(const string& name); + + // Make the node name unique. + string Uniquify(const string& name); + + // Look up how a node name was previously normalized/uniquified. + // Returns empty if name was never seen. + string Lookup(const string& name) const; + + private: + string UniquifyHelper(const string& name) const; + static string Normalize(string name); + + // The normalized/uniquified names already used as + // input names (in signature), output names (in signature), and node names + // (in node_def). + // This is a superset of values in name_mapping_. + std::unordered_set used_names_; + // Mapping from original node name from the graph to the normalized + // and uniqified version of it. + std::unordered_map name_mapping_; +}; + +string NodeNameMapping::Normalize(string name) { + // Convert letters to lowercase and non-alphanumeric characters to '_'. + if (name.empty()) return "unknown"; + const int n = name.size(); + for (int i = 0; i < n; ++i) { + char c = name[i]; + if (isalnum(c)) { + if (isupper(c)) { + name[i] = tolower(c); + } + } else { + name[i] = '_'; + } + } + + // Find the first letter and start with it. + int i = 0; + for (; i < n; ++i) { + if (isalpha(name[i])) break; + } + + // Return "unknown" if none of the name's chars were letters. + return i == n ? "unknown" : name.substr(i); +} + +string NodeNameMapping::UniquifyHelper(const string& name) const { + // If the name hasn't been used yet, use it as-is. + if (used_names_.find(name) == used_names_.end()) return name; + // Add a suffix to name to make it unique. + for (int i = 0;; ++i) { + const string candidate = strings::StrCat(name, "_", i); + if (used_names_.find(candidate) == used_names_.end()) return candidate; + } +} + +string NodeNameMapping::GetIOName(const string& name) { + const string& input_name = UniquifyHelper(Normalize(name)); + // Record that we used this name, but don't add it to name_mapping_ + // since this name is not for a node. + used_names_.insert(input_name); + return input_name; +} + +string NodeNameMapping::Uniquify(const string& name) { + const string uniqued = UniquifyHelper(name); + name_mapping_[name] = uniqued; + used_names_.insert(uniqued); + return uniqued; +} + +string NodeNameMapping::Lookup(const string& name) const { + const auto iter = name_mapping_.find(name); + if (iter == name_mapping_.end()) return string(); + return iter->second; +} + +Status ValidateNoRefOutputs(const Node* node) { + for (int i = 0; i < node->num_outputs(); ++i) { + const DataType& dt = node->output_type(i); + if (IsRefType(dt)) { + return errors::InvalidArgument("Output ", i, " of node '", node->name(), + "' has a reference " + "type ", + DataTypeString(dt)); + } + } + return Status::OK(); +} + +Status FillFunctionBody( + const string& fn_name, const NodeNameMapping& node_names, + const std::vector& body_nodes, + const std::unordered_map& tensor_renaming, + FunctionDef* fdef) { + std::vector in_edges; + std::vector control_edges; + for (const Node* node : body_nodes) { + NodeDef* node_def = fdef->add_node_def(); + // First, copy the node_def as is. We will patch it next. + *node_def = node->def(); + if (!node->assigned_device_name().empty()) { + node_def->set_device(node->assigned_device_name()); + } + node_def->set_name(node_names.Lookup(node->name())); + + // Input names must be set based on nested names in tensor_renaming. + // Clear the flat input names we got from the original node_def + // from the graph. + node_def->clear_input(); + + // Collect regular and control inputs. Regular inputs are indexed + // by the index at which they come into the `node`. Control inputs + // don't follow any order. + in_edges.clear(); + in_edges.resize(node->num_inputs(), nullptr); + control_edges.clear(); + for (const Edge* edge : node->in_edges()) { + if (edge->src()->IsSource()) continue; + if (edge->IsControlEdge()) { + control_edges.push_back(edge); + } else { + in_edges[edge->dst_input()] = edge; + } + } + + // Add regular inputs. + for (size_t i = 0; i < in_edges.size(); ++i) { + const Edge* edge = in_edges[i]; + string original_input_name; + if (edge == nullptr) { + // A backedge might not appear as a regular Edge, but be only present + // in the node_def. Such edges are referred to as requested_inputs(). + if (i >= node->requested_inputs().size()) { + return errors::InvalidArgument( + "Graph to be converted to function appears to be malformed. ", + "Node ", node->name(), " is missing input edge ", i); + } + original_input_name = + ParseTensorName(node->requested_inputs()[i]).ToString(); + } else { + original_input_name = + strings::StrCat(edge->src()->name(), ":", edge->src_output()); + } + + const auto iter = tensor_renaming.find(original_input_name); + if (iter == tensor_renaming.end()) { + return errors::InvalidArgument( + "Input ", i, ", '", original_input_name, "', of node '", + node->name(), "' in function '", fn_name, + "' is not available. You might need to include it in inputs " + "or include its source node in the body"); + } + node_def->add_input(iter->second); + } + + // Add control inputs. + for (const Edge* edge : control_edges) { + // Add this control input only if the src node is in the body. + const string normalized = node_names.Lookup(edge->src()->name()); + // If we did not find a name for the source of control edge, this + // source must be outside of the body. Raise an error. + if (normalized.empty()) { + return errors::InvalidArgument( + "The source of control edge ", edge->DebugString(), + " is not in the body. Encountered while creating function '", + fn_name, "'"); + } + node_def->add_input(strings::StrCat("^", normalized)); + } + } + return Status::OK(); +} + +// Graph to FunctionDef conversion. This code is closely modeled on the Python +// code in third_party/tensorflow/python/framework/function.py. +Status GraphToFunctionDef(const Graph& fn_body, const string& fn_name, + const std::vector& body_nodes, + const std::vector& inputs, + const std::vector& outputs, + const std::vector& output_names, + FunctionDef* fdef) { + fdef->mutable_signature()->set_name(fn_name); + + // Keep track of names we used and how we normalized them. + NodeNameMapping node_names; + + // Mapping from original names of tensors (i.e. ":") to the + // name we used in the function: + // - For input tensors: + // {flat_tensor_name -> normalized_name_of_src_node} + // e.g. {In:3 -> in} + // - For tensors produced by nodes in function's body: + // {flat_tensor_name -> nested_tensor_name} + // e.g. {Add:3 -> add_0:z:1} + std::unordered_map tensor_renaming; + + // Fill inputs in function's signature. + for (size_t i = 0; i < inputs.size(); ++i) { + const Node* node = inputs[i].node; + int idx = inputs[i].index; + OpDef::ArgDef* argdef = fdef->mutable_signature()->add_input_arg(); + argdef->set_type(node->output_type(idx)); + const string& input_name = node_names.GetIOName(node->name()); + argdef->set_name(input_name); + tensor_renaming[strings::StrCat(node->name(), ":", idx)] = input_name; + } + + // Fill outputs in function's signature. + for (size_t i = 0; i < outputs.size(); ++i) { + const Node* node = outputs[i].node; + int idx = outputs[i].index; + OpDef::ArgDef* argdef = fdef->mutable_signature()->add_output_arg(); + argdef->set_type(node->output_type(idx)); + argdef->set_name(node_names.GetIOName(node->name())); + } + + // Populate tensor_renaming and node_names. + // Generate the new output names for every node in the function. + // The NodeDefs in FunctionDefs use a different naming scheme for + // their inputs than the NodeDefs in a graph (see the comment for + // FunctionDef.node_def in function.proto). We do the + // graph tensor name -> function tensor name conversion for every + // possible input (i.e. every node's outputs) and store the result + // in tensor_renaming. + for (const Node* node : body_nodes) { + // Make sure node_name does not collide with an input or output name. + const string& node_name = node_names.Uniquify(node->name()); + // For each output_arg in the op_def, the output_ranges + // map will have [start, end] range of indices that this arg produces + // among all the output tensors of this op. + NameRangeMap output_ranges; + TF_RETURN_IF_ERROR( + NameRangesForNode(*node, node->op_def(), nullptr, &output_ranges)); + for (const auto& output : output_ranges) { + const string& output_name = output.first; + int index_start = output.second.first; + int index_end = output.second.second; + for (int i = index_start; i < index_end; ++i) { + const string& original_name = strings::StrCat(node->name(), ":", i); + const string& new_name = + strings::StrCat(node_name, ":", output_name, ":", i - index_start); + // Record the mapping if this tensor is not already mapped. + // Tensor can be already mapped if it is used as an input. + if (tensor_renaming.find(original_name) == tensor_renaming.end()) { + tensor_renaming[original_name] = new_name; + } + } + } + } + + TF_RETURN_IF_ERROR( + FillFunctionBody(fn_name, node_names, body_nodes, tensor_renaming, fdef)); + + // Remap return values. + for (int r = 0; r < fdef->signature().output_arg_size(); ++r) { + const string& ret_name = fdef->signature().output_arg(r).name(); + + // We convert this flat tensor name to the nested value + // (e.g. `add:z:1`) that we stored in tensor_renaming. + const string& return_value = + strings::StrCat(outputs[r].node->name(), ":", outputs[r].index); + const auto iter = tensor_renaming.find(return_value); + if (iter == tensor_renaming.end()) { + return errors::InvalidArgument( + "TF_Output ", return_value, " is neither in the function body ", + "nor among function inputs. Encountered while creating function '", + fn_name, "'"); + } + (*fdef->mutable_ret())[ret_name] = iter->second; + } + + return Status::OK(); +} + +// Converts `ninputs` and `inputs` into `inputs_tensors` and `input_nodes` and +// does various checks while doing so. `input_nodes` will contain the same +// information as input_tensors just in a different structure to make +// following processing easier. TODO(iga): Simplify this nested structure. +Status ProcessInputs( + const TF_Graph* fn_body, const char* fn_name, int ninputs, + const TF_Output* inputs, std::vector* input_tensors, + std::unordered_map>* input_nodes) + EXCLUSIVE_LOCKS_REQUIRED(fn_body->mu) { + input_tensors->reserve(ninputs); + for (int i = 0; i < ninputs; ++i) { + const Node& node = inputs[i].oper->node; + int idx = inputs[i].index; + + TF_RETURN_WITH_CONTEXT_IF_ERROR( + fn_body->graph.IsValidOutputTensor(&node, idx), + "Encountered while processing input ", i, " into function '", fn_name, + "'"); + TF_RETURN_WITH_CONTEXT_IF_ERROR(ValidateNoRefOutputs(&node), + "Encountered while processing input ", i, + " into function '", fn_name, "'"); + + input_tensors->emplace_back(&node, idx); + + const auto& iter = input_nodes->find(&node); + if (iter == input_nodes->end()) { + input_nodes->insert({&node, {idx}}); + } else { + auto& indices = iter->second; + if (std::find(indices.begin(), indices.end(), idx) != indices.end()) { + return errors::InvalidArgument( + "TF_Output ", node.name(), ":", idx, + " appears more than once in the input list"); + } + indices.push_back(idx); + } + } + return Status::OK(); +} + +// Converts `noutputs` and `outputs` into `outputs_tensors` and does various +// checks while doing so. +Status ProcessOutputs(const TF_Graph* fn_body, const char* fn_name, + int noutputs, const TF_Output* outputs, + std::vector* output_tensors) + EXCLUSIVE_LOCKS_REQUIRED(fn_body->mu) { + output_tensors->reserve(noutputs); + for (int i = 0; i < noutputs; ++i) { + const Node& node = outputs[i].oper->node; + int idx = outputs[i].index; + TF_RETURN_WITH_CONTEXT_IF_ERROR( + fn_body->graph.IsValidOutputTensor(&node, idx), + "Encountered while processing output ", i, " from function '", fn_name, + "'"); + output_tensors->emplace_back(&node, idx); + } + return Status::OK(); +} + +// Populates `body_nodes` with the nodes that will become function's body. +// Performs various checks. +Status ComputeBodyNodes( + const TF_Graph* fn_body, const char* fn_name, int num_opers, + const TF_Operation* const* opers, + const std::unordered_map>& input_nodes, + std::vector* body_nodes) + EXCLUSIVE_LOCKS_REQUIRED(fn_body->mu) { + if (num_opers == -1) { + for (const Node* node : fn_body->graph.op_nodes()) { + const auto& iter = input_nodes.find(node); + if (iter == input_nodes.end()) { + // This node is not referenced in inputs. Add it to the body. + TF_RETURN_WITH_CONTEXT_IF_ERROR(ValidateNoRefOutputs(node), + "Encountered while creating function '", + fn_name, "'"); + body_nodes->push_back(node); + } else { + // This node is referenced in inputs. Currently, we place an + // artificial restriction and require that when num_opers=-1, such + // nodes must have a single output. + if (node->num_outputs() != 1) { + return errors::InvalidArgument( + "When `num_opers` is set to -1, nodes referenced in `inputs` " + "must have a single output. Node ", + node->name(), " has ", node->num_outputs(), + " outputs. Encountered while creating function '", fn_name, "'"); + } + } + } + } else { + body_nodes->reserve(num_opers); + for (int i = 0; i < num_opers; ++i) { + const Node* node = &opers[i]->node; + TF_RETURN_WITH_CONTEXT_IF_ERROR(ValidateNoRefOutputs(node), + "Encountered while creating function '", + fn_name, "'"); + body_nodes->push_back(node); + } + } + return Status::OK(); +} + +} // anonymous namespace +} // namespace tensorflow + +using tensorflow::Node; +using tensorflow::string; + +TF_Function* TF_GraphToFunction(const TF_Graph* fn_body, const char* fn_name, + int num_opers, const TF_Operation* const* opers, + int ninputs, const TF_Output* inputs, + int noutputs, const TF_Output* outputs, + const char* const* output_names, + const TF_FunctionOptions* opts, + TF_Status* status) { + tensorflow::mutex_lock l(*const_cast(&fn_body->mu)); + + // Process inputs. + std::vector input_tensors; + std::unordered_map> input_nodes; + status->status = tensorflow::ProcessInputs(fn_body, fn_name, ninputs, inputs, + &input_tensors, &input_nodes); + if (!status->status.ok()) return nullptr; + + // Process outputs. + std::vector output_tensors; + status->status = tensorflow::ProcessOutputs(fn_body, fn_name, noutputs, + outputs, &output_tensors); + if (!status->status.ok()) return nullptr; + + // Process output names. + std::vector output_names_vec; + if (output_names) { + output_names_vec.reserve(noutputs); + for (int i = 0; i < noutputs; ++i) { + output_names_vec.push_back(string(output_names[i])); + } + } + + // Compute body nodes. + std::vector body_nodes; + status->status = tensorflow::ComputeBodyNodes( + fn_body, fn_name, num_opers, opers, input_nodes, &body_nodes); + if (!status->status.ok()) return nullptr; + + // Do the actual function creation. + TF_Function* tf_function = new TF_Function(); + status->status = tensorflow::GraphToFunctionDef( + fn_body->graph, fn_name, body_nodes, input_tensors, output_tensors, + output_names_vec, tf_function->fdef_lib.add_function()); + if (!status->status.ok()) { + TF_DeleteFunction(tf_function); + return nullptr; + } + return tf_function; +} + +void TF_GraphAddFunction(TF_Graph* g, const TF_Function* function, + TF_Status* status) { + tensorflow::mutex_lock l(g->mu); + + // At the moment, we have only one function and no gradients in fdef_lib. + // This makes the following operation atomic. + // TODO(iga): Add an atomic version of AddFunctionLibrary when we support + // gradients + status->status = g->graph.AddFunctionLibrary(function->fdef_lib); +} + +void TF_FunctionToFunctionDef(TF_Function* func, TF_Buffer* output_func_def, + TF_Status* status) { + DCHECK_EQ(1, func->fdef_lib.function_size()); + status->status = MessageToBuffer(func->fdef_lib.function(0), output_func_def); +} + +void TF_DeleteFunction(TF_Function* function) { delete function; } diff --git a/tensorflow/c/c_api_function_test.cc b/tensorflow/c/c_api_function_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..c9dd38ea15fa49f1fec5f86f9dd2353b1b8398ba --- /dev/null +++ b/tensorflow/c/c_api_function_test.cc @@ -0,0 +1,1039 @@ +/* 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. +==============================================================================*/ + +#include "tensorflow/c/c_api.h" + +#include "tensorflow/c/c_test_util.h" +#include "tensorflow/core/framework/function.pb.h" +#include "tensorflow/core/framework/op_def.pb.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/test.h" + +namespace tensorflow { +namespace { + +// Specification for expected input/output and its type. +// DataType value of DT_INVALID signifies that we don't want to +// check the data type. +typedef std::pair IOSpec; + +std::vector M(const std::initializer_list& names) { + std::vector v; + for (const string& name : names) { + v.push_back(IOSpec(name, DT_INVALID)); + } + return v; +} + +// Specification for an expected edge. +// src is either: +// - input name (as it appears in FunctionDef) +// - name of output tensor (in nested "add:z:0" format) +// dst is either: +// - output name (as it appears in FunctionDef) +// - : (this looks the same as +// output tensor naming, but it the index is actually an input index) +struct EdgeSpec : public std::pair { + typedef std::pair Base; + + // Inherit the set of constructors + using Base::pair; + + string ToString() const { return strings::StrCat(first, "->", second); } +}; + +class CApiFunctionTest : public ::testing::Test { + protected: + CApiFunctionTest() + : s_(TF_NewStatus()), + func_graph_(TF_NewGraph()), + host_graph_(TF_NewGraph()), + func_(nullptr) {} + + void SetUp() override {} + + ~CApiFunctionTest() override { + TF_DeleteFunction(func_); + TF_DeleteGraph(host_graph_); + TF_DeleteGraph(func_graph_); + TF_DeleteStatus(s_); + } + + void Run(const std::vector>& inputs, + TF_Operation* output, int32_t expected_result) { + Run(inputs, {{output, 0}}, {expected_result}); + } + + // Run the host graph, which now contains a function and check that + // outputs are as expected. + // 'T' stands for 'tensor' since the outputs are tensors, not scalars. + void RunT(const std::vector>& inputs, + std::initializer_list outputs, + const std::vector>& expected_results) { + // Create a session for this graph + CSession csession(host_graph_, s_); + ASSERT_EQ(TF_OK, TF_GetCode(s_)) << TF_Message(s_); + + // Run + csession.SetInputs(inputs); + csession.SetOutputs(outputs); + csession.Run(s_); + ASSERT_EQ(TF_OK, TF_GetCode(s_)) << TF_Message(s_); + + // Check results + for (int i = 0; i < expected_results.size(); ++i) { + TF_Tensor* out = csession.output_tensor(i); + ASSERT_TRUE(out != nullptr); + EXPECT_EQ(TF_INT32, TF_TensorType(out)); + EXPECT_EQ(1, TF_NumDims(out)); + CompareInt32Tensor(expected_results[i], out); + } + } + + // Run the host graph, which now contains a function and check that + // outputs are as expected. + void Run(const std::vector>& inputs, + std::initializer_list outputs, + const std::vector& expected_results) { + // Create a session for this graph. + CSession csession(host_graph_, s_); + ASSERT_EQ(TF_OK, TF_GetCode(s_)) << TF_Message(s_); + + csession.SetInputs(inputs); + csession.SetOutputs(outputs); + csession.Run(s_); + ASSERT_EQ(TF_OK, TF_GetCode(s_)) << TF_Message(s_); + + for (int i = 0; i < expected_results.size(); ++i) { + TF_Tensor* out = csession.output_tensor(i); + ASSERT_TRUE(out != nullptr); + EXPECT_EQ(TF_INT32, TF_TensorType(out)); + EXPECT_EQ(0, TF_NumDims(out)); // scalar + ASSERT_EQ(sizeof(int32_t), TF_TensorByteSize(out)); + int32_t* output_contents = static_cast(TF_TensorData(out)); + EXPECT_EQ(expected_results[i], *output_contents); + } + } + + void CompareInt32Tensor(const std::vector& expected, TF_Tensor* t) { + int32_t* data = static_cast(TF_TensorData(t)); + size_t size = TF_TensorByteSize(t); + ASSERT_EQ(expected.size() * sizeof(int32_t), size); + for (int i = 0; i < expected.size(); ++i) { + ASSERT_EQ(expected[i], data[i]) << "Different data at index " << i; + } + } + + std::vector ToOutput(const std::vector ops) { + std::vector out; + for (auto op : ops) { + out.push_back({op, 0}); + } + return out; + } + + void Define(int num_opers, const std::vector& opers, + const std::vector& inputs, + const std::vector& outputs, + const char** output_names, bool expect_failure = false) { + DefineT(num_opers, opers, ToOutput(inputs), ToOutput(outputs), output_names, + expect_failure); + } + + // An explicit `num_opers` is needed so that we can distinguish between the + // case of no operations specified (-1) and the case of an empty set of + // operations specified (0). + void DefineT(int num_opers, const std::vector& opers, + const std::vector& inputs, + const std::vector& outputs, const char** output_names, + bool expect_failure = false) { + ASSERT_EQ(func_, nullptr); + func_ = TF_GraphToFunction(func_graph_, func_name_, num_opers, + num_opers == -1 ? nullptr : opers.data(), + inputs.size(), inputs.data(), outputs.size(), + outputs.data(), output_names, + /*opts=*/nullptr, s_); + if (expect_failure) { + ASSERT_EQ(func_, nullptr); + return; + } + + ASSERT_EQ(TF_OK, TF_GetCode(s_)) << TF_Message(s_); + ASSERT_NE(func_, nullptr); + TF_GraphAddFunction(host_graph_, func_, s_); + ASSERT_EQ(TF_OK, TF_GetCode(s_)) << TF_Message(s_); + } + + TF_Operation* Use(const std::vector& inputs) { + return UseT(ToOutput(inputs)); + } + + TF_Operation* UseT(const std::vector& inputs) { + TF_Operation* op; + UseHelper(inputs, &op); + return op; + } + + // All the *Helper methods are used as a workaround for the restrictions that + // one cannot call ASSERT_* methods in non-void-returning functions (when + // exceptions are disabled during compilation) + void UseHelper(const std::vector& inputs, TF_Operation** op) { + TF_OperationDescription* desc = + TF_NewOperation(host_graph_, func_name_, func_node_name_); + for (auto input : inputs) { + TF_AddInput(desc, input); + } + // Set device to CPU because some ops inside the function might not be + // available on GPU. + TF_SetDevice(desc, "/cpu:0"); + *op = TF_FinishOperation(desc, s_); + ASSERT_EQ(TF_OK, TF_GetCode(s_)) << TF_Message(s_); + ASSERT_NE(*op, nullptr); + } + + FunctionDef fdef() { + tensorflow::FunctionDef fdef; + EXPECT_TRUE(GetFunctionDef(func_, &fdef)); + return fdef; + } + + // logging utility + template + string ToString(const Container& v) { + std::stringstream ss; + ss << "{"; + size_t i = 0; + for (const auto& e : v) { + if (i != 0) { + ss << ", "; + } + ss << e.ToString(); + ++i; + } + ss << "}"; + return ss.str(); + } + + void VerifyFDefNodes(const tensorflow::FunctionDef& fdef, + const std::unordered_set& nodes) { + ASSERT_EQ(nodes.size(), fdef.node_def_size()) + << "Got unexpected number of nodes. Expected: [" + << str_util::Join(nodes, ", ") + << "] Actual nodes in fdef: " << fdef.DebugString(); + for (const NodeDef& node_def : fdef.node_def()) { + ASSERT_TRUE(nodes.find(node_def.name()) != nodes.end()) + << "Got unexpected node: " << node_def.name() + << " in fdef: " << fdef.DebugString(); + } + } + + void VerifyFDefInputs(const tensorflow::FunctionDef& fdef, + const std::vector& inputs) { + const OpDef& signature = fdef.signature(); + ASSERT_EQ(inputs.size(), signature.input_arg_size()); + for (int i = 0; i < inputs.size(); ++i) { + const OpDef::ArgDef& arg = signature.input_arg(i); + const IOSpec& in = inputs[i]; + if (in.second != DT_INVALID) { + ASSERT_EQ(arg.type(), in.second) + << "Got unexpected type for input " << i + << ". fdef: " << fdef.DebugString(); + } + ASSERT_EQ(arg.name(), in.first) << "Got unexpected name for input " << i + << ". fdef: " << fdef.DebugString(); + } + } + + void VerifyFDefOutputs(const tensorflow::FunctionDef& fdef, + const std::vector& outputs) { + const OpDef& signature = fdef.signature(); + ASSERT_EQ(outputs.size(), signature.output_arg_size()); + for (int i = 0; i < outputs.size(); ++i) { + const OpDef::ArgDef& arg = signature.output_arg(i); + const IOSpec& out = outputs[i]; + if (out.second != DT_INVALID) { + ASSERT_EQ(arg.type(), out.second) + << "Got unexpected type for output " << i + << ". fdef: " << fdef.DebugString(); + } + ASSERT_EQ(arg.name(), out.first) << "Got unexpected name for output " << i + << ". fdef: " << fdef.DebugString(); + } + } + + void VerifyFDefEdges( + const tensorflow::FunctionDef& fdef, + const std::vector& e_edges, // expected edges + const std::vector& c_edges, // expected ctrl edges + bool is_exact_edges = true) { + // Build a set of edges from fdef + std::set a_edges; // actual edges + // Get edges from inputs to body nodes and between body nodes + for (const NodeDef& node_def : fdef.node_def()) { + for (int i = 0; i < node_def.input_size(); ++i) { + const string& in = node_def.input(i); + const auto& v = + a_edges.insert({in, strings::StrCat(node_def.name(), ":", i)}); + ASSERT_TRUE(v.second) << "Duplicate edge " << in << " -> " + << strings::StrCat(node_def.name(), ":", i) + << ". fdef: " << fdef.DebugString(); + } + } + // Get edges from body nodes to outputs and from inputs to outputs + for (const OpDef::ArgDef& arg : fdef.signature().output_arg()) { + const auto& iter = fdef.ret().find(arg.name()); + if (iter != fdef.ret().end()) { + const auto& v = a_edges.insert({iter->second, arg.name()}); + ASSERT_TRUE(v.second) << "Duplicate edge " << iter->second << " -> " + << arg.name() << ". fdef: " << fdef.DebugString(); + } else { + const auto& v = a_edges.insert({arg.name(), arg.name()}); + ASSERT_TRUE(v.second) << "Duplicate edge " << arg.name() << " -> " + << arg.name() << ". fdef: " << fdef.DebugString(); + } + } + + // Verify edges + for (const EdgeSpec& e : e_edges) { + ASSERT_TRUE(a_edges.find(e) != a_edges.end()) + << "Failed to find expected edge " << e.ToString() + << " in fdef: " << fdef.DebugString(); + } + + // If caller specified all edges, check that we have seen all + if (is_exact_edges) { + ASSERT_EQ(e_edges.size() + c_edges.size(), a_edges.size()) + << "Expected edges: " << ToString(e_edges) + << " Expected Control edges: " << ToString(c_edges) + << " Actual edges: " << ToString(a_edges) + << " in fdef: " << fdef.DebugString(); + } + } + + void VerifyFDef(const std::unordered_set& nodes, + const std::vector& inputs, + const std::vector& outputs, + const std::vector& e_edges, // expected edges + const std::vector& c_edges, // expected ctrl edges + bool is_exact_edges = true) { + tensorflow::FunctionDef fdef; + ASSERT_TRUE(GetFunctionDef(func_, &fdef)); + VerifyFDefNodes(fdef, nodes); + VerifyFDefInputs(fdef, inputs); + VerifyFDefOutputs(fdef, outputs); + VerifyFDefEdges(fdef, e_edges, c_edges, is_exact_edges); + } + + const char* func_name_ = "MyFunc"; + const char* func_node_name_ = "MyFunc_0"; + TF_Status* s_; + TF_Graph* func_graph_; + TF_Graph* host_graph_; + TF_Function* func_; + + // Workaround for not being able to initialize empty map using {} + std::unordered_set empty_; +}; + +TEST_F(CApiFunctionTest, OneOp_ZeroInputs_OneOutput) { + /* + * constant + * | + * v + */ + // Define + TF_Operation* c = ScalarConst(10, func_graph_, s_, "scalar10"); + Define(-1, {}, {}, {c}, nullptr); + + // Use, run, and verify + TF_Operation* func_op = Use({}); + Run({}, func_op, 10); + VerifyFDef({"scalar10_0"}, {}, {{"scalar10", DT_INT32}}, + {{"scalar10_0:output:0", "scalar10"}}, {}); +} + +TEST_F(CApiFunctionTest, OneOp_OneInput_OneOutput) { + /* + * | + * v + * negate + * | + * v + */ + // Define + TF_Operation* feed = Placeholder(func_graph_, s_); + TF_Operation* neg = Neg(feed, func_graph_, s_); + Define(-1, {}, {feed}, {neg}, nullptr); + + // Use, run, and verify + TF_Operation* func_feed = Placeholder(host_graph_, s_); + TF_Operation* func_op = Use({func_feed}); + Run({{func_feed, Int32Tensor(3)}}, func_op, -3); + VerifyFDef({"neg_0"}, {{"feed", DT_INT32}}, {{"neg", DT_INT32}}, + {{"feed", "neg_0:0"}, {"neg_0:y:0", "neg"}}, {}); +} + +TEST_F(CApiFunctionTest, ZeroOps_Identity) { + /* + * | + * | + * | + * v + */ + // Define + TF_Operation* feed = Placeholder(func_graph_, s_); + Define(-1, {}, {feed}, {feed}, nullptr); + + // Use, run, and verify + TF_Operation* func_feed = Placeholder(host_graph_, s_); + TF_Operation* func_op = Use({func_feed}); + Run({{func_feed, Int32Tensor(3)}}, func_op, 3); + VerifyFDef(empty_, {{"feed", DT_INT32}}, {{"feed_0", DT_INT32}}, + {{"feed", "feed_0"}}, {}); +} + +TEST_F(CApiFunctionTest, ZeroOps_Permutation) { + /* + * | | + * \ / + * \/ + * x + * /\ + * / \ + * | | + * v v + */ + // Define + TF_Operation* feed1 = Placeholder(func_graph_, s_, "feed1"); + TF_Operation* feed2 = Placeholder(func_graph_, s_, "feed2"); + Define(-1, {}, {feed1, feed2}, {feed2, feed1}, nullptr); + + // Use, run, and verify + TF_Operation* two = ScalarConst(2, host_graph_, s_); + TF_Operation* func_feed = Placeholder(host_graph_, s_); + TF_Operation* func_op = Use({two, func_feed}); + Run({{func_feed, Int32Tensor(3)}}, {{func_op, 0}, {func_op, 1}}, {3, 2}); + VerifyFDef(empty_, M({{"feed1"}, {"feed2"}}), M({{"feed2_0"}, {"feed1_0"}}), + {{"feed1", "feed1_0"}, {"feed2", "feed2_0"}}, {}); +} + +TEST_F(CApiFunctionTest, OneOp_TwoInputs_OneOutput) { + /* + * | | + * v v + * add + * | + * v + */ + // Define + TF_Operation* feed1 = Placeholder(func_graph_, s_, "feed1"); + TF_Operation* feed2 = Placeholder(func_graph_, s_, "feed2"); + TF_Operation* add = Add(feed1, feed2, func_graph_, s_); + Define(-1, {}, {feed1, feed2}, {add}, nullptr); + + // Use, run, and verify + TF_Operation* two = ScalarConst(2, host_graph_, s_); + TF_Operation* func_feed = Placeholder(host_graph_, s_); + TF_Operation* func_op = Use({two, func_feed}); + Run({{func_feed, Int32Tensor(3)}}, func_op, 2 + 3); + VerifyFDef( + {"add_0"}, M({{"feed1"}, {"feed2"}}), M({{"add"}}), + {{"feed1", "add_0:0"}, {"feed2", "add_0:1"}, {"add_0:sum:0", "add"}}, {}); +} + +TEST_F(CApiFunctionTest, OneOp_TwoInputs_ZeroOutputs) { + /* + * | | + * v v + * add + * + * (output ignored) + */ + // Define + TF_Operation* feed1 = Placeholder(func_graph_, s_, "feed1"); + TF_Operation* feed2 = Placeholder(func_graph_, s_, "feed2"); + Add(feed1, feed2, func_graph_, s_); + Define(-1, {}, {feed1, feed2}, {}, nullptr); + + // Use, run, and verify + TF_Operation* two = ScalarConst(2, host_graph_, s_); + TF_Operation* func_feed = Placeholder(host_graph_, s_); + Use({two, func_feed}); + VerifyFDef({"add"}, M({{"feed1"}, {"feed2"}}), {}, + {{"feed1", "add:0"}, {"feed2", "add:1"}}, {}); +} + +TEST_F(CApiFunctionTest, TwoOps_ThreeInputs_OneOutput) { + /* + * | | | + * v v / + * add1 / + * | | + * v v + * add2 + * | + * v + */ + // Define + TF_Operation* feed1 = Placeholder(func_graph_, s_, "feed1"); + TF_Operation* feed2 = Placeholder(func_graph_, s_, "feed2"); + TF_Operation* feed3 = Placeholder(func_graph_, s_, "feed3"); + TF_Operation* add1 = Add(feed1, feed2, func_graph_, s_, "add1"); + TF_Operation* add2 = Add(add1, feed3, func_graph_, s_, "add2"); + Define(-1, {}, {feed1, feed2, feed3}, {add2}, nullptr); + + // Use, run, and verify + TF_Operation* two = ScalarConst(2, host_graph_, s_, "two"); + TF_Operation* ten = ScalarConst(10, host_graph_, s_, "ten"); + TF_Operation* func_feed = Placeholder(host_graph_, s_); + TF_Operation* func_op = Use({two, ten, func_feed}); + Run({{func_feed, Int32Tensor(3)}}, func_op, 2 + 10 + 3); + VerifyFDef({"add1", "add2_0"}, M({{"feed1"}, {"feed2"}, {"feed3"}}), + M({{"add2"}}), + {{"feed1", "add1:0"}, + {"feed2", "add1:1"}, + {"add1:sum:0", "add2_0:0"}, + {"feed3", "add2_0:1"}, + {"add2_0:sum:0", "add2"}}, + {}); +} + +TEST_F(CApiFunctionTest, OneOp_TwoInputs_TwoDuplicateOutputs) { + /* + * | | + * v v + * add + * | + * +-+-+ + * | | + * v v + */ + // Define + TF_Operation* feed1 = Placeholder(func_graph_, s_, "feed1"); + TF_Operation* feed2 = Placeholder(func_graph_, s_, "feed2"); + TF_Operation* add = Add(feed1, feed2, func_graph_, s_); + Define(-1, {}, {feed1, feed2}, {add, add}, nullptr); + + // Use, run, and verify + TF_Operation* two = ScalarConst(2, host_graph_, s_); + TF_Operation* func_feed = Placeholder(host_graph_, s_); + TF_Operation* func_op = Use({two, func_feed}); + Run({{func_feed, Int32Tensor(3)}}, {{func_op, 0}, {func_op, 1}}, {5, 5}); + VerifyFDef({"add_1"}, M({{"feed1"}, {"feed2"}}), M({{"add"}, {"add_0"}}), + {{"feed1", "add_1:0"}, + {"feed2", "add_1:1"}, + {"add_1:sum:0", "add"}, + {"add_1:sum:0", "add_0"}}, + {}); +} + +TEST_F(CApiFunctionTest, TwoOps_ThreeInputs_TwoOutputs) { + /* + * | | | + * v v / + * add / + * | | + * +-+ | + * | | | + * | v v + * | add + * | | + * v v + */ + // Define + TF_Operation* feed1 = Placeholder(func_graph_, s_, "feed1"); + TF_Operation* feed2 = Placeholder(func_graph_, s_, "feed2"); + TF_Operation* feed3 = Placeholder(func_graph_, s_, "feed3"); + TF_Operation* add1 = Add(feed1, feed2, func_graph_, s_, "add1"); + TF_Operation* add2 = Add(add1, feed3, func_graph_, s_, "add2"); + Define(-1, {}, {feed1, feed2, feed3}, {add1, add2}, nullptr); + + // Use, run, and verify + TF_Operation* two = ScalarConst(2, host_graph_, s_, "two"); + TF_Operation* ten = ScalarConst(10, host_graph_, s_, "ten"); + TF_Operation* func_feed = Placeholder(host_graph_, s_); + TF_Operation* func_op = Use({two, ten, func_feed}); + Run({{func_feed, Int32Tensor(3)}}, {{func_op, 0}, {func_op, 1}}, {12, 15}); + VerifyFDef({"add1_0", "add2_0"}, M({{"feed1"}, {"feed2"}, {"feed3"}}), + M({{"add1"}, {"add2"}}), + {{"feed1", "add1_0:0"}, + {"feed2", "add1_0:1"}, + {"add1_0:sum:0", "add2_0:0"}, + {"feed3", "add2_0:1"}, + {"add1_0:sum:0", "add1"}, + {"add2_0:sum:0", "add2"}}, + {}); +} + +TEST_F(CApiFunctionTest, FromSubsetOfOps) { + /* + * | | | + * v v / + * add / + * | | + * +---+--+---+ + * Ops used | | | | + * for func | v v | + * | | add | + * +-------> | | | + * | v | + * | | + * +----------+ + */ + // Define + TF_Operation* feed1 = Placeholder(func_graph_, s_, "feed1"); + TF_Operation* feed2 = Placeholder(func_graph_, s_, "feed2"); + TF_Operation* feed3 = Placeholder(func_graph_, s_, "feed3"); + TF_Operation* add1 = Add(feed1, feed2, func_graph_, s_, "add1"); + TF_Operation* add2 = Add(add1, feed3, func_graph_, s_, "add2"); + Define(1, {add2}, {add1, feed3}, {add2}, nullptr); + + // Use, run, and verify + TF_Operation* two = ScalarConst(2, host_graph_, s_, "two"); + TF_Operation* func_feed = Placeholder(host_graph_, s_); + TF_Operation* func_op = Use({two, func_feed}); + Run({{func_feed, Int32Tensor(3)}}, func_op, 2 + 3); + VerifyFDef( + {"add2_0"}, M({{"add1"}, {"feed3"}}), M({{"add2"}}), + {{"add1", "add2_0:0"}, {"feed3", "add2_0:1"}, {"add2_0:sum:0", "add2"}}, + {}); +} + +TEST_F(CApiFunctionTest, UsingOneOutputOfSplit) { + /* + * feed + * | + * +---------+---+ + * | const0 | | + * | | | | + * | v / | + * | split | + * | | | | | + * | v | v | + * | | | + * +------+------+ + * | + * v + * + * Only the second output from split is used as function output + */ + // Define + TF_Operation* feed = Placeholder(func_graph_, s_); + TF_Operation* split = Split3(feed, func_graph_, s_); + DefineT(-1, {}, {{feed, 0}}, {{split, 1}}, nullptr); + + // Use, run, and verify + TF_Operation* func_feed = Placeholder(host_graph_, s_); + TF_Operation* func_op = Use({func_feed}); + RunT({{func_feed, Int32Tensor({1, 2, 3, 4, 5, 6})}}, {{func_op, 0}}, + {{3, 4}}); + VerifyFDef({"split3_const0", "split3_0"}, M({{"feed"}}), M({{"split3"}}), + {{"split3_const0:output:0", "split3_0:0"}, + {"feed", "split3_0:1"}, + {"split3_0:output:1", "split3"}}, + {}); +} + +TEST_F(CApiFunctionTest, UsingTwoOutputsOfSplit) { + /* + * feed + * | + * +---------+---+ + * | const0 | | + * | | | | + * | v / | + * | split | + * | | | | | + * | | v | | + * | | | | + * +---+-----+---+ + * | | + * v v + * + * Second output from split is not used as function output + */ + // Define + TF_Operation* feed = Placeholder(func_graph_, s_); + TF_Operation* split = Split3(feed, func_graph_, s_); + DefineT(-1, {}, {{feed, 0}}, {{split, 0}, {split, 2}}, nullptr); + + // Use, run, and verify + TF_Operation* func_feed = Placeholder(host_graph_, s_); + TF_Operation* func_op = Use({func_feed}); + RunT({{func_feed, Int32Tensor({1, 2, 3, 4, 5, 6})}}, + {{func_op, 0}, {func_op, 1}}, {{1, 2}, {5, 6}}); + VerifyFDef({"split3_const0", "split3_1"}, M({{"feed"}}), + M({{"split3"}, {"split3_0"}}), + {{"split3_const0:output:0", "split3_1:0"}, + {"feed", "split3_1:1"}, + {"split3_1:output:0", "split3"}, + {"split3_1:output:2", "split3_0"}}, + {}); +} + +TEST_F(CApiFunctionTest, UsingTwoOutputsOfSplitAsInputs) { + /* + * | + * v + * split + * | | | + * | v | + * | | + * +---+-----+---+ + * | | | | + * | v v | + * | add | + * | | | + * | | | + * +------+------+ + * | + * v + */ + // Define + TF_Operation* feed = Placeholder(func_graph_, s_); + TF_Operation* split = Split3(feed, func_graph_, s_); + TF_Operation* add = Add({split, 0}, {split, 2}, func_graph_, s_); + ASSERT_EQ(TF_OK, TF_GetCode(s_)) << TF_Message(s_); + DefineT(1, {add}, {{split, 0}, {split, 2}}, {{add, 0}}, nullptr); + + // Use, run, and verify + TF_Operation* two = ScalarConst(2, host_graph_, s_, "two"); + TF_Operation* func_feed = Placeholder(host_graph_, s_); + TF_Operation* func_op = Use({two, func_feed}); + Run({{func_feed, Int32Tensor(3)}}, func_op, 2 + 3); + VerifyFDef( + {"add_0"}, M({{"split3"}, {"split3_0"}}), M({{"add"}}), + {{"split3", "add_0:0"}, {"split3_0", "add_0:1"}, {"add_0:sum:0", "add"}}, + {}); +} + +TEST_F(CApiFunctionTest, NodesUsedInInputsMustHaveSingleOutput) { + /* + * | + * v + * split + * | | | + * | v | + * | | + * input --->| |<--- input + * | | + * v v + * add + * | + * | + * v + */ + // Define + TF_Tensor* tensor_123 = Int32Tensor({1, 2, 3}); + TF_Operation* c = Const(tensor_123, func_graph_, s_, "const_array"); + ASSERT_EQ(TF_OK, TF_GetCode(s_)) << TF_Message(s_); + TF_Operation* split = Split3(c, func_graph_, s_); + TF_Operation* add = Add({split, 0}, {split, 2}, func_graph_, s_); + ASSERT_EQ(TF_OK, TF_GetCode(s_)) << TF_Message(s_); + DefineT(-1, {}, {{split, 0}, {split, 2}}, {{add, 0}}, nullptr, true); + EXPECT_EQ(TF_INVALID_ARGUMENT, TF_GetCode(s_)); + EXPECT_EQ(string("When `num_opers` is set to -1, nodes referenced in " + "`inputs` must have a single output. Node split3 has " + "3 outputs. Encountered while creating function 'MyFunc'"), + string(TF_Message(s_))); + + TF_DeleteTensor(tensor_123); +} + +TEST_F(CApiFunctionTest, FunctionWithWhileLoop) { + // Inputs to the while loop and the function as a whole + TF_Operation* feed1 = Placeholder(func_graph_, s_, "feed1"); + TF_Operation* feed2 = Placeholder(func_graph_, s_, "feed2"); + + // Outputs of the while loop corresponding to the two inputs above + // The first one will the function's output + std::vector outputs; + + // Add while loop to func_graph_ + { + // The inputs to the while loop + std::vector inputs = {{feed1, 0}, {feed2, 0}}; + std::unique_ptr params(new TF_WhileParams( + TF_NewWhile(func_graph_, &inputs[0], inputs.size(), s_))); + ASSERT_EQ(TF_OK, TF_GetCode(s_)) << TF_Message(s_); + params->name = "test_loop"; + + // Initialize outputs so we can easily detect errors/bugs + outputs.resize(2, {nullptr, -1}); + + // Create loop: while (input1 < input2) input1 += input2 + 1 + TF_Operation* less_than = LessThan( + params->cond_inputs[0], params->cond_inputs[1], params->cond_graph, s_); + ASSERT_EQ(TF_OK, TF_GetCode(s_)) << TF_Message(s_); + params->cond_output = {less_than, 0}; + + TF_Operation* add1 = Add(params->body_inputs[0], params->body_inputs[1], + params->body_graph, s_, "add1"); + ASSERT_EQ(TF_OK, TF_GetCode(s_)) << TF_Message(s_); + TF_Operation* one = ScalarConst(1, params->body_graph, s_); + ASSERT_EQ(TF_OK, TF_GetCode(s_)) << TF_Message(s_); + TF_Operation* add2 = Add(add1, one, params->body_graph, s_, "add2"); + ASSERT_EQ(TF_OK, TF_GetCode(s_)) << TF_Message(s_); + params->body_outputs[0] = {add2, 0}; + params->body_outputs[1] = params->body_inputs[1]; + + // Finalize while loop + TF_FinishWhile(params.get(), s_, &outputs[0]); + EXPECT_EQ(TF_OK, TF_GetCode(s_)) << TF_Message(s_); + } + + // Define function, use it in graph, and run + DefineT(-1, {}, {{feed1, 0}, {feed2, 0}}, {outputs[0]}, nullptr); + TF_Operation* five = ScalarConst(5, host_graph_, s_, "five"); + TF_Operation* func_feed = Placeholder(host_graph_, s_); + TF_Operation* func_op = Use({func_feed, five}); + Run({{func_feed, Int32Tensor(2)}}, func_op, 2 /*+=*/ + 5 + 1); + + // Verify input, output, and subset of edges in fdef. + // The subset of edges we verify is a chain between feed1 and output to + // make sure that the correct output is picked. + tensorflow::FunctionDef fdef; + ASSERT_TRUE(GetFunctionDef(func_, &fdef)); + VerifyFDefInputs(fdef, M({{"feed1"}, {"feed2"}})); + VerifyFDefOutputs(fdef, M({{"test_loop_exit"}})); + VerifyFDefEdges(fdef, + {{"feed1", "test_loop/Enter:0"}, + {"test_loop/Enter:output:0", "test_loop/Merge:0"}, + {"test_loop/Merge:output:0", "test_loop/Switch:0"}, + {"test_loop/Switch:output_false:0", "test_loop/Exit:0"}, + {"test_loop/Exit:output:0", "test_loop_exit"}}, + {}, false); +} + +TEST_F(CApiFunctionTest, ControlDependency) { + /* + * | | scalar + * | | . + * v v . <---- control dependency + * add < - + * | + * v + */ + // Define + TF_Operation* feed1 = Placeholder(func_graph_, s_, "feed1"); + TF_Operation* feed2 = Placeholder(func_graph_, s_, "feed2"); + TF_Operation* five = ScalarConst(5, func_graph_, s_); + TF_Operation* add = + AddWithCtrlDependency(feed1, feed2, func_graph_, five, s_); + EXPECT_EQ(TF_OK, TF_GetCode(s_)) << TF_Message(s_); + Define(-1, {}, {feed1, feed2}, {add}, nullptr); + + // Use, run, and verify + TF_Operation* two = ScalarConst(2, host_graph_, s_); + TF_Operation* func_feed = Placeholder(host_graph_, s_); + TF_Operation* func_op = Use({two, func_feed}); + Run({{func_feed, Int32Tensor(3)}}, func_op, 2 + 3); + VerifyFDef( + {"add_0", "scalar"}, M({{"feed1"}, {"feed2"}}), M({{"add"}}), + {{"feed1", "add_0:0"}, {"feed2", "add_0:1"}, {"add_0:sum:0", "add"}}, + {{"scalar", "add_0"}}); +} + +TEST_F(CApiFunctionTest, ControlDependencyOutsideOfBody) { + /* + * | | scalar + * | | . + * v v . <---- control dependency + * add < - + * | + * v + */ + // Define + TF_Operation* feed1 = Placeholder(func_graph_, s_, "feed1"); + TF_Operation* feed2 = Placeholder(func_graph_, s_, "feed2"); + TF_Operation* five = ScalarConst(5, func_graph_, s_); + TF_Operation* add = + AddWithCtrlDependency(feed1, feed2, func_graph_, five, s_); + EXPECT_EQ(TF_OK, TF_GetCode(s_)) << TF_Message(s_); + Define(1, {add}, {feed1, feed2}, {add}, nullptr, true); + EXPECT_EQ(TF_INVALID_ARGUMENT, TF_GetCode(s_)); + EXPECT_EQ(string("The source of control edge [id=3 scalar:-1 -> add:-1] " + "is not in the body. Encountered while creating " + "function 'MyFunc'"), + string(TF_Message(s_))); +} + +TEST_F(CApiFunctionTest, ControlDependencyOutsideOfBody_FromInputNode) { + /* + * | |. + * | | . + * | | . + * v v . <---- control dependency + * add < - + * | + * v + */ + // Define + TF_Operation* feed1 = Placeholder(func_graph_, s_, "feed1"); + TF_Operation* feed2 = Placeholder(func_graph_, s_, "feed2"); + TF_Operation* add = + AddWithCtrlDependency(feed1, feed2, func_graph_, feed1, s_); + EXPECT_EQ(TF_OK, TF_GetCode(s_)) << TF_Message(s_); + Define(-1, {}, {feed1, feed2}, {add}, nullptr, true); + EXPECT_EQ(TF_INVALID_ARGUMENT, TF_GetCode(s_)); + EXPECT_EQ(string("The source of control edge [id=3 feed1:-1 -> add:-1] " + "is not in the body. Encountered while creating " + "function 'MyFunc'"), + string(TF_Message(s_))); +} + +TEST_F(CApiFunctionTest, DuplicateInputsAreNotAllowed) { + /* + * feed + * | + * +++ + * | | + * +---+-+---+ + * | | | | + * | v v | + * | add | + * | | | + * | | | + * +----+----+ + * | + * v + */ + TF_Operation* feed1 = Placeholder(func_graph_, s_, "feed1"); + TF_Operation* add = Add(feed1, feed1, func_graph_, s_); + Define(-1, {}, {feed1, feed1}, {add}, nullptr, true); + EXPECT_EQ(TF_INVALID_ARGUMENT, TF_GetCode(s_)); + EXPECT_EQ( + string("TF_Output feed1:0 appears more than once in the input list"), + string(TF_Message(s_))); +} + +TEST_F(CApiFunctionTest, InvalidInputTensor_HighIndex) { + /* + * | | + * v v + * add + * | + * v + */ + TF_Operation* feed1 = Placeholder(func_graph_, s_, "feed1"); + TF_Operation* feed2 = Placeholder(func_graph_, s_, "feed2"); + TF_Operation* add = Add(feed1, feed2, func_graph_, s_); + DefineT(-1, {}, {{feed1, 0}, {feed2, 2}}, {{add, 0}}, nullptr, true); + EXPECT_EQ(TF_INVALID_ARGUMENT, TF_GetCode(s_)); + EXPECT_EQ(string("Node 'feed2' (type: 'Placeholder', num of outputs: 1) does " + "not have output 2\n\tEncountered while processing " + "input 1 into function 'MyFunc'"), + string(TF_Message(s_))); +} + +TEST_F(CApiFunctionTest, InvalidInputTensor_BadNodePtr) { + /* + * | | + * v v + * add + * | + * v + */ + TF_Operation* feed1 = Placeholder(func_graph_, s_, "feed1"); + TF_Operation* feed2 = Placeholder(func_graph_, s_, "feed2"); + TF_Operation* add = Add(feed1, feed2, func_graph_, s_); + DefineT(-1, {}, {{feed1, 0}, {nullptr, 0}}, {{add, 0}}, nullptr, true); + EXPECT_EQ(TF_INVALID_ARGUMENT, TF_GetCode(s_)); + EXPECT_EQ(string("Node is null\n\tEncountered while processing input 1 " + "into function 'MyFunc'"), + string(TF_Message(s_))); +} + +TEST_F(CApiFunctionTest, InvalidOutputTensor_HighIndex) { + /* + * | | + * v v + * add + * | + * v + */ + TF_Operation* feed1 = Placeholder(func_graph_, s_, "feed1"); + TF_Operation* feed2 = Placeholder(func_graph_, s_, "feed2"); + TF_Operation* add = Add(feed1, feed2, func_graph_, s_); + DefineT(-1, {}, {{feed1, 0}, {feed2, 0}}, {{add, 3}}, nullptr, true); + EXPECT_EQ(TF_INVALID_ARGUMENT, TF_GetCode(s_)); + EXPECT_EQ(string("Node 'add' (type: 'AddN', num of outputs: 1) does " + "not have output 3\n\tEncountered while processing " + "output 0 from function 'MyFunc'"), + string(TF_Message(s_))); +} + +TEST_F(CApiFunctionTest, InvalidOutputTensor_BadNodePtr) { + /* + * | | + * v v + * add + * | + * v + */ + TF_Operation* feed1 = Placeholder(func_graph_, s_, "feed1"); + TF_Operation* feed2 = Placeholder(func_graph_, s_, "feed2"); + Add(feed1, feed2, func_graph_, s_); + DefineT(-1, {}, {{feed1, 0}, {feed2, 0}}, {{nullptr, 3}}, nullptr, true); + EXPECT_EQ(TF_INVALID_ARGUMENT, TF_GetCode(s_)); + EXPECT_EQ(string("Node is null\n\tEncountered while processing output 0 " + "from function 'MyFunc'"), + string(TF_Message(s_))); +} + +TEST_F(CApiFunctionTest, NodeMissingInput) { + /* + * input---> | | <----missing input + * v v + * body----> add + * | + * v + */ + TF_Operation* feed1 = Placeholder(func_graph_, s_, "feed1"); + TF_Operation* feed2 = Placeholder(func_graph_, s_, "feed2"); + TF_Operation* add = Add(feed1, feed2, func_graph_, s_); + DefineT(1, {add}, {{feed1, 0}}, {{add, 0}}, nullptr, true); + EXPECT_EQ(TF_INVALID_ARGUMENT, TF_GetCode(s_)); + EXPECT_EQ(string("Input 1, 'feed2:0', of node 'add' in function 'MyFunc' " + "is not available. You might need to include it in inputs " + "or include its source node in the body"), + string(TF_Message(s_))); +} + +TEST_F(CApiFunctionTest, OutputOpNotInBody) { + /* + * | | + * v v + * add scalar (scalar not included in body) + * | | + * v v (function has two outputs) + */ + // Define + TF_Operation* feed1 = Placeholder(func_graph_, s_, "feed1"); + TF_Operation* feed2 = Placeholder(func_graph_, s_, "feed2"); + TF_Operation* scalar = ScalarConst(2, func_graph_, s_); + TF_Operation* add = Add(feed1, feed2, func_graph_, s_); + Define(1, {add}, {feed1, feed2}, {add, scalar}, nullptr, true); + EXPECT_EQ(TF_INVALID_ARGUMENT, TF_GetCode(s_)); + EXPECT_EQ(string("TF_Output scalar:0 is neither in the function body nor " + "among function inputs. Encountered while creating " + "function 'MyFunc'"), + string(TF_Message(s_))); +} + +} // namespace +} // namespace tensorflow diff --git a/tensorflow/c/c_api_internal.h b/tensorflow/c/c_api_internal.h new file mode 100644 index 0000000000000000000000000000000000000000..68c324f2b992df144db79fb392eb8262a283d250 --- /dev/null +++ b/tensorflow/c/c_api_internal.h @@ -0,0 +1,157 @@ +/* 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. +==============================================================================*/ + +#ifndef TENSORFLOW_C_C_API_INTERNAL_H_ +#define TENSORFLOW_C_C_API_INTERNAL_H_ + +#include "tensorflow/c/c_api.h" + +#include +#include +#include + +#include "tensorflow/core/common_runtime/shape_refiner.h" +#include "tensorflow/core/framework/tensor.h" +#include "tensorflow/core/framework/tensor_shape.h" +#include "tensorflow/core/graph/graph.h" +#include "tensorflow/core/graph/graph_constructor.h" +#include "tensorflow/core/graph/node_builder.h" +#include "tensorflow/core/lib/core/status.h" +#include "tensorflow/core/platform/mutex.h" +#include "tensorflow/core/platform/types.h" +#include "tensorflow/core/public/session.h" + +namespace tensorflow { +class Device; +class DeviceMgr; +} // namespace tensorflow + +// Internal structures used by the C API. These are likely to change and should +// not be depended on. + +struct TF_Status { + tensorflow::Status status; +}; + +struct TF_Tensor { + ~TF_Tensor(); + + TF_DataType dtype; + tensorflow::TensorShape shape; + tensorflow::TensorBuffer* buffer; +}; + +struct TF_SessionOptions { + tensorflow::SessionOptions options; +}; + +struct TF_DeprecatedSession { + tensorflow::Session* session; +}; + +struct TF_Library { + void* lib_handle; + TF_Buffer op_list; +}; + +struct TF_Graph { + TF_Graph(); + + tensorflow::mutex mu; + tensorflow::Graph graph GUARDED_BY(mu); + + // Runs shape inference. + tensorflow::ShapeRefiner refiner GUARDED_BY(mu); + + // Maps from name of an operation to the Node* in 'graph'. + std::unordered_map name_map + GUARDED_BY(mu); + + // TF_Graph may only / must be deleted when + // num_sessions == 0 && delete_requested == true + + // num_sessions incremented by TF_NewSession, and decremented by + // TF_DeleteSession. + int num_sessions GUARDED_BY(mu); + bool delete_requested GUARDED_BY(mu); // set true by TF_DeleteGraph + + // Used to link graphs contained in TF_WhileParams to the parent graph that + // will eventually contain the full while loop. + TF_Graph* parent; + TF_Output* parent_inputs; +}; + +struct TF_OperationDescription { + TF_OperationDescription(TF_Graph* g, const char* op_type, + const char* node_name) + : node_builder(node_name, op_type, g->graph.op_registry()), graph(g) {} + + tensorflow::NodeBuilder node_builder; + TF_Graph* graph; + std::set colocation_constraints; +}; + +struct TF_Operation { + tensorflow::Node node; +}; + +struct TF_Session { + TF_Session(tensorflow::Session* s, TF_Graph* g); + + tensorflow::Session* session; + TF_Graph* graph; + + tensorflow::mutex mu; + int last_num_graph_nodes; + + // NOTE(ashankar): Experimental fields to help keep the + // buffers of a TF_Tensor pinned in device memory. + const tensorflow::DeviceMgr* device_mgr; // Owned by session. + std::vector devices; // Owned by device_mgr. +}; + +struct TF_ImportGraphDefOptions { + tensorflow::ImportGraphDefOptions opts; +}; + +struct TF_DeviceList { + std::vector response; +}; + +struct TF_Function { + // Currently contains a single function and no gradients + tensorflow::FunctionDefLibrary fdef_lib; +}; + +namespace tensorflow { + +class TensorCApi { + public: + static TensorBuffer* Buffer(const Tensor& tensor) { return tensor.buf_; } + static Tensor MakeTensor(TF_DataType type, const TensorShape& shape, + TensorBuffer* buf) { + return Tensor(static_cast(type), shape, buf); + } +}; + +Status TF_TensorToTensor(const TF_Tensor* src, Tensor* dst); + +TF_Tensor* TF_TensorFromTensor(const Tensor& src, TF_Status* status); + +Status MessageToBuffer(const tensorflow::protobuf::Message& in, TF_Buffer* out); + +} // 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 5673f657d3c5b77618c481da614573b9e4a63aba..c4420290099ee10c89792210dad2604328296515 100644 --- a/tensorflow/c/c_api_test.cc +++ b/tensorflow/c/c_api_test.cc @@ -16,13 +16,17 @@ limitations under the License. #include "tensorflow/c/c_api.h" #include +#include #include #include #include + +#include "tensorflow/c/c_test_util.h" #include "tensorflow/cc/saved_model/signature_constants.h" #include "tensorflow/cc/saved_model/tag_constants.h" #include "tensorflow/core/example/example.pb.h" #include "tensorflow/core/example/feature.pb.h" +#include "tensorflow/core/framework/common_shape_fns.h" #include "tensorflow/core/framework/graph.pb_text.h" #include "tensorflow/core/framework/node_def.pb_text.h" #include "tensorflow/core/framework/node_def_util.h" @@ -38,25 +42,15 @@ limitations under the License. #include "tensorflow/core/lib/strings/strcat.h" #include "tensorflow/core/platform/test.h" #include "tensorflow/core/protobuf/meta_graph.pb.h" - -using tensorflow::int32; -using tensorflow::string; -using tensorflow::GraphDef; -using tensorflow::NodeDef; -using tensorflow::Tensor; -using tensorflow::TensorShape; +#include "tensorflow/core/util/equal_graph_def.h" namespace tensorflow { -bool TF_Tensor_DecodeStrings(TF_Tensor* src, Tensor* dst, TF_Status* status); -TF_Tensor* TF_Tensor_EncodeStrings(const Tensor& src); -} // namespace tensorflow +TF_Tensor* TF_TensorFromTensor(const Tensor& src, TF_Status* status); +Status TF_TensorToTensor(const TF_Tensor* src, Tensor* dst); namespace { -typedef std::unique_ptr - unique_tensor_ptr; - -TEST(CAPI, Version) { EXPECT_NE("", string(TF_Version())); } +TEST(CAPI, Version) { EXPECT_STRNE("", TF_Version()); } TEST(CAPI, Status) { TF_Status* s = TF_NewStatus(); @@ -68,7 +62,7 @@ TEST(CAPI, Status) { TF_DeleteStatus(s); } -static void Deallocator(void* data, size_t, void* arg) { +void Deallocator(void* data, size_t, void* arg) { tensorflow::cpu_allocator()->DeallocateRaw(data); *reinterpret_cast(arg) = true; } @@ -105,6 +99,22 @@ TEST(CAPI, AllocateTensor) { TF_DeleteTensor(t); } +TEST(CAPI, MaybeMove) { + const int num_bytes = 6 * sizeof(float); + float* values = + reinterpret_cast(tensorflow::cpu_allocator()->AllocateRaw( + EIGEN_MAX_ALIGN_BYTES, num_bytes)); + int64_t dims[] = {2, 3}; + bool deallocator_called = false; + TF_Tensor* t = TF_NewTensor(TF_FLOAT, dims, 2, values, num_bytes, + &Deallocator, &deallocator_called); + + TF_Tensor* o = TF_TensorMaybeMove(t); + ASSERT_TRUE(o == nullptr); // It is unsafe to move memory TF might not own. + TF_DeleteTensor(t); + EXPECT_TRUE(deallocator_called); +} + TEST(CAPI, LibraryLoadFunctions) { // Load the library. TF_Status* status = TF_NewStatus(); @@ -125,8 +135,9 @@ TEST(CAPI, LibraryLoadFunctions) { TF_DeleteLibraryHandle(lib); } -static void TestEncodeDecode(int line, const std::vector& data) { +void TestEncodeDecode(int line, const std::vector& data) { const tensorflow::int64 n = data.size(); + TF_Status* status = TF_NewStatus(); for (const std::vector& dims : std::vector>{ {n}, {1, n}, {n, 1}, {n / 2, 2}}) { @@ -135,21 +146,20 @@ static void TestEncodeDecode(int line, const std::vector& data) { for (tensorflow::int64 i = 0; i < src.NumElements(); ++i) { src.flat()(i) = data[i]; } - TF_Tensor* dst = TF_Tensor_EncodeStrings(src); + TF_Tensor* dst = TF_TensorFromTensor(src, status); + ASSERT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); // Convert back to a C++ Tensor and ensure we get expected output. - TF_Status* status = TF_NewStatus(); Tensor output; - ASSERT_TRUE(TF_Tensor_DecodeStrings(dst, &output, status)) << line; - ASSERT_EQ(TF_OK, TF_GetCode(status)) << line; + ASSERT_EQ(Status::OK(), TF_TensorToTensor(dst, &output)) << line; ASSERT_EQ(src.NumElements(), output.NumElements()) << line; for (tensorflow::int64 i = 0; i < src.NumElements(); ++i) { ASSERT_EQ(data[i], output.flat()(i)) << line; } - TF_DeleteStatus(status); TF_DeleteTensor(dst); } + TF_DeleteStatus(status); } TEST(CAPI, TensorEncodeDecodeStrings) { @@ -257,176 +267,6 @@ TEST(CAPI, GetAllOpList) { TF_DeleteBuffer(buf); } -static void Int32Deallocator(void* data, size_t, void* arg) { - delete[] static_cast(data); -} - -static TF_Tensor* Int32Tensor(int32 v) { - const int num_bytes = sizeof(int32); - int32* values = new int32[1]; - values[0] = v; - return TF_NewTensor(TF_INT32, nullptr, 0, values, num_bytes, - &Int32Deallocator, nullptr); -} - -TF_Operation* Placeholder(TF_Graph* graph, TF_Status* s, - const char* name = "feed") { - TF_OperationDescription* desc = TF_NewOperation(graph, "Placeholder", name); - TF_SetAttrType(desc, "dtype", TF_INT32); - return TF_FinishOperation(desc, s); -} - -TF_Operation* ScalarConst(int32 v, TF_Graph* graph, TF_Status* s, - const char* name = "scalar") { - unique_tensor_ptr tensor(Int32Tensor(v), TF_DeleteTensor); - TF_OperationDescription* desc = TF_NewOperation(graph, "Const", name); - TF_SetAttrTensor(desc, "value", tensor.get(), s); - if (TF_GetCode(s) != TF_OK) return nullptr; - TF_SetAttrType(desc, "dtype", TF_INT32); - return TF_FinishOperation(desc, s); -} - -TF_Operation* Add(TF_Operation* l, TF_Operation* r, TF_Graph* graph, - TF_Status* s, const char* name = "add") { - TF_OperationDescription* desc = TF_NewOperation(graph, "AddN", name); - TF_Output add_inputs[2] = {{l, 0}, {r, 0}}; - TF_AddInputList(desc, add_inputs, 2); - return TF_FinishOperation(desc, s); -} - -TF_Operation* Add(TF_Output l, TF_Output r, TF_Graph* graph, TF_Status* s, - const char* name = "add") { - TF_OperationDescription* desc = TF_NewOperation(graph, "AddN", name); - TF_Output inputs[2] = {l, r}; - TF_AddInputList(desc, inputs, 2); - return TF_FinishOperation(desc, s); -} - -TF_Operation* Neg(TF_Operation* n, TF_Graph* graph, TF_Status* s) { - TF_OperationDescription* desc = TF_NewOperation(graph, "Neg", "neg"); - TF_Output neg_input = {n, 0}; - TF_AddInput(desc, neg_input); - return TF_FinishOperation(desc, s); -} - -TF_Operation* LessThan(TF_Output l, TF_Output r, TF_Graph* graph, - TF_Status* s) { - TF_OperationDescription* desc = TF_NewOperation(graph, "Less", "less_than"); - TF_AddInput(desc, l); - TF_AddInput(desc, r); - return TF_FinishOperation(desc, s); -} - -bool IsPlaceholder(const NodeDef& node_def) { - if (node_def.op() != "Placeholder" || node_def.name() != "feed") { - return false; - } - bool found_dtype = false; - bool found_shape = false; - for (const auto& attr : node_def.attr()) { - if (attr.first == "dtype") { - if (attr.second.type() == tensorflow::DT_INT32) { - found_dtype = true; - } else { - return false; - } - } else if (attr.first == "shape") { - found_shape = true; - } - } - return found_dtype && found_shape; -} - -bool IsScalarConst(const NodeDef& node_def, int v) { - if (node_def.op() != "Const" || node_def.name() != "scalar") { - return false; - } - bool found_dtype = false; - bool found_value = false; - for (const auto& attr : node_def.attr()) { - if (attr.first == "dtype") { - if (attr.second.type() == tensorflow::DT_INT32) { - found_dtype = true; - } else { - return false; - } - } else if (attr.first == "value") { - if (attr.second.has_tensor() && - attr.second.tensor().int_val_size() == 1 && - attr.second.tensor().int_val(0) == v) { - found_value = true; - } else { - return false; - } - } - } - return found_dtype && found_value; -} - -bool IsAddN(const NodeDef& node_def, int n) { - if (node_def.op() != "AddN" || node_def.name() != "add" || - node_def.input_size() != n) { - return false; - } - bool found_t = false; - bool found_n = false; - for (const auto& attr : node_def.attr()) { - if (attr.first == "T") { - if (attr.second.type() == tensorflow::DT_INT32) { - found_t = true; - } else { - return false; - } - } else if (attr.first == "N") { - if (attr.second.i() == n) { - found_n = true; - } else { - return false; - } - } - } - return found_t && found_n; -} - -bool IsNeg(const NodeDef& node_def, const string& input) { - return node_def.op() == "Neg" && node_def.name() == "neg" && - node_def.input_size() == 1 && node_def.input(0) == input; -} - -bool GetGraphDef(TF_Graph* graph, GraphDef* graph_def) { - TF_Status* s = TF_NewStatus(); - TF_Buffer* buffer = TF_NewBuffer(); - TF_GraphToGraphDef(graph, buffer, s); - bool ret = TF_GetCode(s) == TF_OK; - EXPECT_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s); - if (ret) ret = graph_def->ParseFromArray(buffer->data, buffer->length); - TF_DeleteBuffer(buffer); - TF_DeleteStatus(s); - return ret; -} - -bool GetNodeDef(TF_Operation* oper, NodeDef* node_def) { - TF_Status* s = TF_NewStatus(); - TF_Buffer* buffer = TF_NewBuffer(); - TF_OperationToNodeDef(oper, buffer, s); - bool ret = TF_GetCode(s) == TF_OK; - EXPECT_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s); - if (ret) ret = node_def->ParseFromArray(buffer->data, buffer->length); - TF_DeleteBuffer(buffer); - TF_DeleteStatus(s); - return ret; -} - -bool GetAttrValue(TF_Operation* oper, const char* attr_name, - tensorflow::AttrValue* attr_value, TF_Status* s) { - TF_Buffer* buffer = TF_NewBuffer(); - TF_OperationGetAttrValueProto(oper, attr_name, buffer, s); - bool ret = TF_GetCode(s) == TF_OK; - if (ret) ret = attr_value->ParseFromArray(buffer->data, buffer->length); - TF_DeleteBuffer(buffer); - return ret; -} - TEST(CAPI, SetShape) { TF_Status* s = TF_NewStatus(); TF_Graph* graph = TF_NewGraph(); @@ -805,6 +645,33 @@ TEST(CAPI, ImportGraphDef) { EXPECT_EQ(feed, control_inputs[0]); EXPECT_EQ(feed2, control_inputs[1]); + // Export to a graph def so we can import a graph with control dependencies + TF_DeleteBuffer(graph_def); + graph_def = TF_NewBuffer(); + TF_GraphToGraphDef(graph, graph_def, s); + ASSERT_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s); + + // Import again, with remapped control dependency, into the same graph + TF_DeleteImportGraphDefOptions(opts); + opts = TF_NewImportGraphDefOptions(); + TF_ImportGraphDefOptionsSetPrefix(opts, "imported4"); + TF_ImportGraphDefOptionsRemapControlDependency(opts, "imported/feed", feed); + TF_GraphImportGraphDef(graph, graph_def, opts, s); + ASSERT_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s); + + TF_Operation* scalar4 = + TF_GraphOperationByName(graph, "imported4/imported3/scalar"); + TF_Operation* feed4 = + TF_GraphOperationByName(graph, "imported4/imported2/feed"); + + // Check that imported `imported3/scalar` has remapped control dep from + // original graph and imported control dep + num_control_inputs = TF_OperationGetControlInputs( + scalar4, control_inputs, TF_OperationNumControlInputs(scalar4)); + ASSERT_EQ(2, num_control_inputs); + EXPECT_EQ(feed, control_inputs[0]); + EXPECT_EQ(feed4, control_inputs[1]); + TF_DeleteImportGraphDefOptions(opts); TF_DeleteBuffer(graph_def); @@ -816,114 +683,6 @@ TEST(CAPI, ImportGraphDef) { TF_DeleteStatus(s); } -class CSession { - public: - CSession(TF_Graph* graph, TF_Status* s) { - TF_SessionOptions* opts = TF_NewSessionOptions(); - session_ = TF_NewSession(graph, opts, s); - TF_DeleteSessionOptions(opts); - } - - CSession(TF_Session* session) { session_ = session; } - - ~CSession() { - TF_Status* s = TF_NewStatus(); - CloseAndDelete(s); - EXPECT_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s); - TF_DeleteStatus(s); - } - - void SetInputs(std::vector> inputs) { - DeleteInputValues(); - inputs_.clear(); - for (const auto& p : inputs) { - inputs_.emplace_back(TF_Output{p.first, 0}); - input_values_.emplace_back(p.second); - } - } - - void SetOutputs(std::initializer_list outputs) { - ResetOutputValues(); - outputs_.clear(); - for (TF_Operation* o : outputs) { - outputs_.emplace_back(TF_Output{o, 0}); - } - } - - void SetOutputs(const std::vector& outputs) { - ResetOutputValues(); - outputs_ = outputs; - } - - void SetTargets(std::initializer_list targets) { - targets_.clear(); - for (TF_Operation* t : targets) { - targets_.emplace_back(t); - } - } - - void Run(TF_Status* s) { - if (inputs_.size() != input_values_.size()) { - ADD_FAILURE() << "Call SetInputs() before Run()"; - return; - } - ResetOutputValues(); - output_values_.resize(outputs_.size(), nullptr); - - const TF_Output* inputs_ptr = inputs_.empty() ? nullptr : &inputs_[0]; - TF_Tensor* const* input_values_ptr = - input_values_.empty() ? nullptr : &input_values_[0]; - - const TF_Output* outputs_ptr = outputs_.empty() ? nullptr : &outputs_[0]; - TF_Tensor** output_values_ptr = - output_values_.empty() ? nullptr : &output_values_[0]; - - TF_Operation* const* targets_ptr = - targets_.empty() ? nullptr : &targets_[0]; - - TF_SessionRun(session_, nullptr, inputs_ptr, input_values_ptr, - inputs_.size(), outputs_ptr, output_values_ptr, - outputs_.size(), targets_ptr, targets_.size(), nullptr, s); - - DeleteInputValues(); - } - - void CloseAndDelete(TF_Status* s) { - DeleteInputValues(); - ResetOutputValues(); - if (session_ != nullptr) { - TF_CloseSession(session_, s); - EXPECT_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s); - TF_DeleteSession(session_, s); - session_ = nullptr; - } - } - - TF_Tensor* output_tensor(int i) { return output_values_[i]; } - - private: - void DeleteInputValues() { - for (int i = 0; i < input_values_.size(); ++i) { - TF_DeleteTensor(input_values_[i]); - } - input_values_.clear(); - } - - void ResetOutputValues() { - for (int i = 0; i < output_values_.size(); ++i) { - if (output_values_[i] != nullptr) TF_DeleteTensor(output_values_[i]); - } - output_values_.clear(); - } - - TF_Session* session_; - std::vector inputs_; - std::vector input_values_; - std::vector outputs_; - std::vector output_values_; - std::vector targets_; -}; - TEST(CAPI, Session) { TF_Status* s = TF_NewStatus(); TF_Graph* graph = TF_NewGraph(); @@ -1049,39 +808,201 @@ TEST(CAPI, SessionPRun) { TF_DeleteStatus(s); } -TEST(CAPI, ColocateWith) { - TF_Status* s = TF_NewStatus(); +TEST(CAPI, ShapeInferenceError) { + // TF_FinishOperation should fail if the shape of the added operation cannot + // be inferred. + TF_Status* status = TF_NewStatus(); TF_Graph* graph = TF_NewGraph(); - TF_Operation* feed = Placeholder(graph, s); - ASSERT_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s); + // Create this failure by trying to add two nodes with incompatible shapes + // (A tensor with shape [2] and a tensor with shape [3] cannot be added). + const char data[] = {1, 2, 3}; + const int64_t vec2_dims[] = {2}; + unique_tensor_ptr vec2_tensor( + Int8Tensor(vec2_dims, TF_ARRAYSIZE(vec2_dims), data), TF_DeleteTensor); + TF_Operation* vec2 = Const(vec2_tensor.get(), graph, status, "vec2"); + ASSERT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); + + const int64_t vec3_dims[] = {3}; + unique_tensor_ptr vec3_tensor( + Int8Tensor(vec3_dims, TF_ARRAYSIZE(vec3_dims), data), TF_DeleteTensor); + TF_Operation* vec3 = Const(vec3_tensor.get(), graph, status, "vec3"); + ASSERT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); + + TF_Operation* add = AddNoCheck(vec2, vec3, graph, status); + ASSERT_NE(TF_OK, TF_GetCode(status)); + ASSERT_TRUE(add == nullptr); - TF_Operation* constant = ScalarConst(10, graph, s); - ASSERT_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s); + TF_DeleteGraph(graph); + TF_DeleteStatus(status); +} - TF_OperationDescription* desc = TF_NewOperation(graph, "AddN", "add"); - TF_Output inputs[] = {{feed, 0}, {constant, 0}}; - TF_AddInputList(desc, inputs, TF_ARRAYSIZE(inputs)); - TF_ColocateWith(desc, feed); - TF_Operation* add = TF_FinishOperation(desc, s); - ASSERT_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s); +void StringVectorToArrays(const std::vector& v, + std::unique_ptr* ptrs, + std::unique_ptr* lens) { + ptrs->reset(new const void*[v.size()]); + lens->reset(new size_t[v.size()]); + for (size_t i = 0; i < v.size(); ++i) { + (*ptrs)[i] = v[i].data(); + (*lens)[i] = v[i].size(); + } +} - TF_AttrMetadata m = - TF_OperationGetAttrMetadata(add, tensorflow::kColocationAttrName, s); - EXPECT_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s); - EXPECT_EQ(1, m.is_list); - EXPECT_EQ(1, m.list_size); - EXPECT_EQ(TF_ATTR_STRING, m.type); - void* values[1]; - size_t lens[1]; - std::unique_ptr storage(new char[m.total_size]); - TF_OperationGetAttrStringList(add, tensorflow::kColocationAttrName, values, - lens, 1, storage.get(), m.total_size, s); - EXPECT_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s); - EXPECT_EQ("loc:@feed", string(static_cast(values[0]), lens[0])); +class CApiColocationTest : public ::testing::Test { + protected: + CApiColocationTest() : s_(TF_NewStatus()), graph_(TF_NewGraph()) {} - TF_DeleteGraph(graph); - TF_DeleteStatus(s); + void SetUp() override { + feed1_ = Placeholder(graph_, s_, "feed1"); + ASSERT_EQ(TF_OK, TF_GetCode(s_)) << TF_Message(s_); + + feed2_ = Placeholder(graph_, s_, "feed2"); + ASSERT_EQ(TF_OK, TF_GetCode(s_)) << TF_Message(s_); + + constant_ = ScalarConst(10, graph_, s_); + ASSERT_EQ(TF_OK, TF_GetCode(s_)) << TF_Message(s_); + + desc_ = TF_NewOperation(graph_, "AddN", "add"); + TF_Output inputs[] = {{feed1_, 0}, {constant_, 0}}; + TF_AddInputList(desc_, inputs, TF_ARRAYSIZE(inputs)); + } + + ~CApiColocationTest() override { + TF_DeleteGraph(graph_); + TF_DeleteStatus(s_); + } + + void SetViaStringList(TF_OperationDescription* desc, + const std::vector& list) { + std::unique_ptr list_ptrs; + std::unique_ptr list_lens; + StringVectorToArrays(list, &list_ptrs, &list_lens); + TF_SetAttrStringList(desc, tensorflow::kColocationAttrName, list_ptrs.get(), + list_lens.get(), list.size()); + } + + void SetViaProto(TF_OperationDescription* desc, + const std::vector& list) { + tensorflow::AttrValue attr; + for (const string& v : list) { + attr.mutable_list()->add_s(v); + } + string bytes; + attr.SerializeToString(&bytes); + TF_SetAttrValueProto(desc, tensorflow::kColocationAttrName, bytes.data(), + bytes.size(), s_); + ASSERT_EQ(TF_OK, TF_GetCode(s_)) << TF_Message(s_); + } + + void VerifyCollocation(TF_Operation* op, + const std::vector& expected) { + TF_AttrMetadata m = + TF_OperationGetAttrMetadata(op, tensorflow::kColocationAttrName, s_); + if (expected.empty()) { + ASSERT_EQ(TF_INVALID_ARGUMENT, TF_GetCode(s_)) << TF_Message(s_); + EXPECT_EQ(std::string("Operation has no attr named '_class'."), + std::string(TF_Message(s_))); + return; + } + EXPECT_EQ(TF_OK, TF_GetCode(s_)) << TF_Message(s_); + EXPECT_EQ(1, m.is_list); + EXPECT_EQ(expected.size(), m.list_size); + EXPECT_EQ(TF_ATTR_STRING, m.type); + std::vector values(expected.size()); + std::vector lens(expected.size()); + std::unique_ptr storage(new char[m.total_size]); + TF_OperationGetAttrStringList(op, tensorflow::kColocationAttrName, + values.data(), lens.data(), expected.size(), + storage.get(), m.total_size, s_); + EXPECT_EQ(TF_OK, TF_GetCode(s_)) << TF_Message(s_); + for (int i = 0; i < expected.size(); ++i) { + EXPECT_EQ(expected[i], + string(static_cast(values[i]), lens[i])); + } + } + + void FinishAndVerify(TF_OperationDescription* desc, + const std::vector& expected) { + TF_Operation* op = TF_FinishOperation(desc_, s_); + ASSERT_EQ(TF_OK, TF_GetCode(s_)) << TF_Message(s_); + VerifyCollocation(op, expected); + } + + TF_Status* s_; + TF_Graph* graph_; + TF_Operation* feed1_; + TF_Operation* feed2_; + TF_Operation* constant_; + TF_OperationDescription* desc_; +}; + +TEST_F(CApiColocationTest, ColocateWith) { + TF_ColocateWith(desc_, feed1_); + FinishAndVerify(desc_, {"loc:@feed1"}); +} + +TEST_F(CApiColocationTest, StringList) { + SetViaStringList(desc_, {"loc:@feed1"}); + FinishAndVerify(desc_, {"loc:@feed1"}); +} + +TEST_F(CApiColocationTest, Proto) { + SetViaProto(desc_, {"loc:@feed1"}); + FinishAndVerify(desc_, {"loc:@feed1"}); +} + +TEST_F(CApiColocationTest, ColocateWith_StringList) { + TF_ColocateWith(desc_, feed1_); + SetViaStringList(desc_, {"loc:@feed2"}); + FinishAndVerify(desc_, {"loc:@feed2"}); +} + +TEST_F(CApiColocationTest, ColocateWith_Proto) { + TF_ColocateWith(desc_, feed1_); + SetViaProto(desc_, {"loc:@feed2"}); + FinishAndVerify(desc_, {"loc:@feed2"}); +} + +TEST_F(CApiColocationTest, StringList_ColocateWith) { + SetViaStringList(desc_, {"loc:@feed2"}); + TF_ColocateWith(desc_, feed1_); + FinishAndVerify(desc_, {"loc:@feed1", "loc:@feed2"}); +} + +TEST_F(CApiColocationTest, Proto_ColocateWith) { + SetViaProto(desc_, {"loc:@feed2"}); + TF_ColocateWith(desc_, feed1_); + FinishAndVerify(desc_, {"loc:@feed1", "loc:@feed2"}); +} + +TEST_F(CApiColocationTest, ColocateWith_ColocateWith) { + TF_ColocateWith(desc_, feed1_); + TF_ColocateWith(desc_, feed2_); + FinishAndVerify(desc_, {"loc:@feed1", "loc:@feed2"}); +} + +TEST_F(CApiColocationTest, Proto_StringList) { + SetViaProto(desc_, {"loc:@feed1"}); + SetViaStringList(desc_, {"loc:@feed2"}); + FinishAndVerify(desc_, {"loc:@feed2"}); +} + +TEST_F(CApiColocationTest, StringList_Proto) { + SetViaStringList(desc_, {"loc:@feed1"}); + SetViaProto(desc_, {"loc:@feed2"}); + FinishAndVerify(desc_, {"loc:@feed2"}); +} + +TEST_F(CApiColocationTest, ClearViaStringList) { + TF_ColocateWith(desc_, feed1_); + SetViaStringList(desc_, {}); + FinishAndVerify(desc_, {}); +} + +TEST_F(CApiColocationTest, ClearViaProto) { + TF_ColocateWith(desc_, feed1_); + SetViaProto(desc_, {}); + FinishAndVerify(desc_, {}); } TEST(CAPI, SavedModel) { @@ -1129,7 +1050,8 @@ TEST(CAPI, SavedModel) { TF_Operation* input_op = TF_GraphOperationByName(graph, input_op_name.c_str()); ASSERT_TRUE(input_op != nullptr); - csession.SetInputs({{input_op, TF_Tensor_EncodeStrings(input)}}); + csession.SetInputs({{input_op, TF_TensorFromTensor(input, s)}}); + ASSERT_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s); const tensorflow::string output_op_name = tensorflow::ParseTensorName(output_name).first.ToString(); @@ -1180,338 +1102,293 @@ TEST(CAPI, SavedModelNullArgsAreValid) { TF_DeleteStatus(s); } -class CApiWhileLoopTest : public ::testing::Test { +REGISTER_OP("TestOpWithNoGradient") + .Input("x: T") + .Output("y: T") + .Attr("T: {float, double}") + .Doc(R"doc( +Test op with no grad registered. + +x: input +y: output +)doc") + .SetShapeFn(tensorflow::shape_inference::UnknownShape); + +class CApiGradientsTest : public ::testing::Test { protected: - CApiWhileLoopTest() : s_(TF_NewStatus()), graph_(TF_NewGraph()) {} + CApiGradientsTest() + : s_(TF_NewStatus()), + graph_(TF_NewGraph()), + expected_graph_(TF_NewGraph()) {} - ~CApiWhileLoopTest() override { + ~CApiGradientsTest() override { TF_DeleteGraph(graph_); + TF_DeleteGraph(expected_graph_); TF_DeleteStatus(s_); } - void Init(int ninputs) { - DCHECK(inputs_.empty()); - DCHECK_GT(ninputs, 0); + void TestGradientsSuccess(bool grad_inputs_provided) { + TF_Output inputs[2]; + TF_Output outputs[1]; + TF_Output grad_outputs[2]; + TF_Output expected_grad_outputs[2]; - for (int i = 0; i < ninputs; ++i) { - TF_Operation* placeholder = Placeholder( - graph_, s_, ::tensorflow::strings::StrCat("p", i).c_str()); - DCHECK_EQ(TF_OK, TF_GetCode(s_)) << TF_Message(s_); - inputs_.push_back({placeholder, 0}); - } + BuildSuccessGraph(inputs, outputs); + BuildExpectedGraph(grad_inputs_provided, expected_grad_outputs); - original_graph_description_ = GraphDebugString(); + AddGradients(grad_inputs_provided, inputs, 2, outputs, 1, grad_outputs); - params_.reset(new TF_WhileParams( - TF_NewWhile(graph_, &inputs_[0], inputs_.size(), s_))); - ASSERT_EQ(TF_OK, TF_GetCode(s_)) << TF_Message(s_); - ASSERT_EQ(original_graph_description_, GraphDebugString()) - << "TF_NewWhile() altered graph"; + EXPECT_EQ(TF_OK, TF_GetCode(s_)) << TF_Message(s_); - params_->name = "test_loop"; + // Compare that the graphs match. + GraphDef expected_gdef; + GraphDef gdef; + EXPECT_TRUE(GetGraphDef(expected_graph_, &expected_gdef)); + EXPECT_TRUE(GetGraphDef(graph_, &gdef)); + TF_EXPECT_GRAPH_EQ(expected_gdef, gdef); - // Initialize outputs_ so we can easily detect errors/bugs - outputs_.resize(ninputs, {nullptr, -1}); + // Compare that the output of the gradients of both graphs match. + RunGraphsAndCompareOutputs(grad_outputs, expected_grad_outputs); } - void ExpectOK() { - TF_FinishWhile(params_.get(), s_, &outputs_[0]); - EXPECT_EQ(TF_OK, TF_GetCode(s_)) << TF_Message(s_); - } + void TestGradientsError(bool grad_inputs_provided) { + TF_Output inputs[1]; + TF_Output outputs[1]; + TF_Output grad_outputs[1]; + + BuildErrorGraph(inputs, outputs); - void ExpectError(TF_Code expected_code, const string& expected_msg) { - TF_FinishWhile(params_.get(), s_, &outputs_[0]); - EXPECT_EQ(expected_code, TF_GetCode(s_)); + AddGradients(grad_inputs_provided, inputs, 1, outputs, 1, grad_outputs); + + string expected_msg = + "No gradient defined for op: TestOpWithNoGradient. Please see " + "https://www.tensorflow.org/code/" + "tensorflow/cc/gradients/README.md" + " for instructions on how to add C++ gradients."; EXPECT_EQ(expected_msg, TF_Message(s_)); - // TODO(skyewm): this assert is currently broken. Fix or remove guarantee. - // ASSERT_EQ(original_graph_description_, GraphDebugString()) << - // "TF_FinishWhile() altered graph on error"; } - void Run(std::initializer_list input_values) { - DCHECK_EQ(inputs_.size(), input_values.size()); - std::vector> inputs(inputs_.size()); - int i = 0; - for (int v : input_values) { - inputs[i] = {inputs_[i].oper, Int32Tensor(v)}; - ++i; - } - csession_.reset(new CSession(graph_, s_)); - csession_->SetInputs(inputs); - csession_->SetOutputs(outputs_); - csession_->Run(s_); + // Run the graph and ensure that the gradient values are as expected. + void RunGraphsAndCompareOutputs(TF_Output* grad_outputs, + TF_Output* expected_grad_outputs) { + std::unique_ptr csession(new CSession(graph_, s_)); + std::unique_ptr expected_csession( + new CSession(expected_graph_, s_)); + + std::vector grad_outputs_vec; + grad_outputs_vec.assign(grad_outputs, grad_outputs + 2); + csession->SetOutputs(grad_outputs_vec); + csession->Run(s_); ASSERT_EQ(TF_OK, TF_GetCode(s_)) << TF_Message(s_); - } + TF_Tensor* out0 = csession->output_tensor(0); + TF_Tensor* out1 = csession->output_tensor(1); + + std::vector expected_grad_outputs_vec; + expected_grad_outputs_vec.assign(expected_grad_outputs, + expected_grad_outputs + 2); + expected_csession->SetOutputs(expected_grad_outputs_vec); + expected_csession->Run(s_); + ASSERT_EQ(TF_OK, TF_GetCode(s_)) << TF_Message(s_); + TF_Tensor* expected_out0 = expected_csession->output_tensor(0); + TF_Tensor* expected_out1 = expected_csession->output_tensor(1); - void ExpectOutputValue(int idx, int expected_value) { - TF_Tensor* out = csession_->output_tensor(idx); - ASSERT_TRUE(out != nullptr); - EXPECT_EQ(TF_INT32, TF_TensorType(out)); - EXPECT_EQ(0, TF_NumDims(out)); - ASSERT_EQ(sizeof(int32), TF_TensorByteSize(out)); - int32* data = static_cast(TF_TensorData(out)); - EXPECT_EQ(expected_value, *data); + CompareTensors(out0, expected_out0); + CompareTensors(out1, expected_out1); } - // Create a valid conditonal graph. Useful for testing unrelated errors. - void CreateCondGraph() { - TF_Operation* one = ScalarConst(1, params_->cond_graph, s_); - TF_Operation* less_than = - LessThan(params_->cond_inputs[0], {one, 0}, params_->cond_graph, s_); - DCHECK_EQ(TF_OK, TF_GetCode(s_)) << TF_Message(s_); - params_->cond_output = {less_than, 0}; + void CompareTensors(TF_Tensor* a, TF_Tensor* b) { + float* a_data = static_cast(TF_TensorData(a)); + float* b_data = static_cast(TF_TensorData(b)); + EXPECT_EQ(*a_data, *b_data); } - string GraphDebugString() const { - TF_Buffer* buf = TF_NewBuffer(); - TF_GraphToGraphDef(graph_, buf, s_); - DCHECK_EQ(TF_OK, TF_GetCode(s_)) << TF_Message(s_); - GraphDef def; - bool success = def.ParseFromArray(buf->data, buf->length); - DCHECK(success); - TF_DeleteBuffer(buf); - return def.DebugString(); + void AddGradients(bool grad_inputs_provided, TF_Output* inputs, int ninputs, + TF_Output* outputs, int noutputs, TF_Output* grad_outputs) { + if (grad_inputs_provided) { + TF_Output grad_inputs[1]; + const float grad_inputs_val[] = {1.0, 1.0, 1.0, 1.0}; + TF_Operation* grad_inputs_op = + FloatConst2x2(graph_, s_, grad_inputs_val, "GradInputs"); + grad_inputs[0] = TF_Output{grad_inputs_op, 0}; + TF_AddGradients(graph_, outputs, noutputs, inputs, ninputs, grad_inputs, + s_, grad_outputs); + } else { + TF_AddGradients(graph_, outputs, noutputs, inputs, ninputs, nullptr, s_, + grad_outputs); + } } - TF_Status* s_; - TF_Graph* graph_; - std::vector inputs_; // The inputs to the while loop - std::vector outputs_; // The final outputs of the while loop - std::unique_ptr params_; - std::unique_ptr csession_; - - private: - // Used to verify that errors don't change graph_ - string original_graph_description_; -}; - -TEST_F(CApiWhileLoopTest, BasicLoop) { - Init(2); - - // Validate TF_WhileParams returned by TF_NewWhile() - EXPECT_TRUE(params_->body_graph != nullptr); - EXPECT_TRUE(params_->cond_graph != nullptr); - - EXPECT_EQ(params_->ninputs, 2); - - ASSERT_TRUE(params_->cond_inputs != nullptr); - ASSERT_TRUE(params_->cond_inputs[0].oper != nullptr); - EXPECT_TRUE(params_->cond_inputs[1].oper != nullptr); - - ASSERT_TRUE(params_->body_inputs != nullptr); - EXPECT_TRUE(params_->body_inputs[0].oper != nullptr); - EXPECT_TRUE(params_->body_inputs[1].oper != nullptr); - - ASSERT_TRUE(params_->body_outputs != nullptr); - - // Create loop: while (input1 < input2) input1 += input2 + 1 - TF_Operation* less_than = - LessThan(params_->cond_inputs[0], params_->cond_inputs[1], - params_->cond_graph, s_); - ASSERT_EQ(TF_OK, TF_GetCode(s_)) << TF_Message(s_); - params_->cond_output = {less_than, 0}; - - TF_Operation* add1 = Add(params_->body_inputs[0], params_->body_inputs[1], - params_->body_graph, s_, "add1"); - ASSERT_EQ(TF_OK, TF_GetCode(s_)) << TF_Message(s_); - TF_Operation* one = ScalarConst(1, params_->body_graph, s_); - ASSERT_EQ(TF_OK, TF_GetCode(s_)) << TF_Message(s_); - TF_Operation* add2 = Add(add1, one, params_->body_graph, s_, "add2"); - ASSERT_EQ(TF_OK, TF_GetCode(s_)) << TF_Message(s_); - params_->body_outputs[0] = {add2, 0}; - params_->body_outputs[1] = params_->body_inputs[1]; - - // Finalize while loop - ExpectOK(); - - // Validate while loop outputs returned by TF_FinishWhile() - EXPECT_TRUE(outputs_[0].oper != nullptr); - EXPECT_GE(outputs_[0].index, 0); - EXPECT_TRUE(outputs_[1].oper != nullptr); - EXPECT_GE(outputs_[1].index, 0); - - // Run the graph - Run({-9, 2}); - ExpectOutputValue(0, 3); - ExpectOutputValue(1, 2); -} - -TEST_F(CApiWhileLoopTest, NestedLoop) { - Init(2); - // Create nested loop: - // while (input1 < 6) { - // inner_input1 = input1 - // while (inner_input1 < 3) { - // input2 += 1 - // inner_input1 += 2 - // } - // input1 += input2 - // } - // - // Expected execution with initial values input1 = input2 = 0: - // - // outer inner inner_ - // step# step# input1 input2 input1 - // ------------------------------------ - // 0 0 0 0 0 - // 0 1 0 1 2 - // 0 2 0 2 4 - // 0 - 2 2 - - // 1 0 2 2 2 - // 1 1 2 3 4 - // 1 - 5 3 - - // 2 0 5 3 5 - // 2 - 8 3 - - - // Create outer cond graph - TF_Operation* six = ScalarConst(6, params_->cond_graph, s_); - ASSERT_EQ(TF_OK, TF_GetCode(s_)) << TF_Message(s_); - TF_Operation* less_than = - LessThan(params_->cond_inputs[0], {six, 0}, params_->cond_graph, s_); - ASSERT_EQ(TF_OK, TF_GetCode(s_)) << TF_Message(s_); - params_->cond_output = {less_than, 0}; - - // Create outer body graph - // Init inner graph - TF_Output inner_inputs[] = {params_->body_inputs[0], params_->body_inputs[1]}; - TF_WhileParams inner_params = - TF_NewWhile(params_->body_graph, inner_inputs, 2, s_); - ASSERT_EQ(TF_OK, TF_GetCode(s_)) << TF_Message(s_); - inner_params.name = "inner_loop"; + void BuildErrorGraph(TF_Output* inputs, TF_Output* outputs) { + const float const0_val[] = {1.0, 2.0, 3.0, 4.0}; + TF_Operation* const0 = FloatConst2x2(graph_, s_, const0_val, "Const_0"); + TF_Operation* nograd = NoGradientOp(graph_, s_, const0, "NoGrad"); + inputs[0] = TF_Output{const0, 0}; + outputs[0] = TF_Output{nograd, 0}; + EXPECT_EQ(TF_OK, TF_GetCode(s_)) << TF_Message(s_); + } - // Create inner cond graph - TF_Operation* three = ScalarConst(3, inner_params.cond_graph, s_); - ASSERT_EQ(TF_OK, TF_GetCode(s_)) << TF_Message(s_); - TF_Operation* inner_less_than = LessThan( - inner_params.cond_inputs[0], {three, 0}, inner_params.cond_graph, s_); - ASSERT_EQ(TF_OK, TF_GetCode(s_)) << TF_Message(s_); - inner_params.cond_output = {inner_less_than, 0}; + void BuildSuccessGraph(TF_Output* inputs, TF_Output* outputs) { + // Construct the following graph: + // | + // z| + // | + // MatMul + // / \ + // ^ ^ + // | | + // x| y| + // | | + // | | + // Const_0 Const_1 + // + const float const0_val[] = {1.0, 2.0, 3.0, 4.0}; + const float const1_val[] = {1.0, 0.0, 0.0, 1.0}; + TF_Operation* const0 = FloatConst2x2(graph_, s_, const0_val, "Const_0"); + TF_Operation* const1 = FloatConst2x2(graph_, s_, const1_val, "Const_1"); + TF_Operation* matmul = MatMul(graph_, s_, const0, const1, "MatMul"); + inputs[0] = TF_Output{const0, 0}; + inputs[1] = TF_Output{const1, 0}; + outputs[0] = TF_Output{matmul, 0}; + EXPECT_EQ(TF_OK, TF_GetCode(s_)) << TF_Message(s_); + } - // Create inner body graph - TF_Operation* one = ScalarConst(1, inner_params.body_graph, s_, "one"); - ASSERT_EQ(TF_OK, TF_GetCode(s_)) << TF_Message(s_); - TF_Operation* two = ScalarConst(2, inner_params.body_graph, s_, "two"); - ASSERT_EQ(TF_OK, TF_GetCode(s_)) << TF_Message(s_); + void BuildExpectedGraph(bool grad_inputs_provided, + TF_Output* expected_grad_outputs) { + // The expected graph looks like this if grad_inputs_provided. + // If grad_inputs_provided is false, Const_0 will be a OnesLike op. + // ^ ^ + // dy| dx| // MatMul Gradient Graph + // | | + // MatMul_2 MatMul_1 + // ^ ^ ^ ^ + // | |----------| | + // | ^ | + // | dz| | + // | | | + // | Const_3 | + // | | + // | ^ | + // | z| | // MatMul Forward Graph + // | | | + // | MatMul | + // | / \ | + // | ^ ^ | + // | | | | + // |---x| y|----| + // | | + // | | + // Const_0 Const_1 + // + const float const0_val[] = {1.0, 2.0, 3.0, 4.0}; + const float const1_val[] = {1.0, 0.0, 0.0, 1.0}; + TF_Operation* const0 = + FloatConst2x2(expected_graph_, s_, const0_val, "Const_0"); + TF_Operation* const1 = + FloatConst2x2(expected_graph_, s_, const1_val, "Const_1"); + TF_Operation* matmul = + MatMul(expected_graph_, s_, const0, const1, "MatMul"); + + TF_Operation* const3; + if (grad_inputs_provided) { + const float const3_val[] = {1.0, 1.0, 1.0, 1.0}; + const3 = FloatConst2x2(expected_graph_, s_, const3_val, "GradInputs"); + } else { + const3 = OnesLike(expected_graph_, s_, matmul, "gradients/OnesLike"); + } - TF_Operation* input2_add = - Add(inner_params.body_inputs[1].oper, one, inner_params.body_graph, s_); - ASSERT_EQ(TF_OK, TF_GetCode(s_)) << TF_Message(s_); - inner_params.body_outputs[1] = {input2_add, 0}; + TF_Operation* matmul1 = MatMul(expected_graph_, s_, const3, const1, + "gradients/MatMul", false, true); + TF_Operation* matmul2 = MatMul(expected_graph_, s_, const0, const3, + "gradients/MatMul_1", true, false); + expected_grad_outputs[0] = {matmul1, 0}; + expected_grad_outputs[1] = {matmul2, 0}; + } - TF_Operation* inner_input1_add = Add(inner_params.body_inputs[0].oper, two, - inner_params.body_graph, s_, "add2"); - ASSERT_EQ(TF_OK, TF_GetCode(s_)) << TF_Message(s_); - inner_params.body_outputs[0] = {inner_input1_add, 0}; + TF_Tensor* FloatTensor2x2(const float* values) { + const int64_t dims[2] = {2, 2}; + TF_Tensor* t = TF_AllocateTensor(TF_FLOAT, dims, 2, sizeof(float) * 4); + memcpy(TF_TensorData(t), values, sizeof(float) * 4); + return t; + } - // Finalize inner graph - TF_Output inner_outputs[2] = {{nullptr, -1}}; - TF_FinishWhile(&inner_params, s_, inner_outputs); - ASSERT_EQ(TF_OK, TF_GetCode(s_)) << TF_Message(s_); + TF_Operation* FloatConst2x2(TF_Graph* graph, TF_Status* s, + const float* values, const char* name) { + unique_tensor_ptr tensor(FloatTensor2x2(values), TF_DeleteTensor); + TF_OperationDescription* desc = TF_NewOperation(graph, "Const", name); + TF_SetAttrTensor(desc, "value", tensor.get(), s); + if (TF_GetCode(s) != TF_OK) return nullptr; + TF_SetAttrType(desc, "dtype", TF_FLOAT); + TF_Operation* op = TF_FinishOperation(desc, s); + EXPECT_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s); + return op; + } - TF_Operation* input1_add = - Add(params_->body_inputs[0], inner_outputs[1], params_->body_graph, s_); - ASSERT_EQ(TF_OK, TF_GetCode(s_)) << TF_Message(s_); - params_->body_outputs[0] = {input1_add, 0}; - - params_->body_outputs[1] = inner_outputs[1]; - - // Finalize outer graph - ExpectOK(); - - // Check for a few expected nodes - const char* node_name = "test_loop/cond/scalar"; - EXPECT_TRUE(TF_GraphOperationByName(graph_, node_name) != nullptr); - node_name = "test_loop/body/add"; - EXPECT_TRUE(TF_GraphOperationByName(graph_, node_name) != nullptr); - node_name = "test_loop/body/inner_loop/body/one"; - EXPECT_TRUE(TF_GraphOperationByName(graph_, node_name) != nullptr); - node_name = "test_loop/body/inner_loop/cond/less_than"; - EXPECT_TRUE(TF_GraphOperationByName(graph_, node_name) != nullptr); - - // Run the graph - Run({0, 0}); - ExpectOutputValue(0, 8); - ExpectOutputValue(1, 3); -} + TF_Operation* MatMul(TF_Graph* graph, TF_Status* s, TF_Operation* l, + TF_Operation* r, const char* name, + bool transpose_a = false, bool transpose_b = false) { + TF_OperationDescription* desc = TF_NewOperation(graph, "MatMul", name); + if (transpose_a) { + TF_SetAttrBool(desc, "transpose_a", 1); + } + if (transpose_b) { + TF_SetAttrBool(desc, "transpose_b", 1); + } + TF_AddInput(desc, {l, 0}); + TF_AddInput(desc, {r, 0}); + TF_Operation* op = TF_FinishOperation(desc, s); + EXPECT_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s); + return op; + } -TEST_F(CApiWhileLoopTest, BadCondOutput) { - Init(1); - params_->body_outputs[0] = params_->body_inputs[0]; - ExpectError(TF_INVALID_ARGUMENT, - "TF_WhileParams `cond_output` field isn't set"); -} + TF_Operation* OnesLike(TF_Graph* graph, TF_Status* s, TF_Operation* in, + const char* name) { + TF_OperationDescription* desc = TF_NewOperation(graph, "OnesLike", name); + TF_AddInput(desc, {in, 0}); + TF_Operation* op = TF_FinishOperation(desc, s); + EXPECT_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s); + return op; + } -TEST_F(CApiWhileLoopTest, BadBodyOutput) { - Init(1); - CreateCondGraph(); - ExpectError(TF_INVALID_ARGUMENT, - "TF_WhileParams `body_outputs[0]` field isn't set"); -} + TF_Operation* NoGradientOp(TF_Graph* graph, TF_Status* s, TF_Operation* in, + const char* name) { + TF_OperationDescription* desc = + TF_NewOperation(graph, "TestOpWithNoGradient", name); + TF_AddInput(desc, {in, 0}); + TF_Operation* op = TF_FinishOperation(desc, s); + EXPECT_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s); + return op; + } -TEST_F(CApiWhileLoopTest, NullName) { - Init(1); - CreateCondGraph(); - params_->body_outputs[0] = params_->body_inputs[0]; - params_->name = nullptr; - ExpectError(TF_INVALID_ARGUMENT, "TF_WhileParams `name` field is null"); -} + TF_Status* s_; + TF_Graph* graph_; + TF_Graph* expected_graph_; +}; -TEST_F(CApiWhileLoopTest, WrongGraph) { - Init(1); - CreateCondGraph(); - // Set body output to output from outer graph - params_->body_outputs[0] = inputs_[0]; - // TODO(skyewm): improve error message - ExpectError(TF_INVALID_ARGUMENT, - "Requested return node 'p0' not found in graph def"); -} +TEST_F(CApiGradientsTest, Gradients_GradInputs) { TestGradientsSuccess(true); } -TEST_F(CApiWhileLoopTest, BadTypes) { - Init(1); - CreateCondGraph(); - // Op that has a float input + output - TF_OperationDescription* desc = TF_NewOperation( - params_->body_graph, "FakeQuantWithMinMaxArgs", "float_op"); - TF_AddInput(desc, params_->body_inputs[0]); - TF_FinishOperation(desc, s_); - ASSERT_EQ(TF_INVALID_ARGUMENT, TF_GetCode(s_)); - string msg(TF_Message(s_)); - EXPECT_NE(msg.find("Input 'inputs' passed int32 expected float while " - "building NodeDef 'float_op'"), - msg.npos); - TF_AbortWhile(params_.get()); +TEST_F(CApiGradientsTest, Gradients_NoGradInputs) { + TestGradientsSuccess(false); } -// Create a tensor with values of type TF_INT8 provided by `values`. -TF_Tensor* Int8Tensor(const int64_t* dims, int num_dims, const char* values) { - int64_t num_values = 1; - for (int i = 0; i < num_dims; ++i) { - num_values *= dims[i]; - } - TF_Tensor* t = - TF_AllocateTensor(TF_INT8, dims, num_dims, sizeof(char) * num_values); - memcpy(TF_TensorData(t), values, sizeof(char) * num_values); - return t; +TEST_F(CApiGradientsTest, OpWithNoGradientRegistered_GradInputs) { + TestGradientsError(true); } -void StringVectorToArrays(const std::vector& v, - std::unique_ptr* ptrs, - std::unique_ptr* lens) { - ptrs->reset(new const void*[v.size()]); - lens->reset(new size_t[v.size()]); - for (size_t i = 0; i < v.size(); ++i) { - (*ptrs)[i] = v[i].data(); - (*lens)[i] = v[i].size(); - } +TEST_F(CApiGradientsTest, OpWithNoGradientRegistered_NoGradInputs) { + TestGradientsError(false); } // REGISTER_OP for CApiTestAttributesTest test cases. // Registers two ops, each with a single attribute called 'v'. // The attribute in one op will have a type 'type', the other // will have list(type). -#define ATTR_TEST_REGISTER_OP(type) \ - REGISTER_OP("CApiAttributesTestOp" #type).Attr("v: " #type); \ - REGISTER_OP("CApiAttributesTestOpList" #type).Attr("v: list(" #type ")") +#define ATTR_TEST_REGISTER_OP(type) \ + REGISTER_OP("CApiAttributesTestOp" #type) \ + .Attr("v: " #type) \ + .SetShapeFn(tensorflow::shape_inference::UnknownShape); \ + REGISTER_OP("CApiAttributesTestOpList" #type) \ + .Attr("v: list(" #type ")") \ + .SetShapeFn(tensorflow::shape_inference::UnknownShape) ATTR_TEST_REGISTER_OP(string); ATTR_TEST_REGISTER_OP(int); ATTR_TEST_REGISTER_OP(float); @@ -1881,6 +1758,39 @@ TEST_F(CApiAttributesTest, Tensor) { TF_DeleteTensor(value); } +TEST_F(CApiAttributesTest, StringTensor) { + // Create the string-Tensor "atttribute" value. + char encoded[] = { + 0, 0, 0, 0, 0, 0, 0, 0, // array[uint64] offsets + 1, // varint encoded string length + 'A', + }; + auto deallocator = [](void* data, size_t len, void* arg) {}; + unique_tensor_ptr t_in(TF_NewTensor(TF_STRING, nullptr, 0, &encoded[0], + sizeof(encoded), deallocator, nullptr), + TF_DeleteTensor); + + // Create a TF_Operation with the attribute t_in + auto desc = init("tensor"); + TF_SetAttrTensor(desc, "v", t_in.get(), s_); + ASSERT_EQ(TF_OK, TF_GetCode(s_)) << TF_Message(s_); + + auto oper = TF_FinishOperation(desc, s_); + ASSERT_EQ(TF_OK, TF_GetCode(s_)) << TF_Message(s_); + + // Fetch the attribute back. + EXPECT_TF_META("v", -1, TF_ATTR_TENSOR, -1); + TF_Tensor* t_out = nullptr; + TF_OperationGetAttrTensor(oper, "v", &t_out, s_); + ASSERT_EQ(TF_OK, TF_GetCode(s_)) << TF_Message(s_); + EXPECT_EQ(TF_STRING, TF_TensorType(t_out)); + EXPECT_EQ(0, TF_NumDims(t_out)); + ASSERT_EQ(TF_TensorByteSize(t_in.get()), TF_TensorByteSize(t_out)); + EXPECT_EQ(0, memcmp(TF_TensorData(t_in.get()), TF_TensorData(t_out), + TF_TensorByteSize(t_out))); + TF_DeleteTensor(t_out); +} + TEST_F(CApiAttributesTest, TensorList) { const char tensor1[] = {5, 7}; const int64_t dims1[] = {1, 2}; @@ -1892,7 +1802,8 @@ TEST_F(CApiAttributesTest, TensorList) { auto desc = init("list(tensor)"); TF_Tensor* tmp[] = { - Int8Tensor(dims1, ndims1, tensor1), Int8Tensor(dims2, ndims2, tensor2), + Int8Tensor(dims1, ndims1, tensor1), + Int8Tensor(dims2, ndims2, tensor2), }; TF_SetAttrTensorList(desc, "v", tmp, TF_ARRAYSIZE(tmp), s_); for (int i = 0; i < TF_ARRAYSIZE(tmp); ++i) { @@ -1944,12 +1855,14 @@ TEST_F(CApiAttributesTest, Errors) { TF_OperationGetAttrString(oper, "v", nullptr, 0, s_); EXPECT_EQ(TF_INVALID_ARGUMENT, TF_GetCode(s_)) << TF_Message(s_); } + #undef EXPECT_TF_META +} // namespace +} // namespace tensorflow + // TODO(josh11b): Test: // * TF_SetDevice(desc, "/job:worker"); // * control inputs / outputs // * targets // * TF_DeleteGraph() before TF_DeleteSession() - -} // namespace diff --git a/tensorflow/c/c_test_util.cc b/tensorflow/c/c_test_util.cc new file mode 100644 index 0000000000000000000000000000000000000000..9cd978c97eada2123950da6271886ee20f918d5f --- /dev/null +++ b/tensorflow/c/c_test_util.cc @@ -0,0 +1,419 @@ +/* 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/c_test_util.h" + +#include "tensorflow/core/framework/function.pb.h" +#include "tensorflow/core/framework/tensor.pb.h" +#include "tensorflow/core/lib/strings/strcat.h" +#include "tensorflow/core/platform/logging.h" + +using tensorflow::GraphDef; +using tensorflow::NodeDef; + +static void Int32Deallocator(void* data, size_t, void* arg) { + delete[] static_cast(data); +} + +TF_Tensor* Int8Tensor(const int64_t* dims, int num_dims, const char* values) { + int64_t num_values = 1; + for (int i = 0; i < num_dims; ++i) { + num_values *= dims[i]; + } + TF_Tensor* t = + TF_AllocateTensor(TF_INT8, dims, num_dims, sizeof(char) * num_values); + memcpy(TF_TensorData(t), values, sizeof(char) * num_values); + return t; +} + +TF_Tensor* Int32Tensor(const int64_t* dims, int num_dims, + const int32_t* values) { + int64_t num_values = 1; + for (int i = 0; i < num_dims; ++i) { + num_values *= dims[i]; + } + TF_Tensor* t = + TF_AllocateTensor(TF_INT32, dims, num_dims, sizeof(int32_t) * num_values); + memcpy(TF_TensorData(t), values, sizeof(int32_t) * num_values); + return t; +} + +TF_Tensor* Int32Tensor(const std::vector& values) { + int64_t dims = values.size(); + return Int32Tensor(&dims, 1, values.data()); +} + +TF_Tensor* Int32Tensor(int32_t v) { + const int num_bytes = sizeof(int32_t); + int32_t* values = new int32_t[1]; + values[0] = v; + return TF_NewTensor(TF_INT32, nullptr, 0, values, num_bytes, + &Int32Deallocator, nullptr); +} + +// All the *Helper methods are used as a workaround for the restrictions that +// one cannot call ASSERT_* methods in non-void-returning functions (when +// exceptions are disabled during compilation) +void PlaceholderHelper(TF_Graph* graph, TF_Status* s, const char* name, + TF_Operation** op) { + TF_OperationDescription* desc = TF_NewOperation(graph, "Placeholder", name); + TF_SetAttrType(desc, "dtype", TF_INT32); + *op = TF_FinishOperation(desc, s); + ASSERT_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s); + ASSERT_NE(*op, nullptr); +} + +TF_Operation* Placeholder(TF_Graph* graph, TF_Status* s, const char* name) { + TF_Operation* op; + PlaceholderHelper(graph, s, name, &op); + return op; +} + +void ConstHelper(TF_Tensor* t, TF_Graph* graph, TF_Status* s, const char* name, + TF_Operation** op) { + TF_OperationDescription* desc = TF_NewOperation(graph, "Const", name); + TF_SetAttrTensor(desc, "value", t, s); + ASSERT_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s); + TF_SetAttrType(desc, "dtype", TF_TensorType(t)); + *op = TF_FinishOperation(desc, s); + ASSERT_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s); + ASSERT_NE(*op, nullptr); +} + +TF_Operation* Const(TF_Tensor* t, TF_Graph* graph, TF_Status* s, + const char* name) { + TF_Operation* op; + ConstHelper(t, graph, s, name, &op); + return op; +} + +TF_Operation* ScalarConst(int32_t v, TF_Graph* graph, TF_Status* s, + const char* name) { + unique_tensor_ptr tensor(Int32Tensor(v), TF_DeleteTensor); + return Const(tensor.get(), graph, s, name); +} + +void AddHelper(TF_Operation* l, TF_Operation* r, TF_Graph* graph, TF_Status* s, + const char* name, TF_Operation** op, bool check) { + TF_OperationDescription* desc = TF_NewOperation(graph, "AddN", name); + TF_Output add_inputs[2] = {{l, 0}, {r, 0}}; + TF_AddInputList(desc, add_inputs, 2); + *op = TF_FinishOperation(desc, s); + if (check) { + ASSERT_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s); + ASSERT_NE(*op, nullptr); + } +} + +TF_Operation* Add(TF_Operation* l, TF_Operation* r, TF_Graph* graph, + TF_Status* s, const char* name) { + TF_Operation* op; + AddHelper(l, r, graph, s, name, &op, true); + return op; +} + +TF_Operation* AddNoCheck(TF_Operation* l, TF_Operation* r, TF_Graph* graph, + TF_Status* s, const char* name) { + TF_Operation* op; + AddHelper(l, r, graph, s, name, &op, false); + return op; +} + +TF_Operation* AddWithCtrlDependency(TF_Operation* l, TF_Operation* r, + TF_Graph* graph, TF_Operation* ctrl_op, + TF_Status* s, const char* name) { + TF_OperationDescription* desc = TF_NewOperation(graph, "AddN", name); + TF_Output add_inputs[2] = {{l, 0}, {r, 0}}; + TF_AddInputList(desc, add_inputs, 2); + TF_AddControlInput(desc, ctrl_op); + return TF_FinishOperation(desc, s); +} + +TF_Operation* Add(TF_Output l, TF_Output r, TF_Graph* graph, TF_Status* s, + const char* name) { + TF_OperationDescription* desc = TF_NewOperation(graph, "AddN", name); + TF_Output inputs[2] = {l, r}; + TF_AddInputList(desc, inputs, 2); + return TF_FinishOperation(desc, s); +} + +void NegHelper(TF_Operation* n, TF_Graph* graph, TF_Status* s, + TF_Operation** op) { + TF_OperationDescription* desc = TF_NewOperation(graph, "Neg", "neg"); + TF_Output neg_input = {n, 0}; + TF_AddInput(desc, neg_input); + *op = TF_FinishOperation(desc, s); + ASSERT_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s); + ASSERT_NE(*op, nullptr); +} + +TF_Operation* Neg(TF_Operation* n, TF_Graph* graph, TF_Status* s) { + TF_Operation* op; + NegHelper(n, graph, s, &op); + return op; +} + +TF_Operation* LessThan(TF_Output l, TF_Output r, TF_Graph* graph, + TF_Status* s) { + TF_OperationDescription* desc = TF_NewOperation(graph, "Less", "less_than"); + TF_AddInput(desc, l); + TF_AddInput(desc, r); + return TF_FinishOperation(desc, s); +} + +void Split3Helper(TF_Operation* input, TF_Graph* graph, TF_Status* s, + const char* name, TF_Operation** op) { + TF_Operation* zero = ScalarConst( + 0, graph, s, ::tensorflow::strings::StrCat(name, "_const0").c_str()); + TF_OperationDescription* desc = TF_NewOperation(graph, "Split", name); + TF_AddInput(desc, {zero, 0}); + TF_AddInput(desc, {input, 0}); + TF_SetAttrInt(desc, "num_split", 3); + TF_SetAttrType(desc, "T", TF_INT32); + // Set device to CPU since there is no version of split for int32 on GPU + // TODO(iga): Convert all these helpers and tests to use floats because + // they are usually available on GPUs. After doing this, remove TF_SetDevice + // call in c_api_function_test.cc + TF_SetDevice(desc, "/cpu:0"); + *op = TF_FinishOperation(desc, s); + ASSERT_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s); + ASSERT_NE(*op, nullptr); +} + +TF_Operation* Split3(TF_Operation* input, TF_Graph* graph, TF_Status* s, + const char* name) { + TF_Operation* op; + Split3Helper(input, graph, s, name, &op); + return op; +} + +bool IsPlaceholder(const tensorflow::NodeDef& node_def) { + if (node_def.op() != "Placeholder" || node_def.name() != "feed") { + return false; + } + bool found_dtype = false; + bool found_shape = false; + for (const auto& attr : node_def.attr()) { + if (attr.first == "dtype") { + if (attr.second.type() == tensorflow::DT_INT32) { + found_dtype = true; + } else { + return false; + } + } else if (attr.first == "shape") { + found_shape = true; + } + } + return found_dtype && found_shape; +} + +bool IsScalarConst(const tensorflow::NodeDef& node_def, int v) { + if (node_def.op() != "Const" || node_def.name() != "scalar") { + return false; + } + bool found_dtype = false; + bool found_value = false; + for (const auto& attr : node_def.attr()) { + if (attr.first == "dtype") { + if (attr.second.type() == tensorflow::DT_INT32) { + found_dtype = true; + } else { + return false; + } + } else if (attr.first == "value") { + if (attr.second.has_tensor() && + attr.second.tensor().int_val_size() == 1 && + attr.second.tensor().int_val(0) == v) { + found_value = true; + } else { + return false; + } + } + } + return found_dtype && found_value; +} + +bool IsAddN(const tensorflow::NodeDef& node_def, int n) { + if (node_def.op() != "AddN" || node_def.name() != "add" || + node_def.input_size() != n) { + return false; + } + bool found_t = false; + bool found_n = false; + for (const auto& attr : node_def.attr()) { + if (attr.first == "T") { + if (attr.second.type() == tensorflow::DT_INT32) { + found_t = true; + } else { + return false; + } + } else if (attr.first == "N") { + if (attr.second.i() == n) { + found_n = true; + } else { + return false; + } + } + } + return found_t && found_n; +} + +bool IsNeg(const tensorflow::NodeDef& node_def, const string& input) { + return node_def.op() == "Neg" && node_def.name() == "neg" && + node_def.input_size() == 1 && node_def.input(0) == input; +} + +bool GetGraphDef(TF_Graph* graph, tensorflow::GraphDef* graph_def) { + TF_Status* s = TF_NewStatus(); + TF_Buffer* buffer = TF_NewBuffer(); + TF_GraphToGraphDef(graph, buffer, s); + bool ret = TF_GetCode(s) == TF_OK; + EXPECT_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s); + if (ret) ret = graph_def->ParseFromArray(buffer->data, buffer->length); + TF_DeleteBuffer(buffer); + TF_DeleteStatus(s); + return ret; +} + +bool GetNodeDef(TF_Operation* oper, tensorflow::NodeDef* node_def) { + TF_Status* s = TF_NewStatus(); + TF_Buffer* buffer = TF_NewBuffer(); + TF_OperationToNodeDef(oper, buffer, s); + bool ret = TF_GetCode(s) == TF_OK; + EXPECT_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s); + if (ret) ret = node_def->ParseFromArray(buffer->data, buffer->length); + TF_DeleteBuffer(buffer); + TF_DeleteStatus(s); + return ret; +} + +bool GetFunctionDef(TF_Function* func, tensorflow::FunctionDef* func_def) { + TF_Status* s = TF_NewStatus(); + TF_Buffer* buffer = TF_NewBuffer(); + TF_FunctionToFunctionDef(func, buffer, s); + bool ret = TF_GetCode(s) == TF_OK; + EXPECT_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s); + if (ret) ret = func_def->ParseFromArray(buffer->data, buffer->length); + TF_DeleteBuffer(buffer); + TF_DeleteStatus(s); + return ret; +} + +bool GetAttrValue(TF_Operation* oper, const char* attr_name, + tensorflow::AttrValue* attr_value, TF_Status* s) { + TF_Buffer* buffer = TF_NewBuffer(); + TF_OperationGetAttrValueProto(oper, attr_name, buffer, s); + bool ret = TF_GetCode(s) == TF_OK; + if (ret) ret = attr_value->ParseFromArray(buffer->data, buffer->length); + TF_DeleteBuffer(buffer); + return ret; +} + +CSession::CSession(TF_Graph* graph, TF_Status* s) { + TF_SessionOptions* opts = TF_NewSessionOptions(); + session_ = TF_NewSession(graph, opts, s); + TF_DeleteSessionOptions(opts); +} + +CSession::CSession(TF_Session* session) : session_(session) {} + +CSession::~CSession() { + TF_Status* s = TF_NewStatus(); + CloseAndDelete(s); + EXPECT_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s); + TF_DeleteStatus(s); +} + +void CSession::SetInputs( + std::vector> inputs) { + DeleteInputValues(); + inputs_.clear(); + for (const auto& p : inputs) { + inputs_.emplace_back(TF_Output{p.first, 0}); + input_values_.emplace_back(p.second); + } +} + +void CSession::SetOutputs(std::initializer_list outputs) { + ResetOutputValues(); + outputs_.clear(); + for (TF_Operation* o : outputs) { + outputs_.emplace_back(TF_Output{o, 0}); + } + output_values_.resize(outputs_.size()); +} + +void CSession::SetOutputs(const std::vector& outputs) { + ResetOutputValues(); + outputs_ = outputs; + output_values_.resize(outputs_.size()); +} + +void CSession::SetTargets(std::initializer_list targets) { + targets_.clear(); + for (TF_Operation* t : targets) { + targets_.emplace_back(t); + } +} + +void CSession::Run(TF_Status* s) { + if (inputs_.size() != input_values_.size()) { + ADD_FAILURE() << "Call SetInputs() before Run()"; + return; + } + ResetOutputValues(); + output_values_.resize(outputs_.size(), nullptr); + + const TF_Output* inputs_ptr = inputs_.empty() ? nullptr : &inputs_[0]; + TF_Tensor* const* input_values_ptr = + input_values_.empty() ? nullptr : &input_values_[0]; + + const TF_Output* outputs_ptr = outputs_.empty() ? nullptr : &outputs_[0]; + TF_Tensor** output_values_ptr = + output_values_.empty() ? nullptr : &output_values_[0]; + + TF_Operation* const* targets_ptr = targets_.empty() ? nullptr : &targets_[0]; + + TF_SessionRun(session_, nullptr, inputs_ptr, input_values_ptr, inputs_.size(), + outputs_ptr, output_values_ptr, outputs_.size(), targets_ptr, + targets_.size(), nullptr, s); + + DeleteInputValues(); +} + +void CSession::CloseAndDelete(TF_Status* s) { + DeleteInputValues(); + ResetOutputValues(); + if (session_ != nullptr) { + TF_CloseSession(session_, s); + EXPECT_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s); + TF_DeleteSession(session_, s); + session_ = nullptr; + } +} + +void CSession::DeleteInputValues() { + for (size_t i = 0; i < input_values_.size(); ++i) { + TF_DeleteTensor(input_values_[i]); + } + input_values_.clear(); +} + +void CSession::ResetOutputValues() { + for (size_t i = 0; i < output_values_.size(); ++i) { + if (output_values_[i] != nullptr) TF_DeleteTensor(output_values_[i]); + } + output_values_.clear(); +} diff --git a/tensorflow/c/c_test_util.h b/tensorflow/c/c_test_util.h new file mode 100644 index 0000000000000000000000000000000000000000..a927739d462edfd25c9652cedbd0ab506991af45 --- /dev/null +++ b/tensorflow/c/c_test_util.h @@ -0,0 +1,122 @@ +/* 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 THIRD_PARTY_TENSORFLOW_C_C_TEST_UTIL_H_ +#define THIRD_PARTY_TENSORFLOW_C_C_TEST_UTIL_H_ + +#include "tensorflow/c/c_api.h" + +#include +#include "tensorflow/core/framework/attr_value.pb.h" +#include "tensorflow/core/framework/graph.pb.h" +#include "tensorflow/core/framework/node_def.pb.h" +#include "tensorflow/core/framework/types.pb.h" +#include "tensorflow/core/platform/test.h" + +using ::tensorflow::string; + +typedef std::unique_ptr + unique_tensor_ptr; + +// Create a tensor with values of type TF_INT8 provided by `values`. +TF_Tensor* Int8Tensor(const int64_t* dims, int num_dims, const char* values); + +// Create a tensor with values of type TF_INT32 provided by `values`. +TF_Tensor* Int32Tensor(const int64_t* dims, int num_dims, + const int32_t* values); + +// Create 1 dimensional tensor with values from `values` +TF_Tensor* Int32Tensor(const std::vector& values); + +TF_Tensor* Int32Tensor(int32_t v); + +TF_Operation* Placeholder(TF_Graph* graph, TF_Status* s, + const char* name = "feed"); + +TF_Operation* Const(TF_Tensor* t, TF_Graph* graph, TF_Status* s, + const char* name = "const"); + +TF_Operation* ScalarConst(int32_t v, TF_Graph* graph, TF_Status* s, + const char* name = "scalar"); + +TF_Operation* Add(TF_Operation* l, TF_Operation* r, TF_Graph* graph, + TF_Status* s, const char* name = "add"); + +TF_Operation* AddNoCheck(TF_Operation* l, TF_Operation* r, TF_Graph* graph, + TF_Status* s, const char* name = "add"); + +TF_Operation* AddWithCtrlDependency(TF_Operation* l, TF_Operation* r, + TF_Graph* graph, TF_Operation* ctrl_op, + TF_Status* s, const char* name = "add"); + +TF_Operation* Add(TF_Output l, TF_Output r, TF_Graph* graph, TF_Status* s, + const char* name = "add"); + +TF_Operation* Neg(TF_Operation* n, TF_Graph* graph, TF_Status* s); + +TF_Operation* LessThan(TF_Output l, TF_Output r, TF_Graph* graph, TF_Status* s); + +// Split `input` along the first dimention into 3 tensors +TF_Operation* Split3(TF_Operation* input, TF_Graph* graph, TF_Status* s, + const char* name = "split3"); + +bool IsPlaceholder(const tensorflow::NodeDef& node_def); + +bool IsScalarConst(const tensorflow::NodeDef& node_def, int v); + +bool IsAddN(const tensorflow::NodeDef& node_def, int n); + +bool IsNeg(const tensorflow::NodeDef& node_def, const string& input); + +bool GetGraphDef(TF_Graph* graph, tensorflow::GraphDef* graph_def); + +bool GetNodeDef(TF_Operation* oper, tensorflow::NodeDef* node_def); + +bool GetFunctionDef(TF_Function* func, tensorflow::FunctionDef* func_def); + +bool GetAttrValue(TF_Operation* oper, const char* attr_name, + tensorflow::AttrValue* attr_value, TF_Status* s); + +class CSession { + public: + CSession(TF_Graph* graph, TF_Status* s); + explicit CSession(TF_Session* session); + + ~CSession(); + + void SetInputs(std::vector> inputs); + void SetOutputs(std::initializer_list outputs); + void SetOutputs(const std::vector& outputs); + void SetTargets(std::initializer_list targets); + + void Run(TF_Status* s); + + void CloseAndDelete(TF_Status* s); + + TF_Tensor* output_tensor(int i) { return output_values_[i]; } + + private: + void DeleteInputValues(); + void ResetOutputValues(); + + TF_Session* session_; + std::vector inputs_; + std::vector input_values_; + std::vector outputs_; + std::vector output_values_; + std::vector targets_; +}; + +#endif // THIRD_PARTY_TENSORFLOW_C_C_TEST_UTIL_H_ diff --git a/tensorflow/c/eager/BUILD b/tensorflow/c/eager/BUILD new file mode 100644 index 0000000000000000000000000000000000000000..85c4e4fd93cec8bc77f11ea5e1ac0ad5ade89701 --- /dev/null +++ b/tensorflow/c/eager/BUILD @@ -0,0 +1,85 @@ +# Experimental extensions to the C API for eager execution of kernels. +licenses(["notice"]) # Apache 2.0 + +load( + "//tensorflow:tensorflow.bzl", + "tf_cc_test", + "tf_copts", + "tf_cuda_library", +) + +tf_cuda_library( + name = "c_api", + srcs = ["c_api.cc"], + hdrs = ["c_api.h"], + copts = tf_copts(), + visibility = ["//visibility:public"], + deps = select({ + "//tensorflow:android": [ + ":c_api_internal", + "//tensorflow/core:android_tensorflow_lib_lite", + ], + "//conditions:default": [ + ":runtime", + "//tensorflow/c:c_api", + "//tensorflow/c:c_api_internal", + "//tensorflow/core:core_cpu_internal", + "//tensorflow/core:framework", + "//tensorflow/core:framework_internal", + "//tensorflow/core:lib", + "//tensorflow/core:lib_internal", + "//tensorflow/core:protos_all_cc", + ], + }), +) + +tf_cc_test( + name = "c_api_test", + srcs = ["c_api_test.cc"], + deps = [ + ":c_api", + "//tensorflow/core:lib", + "//tensorflow/core:protos_all_cc", + "//tensorflow/core:test", + "//tensorflow/core:test_main", + ], +) + +tf_cuda_library( + name = "runtime", + srcs = ["runtime.cc"], + hdrs = ["runtime.h"], + copts = tf_copts(), + visibility = ["//tensorflow:internal"], + deps = select({ + "//tensorflow:android": [ + ":c_api_internal", + "//tensorflow/core:android_tensorflow_lib_lite", + ], + "//conditions:default": [ + "//tensorflow/c:c_api", + "//tensorflow/core:core_cpu", + "//tensorflow/core:core_cpu_internal", + "//tensorflow/core:framework", + "//tensorflow/core:framework_internal", + "//tensorflow/core:lib", + "//tensorflow/core:lib_internal", + "//tensorflow/core:protos_all_cc", + ], + }), +) + +tf_cc_test( + name = "runtime_test", + srcs = ["runtime_test.cc"], + deps = [ + ":runtime", + "//tensorflow/cc:cc_ops", + "//tensorflow/cc:client_session", + "//tensorflow/cc:ops", + "//tensorflow/cc:scope", + "//tensorflow/core:core_cpu_internal", + "//tensorflow/core:test", + "//tensorflow/core:test_main", + ], +) diff --git a/tensorflow/c/eager/c_api.cc b/tensorflow/c/eager/c_api.cc new file mode 100644 index 0000000000000000000000000000000000000000..e70539ceefa1e9b3b70be0ac2dd8acb431ed8caa --- /dev/null +++ b/tensorflow/c/eager/c_api.cc @@ -0,0 +1,552 @@ +/* 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.h" + +#include +#include +#include +#include +#include + +#include "tensorflow/c/c_api.h" +#include "tensorflow/c/c_api_internal.h" +#include "tensorflow/c/eager/runtime.h" +#include "tensorflow/core/common_runtime/device_factory.h" +#include "tensorflow/core/common_runtime/device_mgr.h" +#include "tensorflow/core/common_runtime/function.h" +#include "tensorflow/core/common_runtime/rendezvous_mgr.h" +#include "tensorflow/core/framework/rendezvous.h" +#include "tensorflow/core/framework/tensor_shape.pb.h" +#include "tensorflow/core/framework/types.h" +#include "tensorflow/core/lib/core/refcount.h" +#include "tensorflow/core/lib/gtl/map_util.h" +#include "tensorflow/core/lib/gtl/stl_util.h" +#include "tensorflow/core/platform/mutex.h" +#include "tensorflow/core/platform/thread_annotations.h" +#include "tensorflow/core/public/version.h" + +using tensorflow::int64; +using tensorflow::string; + +namespace { +bool IsCPU(tensorflow::Device* d) { + return d == nullptr || d->tensorflow_gpu_device_info() == nullptr; +} + +string DeviceName(tensorflow::Device* d) { + return (d == nullptr) ? "cpu:0" : d->name(); +} +} // namespace + +struct TFE_Context { + explicit TFE_Context(TF_Session* s) : session(s) {} + + // TFE_Context is an extension of TF_Session. And TF_Session needs a TF_Graph. + TF_Session* session; + tensorflow::Rendezvous* rendezvous; + + tensorflow::mutex functions_mu; + tensorflow::FunctionLibraryDefinition func_lib_def GUARDED_BY(functions_mu){ + tensorflow::OpRegistry::Global(), {}}; + + // One FunctionLibraryRuntime per device. + // func_libs[i] is the FunctionLibraryRuntime corresponding to + // session->devices[i]. + std::unique_ptr pflr; + + std::unordered_map + kernel_cache; + + tensorflow::FunctionLibraryRuntime* func_lib(tensorflow::Device* d) { + return pflr->GetFLR(d->name()); + } + + const std::vector& devices() { return session->devices; } +}; + +struct TFE_TensorHandle { + TFE_TensorHandle(const tensorflow::Tensor& t, tensorflow::Device* d) + : t(t), d(d) {} + + tensorflow::Tensor t; + // TODO(ashankar): d == nullptr iff local CPU + // This was expedient, but perhaps worth revisiting ('d' should always be a + // valid pointer?) + // This can be done if TFE_NewOp() and the TFE_TensorHandle constructors are + // provided with the appropriate TFE_Context. + // + // TODO(ashankar): Reference count TFE_Context to ensure that 'd' of a + // TFE_TensorHandle does not outlive the TFE_Context from which it came? + tensorflow::Device* d; +}; + +struct TFE_Op { + TFE_Op(TFE_Context* ctx, const char* op, const tensorflow::AttrTypeMap* t) + : ctx(ctx), name(op), attrs(op), attr_types(t), device(nullptr) {} + + bool const is_function() const { return attr_types == nullptr; } + + TFE_Context* ctx; // Must outlive the TFE_Op. + const char* name; + tensorflow::AttrBuilder attrs; + const tensorflow::AttrTypeMap* attr_types; + std::vector inputs; + std::vector input_devices; + tensorflow::Device* device; +}; + +extern "C" { + +TFE_Context* TFE_NewContext(const TF_SessionOptions* opts, TF_Status* status) { + TF_Graph* graph = TF_NewGraph(); + TF_Session* session = TF_NewSession(graph, opts, status); + if (status->status.ok()) { + if (session->device_mgr == nullptr || session->devices.empty()) { + status->status = tensorflow::errors::InvalidArgument( + "Provided TF_SessionOptions are not compatible with eager execution " + "(perhaps the TF_SessionOptions alluded to session execution in a " + "remote address space?)"); + } + } + if (!status->status.ok()) { + TF_DeleteGraph(graph); + return nullptr; + } + + TFE_Context* ret = new TFE_Context(session); + ret->pflr.reset(new tensorflow::ProcessFunctionLibraryRuntime( + ret->session->device_mgr, opts->options.env, TF_GRAPH_DEF_VERSION, + &ret->func_lib_def, {})); + ret->rendezvous = + new tensorflow::IntraProcessRendezvous(ret->session->device_mgr); + + return ret; +} + +void TFE_DeleteContext(TFE_Context* ctx, TF_Status* status) { + status->status = tensorflow::Status::OK(); + tensorflow::gtl::STLDeleteValues(&ctx->kernel_cache); + TF_Graph* graph = ctx->session->graph; + TF_DeleteSession(ctx->session, status); + TF_DeleteGraph(graph); + ctx->rendezvous->Unref(); + delete ctx; +} + +TF_DeviceList* TFE_ContextListDevices(TFE_Context* ctx, TF_Status* status) { + return TF_SessionListDevices(ctx->session, status); +} + +TFE_TensorHandle* TFE_NewTensorHandle(TF_Tensor* t, TF_Status* status) { + tensorflow::Tensor tensor; + status->status = tensorflow::TF_TensorToTensor(t, &tensor); + if (!status->status.ok()) return nullptr; + return new TFE_TensorHandle(tensor, nullptr); +} + +void TFE_DeleteTensorHandle(TFE_TensorHandle* h) { delete h; } + +TF_DataType TFE_TensorHandleDataType(TFE_TensorHandle* h) { + return static_cast(h->t.dtype()); +} + +int TFE_TensorHandleNumDims(TFE_TensorHandle* h) { return h->t.dims(); } + +int64_t TFE_TensorHandleDim(TFE_TensorHandle* h, int dim_index) { + return h->t.dim_size(dim_index); +} + +const char* TFE_TensorHandleDeviceName(TFE_TensorHandle* h) { + // This might be a bit confusing as a tensor on CPU can sometimes return + // "CPU:0" and sometimes "/job:localhost/replica:0/task:0/cpu:0". + // TODO(ashankar): Figure out which one would be nicer. + return (h->d == nullptr) ? "CPU:0" : h->d->name().c_str(); +} + +TF_Tensor* TFE_TensorHandleResolve(TFE_TensorHandle* h, TF_Status* status) { + if (!IsCPU(h->d)) { + TF_SetStatus(status, TF_UNIMPLEMENTED, + tensorflow::strings::StrCat( + "TFE_TensorHandle can be resolved iff it is on CPU (this " + "handle is on ", + h->d->name(), + "). Consider using TFE_TensorHandleCopyToDevice to get a " + "copy of the tensor on CPU") + .c_str()); + return nullptr; + } + return tensorflow::TF_TensorFromTensor(h->t, status); +} + +TFE_TensorHandle* TFE_TensorHandleCopyToDevice(TFE_TensorHandle* h, + TFE_Context* ctx, + const char* device_name, + TF_Status* status) { + tensorflow::Device* dstd = ctx->devices()[0]; + if (device_name != nullptr && strlen(device_name) > 0) { + status->status = ctx->session->device_mgr->LookupDevice(device_name, &dstd); + if (!status->status.ok()) return nullptr; + } + + tensorflow::Device* srcd = h->d == nullptr ? ctx->devices()[0] : h->d; + bool is_same_device = + (srcd == dstd) || (DeviceName(srcd) == DeviceName(dstd)); + const bool dst_cpu = IsCPU(dstd); + if (is_same_device) { + return new TFE_TensorHandle(h->t, dst_cpu ? nullptr : dstd); + } + const bool src_cpu = IsCPU(srcd); + if (src_cpu == dst_cpu) { + TF_SetStatus( + status, TF_INVALID_ARGUMENT, + tensorflow::strings::StrCat( + "TFE_TensorHandleCopyToDevice requires either the source " + "TFE_TensorHandle be on or the destination device be on CPU " + "or be the same (they are ", + DeviceName(srcd), " and ", DeviceName(dstd), " in this call)") + .c_str()); + return nullptr; + } + tensorflow::Tensor* src = &(h->t); + if (src_cpu) { + tensorflow::Tensor dst( + dstd->GetAllocator(tensorflow::AllocatorAttributes()), src->dtype(), + src->shape()); + tensorflow::Notification n; + dstd->tensorflow_gpu_device_info()->default_context->CopyCPUTensorToDevice( + src, dstd, &dst, [status, &n](const tensorflow::Status& s) { + status->status = s; + n.Notify(); + }); + n.WaitForNotification(); + return (TF_GetCode(status) == TF_OK) ? new TFE_TensorHandle(dst, dstd) + : nullptr; + } + CHECK(dst_cpu); + tensorflow::Tensor dst(src->dtype(), src->shape()); + tensorflow::Notification n; + // TODO(ashankar): The Sync() call below may be more aggressive than + // necessary. It is based on knowledge of implementation details - that + // GPU devices are implemented using 3 streams - one for host->device copies, + // one for device->host copies and one for sending operations to the GPU. + // With that setup, Sync()ing across all 3 streams should be sufficient + // but more than necessary (since it waits for operations that might have + // nothing to do with this tensor to complete). + status->status = srcd->Sync(); + if (!status->status.ok()) return nullptr; + srcd->tensorflow_gpu_device_info()->default_context->CopyDeviceTensorToCPU( + src, "IGNORE_MY_TENSOR_NAME", srcd, &dst, + [status, &n](const tensorflow::Status& s) { + status->status = s; + n.Notify(); + }); + n.WaitForNotification(); + return (TF_GetCode(status) == TF_OK) ? new TFE_TensorHandle(dst, nullptr) + : nullptr; +} + +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) { + tensorflow::mutex_lock l(ctx->functions_mu); + if (ctx->func_lib_def.Find(name) != nullptr) { + status->status = tensorflow::Status::OK(); + return new TFE_Op(ctx, name, nullptr); + } + } + return nullptr; +} + +void TFE_DeleteOp(TFE_Op* op) { delete op; } + +static void TFE_OpSetDeviceHelper(TFE_Op* op, tensorflow::Device* device, + TF_Status* status) { + // Questionable heuristic: Place the op on the same device as the first input + // placed outside of host memory? + if (IsCPU(op->device) && !IsCPU(device)) { + op->device = device; + } +} + +void TFE_OpSetDevice(TFE_Op* op, TFE_Context* ctx, const char* device_name, + TF_Status* status) { + tensorflow::Device* d = nullptr; + if (device_name != nullptr && strlen(device_name) > 0) { + status->status = ctx->session->device_mgr->LookupDevice(device_name, &d); + if (!status->status.ok()) return; + } + TFE_OpSetDeviceHelper(op, d, status); +} + +void TFE_OpAddInput(TFE_Op* op, TFE_TensorHandle* h, TF_Status* status) { + TFE_OpSetDeviceHelper(op, h->d, status); + if (!status->status.ok()) return; + op->inputs.push_back(h->t); + op->input_devices.push_back(h->d); + op->attrs.NumInputs(op->inputs.size()); +} + +TF_AttrType TFE_OpGetAttrType(TFE_Op* op, const char* attr_name, + unsigned char* is_list, TF_Status* status) { + TF_AttrType ret; + if (op->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->attr_types, attr_name, &ret, is_list); + return ret; +} + +void TFE_OpSetAttrString(TFE_Op* op, const char* attr_name, const char* value) { + op->attrs.Set(attr_name, value); +} + +void TFE_OpSetAttrInt(TFE_Op* op, const char* attr_name, int64_t value) { + op->attrs.Set(attr_name, static_cast(value)); +} + +void TFE_OpSetAttrFloat(TFE_Op* op, const char* attr_name, float value) { + op->attrs.Set(attr_name, value); +} + +void TFE_OpSetAttrBool(TFE_Op* op, const char* attr_name, unsigned char value) { + op->attrs.Set(attr_name, (value == 0) ? false : true); +} + +void TFE_OpSetAttrType(TFE_Op* op, const char* attr_name, TF_DataType value) { + op->attrs.Set(attr_name, static_cast(value)); +} + +void TFE_OpSetAttrShape(TFE_Op* op, const char* attr_name, const int64_t* dims, + const int num_dims, TF_Status* out_status) { + if (num_dims > tensorflow::TensorShape::MaxDimensions()) { + TF_SetStatus(out_status, TF_INVALID_ARGUMENT, + tensorflow::strings::StrCat( + "Value specified for `", attr_name, "` has ", num_dims, + " dimensions which is over the limit of ", + tensorflow::TensorShape::MaxDimensions(), ".") + .c_str()); + return; + } + tensorflow::TensorShapeProto proto; + if (num_dims < 0) { + proto.set_unknown_rank(true); + } else { + for (int d = 0; d < num_dims; ++d) { + proto.add_dim()->set_size(dims[d]); + } + } + op->attrs.Set(attr_name, proto); +} + +#define TFE_OP_SET_ATTR_LIST(fn, type) \ + void fn(TFE_Op* op, const char* attr_name, const type* values, \ + int num_values) { \ + op->attrs.Set(attr_name, tensorflow::gtl::ArraySlice( \ + values, num_values)); \ + } +TFE_OP_SET_ATTR_LIST(TFE_OpSetAttrStringList, char*) +TFE_OP_SET_ATTR_LIST(TFE_OpSetAttrFloatList, float) +#undef TFE_OP_SET_ATTR_LIST + +void TFE_OpSetAttrIntList(TFE_Op* op, const char* attr_name, + const int64_t* values, int num_values) { + op->attrs.Set(attr_name, + tensorflow::gtl::ArraySlice( + reinterpret_cast(values), num_values)); +} + +void TFE_OpSetAttrTypeList(TFE_Op* op, const char* attr_name, + const TF_DataType* values, int num_values) { + op->attrs.Set( + attr_name, + tensorflow::gtl::ArraySlice( + reinterpret_cast(values), num_values)); +} + +void TFE_OpSetAttrBoolList(TFE_Op* op, const char* attr_name, + const unsigned char* values, int num_values) { + std::unique_ptr b(new bool[num_values]); + for (int i = 0; i < num_values; ++i) { + b[i] = values[i]; + } + op->attrs.Set(attr_name, + tensorflow::gtl::ArraySlice(b.get(), num_values)); +} + +void TFE_OpSetAttrShapeList(TFE_Op* op, const char* attr_name, + const int64_t** dims, const int* num_dims, + int num_values, TF_Status* out_status) { + std::unique_ptr proto( + new tensorflow::TensorShapeProto[num_values]); + for (int i = 0; i < num_values; ++i) { + const auto num_dims_i = num_dims[i]; + + if (num_dims_i > tensorflow::TensorShape::MaxDimensions()) { + TF_SetStatus(out_status, TF_INVALID_ARGUMENT, + tensorflow::strings::StrCat( + "Value specified for `", attr_name, "` has ", num_dims_i, + " dimensions which is over the limit of ", + tensorflow::TensorShape::MaxDimensions(), ".") + .c_str()); + return; + } + if (num_dims_i < 0) { + proto[i].set_unknown_rank(true); + } else { + const int64_t* dims_i = dims[i]; + auto proto_i = &proto[i]; + for (int d = 0; d < num_dims_i; ++d) { + proto_i->add_dim()->set_size(dims_i[d]); + } + } + } + op->attrs.Set(attr_name, + tensorflow::gtl::ArraySlice( + proto.get(), num_values)); +} + +namespace { + +tensorflow::Status ValidateInputTypeAndPlacement( + tensorflow::Device* host_device, tensorflow::Device* op_device, TFE_Op* op, + const tensorflow::OpKernel* kernel) { + const tensorflow::MemoryTypeVector& memtypes = kernel->input_memory_types(); + if (memtypes.size() != op->inputs.size()) { + return tensorflow::errors::InvalidArgument( + "expected ", memtypes.size(), " inputs, got ", op->inputs.size()); + } + for (int i = 0; i < op->inputs.size(); ++i) { + const tensorflow::Device* expected_device = + memtypes[i] == tensorflow::HOST_MEMORY ? host_device : op_device; + const tensorflow::Device* actual_device = + op->input_devices[i] == nullptr ? host_device : op->input_devices[i]; + if (expected_device != actual_device) { + return tensorflow::errors::InvalidArgument( + "cannot compute ", op->name, " as input #", i, + " was expected to be on ", expected_device->name(), + " but is actually on ", actual_device->name(), + " (operation running on ", op_device->name(), ")"); + } + if (op->inputs[i].dtype() != kernel->input_type(i)) { + return tensorflow::errors::InvalidArgument( + "cannot compute ", op->name, " as input #", i, + " was expected to be a ", + tensorflow::DataTypeString(kernel->input_type(i)), + " tensor but is a ", + tensorflow::DataTypeString(op->inputs[i].dtype()), " tensor"); + } + } + return tensorflow::Status::OK(); +} +} // namespace + +void TFE_Execute(TFE_Op* op, TFE_TensorHandle** retvals, int* num_retvals, + TF_Status* status) { + TFE_Context* ctx = op->ctx; + // TODO(ashankar): ASSUMPTION: ctx->devices()[0] is always CPU + tensorflow::Device* device = + (op->device == nullptr) ? ctx->devices()[0] : op->device; + std::vector outputs(1); + const tensorflow::MemoryTypeVector* output_memory_types = nullptr; + tensorflow::Fprint128 cache_key = op->attrs.CacheKey(device->name()); + tensorflow::KernelAndDevice* kernel = + tensorflow::gtl::FindPtrOrNull(ctx->kernel_cache, cache_key); + if (kernel == nullptr) { + const tensorflow::NodeDef& ndef = op->attrs.BuildNodeDef(); + kernel = new tensorflow::KernelAndDevice(ctx->rendezvous); + if (!op->is_function()) { + status->status = + tensorflow::KernelAndDevice::InitOp(device, ndef, kernel); + } else { + // Knowledge of the implementation of InitFn (and in-turn + // FunctionLibraryRuntime::CreateKernel) tells us that ctx->func_lib_def + // will be accessed, so grab on to the lock. + // See WARNING comment below - would be nice to rework to avoid this + // subtlety. + tensorflow::mutex_lock l(ctx->functions_mu); + status->status = tensorflow::KernelAndDevice::InitFn( + ndef, ctx->func_lib(device), kernel); + } + if (!status->status.ok()) { + return; + } + tensorflow::gtl::InsertOrUpdate(&(ctx->kernel_cache), cache_key, kernel); + } + status->status = ValidateInputTypeAndPlacement(ctx->devices()[0], device, op, + kernel->kernel()); + output_memory_types = &kernel->kernel()->output_memory_types(); + if (!status->status.ok()) { + return; + } + // WARNING: kernel->Run utilizes the FunctionLibraryRuntime + // (ctx->func_lib(device)), which in turn holds a pointer to func_lib_def, + // which is GUARDED_BY(ctx->functions_mu). But knowledge of the implementation + // of FunctionLibraryRuntime tells use that func_lib_def is not accessed by + // FunctionLibraryRuntime::Run(), so there is no thread-safety concern here. + // This is quite subtle. Re-work things to make this better? (Would it make + // sense for FunctionLibraryRuntime to ensure thread-safe access to + // FunctionLibraryDefinition?). + status->status = kernel->Run(&op->inputs, &outputs); + if (!status->status.ok()) return; + *num_retvals = std::min(*num_retvals, outputs.size()); + for (int i = 0; i < *num_retvals; ++i) { + tensorflow::Device* d = IsCPU(device) ? nullptr : device; + if (d != nullptr && output_memory_types != nullptr && + (*output_memory_types)[i] == tensorflow::HOST_MEMORY) { + d = nullptr; + } + retvals[i] = new TFE_TensorHandle(outputs[i], d); + } +} + +void TFE_ContextAddFunctionDef(TFE_Context* ctx, + const char* serialized_function_def, size_t size, + TF_Status* status) { + tensorflow::FunctionDef function_def; + if (!function_def.ParseFromArray(serialized_function_def, size)) { + status->status = + tensorflow::errors::InvalidArgument("Invalid FunctionDef proto"); + return; + } + tensorflow::mutex_lock l(ctx->functions_mu); + status->status = ctx->func_lib_def.AddFunctionDef(function_def); +} + +} // extern "C" + +TFE_TensorHandle* TFE_NewTensorHandle(const tensorflow::Tensor& t) { + return new TFE_TensorHandle(t, nullptr); +} + +const tensorflow::Tensor* TFE_TensorHandleUnderlyingTensorInHostMemory( + TFE_TensorHandle* h, TF_Status* status) { + if (h->d != nullptr) { + status->status = tensorflow::errors::FailedPrecondition( + "TFE_TensorHandle is placed in device (not host) memory. Cannot return " + "a tensorflow::Tensor"); + return nullptr; + } + return &h->t; +} diff --git a/tensorflow/c/eager/c_api.h b/tensorflow/c/eager/c_api.h new file mode 100644 index 0000000000000000000000000000000000000000..88a0dd343f2efd70c3f0da693114d9ceb99125e9 --- /dev/null +++ b/tensorflow/c/eager/c_api.h @@ -0,0 +1,159 @@ +/* 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_C_EAGER_C_API_H_ +#define TENSORFLOW_C_EAGER_C_API_H_ + +// C API extensions to experiment with eager execution of kernels. + +#include "tensorflow/c/c_api.h" + +#ifdef __cplusplus +extern "C" { +#endif + +// "Context" under which operations/functions are executed. It encapsulates +// things like the available devices, resource manager etc. +// +// TODO(ashankar): Merge with TF_Session? +typedef struct TFE_Context TFE_Context; + +extern TFE_Context* TFE_NewContext(const TF_SessionOptions* opts, + TF_Status* status); +extern void TFE_DeleteContext(TFE_Context* ctx, TF_Status* status); +extern TF_DeviceList* TFE_ContextListDevices(TFE_Context* ctx, + TF_Status* status); + +// A handle to a tensor on a device. +// +// Like a TF_Tensor, a TFE_TensorHandle refers to a tensor with a value, shape, +// type etc. Unlike a TF_Tensor, a TFE_TensorHandle may refer to such tensors +// placed in memory of different devices or remote address spaces. +typedef struct TFE_TensorHandle TFE_TensorHandle; + +extern TFE_TensorHandle* TFE_NewTensorHandle(TF_Tensor* t, TF_Status* status); +extern void TFE_DeleteTensorHandle(TFE_TensorHandle* h); +extern TF_DataType TFE_TensorHandleDataType(TFE_TensorHandle* h); +extern int TFE_TensorHandleNumDims(TFE_TensorHandle* h); +extern int64_t TFE_TensorHandleDim(TFE_TensorHandle* h, int dim_index); +extern const char* TFE_TensorHandleDeviceName(TFE_TensorHandle* h); +extern TF_Tensor* TFE_TensorHandleResolve(TFE_TensorHandle* h, + TF_Status* status); + +// Create a new TFE_TensorHandle with the same contents as 'h' but placed +// in the memory of the device name 'device_name'. +// If source and destination are the same device, then this creates a new handle +// that shares the underlying buffer. Otherwise, it currently requires at least +// one of the source or destination devices to be CPU (i.e., for the source or +// destination tensor to be placed in host memory). +extern TFE_TensorHandle* TFE_TensorHandleCopyToDevice(TFE_TensorHandle* h, + TFE_Context* ctx, + const char* device_name, + TF_Status* status); + +// Description of the TensorFlow op to execute. +// +// Assumes that the provided 'ctx' outlives the returned TFE_Op, i.e., +// TFE_DeleteOp() is called before TFE_DeleteContext(). +// +// Very similar to TF_OperationDescription with some differences: +// (1) TF_Output or TFE_TensorHandle* as arguments to TF_AddInput, +// TF_AddInputList +// (2) TF_ColocateWith, TF_AddControlInput etc. do not make sense. +// (3) Implementation detail: Avoid use of NodeBuilder/NodeDefBuilder since +// the additional sanity checks there seem unnecessary; +typedef struct TFE_Op TFE_Op; + +extern TFE_Op* TFE_NewOp(TFE_Context* ctx, const char* op_or_function_name, + TF_Status* status); +extern void TFE_DeleteOp(TFE_Op* op); + +// TODO(ashankar): TFE_OpSetDevice and TFE_Execute should not have a TFE_Context +// parameter. Instead, the TFE_Context should be captured when creating the +// TFE_Op. +extern void TFE_OpSetDevice(TFE_Op* op, TFE_Context* ctx, + const char* device_name, TF_Status* status); + +extern void TFE_OpAddInput(TFE_Op* op, TFE_TensorHandle* h, TF_Status* status); + +extern TF_AttrType TFE_OpGetAttrType(TFE_Op* op, const char* attr_name, + unsigned char* is_list, TF_Status* status); + +extern void TFE_OpSetAttrString(TFE_Op* op, const char* attr_name, + const char* value); +extern void TFE_OpSetAttrInt(TFE_Op* op, const char* attr_name, int64_t value); +extern void TFE_OpSetAttrFloat(TFE_Op* op, const char* attr_name, float value); +extern void TFE_OpSetAttrBool(TFE_Op* op, const char* attr_name, + unsigned char value); +extern void TFE_OpSetAttrType(TFE_Op* op, const char* attr_name, + TF_DataType value); +// If the number of dimensions is unknown, `num_dims` must be set to +// -1 and `dims` can be null. If a dimension is unknown, the +// corresponding entry in the `dims` array must be -1. +extern void TFE_OpSetAttrShape(TFE_Op* op, const char* attr_name, + const int64_t* dims, const int num_dims, + TF_Status* out_status); + +extern void TFE_OpSetAttrStringList(TFE_Op* op, const char* attr_name, + const char** value, int num_values); +extern void TFE_OpSetAttrIntList(TFE_Op* op, const char* attr_name, + const int64_t* values, int num_values); +extern void TFE_OpSetAttrFloatList(TFE_Op* op, const char* attr_name, + const float* values, int num_values); +extern void TFE_OpSetAttrBoolList(TFE_Op* op, const char* attr_name, + const unsigned char* values, int num_values); +extern void TFE_OpSetAttrTypeList(TFE_Op* op, const char* attr_name, + const TF_DataType* values, int num_values); +extern void TFE_OpSetAttrShapeList(TFE_Op* op, const char* attr_name, + const int64_t** dims, const int* num_dims, + int num_values, TF_Status* out_status); + +// Execute the operation defined by 'op' and return handles to computed +// tensors in 'retvals'. +// +// 'retvals' must point to a pre-allocated array of TFE_TensorHandle* +// and '*num_retvals' should be set to the size of this array. +// +// On return, 'num_retvals' will be set to the actual number of outputs +// returned by the operation. +extern void TFE_Execute(TFE_Op* op, TFE_TensorHandle** retvals, + int* num_retvals, TF_Status* status); + +// Add a function (serialized FunctionDef protocol buffer) to ctx so +// that it can be invoked using TFE_Execute. +extern void TFE_ContextAddFunctionDef(TFE_Context* ctx, + const char* serialized_function_def, + size_t size, TF_Status* status); + +#ifdef __cplusplus +} /* end extern "C" */ +#endif + +#ifdef __cplusplus +// A workaround to ease conversion to and from numpy objects and +// TFE_TensorHandle's. +// +// TODO(ashankar): Figure out an alternative scheme that precludes the need for +// these API-boundary breaking methods. +namespace tensorflow { +class Tensor; +} // namespace tensorflow + +const tensorflow::Tensor* TFE_TensorHandleUnderlyingTensorInHostMemory( + TFE_TensorHandle* h, TF_Status* status); +TFE_TensorHandle* TFE_NewTensorHandle(const tensorflow::Tensor& t); +#endif + +#endif // TENSORFLOW_C_EAGER_C_API_H_ diff --git a/tensorflow/c/eager/c_api_test.cc b/tensorflow/c/eager/c_api_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..d19583a3abed0b080793e0ed7c34295b6f105de5 --- /dev/null +++ b/tensorflow/c/eager/c_api_test.cc @@ -0,0 +1,483 @@ +/* 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.h" + +#include +#include "tensorflow/core/framework/function.pb.h" +#include "tensorflow/core/lib/strings/strcat.h" +#include "tensorflow/core/platform/logging.h" +#include "tensorflow/core/platform/macros.h" +#include "tensorflow/core/platform/protobuf.h" +#include "tensorflow/core/platform/test.h" +#include "tensorflow/core/platform/test_benchmark.h" + +using tensorflow::string; + +namespace { + +TFE_TensorHandle* TestMatrixTensorHandle() { + int64_t dims[] = {2, 2}; + float data[] = {1.0f, 2.0f, 3.0f, 4.0f}; + TF_Tensor* t = TF_AllocateTensor( + TF_FLOAT, &dims[0], sizeof(dims) / sizeof(int64_t), sizeof(data)); + memcpy(TF_TensorData(t), &data[0], TF_TensorByteSize(t)); + TF_Status* status = TF_NewStatus(); + TFE_TensorHandle* th = TFE_NewTensorHandle(t, status); + CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); + TF_DeleteTensor(t); + return th; +} + +TFE_Op* MatMulOp(TFE_Context* ctx, TFE_TensorHandle* a, TFE_TensorHandle* b) { + TF_Status* status = TF_NewStatus(); + + TFE_Op* op = TFE_NewOp(ctx, "MatMul", 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); + TFE_OpAddInput(op, b, status); + CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); + TF_DeleteStatus(status); + TFE_OpSetAttrBool(op, "transpose_a", 0); + TFE_OpSetAttrBool(op, "transpose_b", 0); + TFE_OpSetAttrType(op, "T", TFE_TensorHandleDataType(a)); + + return op; +} + +void BM_InitOp(int iters) { + tensorflow::testing::StopTiming(); + TF_Status* status = TF_NewStatus(); + TF_SessionOptions* opts = TF_NewSessionOptions(); + TFE_Context* ctx = TFE_NewContext(opts, status); + CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); + TF_DeleteSessionOptions(opts); + + TFE_TensorHandle* m = TestMatrixTensorHandle(); + tensorflow::testing::StartTiming(); + for (int i = 0; i < iters; ++i) { + TFE_Op* matmul = MatMulOp(ctx, m, m); + TFE_DeleteOp(matmul); + } + tensorflow::testing::StopTiming(); + TFE_DeleteTensorHandle(m); + TFE_DeleteContext(ctx, status); + CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); + TF_DeleteStatus(status); +} +BENCHMARK(BM_InitOp); + +void BM_Execute(int iters) { + tensorflow::testing::StopTiming(); + TF_Status* status = TF_NewStatus(); + TF_SessionOptions* opts = TF_NewSessionOptions(); + TFE_Context* ctx = TFE_NewContext(opts, status); + CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); + TF_DeleteSessionOptions(opts); + + TFE_TensorHandle* m = TestMatrixTensorHandle(); + TFE_Op* matmul = MatMulOp(ctx, m, m); + TFE_TensorHandle* retvals[1]; + int num_retvals = 1; + tensorflow::testing::StartTiming(); + for (int i = 0; i < iters; ++i) { + TFE_Execute(matmul, &retvals[0], &num_retvals, status); + CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); + } + tensorflow::testing::StopTiming(); + TFE_DeleteOp(matmul); + TFE_DeleteTensorHandle(m); + TFE_DeleteContext(ctx, status); + CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); + TF_DeleteStatus(status); +} +BENCHMARK(BM_Execute); + +TEST(CAPI, Context) { + TF_Status* status = TF_NewStatus(); + TF_SessionOptions* opts = TF_NewSessionOptions(); + TFE_Context* ctx = TFE_NewContext(opts, status); + TF_DeleteSessionOptions(opts); + + TF_DeviceList* devices = TFE_ContextListDevices(ctx, status); + EXPECT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); + + TFE_DeleteContext(ctx, status); + EXPECT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); + + const int num_devices = TF_DeviceListCount(devices); + EXPECT_GE(num_devices, 1) << "At least one CPU device should exist"; + for (int i = 0; i < num_devices; ++i) { + EXPECT_NE("", TF_DeviceListName(devices, i, status)) << i; + EXPECT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); + } + TF_DeleteDeviceList(devices); + TF_DeleteStatus(status); +} + +TEST(CAPI, TensorHandle) { + TFE_TensorHandle* h = TestMatrixTensorHandle(); + EXPECT_EQ(TF_FLOAT, TFE_TensorHandleDataType(h)); + + std::unique_ptr status( + TF_NewStatus(), TF_DeleteStatus); + TF_Tensor* t = TFE_TensorHandleResolve(h, status.get()); + ASSERT_EQ(16, TF_TensorByteSize(t)); + float data[4] = {0}; + memcpy(&data[0], TF_TensorData(t), TF_TensorByteSize(t)); + EXPECT_EQ(1.0, data[0]); + EXPECT_EQ(2.0, data[1]); + EXPECT_EQ(3.0, data[2]); + EXPECT_EQ(4.0, data[3]); + TF_DeleteTensor(t); + TFE_DeleteTensorHandle(h); +} + +TEST(CAPI, TensorHandleCopyBetweenDevices) { + std::unique_ptr status( + TF_NewStatus(), TF_DeleteStatus); + TF_SessionOptions* opts = TF_NewSessionOptions(); + TFE_Context* ctx = TFE_NewContext(opts, status.get()); + TF_DeleteSessionOptions(opts); + ASSERT_EQ(TF_OK, TF_GetCode(status.get())) << TF_Message(status.get()); + + TFE_TensorHandle* hcpu = TestMatrixTensorHandle(); + TF_Tensor* t = TFE_TensorHandleResolve(hcpu, status.get()); + ASSERT_EQ(TF_OK, TF_GetCode(status.get())) << TF_Message(status.get()); + + 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); + + const char* kCPUDevice = "CPU:0"; + for (int i = 0; i < num_devices; ++i) { + const string name(TF_DeviceListName(devices, i, status.get())); + if (TF_GetCode(status.get()) != TF_OK) { + ADD_FAILURE() << i << " -- " << TF_Message(status.get()); + continue; + } + auto tag = tensorflow::strings::StrCat("Device #", i, " (", name, ")"); + // Copy to device + TFE_TensorHandle* hdevice = + TFE_TensorHandleCopyToDevice(hcpu, ctx, name.c_str(), status.get()); + if (TF_GetCode(status.get()) != TF_OK) { + ADD_FAILURE() << tag << " -- " << TF_Message(status.get()); + continue; + } + // Copy from device to the same device. + TFE_TensorHandle* hdevice2 = + TFE_TensorHandleCopyToDevice(hdevice, ctx, name.c_str(), status.get()); + if (TF_GetCode(status.get()) != TF_OK) { + ADD_FAILURE() << tag << " -- " << TF_Message(status.get()); + continue; + } + TFE_DeleteTensorHandle(hdevice); + // Copy back to CPU + TFE_TensorHandle* hcopy = + TFE_TensorHandleCopyToDevice(hdevice2, ctx, kCPUDevice, status.get()); + if (TF_GetCode(status.get()) != TF_OK) { + ADD_FAILURE() << tag << " -- " << TF_Message(status.get()); + continue; + } + TFE_DeleteTensorHandle(hdevice2); + + // Ensure that the contents are the same! + TF_Tensor* tcopy = TFE_TensorHandleResolve(hcopy, status.get()); + TFE_DeleteTensorHandle(hcopy); + if (TF_GetCode(status.get()) != TF_OK) { + ADD_FAILURE() << tag; + continue; + } + EXPECT_EQ(TF_TensorByteSize(t), TF_TensorByteSize(tcopy)) << tag; + EXPECT_EQ( + 0, memcmp(TF_TensorData(t), TF_TensorData(tcopy), TF_TensorByteSize(t))) + << tag; + TF_DeleteTensor(tcopy); + } + + TF_DeleteDeviceList(devices); + TF_DeleteTensor(t); + TFE_DeleteTensorHandle(hcpu); + TFE_DeleteContext(ctx, status.get()); + EXPECT_EQ(TF_OK, TF_GetCode(status.get())) << TF_Message(status.get()); +} + +TEST(CAPI, Execute) { + TF_Status* status = TF_NewStatus(); + TF_SessionOptions* opts = TF_NewSessionOptions(); + TFE_Context* ctx = TFE_NewContext(opts, status); + CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); + TF_DeleteSessionOptions(opts); + + TFE_TensorHandle* m = TestMatrixTensorHandle(); + TFE_Op* matmul = MatMulOp(ctx, m, m); + TFE_TensorHandle* retvals[2] = {nullptr}; + int num_retvals = 2; // Should be reduced to 1 by the TFE_Execute call. + TFE_Execute(matmul, &retvals[0], &num_retvals, status); + EXPECT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); + TFE_DeleteOp(matmul); + TFE_DeleteTensorHandle(m); + TFE_DeleteContext(ctx, status); + ASSERT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); + ASSERT_EQ(1, num_retvals); + + TF_Tensor* t = TFE_TensorHandleResolve(retvals[0], status); + TFE_DeleteTensorHandle(retvals[0]); + 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); +} + +string MatMulFunction() { + tensorflow::FunctionDef def; + CHECK(tensorflow::protobuf::TextFormat::ParseFromString( + " signature {" + " name: 'MatMulFunction'" + " input_arg {" + " name: 'a'" + " type: DT_FLOAT" + " }" + " output_arg {" + " name: 'm'" + " type: DT_FLOAT" + " }" + " }" + " node_def {" + " name: 'matmul'" + " op: 'MatMul'" + " input: 'a'" + " input: 'a'" + " attr {" + " key: 'T'" + " value {" + " type: DT_FLOAT" + " }" + " }" + " }" + " ret {" + " key: 'm'" + " value: 'matmul:product'" + " }", + &def)); + return def.SerializeAsString(); +} + +TEST(CAPI, FunctionDefAndExecute) { + TF_Status* status = TF_NewStatus(); + TF_SessionOptions* opts = TF_NewSessionOptions(); + TFE_Context* ctx = TFE_NewContext(opts, status); + CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); + TF_DeleteSessionOptions(opts); + + string function_def = MatMulFunction(); + TFE_ContextAddFunctionDef(ctx, function_def.data(), function_def.size(), + status); + CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); + + TFE_TensorHandle* m = TestMatrixTensorHandle(); + TFE_TensorHandle* retval[1] = {nullptr}; + int num_retvals = 1; + TFE_Op* op = TFE_NewOp(ctx, "MatMulFunction", status); + ASSERT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); + TFE_OpAddInput(op, m, status); + ASSERT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); + TFE_Execute(op, &retval[0], &num_retvals, status); + ASSERT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); + ASSERT_EQ(1, num_retvals); + TFE_DeleteOp(op); + TFE_DeleteTensorHandle(m); + TF_Tensor* t = TFE_TensorHandleResolve(retval[0], status); + TFE_DeleteTensorHandle(retval[0]); + 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_DeleteContext(ctx, status); + EXPECT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); + TF_DeleteStatus(status); +} + +void BM_ExecuteFunction(int iters) { + tensorflow::testing::StopTiming(); + TF_Status* status = TF_NewStatus(); + TF_SessionOptions* opts = TF_NewSessionOptions(); + TFE_Context* ctx = TFE_NewContext(opts, status); + CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); + TF_DeleteSessionOptions(opts); + + string function_def = MatMulFunction(); + TFE_ContextAddFunctionDef(ctx, function_def.data(), function_def.size(), + status); + CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); + + TFE_TensorHandle* m = TestMatrixTensorHandle(); + TFE_Op* matmul = TFE_NewOp(ctx, "MatMulFunction", status); + CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); + TFE_OpAddInput(matmul, m, status); + CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); + TFE_TensorHandle* retval[1] = {nullptr}; + int num_retvals = 1; + tensorflow::testing::StartTiming(); + for (int i = 0; i < iters; ++i) { + TFE_Execute(matmul, &retval[0], &num_retvals, status); + CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); + } + tensorflow::testing::StopTiming(); + TFE_DeleteTensorHandle(m); + TFE_DeleteTensorHandle(retval[0]); + TFE_DeleteContext(ctx, status); + EXPECT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); + TF_DeleteStatus(status); +} +BENCHMARK(BM_ExecuteFunction); + +TFE_TensorHandle* CreateVariable(TFE_Context* ctx, float value, + TF_Status* status) { + // Create the variable handle. + TFE_Op* op = TFE_NewOp(ctx, "VarHandleOp", status); + if (TF_GetCode(status) != TF_OK) return nullptr; + TFE_OpSetAttrType(op, "dtype", TF_FLOAT); + TFE_OpSetAttrShape(op, "shape", {}, 0, status); + TFE_OpSetAttrString(op, "container", ""); + TFE_OpSetAttrString(op, "shared_name", ""); + if (TF_GetCode(status) != TF_OK) return nullptr; + TFE_TensorHandle* var_handle = nullptr; + int num_retvals = 1; + TFE_Execute(op, &var_handle, &num_retvals, status); + TFE_DeleteOp(op); + if (TF_GetCode(status) != TF_OK) return nullptr; + CHECK_EQ(1, num_retvals); + + // Assign 'value' to it. + op = TFE_NewOp(ctx, "AssignVariableOp", status); + if (TF_GetCode(status) != TF_OK) return nullptr; + TFE_OpSetAttrType(op, "dtype", TF_FLOAT); + TFE_OpAddInput(op, var_handle, status); + + // Convert 'value' to a TF_Tensor then a TFE_TensorHandle. + std::unique_ptr t( + TF_AllocateTensor(TF_FLOAT, nullptr, 0, sizeof(value)), TF_DeleteTensor); + memcpy(TF_TensorData(t.get()), &value, TF_TensorByteSize(t.get())); + + std::unique_ptr + value_handle(TFE_NewTensorHandle(t.get(), status), TFE_DeleteTensorHandle); + if (TF_GetCode(status) != TF_OK) return nullptr; + + TFE_OpAddInput(op, value_handle.get(), status); + if (TF_GetCode(status) != TF_OK) return nullptr; + + num_retvals = 0; + TFE_Execute(op, nullptr, &num_retvals, status); + TFE_DeleteOp(op); + if (TF_GetCode(status) != TF_OK) return nullptr; + CHECK_EQ(0, num_retvals); + + return var_handle; +} + +TEST(CAPI, Variables) { + // Variables use resource handles, so this is really a test for resource + // tensor handling. + TF_Status* status = TF_NewStatus(); + TF_SessionOptions* opts = TF_NewSessionOptions(); + TFE_Context* ctx = TFE_NewContext(opts, status); + ASSERT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); + TF_DeleteSessionOptions(opts); + + TFE_TensorHandle* var_handle = CreateVariable(ctx, 12.0, status); + ASSERT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); + + TFE_Op* op = TFE_NewOp(ctx, "ReadVariableOp", status); + ASSERT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); + TFE_OpSetAttrType(op, "dtype", TF_FLOAT); + TFE_OpAddInput(op, var_handle, status); + ASSERT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); + int num_retvals = 1; + TFE_TensorHandle* value_handle = nullptr; + TFE_Execute(op, &value_handle, &num_retvals, status); + TFE_DeleteOp(op); + + ASSERT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); + ASSERT_EQ(1, num_retvals); + EXPECT_EQ(TF_FLOAT, TFE_TensorHandleDataType(value_handle)); + EXPECT_EQ(0, TFE_TensorHandleNumDims(value_handle)); + float value = 0.0f; + TF_Tensor* t = TFE_TensorHandleResolve(value_handle, status); + ASSERT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); + ASSERT_EQ(sizeof(float), TF_TensorByteSize(t)); + memcpy(&value, TF_TensorData(t), sizeof(float)); + TF_DeleteTensor(t); + EXPECT_EQ(12.0, value); + + TFE_DeleteTensorHandle(var_handle); + TFE_DeleteTensorHandle(value_handle); + TFE_DeleteContext(ctx, status); + CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); + TF_DeleteStatus(status); +} + +void BM_ReadVariable(int iters) { + tensorflow::testing::StopTiming(); + TF_Status* status = TF_NewStatus(); + TF_SessionOptions* opts = TF_NewSessionOptions(); + TFE_Context* ctx = TFE_NewContext(opts, status); + CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); + TF_DeleteSessionOptions(opts); + + TFE_TensorHandle* var_handle = CreateVariable(ctx, 5.0, status); + CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); + + TFE_Op* op = TFE_NewOp(ctx, "ReadVariableOp", status); + CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); + TFE_OpSetAttrType(op, "dtype", TF_FLOAT); + TFE_OpAddInput(op, var_handle, status); + CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); + + int num_retvals = 1; + TFE_TensorHandle* h = nullptr; + tensorflow::testing::StartTiming(); + for (int i = 0; i < iters; ++i) { + TFE_Execute(op, &h, &num_retvals, status); + CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); + CHECK_EQ(1, num_retvals); + CHECK(h); + CHECK_EQ(TF_FLOAT, TFE_TensorHandleDataType(h)); + CHECK_EQ(0, TFE_TensorHandleNumDims(h)); + h = nullptr; + } + tensorflow::testing::StopTiming(); + TFE_DeleteOp(op); + + TFE_DeleteTensorHandle(var_handle); + TFE_DeleteContext(ctx, status); + CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); + TF_DeleteStatus(status); +} +BENCHMARK(BM_ReadVariable); + +} // namespace diff --git a/tensorflow/c/eager/runtime.cc b/tensorflow/c/eager/runtime.cc new file mode 100644 index 0000000000000000000000000000000000000000..b6d53872c97cfb809ff3beb4110d5cad64790f46 --- /dev/null +++ b/tensorflow/c/eager/runtime.cc @@ -0,0 +1,291 @@ +/* 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/runtime.h" + +#include "tensorflow/core/common_runtime/device_factory.h" +#include "tensorflow/core/common_runtime/rendezvous_mgr.h" +#include "tensorflow/core/framework/allocator.h" +#include "tensorflow/core/framework/node_def.pb.h" +#include "tensorflow/core/lib/core/errors.h" +#include "tensorflow/core/lib/gtl/map_util.h" +#include "tensorflow/core/lib/gtl/stl_util.h" +#include "tensorflow/core/platform/fingerprint.h" +#include "tensorflow/core/platform/mutex.h" +#include "tensorflow/core/public/version.h" +#include "tensorflow/core/util/tensor_slice_reader_cache.h" + +namespace tensorflow { +namespace { + +mutex g_op_name_to_attr_type_map_lock(LINKER_INITIALIZED); + +std::unordered_map* OpNameToAttrTypeMap() { + static auto* const m = new std::unordered_map; + return m; +} + +const uint32 kIsList = 1U << 31; + +} // namespace + +Status AttrTypeMapForOp(const char* op_name, const AttrTypeMap** out) { + mutex_lock l(g_op_name_to_attr_type_map_lock); + *out = gtl::FindPtrOrNull(*OpNameToAttrTypeMap(), op_name); + if (*out != nullptr) return Status::OK(); + const OpRegistrationData* op_reg_data = nullptr; + Status s = OpRegistry::Global()->LookUp(op_name, &op_reg_data); + if (!s.ok()) return s; + std::unique_ptr m(new AttrTypeMap); + // TODO(agarwal): Avoid having to create this "registry" at runtime, + // perhaps can be done at op registration time? + for (const auto& attr : op_reg_data->op_def.attr()) { + string type = attr.type(); + const bool is_list = (type.length() > 6 && type.compare(0, 4, "list") == 0); + if (is_list) { + type = type.substr(5, type.length() - 6); + } + uint32 t = is_list ? kIsList : 0; + if (type == "string") { + t |= TF_ATTR_STRING; + } else if (type == "int") { + t |= TF_ATTR_INT; + } else if (type == "float") { + t |= TF_ATTR_FLOAT; + } else if (type == "bool") { + t |= TF_ATTR_BOOL; + } else if (type == "type") { + t |= TF_ATTR_TYPE; + } else if (type == "shape") { + t |= TF_ATTR_SHAPE; + } else if (type == "tensor") { + t |= TF_ATTR_TENSOR; + } else { + return errors::Unimplemented( + "TODO(agarwal): Enable support for ops with attributes of type '", + type, "'"); + } + gtl::InsertIfNotPresent(m.get(), attr.name(), t); + } + *out = m.get(); + (*OpNameToAttrTypeMap())[op_name] = m.release(); + return Status::OK(); +} + +Status AttrTypeByName(const AttrTypeMap* m, const string& attr_name, + TF_AttrType* out, unsigned char* is_list) { + CHECK(m); + auto* t = gtl::FindOrNull(*m, attr_name); + if (t == nullptr) { + return errors::InvalidArgument("Attribute '", attr_name, + "' does not exist for this operation"); + } + *out = static_cast(*t & ~kIsList); + if (*t & kIsList) { + *is_list = 1; + } else { + *is_list = 0; + } + return Status::OK(); +} + +#define DEFINE_SET_ATTR(value_type, value_field) \ + template <> \ + AttrBuilder& AttrBuilder::Set(StringPiece attr_name, value_type&& value) { \ + value_field.push_back(std::make_pair(attr_name, value)); \ + return *this; \ + } + +DEFINE_SET_ATTR(StringPiece, string_attrs_); +DEFINE_SET_ATTR(float, float_attrs_); +DEFINE_SET_ATTR(int, int_attrs_); +DEFINE_SET_ATTR(bool, bool_attrs_); +DEFINE_SET_ATTR(tensorflow::DataType, type_attrs_); + +#undef DEFINE_SET_ATTR + +AttrBuilder& AttrBuilder::NumInputs(int n) { + DCHECK(!node_def_finalized_) << "Calling NumInputs after BuildNodeDef."; + num_inputs_ = n; + return *this; +} + +const NodeDef& AttrBuilder::BuildNodeDef() { + if (node_def_finalized_) return *node_def_; + MayBeInitializeNodeDef(); + for (int i = 0; i < num_inputs_; ++i) { + node_def_->add_input("dummy_input"); + } + for (const auto& p : string_attrs_) { + SetInNodeDef(p.first, p.second); + } + for (const auto& p : int_attrs_) { + SetInNodeDef(p.first, p.second); + } + for (const auto& p : float_attrs_) { + SetInNodeDef(p.first, p.second); + } + for (const auto& p : bool_attrs_) { + SetInNodeDef(p.first, p.second); + } + for (const auto& p : type_attrs_) { + SetInNodeDef(p.first, p.second); + } + node_def_finalized_ = true; + return *node_def_; +} + +namespace { +inline tensorflow::Fprint128 FingerprintCat128(const tensorflow::Fprint128& a, + const tensorflow::Fprint128& b) { + return {tensorflow::FingerprintCat64(a.low64, b.low64), + tensorflow::FingerprintCat64(a.low64, b.low64)}; +} + +void CombineUnordered(const tensorflow::Fprint128& a, + tensorflow::Fprint128* b) { + b->low64 += a.low64; + b->high64 += a.high64; +} + +inline tensorflow::Fprint128 CacheKeyHelper(const StringPiece& s, + const tensorflow::Fprint128& b) { + // TODO(agarwal): avoid ToString(). + tensorflow::Fprint128 a = tensorflow::Fingerprint128(s.ToString()); + return FingerprintCat128(a, b); +} + +inline tensorflow::Fprint128 CacheKeyHelper(const StringPiece& s, uint64 b) { + return CacheKeyHelper(s, {b, b}); +} + +} // namespace + +tensorflow::Fprint128 AttrBuilder::CacheKey(const string& device) const { + tensorflow::Fprint128 f = tensorflow::Fingerprint128(op_name_); + f = tensorflow::FingerprintCat128(f, tensorflow::Fingerprint128(device)); + if (node_def_ != nullptr) { + // Some attributes are directly written to node_def_ instead of being + // stored explicitly. + string value; + for (const auto& attr : node_def_->attr()) { + attr.second.SerializeToString(&value); + CombineUnordered( + CacheKeyHelper(attr.first, tensorflow::Fingerprint128(value)), &f); + } + // Note that node_def_ may be created but not finalized. This can happen + // when the creation was triggered by a call to Set, but BuildNodeDef has + // not been called. + if (node_def_finalized_) return f; + } + for (const auto& p : string_attrs_) { + // TODO(agarwal): avoid ToString(). + CombineUnordered(CacheKeyHelper(p.first, tensorflow::Fingerprint128( + p.second.ToString())), + &f); + } + for (const auto& p : int_attrs_) { + CombineUnordered(CacheKeyHelper(p.first, static_cast(p.second)), + &f); + } + static std::hash float_hasher; + for (const auto& p : float_attrs_) { + CombineUnordered( + CacheKeyHelper(p.first, static_cast(float_hasher(p.second))), + &f); + } + for (const auto& p : bool_attrs_) { + CombineUnordered(CacheKeyHelper(p.first, p.second ? 1u : 0u), &f); + } + for (const auto& p : type_attrs_) { + CombineUnordered(CacheKeyHelper(p.first, static_cast(p.second)), + &f); + } + return f; +} + +void AttrBuilder::MayBeInitializeNodeDef() { + if (node_def_ == nullptr) { + node_def_.reset(new NodeDef()); + node_def_->set_name(op_name_); + node_def_->set_op(op_name_); + } +} + +// static +Status KernelAndDevice::InitOp(Device* device, const NodeDef& ndef, + KernelAndDevice* out) { + OpKernel* k = nullptr; + Status s = CreateOpKernel(device->device_type().c_str(), device, + device->GetAllocator(AllocatorAttributes()), + nullptr, ndef, TF_GRAPH_DEF_VERSION, &k); + out->device_ = device; + out->kernel_.reset(k); + out->flib_ = nullptr; + return s; +} + +// static +Status KernelAndDevice::InitFn(const NodeDef& ndef, + FunctionLibraryRuntime* flib, + KernelAndDevice* out) { + OpKernel* k = nullptr; + Status s = flib->CreateKernel(ndef, &k); + out->device_ = flib->device(); + out->kernel_.reset(k); + out->flib_ = flib; + return s; +} + +Status KernelAndDevice::Run(std::vector* input_tensors, + std::vector* output_tensors) { + gtl::InlinedVector inputs; + for (Tensor& t : *input_tensors) { + inputs.push_back(TensorValue(&t)); + } + + std::vector out_attrs(kernel_->num_outputs()); + for (size_t i = 0; i < out_attrs.size(); ++i) { + out_attrs[i].set_on_host(kernel_->output_memory_types()[i] == + tensorflow::HOST_MEMORY); + } + + OpKernelContext::Params params; + params.device = device_; + params.frame_iter = FrameAndIter(0, 0); + params.inputs = &inputs; + params.op_kernel = kernel_.get(); + params.resource_manager = device_->resource_manager(); + params.output_attr_array = gtl::vector_as_array(&out_attrs); + params.function_library = flib_; + params.slice_reader_cache = &slice_reader_cache_; + params.rendezvous = rendez_; + // TODO(apassos): use a thread pool. + std::function)> runner = + [](std::function f) { f(); }; + params.runner = &runner; + + OpKernelContext context(¶ms); + device_->Compute(kernel_.get(), &context); + if (!context.status().ok()) return context.status(); + + output_tensors->clear(); + for (int i = 0; i < context.num_outputs(); ++i) { + output_tensors->push_back(Tensor(*context.mutable_output(i))); + } + return Status::OK(); +} + +} // namespace tensorflow diff --git a/tensorflow/c/eager/runtime.h b/tensorflow/c/eager/runtime.h new file mode 100644 index 0000000000000000000000000000000000000000..bb098f74013a697ac1552f223b101bcfce9bd2ec --- /dev/null +++ b/tensorflow/c/eager/runtime.h @@ -0,0 +1,195 @@ +/* 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_C_EAGER_RUNTIME_H_ +#define TENSORFLOW_C_EAGER_RUNTIME_H_ + +// Support for eager execution of TensorFlow kernels. + +#include +#include + +#include "tensorflow/c/c_api.h" +#include "tensorflow/core/common_runtime/device.h" +#include "tensorflow/core/framework/node_def.pb.h" +#include "tensorflow/core/framework/op_kernel.h" +#include "tensorflow/core/framework/types.h" +#include "tensorflow/core/lib/core/status.h" +#include "tensorflow/core/lib/gtl/inlined_vector.h" +#include "tensorflow/core/platform/fingerprint.h" +#include "tensorflow/core/util/tensor_slice_reader_cache.h" + +namespace tensorflow { + +// Maps attribute name to an encoding of the type of the attribute value. +// If the type is not a list type, the value is the same as the TF_AttrType type +// of the value. Else, the highest order bit is on, and the rest of the bits +// represent the TF_AttrType type of the values in the list. +typedef std::unordered_map AttrTypeMap; + +// Returns the AttrTypeMap for the TensorFlow operation named op_name. +Status AttrTypeMapForOp(const char* op_name, const AttrTypeMap** out); + +// Looks for 'attr_name' in 'm' and sets 'out' and 'is_list'. +Status AttrTypeByName(const AttrTypeMap* m, const string& attr_name, + TF_AttrType* out, unsigned char* is_list); + +// KernelAndDevice::Init needs a NodeDef only to pass the attribute map through. +// An AttrBuilder is a convenience class to help with that - providing a smaller +// interface than NodeDefBuilder and avoiding expensive (unnecessary?) sanity +// checks (like number of inputs matching the OpDef - we only care about +// attributes here). +// +// TODO(ashankar): Take a closer look at checks in NodeDefBuilder and see which +// ones make sense to replicate. + +// This is a helper class for creating a NodeDef. Additionally, this class +// allows computing a cache key based on fingerprinting the attributes of this +// NodeDef. +// +// Example usage: +// AttrBuilder a; +// a.NumInputs(2); +// a.Set("T", TF_FLOAT); +// uint64 cache_key = a.CacheKey("cpu:0"); +// const NodeDef& n = a.BuildNodeDef(); +// +// Note that all calls to Set and NumInputs should happen before calling +// BuildNodeDef. Also, calls to NumInputs or Set between multiple invocations +// to CacheKey may cause different values to be returned by CacheKey. +// +// For performance reasons, the class internally delays the actual construction +// of the NodeDef till BuildNodeDef is called, or Set is called with certain +// uncommon types (see template specializations of Set to see which types +// trigger a NodeDef creation). +class AttrBuilder { + public: + explicit AttrBuilder(const char* op) + : op_name_(op), + num_inputs_(0), + node_def_(nullptr), + node_def_finalized_(false) {} + + // Needed to work around call to ValidateNodeDef in CreateOpKernel. + AttrBuilder& NumInputs(int n); + + template + AttrBuilder& Set(StringPiece attr_name, T&& value) { + MayBeInitializeNodeDef(); + return SetInNodeDef(attr_name, value); + } + + tensorflow::Fprint128 CacheKey(const string& device) const; + + const NodeDef& BuildNodeDef(); + + private: + template + using AttrVec = tensorflow::gtl::InlinedVector, 2>; + + void MayBeInitializeNodeDef(); + + template + AttrBuilder& SetInNodeDef(StringPiece attr_name, T&& value) { + DCHECK(!node_def_finalized_) << "Calling SetInNodeDef after BuildNodeDef."; + // Copied from NodeDefBuilder::Attr + const AttrValue* found = AttrSlice(*node_def_).Find(attr_name); + if (found == nullptr) { + AddNodeAttr(attr_name, std::forward(value), node_def_.get()); + } else { + AttrValue attr_value; + SetAttrValue(std::forward(value), &attr_value); + // TODO(ashankar): Do what is done in + // NodeDefBuilder::CheckInconsistency(attr_name, *found, attr_value); + } + return *this; + } + + AttrVec string_attrs_; + AttrVec int_attrs_; + AttrVec float_attrs_; + AttrVec bool_attrs_; + AttrVec type_attrs_; + string op_name_; + int num_inputs_; + std::unique_ptr node_def_; + bool node_def_finalized_; +}; // namespace tensorflow + +template <> +AttrBuilder& AttrBuilder::Set(StringPiece attr_name, StringPiece&& value); +template <> +AttrBuilder& AttrBuilder::Set(StringPiece attr_name, int&& value); +template <> +AttrBuilder& AttrBuilder::Set(StringPiece attr_name, float&& value); +template <> +AttrBuilder& AttrBuilder::Set(StringPiece attr_name, bool&& value); +template <> +AttrBuilder& AttrBuilder::Set(StringPiece attr_name, + tensorflow::DataType&& value); + +// KernelAndDevice encapsulates an instantiated kernel and the device it is on. +// +// Also see: +// https://www.tensorflow.org/code/tensorflow/core/common_runtime/kernel_benchmark_testlib.h +// and +// https://www.tensorflow.org/code/tensorflow/core/kernels/ops_testutil.h +class KernelAndDevice { + public: + // Populates 'out' with a kernel appropriate for 'ndef'. + // + // Assumes that 'ndef' refers to a primitive op (as opposed to a function). + static Status InitOp(Device* device, const NodeDef& ndef, + KernelAndDevice* out); + + // Like InitOp but for functions defined in flib (i.e., ndef.op() refers to a + // TensorFlow function in the FunctionLibraryRuntime). + // + // The provided FunctionLibraryRuntime MUST outlive all calls to + // Run() on the returned KernelAndDevice. + // + // TODO(ashankar): There shouldn't be a need for a separate InitOp and InitFn. + // The implementation of InitFn should work for both because + // FunctionLibraryRuntime::CreateKernel will create a primitive op kernel if + // appropriate. However, for now we keep them separate because I haven't + // figured out thread-safety concerns around FunctionLibraryRuntime (in + // particular, how the underlying FunctionLibraryDefinition might be mutated + // by another thread as new functions are registered with it). + // Conservatively, thread-safe usage of the FunctionLibraryRuntime is pushed + // on to the caller (see locking in c_api.cc) for now. But I really should + // dig into this so that both InitOp and InitFn can be collapsed to + // FunctionLibraryRuntime::CreateKernel. + static Status InitFn(const NodeDef& ndef, FunctionLibraryRuntime* flib, + KernelAndDevice* out); + + KernelAndDevice(tensorflow::Rendezvous* rendez) + : device_(nullptr), flib_(nullptr), rendez_(rendez) {} + + // TODO(ashankar): Handle list-valued inputs. + Status Run(std::vector* inputs, std::vector* outputs); + + const OpKernel* kernel() const { return kernel_.get(); } + + private: + std::unique_ptr kernel_; + tensorflow::Device* device_; + tensorflow::FunctionLibraryRuntime* flib_; + tensorflow::checkpoint::TensorSliceReaderCacheWrapper slice_reader_cache_; + tensorflow::Rendezvous* rendez_; +}; + +} // namespace tensorflow + +#endif // TENSORFLOW_C_EAGER_RUNTIME_H_ diff --git a/tensorflow/c/eager/runtime_test.cc b/tensorflow/c/eager/runtime_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..f9bfce3858064312788ec796786ace2bd4bd272b --- /dev/null +++ b/tensorflow/c/eager/runtime_test.cc @@ -0,0 +1,167 @@ +/* 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/runtime.h" + +#include +#include + +#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/core/common_runtime/device_factory.h" +#include "tensorflow/core/platform/test.h" +#include "tensorflow/core/platform/test_benchmark.h" + +namespace tensorflow { +namespace { + +Device* CPUDevice() { + return DeviceFactory::NewDevice("CPU", {}, "/job:a/replica:0/task:0"); +} + +TEST(AttrTypeMap, Lookup) { + const AttrTypeMap* m = nullptr; + Status s = AttrTypeMapForOp("ThisOpCannotPossiblyExist", &m); + EXPECT_FALSE(s.ok()); + s = AttrTypeMapForOp("MatMul", &m); + ASSERT_TRUE(s.ok()) << s; + + TF_AttrType t; + unsigned char is_list = 1; + s = AttrTypeByName(m, "ThisAttribyteCannotPossiblyExist", &t, &is_list); + EXPECT_FALSE(s.ok()); + EXPECT_NE(is_list, 0); + s = AttrTypeByName(m, "transpose_a", &t, &is_list); + ASSERT_TRUE(s.ok()) << s; + EXPECT_EQ(TF_ATTR_BOOL, t); + EXPECT_EQ(is_list, 0); + + s = AttrTypeMapForOp("Squeeze", &m); + ASSERT_TRUE(s.ok()) << s; + s = AttrTypeByName(m, "squeeze_dims", &t, &is_list); + ASSERT_TRUE(s.ok()) << s; + EXPECT_EQ(TF_ATTR_INT, t); + EXPECT_NE(is_list, 0); +} + +TEST(KernelAndDevice, Run) { + Tensor t(Input({{1.0f, 2.0f}, {3.0f, 4.0f}}).tensor()); + std::vector inputs; + inputs.push_back(t); + inputs.push_back(t); + NodeDef ndef(AttrBuilder("MatMul") + .Set("T", DT_FLOAT) + .Set("transpose_a", false) + .Set("transpose_b", false) + .NumInputs(inputs.size()) + .BuildNodeDef()); + std::unique_ptr device(CPUDevice()); + KernelAndDevice kernel(nullptr); + Status s = KernelAndDevice::InitOp(device.get(), ndef, &kernel); + ASSERT_TRUE(s.ok()) << s; + std::vector outputs; + s = kernel.Run(&inputs, &outputs); + ASSERT_TRUE(s.ok()) << s; + ASSERT_EQ(1, outputs.size()); + const Tensor& out = outputs[0]; + EXPECT_EQ(7, out.matrix()(0, 0)); + EXPECT_EQ(10, out.matrix()(0, 1)); + EXPECT_EQ(15, out.matrix()(1, 0)); + EXPECT_EQ(22, out.matrix()(1, 1)); +} + +void BM_CreateGraph(int iters) { + for (int i = 0; i < iters; ++i) { + Scope root = Scope::NewRootScope(); + auto C = ops::Const(root, {{1.0, 2.0}, {3.0, 4.0}}); + auto M = ops::MatMul(root, C, C); + TF_CHECK_OK(root.status()); + } +} +BENCHMARK(BM_CreateGraph); + +void BM_RunGraph(int iters) { + tensorflow::testing::StopTiming(); + Scope root = Scope::NewRootScope(); + auto C = ops::Const(root, {{1.0, 2.0}, {3.0, 4.0}}); + auto M = ops::MatMul(root, C, C); + SessionOptions opts; + opts.config.set_inter_op_parallelism_threads(1); + opts.config.set_intra_op_parallelism_threads(1); + ClientSession sess(root, opts); + std::vector outputs; + tensorflow::testing::StartTiming(); + for (int i = 0; i < iters; ++i) { + outputs.clear(); + TF_CHECK_OK(sess.Run({M}, &outputs)); + } +} +BENCHMARK(BM_RunGraph); + +void BM_CreateAndDestroySession(int iters) { + tensorflow::testing::StopTiming(); + Scope root = Scope::NewRootScope(); + auto C = ops::Const(root, {{1.0, 2.0}, {3.0, 4.0}}); + auto M = ops::MatMul(root, C, C); + tensorflow::testing::StartTiming(); + for (int i = 0; i < iters; ++i) { + ClientSession sess(root); + } +} +BENCHMARK(BM_CreateAndDestroySession); + +void BM_KernelAndDeviceInit(int iters) { + tensorflow::testing::StopTiming(); + NodeDef ndef(AttrBuilder("MatMul") + .Set("T", DT_FLOAT) + .Set("transpose_a", false) + .Set("transpose_b", false) + .NumInputs(2) + .BuildNodeDef()); + std::unique_ptr device(CPUDevice()); + KernelAndDevice k(nullptr); + tensorflow::testing::StartTiming(); + for (int i = 0; i < iters; ++i) { + TF_CHECK_OK(KernelAndDevice::InitOp(device.get(), ndef, &k)); + } +} +BENCHMARK(BM_KernelAndDeviceInit); + +void BM_KernelAndDeviceRun(int iters) { + tensorflow::testing::StopTiming(); + Tensor t(Input({{1.0f, 2.0f}, {3.0f, 4.0f}}).tensor()); + std::vector inputs; + inputs.push_back(t); + inputs.push_back(t); + std::vector outputs; + NodeDef ndef(AttrBuilder("MatMul") + .Set("T", DT_FLOAT) + .Set("transpose_a", false) + .Set("transpose_b", false) + .NumInputs(inputs.size()) + .BuildNodeDef()); + std::unique_ptr device(CPUDevice()); + KernelAndDevice kernel(nullptr); + TF_CHECK_OK(KernelAndDevice::InitOp(device.get(), ndef, &kernel)); + tensorflow::testing::StartTiming(); + for (int i = 0; i < iters; ++i) { + TF_CHECK_OK(kernel.Run(&inputs, &outputs)); + } +} +BENCHMARK(BM_KernelAndDeviceRun); +} // namespace +} // namespace tensorflow diff --git a/tensorflow/c/exported_symbols.lds b/tensorflow/c/exported_symbols.lds new file mode 100644 index 0000000000000000000000000000000000000000..41f0637c99c8842f4b52701e2f9b55cfdac46309 --- /dev/null +++ b/tensorflow/c/exported_symbols.lds @@ -0,0 +1,2 @@ +_TF_* +_TFE_* diff --git a/tensorflow/c/generate-pc.sh b/tensorflow/c/generate-pc.sh new file mode 100755 index 0000000000000000000000000000000000000000..02a6a58b6153bb78c684f9290ef95900f96e9357 --- /dev/null +++ b/tensorflow/c/generate-pc.sh @@ -0,0 +1,67 @@ +#!/usr/bin/env 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. +# ============================================================================== + +TF_PREFIX='/usr/local' + +usage() { + echo "Usage: $0 OPTIONS" + echo -e "-p, --prefix\tset installation prefix (default: /usr/local)" + echo -e "-v, --version\tset TensorFlow version" + echo -e "-h, --help\tdisplay this message" +} + +[ $# == 0 ] && usage && exit 0 + +# read the options +ARGS=$(getopt -o p:v:h --long prefix:,version:,help -n $0 -- "$@") +eval set -- "$ARGS" + +# extract options and their arguments into variables. +while true ; do + case "$1" in + -h|--help) usage ; exit ;; + -p|--prefix) + case "$2" in + "") shift 2 ;; + *) TF_PREFIX=$2 ; shift 2 ;; + esac ;; + -v|--version) + case "$2" in + "") shift 2 ;; + *) TF_VERSION=$2 ; shift 2 ;; + esac ;; + --) shift ; break ;; + *) echo "Internal error! Try '$0 --help' for more information." ; exit 1 ;; + esac +done + +[ -z $TF_VERSION ] && echo "Specify a version using -v or --version" && exit 1 + +echo "Generating pkgconfig file for TensorFlow $TF_VERSION in $TF_PREFIX" + +cat << EOF > tensorflow.pc +prefix=${TF_PREFIX} +exec_prefix=\${prefix} +libdir=\${exec_prefix}/lib +includedir=\${prefix}/include + +Name: TensorFlow +Version: ${TF_VERSION} +Description: Library for computation using data flow graphs for scalable machine learning +Requires: +Libs: -L\${libdir} -ltensorflow +Cflags: -I\${includedir} +EOF diff --git a/tensorflow/c/python_api.cc b/tensorflow/c/python_api.cc new file mode 100644 index 0000000000000000000000000000000000000000..adca6c762526a85f015560efb22d3de185e2ae6c --- /dev/null +++ b/tensorflow/c/python_api.cc @@ -0,0 +1,33 @@ +/* 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/python_api.h" + +#include "tensorflow/c/c_api_internal.h" + +namespace tensorflow { + +void AddControlInput(TF_Graph* graph, TF_Operation* op, TF_Operation* input) { + // TODO(skyewm): make sure cycles are prevented + mutex_lock l(graph->mu); + graph->graph.AddControlEdge(&input->node, &op->node); +} + +void SetRequestedDevice(TF_Graph* graph, TF_Operation* op, const char* device) { + mutex_lock l(graph->mu); + op->node.set_requested_device(device); +} + +} // namespace tensorflow diff --git a/tensorflow/c/python_api.h b/tensorflow/c/python_api.h new file mode 100644 index 0000000000000000000000000000000000000000..e1a55d7755a76c778bf6a8120a8cf81adb6941dc --- /dev/null +++ b/tensorflow/c/python_api.h @@ -0,0 +1,32 @@ +/* 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 THIRD_PARTY_TENSORFLOW_C_PYTHON_API_H_ +#define THIRD_PARTY_TENSORFLOW_C_PYTHON_API_H_ + +#include "tensorflow/c/c_api.h" + +// These functions can be removed without notice. They exist to facilitate some +// refactoring of graph construction code in the Python API. + +namespace tensorflow { + +void AddControlInput(TF_Graph* graph, TF_Operation* op, TF_Operation* input); + +void SetRequestedDevice(TF_Graph* graph, TF_Operation* op, const char* device); + +} // namespace tensorflow + +#endif // THIRD_PARTY_TENSORFLOW_C_PYTHON_API_H_ diff --git a/tensorflow/c/tf_status_helper.cc b/tensorflow/c/tf_status_helper.cc index 747fd672f08605171da4556e91fac4128220aac9..eaaed89500b94e9f3aa82396ca641827215cb283 100644 --- a/tensorflow/c/tf_status_helper.cc +++ b/tensorflow/c/tf_status_helper.cc @@ -14,6 +14,7 @@ limitations under the License. ==============================================================================*/ #include "tensorflow/c/tf_status_helper.h" +#include "tensorflow/c/c_api_internal.h" namespace tensorflow { @@ -79,4 +80,8 @@ void Set_TF_Status_from_Status(TF_Status* tf_status, const Status& status) { } } +Status StatusFromTF_Status(const TF_Status* tf_status) { + return tf_status->status; +} + } // namespace tensorflow diff --git a/tensorflow/c/tf_status_helper.h b/tensorflow/c/tf_status_helper.h index 4bc56f9cb40e59bb0c1dcc6b88cd23c3949ce8cd..86e687df205617018d94c19ac34fdc3bf54dcc6f 100644 --- a/tensorflow/c/tf_status_helper.h +++ b/tensorflow/c/tf_status_helper.h @@ -24,6 +24,9 @@ namespace tensorflow { // Set the attribute of "tf_status" from the attributes of "status". void Set_TF_Status_from_Status(TF_Status* tf_status, const Status& status); +// Returns a "status" from "tf_status". +Status StatusFromTF_Status(const TF_Status* tf_status); + } // namespace tensorflow #endif // TENSORFLOW_C_TF_STATUS_HELPER_H diff --git a/tensorflow/c/version_script.lds b/tensorflow/c/version_script.lds new file mode 100644 index 0000000000000000000000000000000000000000..455bd7362bb36d30af421a17f0e2f8e9ba66e02b --- /dev/null +++ b/tensorflow/c/version_script.lds @@ -0,0 +1,9 @@ +VERS_1.0 { + # Export symbols in c_api.h. + global: + TF_*; + + # Hide everything else. + local: + *; +}; diff --git a/tensorflow/c/while_loop_test.cc b/tensorflow/c/while_loop_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..ce4f86bb25b9bf2579f08bb9af98bf6debcbd4c1 --- /dev/null +++ b/tensorflow/c/while_loop_test.cc @@ -0,0 +1,339 @@ +/* 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/c_api.h" + +#include "tensorflow/c/c_test_util.h" +#include "tensorflow/core/lib/strings/strcat.h" +#include "tensorflow/core/platform/logging.h" +#include "tensorflow/core/platform/test.h" + +using tensorflow::GraphDef; + +namespace { + +class CApiWhileLoopTest : public ::testing::Test { + protected: + CApiWhileLoopTest() : s_(TF_NewStatus()), graph_(TF_NewGraph()) {} + + ~CApiWhileLoopTest() override { + TF_DeleteGraph(graph_); + TF_DeleteStatus(s_); + } + + void Init(int ninputs) { + DCHECK(inputs_.empty()); + DCHECK_GT(ninputs, 0); + + for (int i = 0; i < ninputs; ++i) { + TF_Operation* placeholder = Placeholder( + graph_, s_, ::tensorflow::strings::StrCat("p", i).c_str()); + DCHECK_EQ(TF_OK, TF_GetCode(s_)) << TF_Message(s_); + inputs_.push_back({placeholder, 0}); + } + + original_graph_description_ = GraphDebugString(); + + params_.reset(new TF_WhileParams( + TF_NewWhile(graph_, &inputs_[0], inputs_.size(), s_))); + ASSERT_EQ(TF_OK, TF_GetCode(s_)) << TF_Message(s_); + ASSERT_EQ(original_graph_description_, GraphDebugString()) + << "TF_NewWhile() altered graph"; + + params_->name = "test_loop"; + + // Initialize outputs_ so we can easily detect errors/bugs + outputs_.resize(ninputs, {nullptr, -1}); + } + + void ExpectOK() { + TF_FinishWhile(params_.get(), s_, &outputs_[0]); + EXPECT_EQ(TF_OK, TF_GetCode(s_)) << TF_Message(s_); + } + + void ExpectError(TF_Code expected_code, const string& expected_msg) { + TF_FinishWhile(params_.get(), s_, &outputs_[0]); + EXPECT_EQ(expected_code, TF_GetCode(s_)); + EXPECT_EQ(expected_msg, TF_Message(s_)); + // TODO(skyewm): this assert is currently broken. Fix or remove guarantee. + // ASSERT_EQ(original_graph_description_, GraphDebugString()) << + // "TF_FinishWhile() altered graph on error"; + } + + void Run(std::initializer_list input_values) { + DCHECK_EQ(inputs_.size(), input_values.size()); + std::vector> inputs(inputs_.size()); + int i = 0; + for (int v : input_values) { + inputs[i] = {inputs_[i].oper, Int32Tensor(v)}; + ++i; + } + csession_.reset(new CSession(graph_, s_)); + csession_->SetInputs(inputs); + csession_->SetOutputs(outputs_); + csession_->Run(s_); + ASSERT_EQ(TF_OK, TF_GetCode(s_)) << TF_Message(s_); + } + + void ExpectOutputValue(int idx, int expected_value) { + TF_Tensor* out = csession_->output_tensor(idx); + ASSERT_TRUE(out != nullptr); + EXPECT_EQ(TF_INT32, TF_TensorType(out)); + EXPECT_EQ(0, TF_NumDims(out)); + ASSERT_EQ(sizeof(int32_t), TF_TensorByteSize(out)); + int32_t* data = static_cast(TF_TensorData(out)); + EXPECT_EQ(expected_value, *data); + } + + // Create a valid conditional graph. Useful for testing unrelated errors. + void CreateCondGraph() { + TF_Operation* one = ScalarConst(1, params_->cond_graph, s_); + TF_Operation* less_than = + LessThan(params_->cond_inputs[0], {one, 0}, params_->cond_graph, s_); + DCHECK_EQ(TF_OK, TF_GetCode(s_)) << TF_Message(s_); + params_->cond_output = {less_than, 0}; + } + + string GraphDebugString() const { + TF_Buffer* buf = TF_NewBuffer(); + TF_GraphToGraphDef(graph_, buf, s_); + DCHECK_EQ(TF_OK, TF_GetCode(s_)) << TF_Message(s_); + GraphDef def; + bool success = def.ParseFromArray(buf->data, buf->length); + DCHECK(success); + TF_DeleteBuffer(buf); + return def.DebugString(); + } + + TF_Status* s_; + TF_Graph* graph_; + std::vector inputs_; // The inputs to the while loop + std::vector outputs_; // The final outputs of the while loop + std::unique_ptr params_; + std::unique_ptr csession_; + + private: + // Used to verify that errors don't change graph_ + string original_graph_description_; +}; + +TEST_F(CApiWhileLoopTest, BasicLoop) { + Init(2); + + // Validate TF_WhileParams returned by TF_NewWhile() + EXPECT_TRUE(params_->body_graph != nullptr); + EXPECT_TRUE(params_->cond_graph != nullptr); + + EXPECT_EQ(params_->ninputs, 2); + + ASSERT_TRUE(params_->cond_inputs != nullptr); + ASSERT_TRUE(params_->cond_inputs[0].oper != nullptr); + EXPECT_TRUE(params_->cond_inputs[1].oper != nullptr); + + ASSERT_TRUE(params_->body_inputs != nullptr); + EXPECT_TRUE(params_->body_inputs[0].oper != nullptr); + EXPECT_TRUE(params_->body_inputs[1].oper != nullptr); + + ASSERT_TRUE(params_->body_outputs != nullptr); + + // Create loop: while (input1 < input2) input1 += input2 + 1 + TF_Operation* less_than = + LessThan(params_->cond_inputs[0], params_->cond_inputs[1], + params_->cond_graph, s_); + ASSERT_EQ(TF_OK, TF_GetCode(s_)) << TF_Message(s_); + params_->cond_output = {less_than, 0}; + + TF_Operation* add1 = Add(params_->body_inputs[0], params_->body_inputs[1], + params_->body_graph, s_, "add1"); + ASSERT_EQ(TF_OK, TF_GetCode(s_)) << TF_Message(s_); + TF_Operation* one = ScalarConst(1, params_->body_graph, s_); + ASSERT_EQ(TF_OK, TF_GetCode(s_)) << TF_Message(s_); + TF_Operation* add2 = Add(add1, one, params_->body_graph, s_, "add2"); + ASSERT_EQ(TF_OK, TF_GetCode(s_)) << TF_Message(s_); + params_->body_outputs[0] = {add2, 0}; + params_->body_outputs[1] = params_->body_inputs[1]; + + // Finalize while loop + ExpectOK(); + + // Validate while loop outputs returned by TF_FinishWhile() + EXPECT_TRUE(outputs_[0].oper != nullptr); + EXPECT_GE(outputs_[0].index, 0); + EXPECT_TRUE(outputs_[1].oper != nullptr); + EXPECT_GE(outputs_[1].index, 0); + + // Check that cond and body inputs are not present + for (int i = 0; i < params_->ninputs; ++i) { + string cond_name = + ::tensorflow::strings::StrCat(params_->name, "/cond/cond_input", i); + string body_name = + ::tensorflow::strings::StrCat(params_->name, "/body/body_input", i); + EXPECT_TRUE(TF_GraphOperationByName(graph_, cond_name.c_str()) == nullptr); + EXPECT_TRUE(TF_GraphOperationByName(graph_, body_name.c_str()) == nullptr); + } + + // Run the graph + Run({-9, 2}); + ExpectOutputValue(0, 3); + ExpectOutputValue(1, 2); +} + +TEST_F(CApiWhileLoopTest, NestedLoop) { + Init(2); + // Create nested loop: + // while (input1 < 6) { + // inner_input1 = input1 + // while (inner_input1 < 3) { + // input2 += 1 + // inner_input1 += 2 + // } + // input1 += input2 + // } + // + // Expected execution with initial values input1 = input2 = 0: + // + // outer inner inner_ + // step# step# input1 input2 input1 + // ------------------------------------ + // 0 0 0 0 0 + // 0 1 0 1 2 + // 0 2 0 2 4 + // 0 - 2 2 - + // 1 0 2 2 2 + // 1 1 2 3 4 + // 1 - 5 3 - + // 2 0 5 3 5 + // 2 - 8 3 - + + // Create outer cond graph + TF_Operation* six = ScalarConst(6, params_->cond_graph, s_); + ASSERT_EQ(TF_OK, TF_GetCode(s_)) << TF_Message(s_); + TF_Operation* less_than = + LessThan(params_->cond_inputs[0], {six, 0}, params_->cond_graph, s_); + ASSERT_EQ(TF_OK, TF_GetCode(s_)) << TF_Message(s_); + params_->cond_output = {less_than, 0}; + + // Create outer body graph + // Init inner graph + TF_Output inner_inputs[] = {params_->body_inputs[0], params_->body_inputs[1]}; + TF_WhileParams inner_params = + TF_NewWhile(params_->body_graph, inner_inputs, 2, s_); + ASSERT_EQ(TF_OK, TF_GetCode(s_)) << TF_Message(s_); + inner_params.name = "inner_loop"; + + // Create inner cond graph + TF_Operation* three = ScalarConst(3, inner_params.cond_graph, s_); + ASSERT_EQ(TF_OK, TF_GetCode(s_)) << TF_Message(s_); + TF_Operation* inner_less_than = LessThan( + inner_params.cond_inputs[0], {three, 0}, inner_params.cond_graph, s_); + ASSERT_EQ(TF_OK, TF_GetCode(s_)) << TF_Message(s_); + inner_params.cond_output = {inner_less_than, 0}; + + // Create inner body graph + TF_Operation* one = ScalarConst(1, inner_params.body_graph, s_, "one"); + ASSERT_EQ(TF_OK, TF_GetCode(s_)) << TF_Message(s_); + TF_Operation* two = ScalarConst(2, inner_params.body_graph, s_, "two"); + ASSERT_EQ(TF_OK, TF_GetCode(s_)) << TF_Message(s_); + + TF_Operation* input2_add = + Add(inner_params.body_inputs[1].oper, one, inner_params.body_graph, s_); + ASSERT_EQ(TF_OK, TF_GetCode(s_)) << TF_Message(s_); + inner_params.body_outputs[1] = {input2_add, 0}; + + TF_Operation* inner_input1_add = Add(inner_params.body_inputs[0].oper, two, + inner_params.body_graph, s_, "add2"); + ASSERT_EQ(TF_OK, TF_GetCode(s_)) << TF_Message(s_); + inner_params.body_outputs[0] = {inner_input1_add, 0}; + + // Finalize inner graph + TF_Output inner_outputs[2] = {{nullptr, -1}}; + TF_FinishWhile(&inner_params, s_, inner_outputs); + ASSERT_EQ(TF_OK, TF_GetCode(s_)) << TF_Message(s_); + + TF_Operation* input1_add = + Add(params_->body_inputs[0], inner_outputs[1], params_->body_graph, s_); + ASSERT_EQ(TF_OK, TF_GetCode(s_)) << TF_Message(s_); + params_->body_outputs[0] = {input1_add, 0}; + + params_->body_outputs[1] = inner_outputs[1]; + + // Finalize outer graph + ExpectOK(); + + // Check for a few expected nodes + const char* node_name = "test_loop/cond/scalar"; + EXPECT_TRUE(TF_GraphOperationByName(graph_, node_name) != nullptr); + node_name = "test_loop/body/add"; + EXPECT_TRUE(TF_GraphOperationByName(graph_, node_name) != nullptr); + node_name = "test_loop/body/inner_loop/body/one"; + EXPECT_TRUE(TF_GraphOperationByName(graph_, node_name) != nullptr); + node_name = "test_loop/body/inner_loop/cond/less_than"; + EXPECT_TRUE(TF_GraphOperationByName(graph_, node_name) != nullptr); + + // Run the graph + Run({0, 0}); + ExpectOutputValue(0, 8); + ExpectOutputValue(1, 3); +} + +TEST_F(CApiWhileLoopTest, BadCondOutput) { + Init(1); + params_->body_outputs[0] = params_->body_inputs[0]; + ExpectError(TF_INVALID_ARGUMENT, + "TF_WhileParams `cond_output` field isn't set"); +} + +TEST_F(CApiWhileLoopTest, BadBodyOutput) { + Init(1); + CreateCondGraph(); + ExpectError(TF_INVALID_ARGUMENT, + "TF_WhileParams `body_outputs[0]` field isn't set"); +} + +TEST_F(CApiWhileLoopTest, NullName) { + Init(1); + CreateCondGraph(); + params_->body_outputs[0] = params_->body_inputs[0]; + params_->name = nullptr; + ExpectError(TF_INVALID_ARGUMENT, "TF_WhileParams `name` field is null"); +} + +TEST_F(CApiWhileLoopTest, WrongGraph) { + Init(1); + CreateCondGraph(); + // Set body output to output from outer graph + params_->body_outputs[0] = inputs_[0]; + // TODO(skyewm): improve error message + ExpectError(TF_INVALID_ARGUMENT, + "Requested return node 'p0' not found in graph def"); +} + +TEST_F(CApiWhileLoopTest, BadTypes) { + Init(1); + CreateCondGraph(); + // Op that has a float input + output + TF_OperationDescription* desc = TF_NewOperation( + params_->body_graph, "FakeQuantWithMinMaxArgs", "float_op"); + TF_AddInput(desc, params_->body_inputs[0]); + TF_FinishOperation(desc, s_); + ASSERT_EQ(TF_INVALID_ARGUMENT, TF_GetCode(s_)); + string msg(TF_Message(s_)); + EXPECT_NE(msg.find("Input 'inputs' passed int32 expected float while " + "building NodeDef 'float_op'"), + msg.npos); + TF_AbortWhile(params_.get()); +} + +} // namespace diff --git a/tensorflow/cc/BUILD b/tensorflow/cc/BUILD index aaebdded9a5b3232bec824b0768a536e36349204..d9071ba6e460b01746026db3215ff21ee52ac1b1 100644 --- a/tensorflow/cc/BUILD +++ b/tensorflow/cc/BUILD @@ -34,6 +34,7 @@ cc_library( tf_cc_test( name = "framework_gradients_test", + size = "small", srcs = ["framework/gradients_test.cc"], deps = [ ":cc_ops", @@ -44,6 +45,7 @@ tf_cc_test( "//tensorflow/core:all_kernels", "//tensorflow/core:framework", "//tensorflow/core:framework_internal", + "//tensorflow/core:protos_all_cc", "//tensorflow/core:test", "//tensorflow/core:test_main", "//tensorflow/core:testlib", @@ -57,11 +59,9 @@ cc_library( deps = [ ":cc_ops", ":client_session", - ":grad_op_registry", ":gradients", ":ops", ":scope", - "//tensorflow/core:core_cpu", "//tensorflow/core:framework", "//tensorflow/core:lib", "//tensorflow/core:lib_internal", @@ -70,6 +70,7 @@ cc_library( tf_cc_test( name = "framework_gradient_checker_test", + size = "small", srcs = ["framework/gradient_checker_test.cc"], deps = [ ":cc_ops", @@ -91,6 +92,7 @@ cc_library( deps = [ ":array_grad", ":math_grad", + ":nn_grad", ], ) @@ -122,7 +124,10 @@ cc_library_with_android_deps( cc_library_with_android_deps( name = "scope", - srcs = ["framework/scope.cc"], + srcs = [ + "framework/scope.cc", + "framework/scope_internal.h", + ], hdrs = ["framework/scope.h"], android_deps = ["//tensorflow/core:android_tensorflow_lib"], common_deps = [ @@ -136,8 +141,18 @@ cc_library_with_android_deps( ], ) +cc_library_with_android_deps( + name = "scope_internal", + hdrs = ["framework/scope_internal.h"], + common_deps = [ + ":scope", + ], + deps = [], +) + tf_cc_test( name = "framework_scope_test", + size = "small", srcs = ["framework/scope_test.cc"], deps = [ ":ops", @@ -167,6 +182,7 @@ cc_library_with_android_deps( tf_cc_test( name = "client_client_session_test", + size = "small", srcs = ["client/client_session_test.cc"], deps = [ ":cc_ops", @@ -201,6 +217,7 @@ cc_library_with_android_deps( tf_cc_test( name = "ops_const_op_test", + size = "small", srcs = ["ops/const_op_test.cc"], deps = [ ":const_op", @@ -211,6 +228,26 @@ tf_cc_test( ], ) +cc_library_with_android_deps( + name = "while_loop", + srcs = ["ops/while_loop.cc"], + hdrs = ["ops/while_loop.h"], + android_deps = [ + "//tensorflow/core:android_tensorflow_lib", + ], + common_deps = [ + ":cc_ops", + ":cc_ops_internal", + ":ops", + ":scope", + ":scope_internal", + ], + deps = [ + "//tensorflow/core:core_cpu", + "//tensorflow/core:framework", + ], +) + cc_library( name = "grad_op_registry", srcs = ["framework/grad_op_registry.cc"], @@ -231,10 +268,12 @@ cc_library( ":gradients", "//tensorflow/core:lib_proto_parsing", ], + alwayslink = 1, ) tf_cc_test( name = "gradients_array_grad_test", + size = "small", srcs = ["gradients/array_grad_test.cc"], deps = [ ":array_grad", @@ -255,21 +294,22 @@ cc_library( srcs = ["gradients/math_grad.cc"], deps = [ ":cc_ops", + ":cc_ops_internal", ":grad_op_registry", - ":ops", - ":scope", - "//tensorflow/core:core_cpu", - "//tensorflow/core:framework", + ":gradients", ], + alwayslink = 1, ) tf_cc_test( name = "gradients_math_grad_test", + size = "small", srcs = ["gradients/math_grad_test.cc"], deps = [ ":cc_ops", ":grad_op_registry", ":grad_testutil", + ":gradient_checker", ":math_grad", ":testutil", "//tensorflow/core:lib_internal", @@ -286,15 +326,13 @@ cc_library( ":cc_ops", ":cc_ops_internal", ":grad_op_registry", - ":ops", - ":scope", - "//tensorflow/core:core_cpu", - "//tensorflow/core:framework", ], + alwayslink = 1, ) tf_cc_test( name = "gradients_nn_grad_test", + size = "small", srcs = ["gradients/nn_grad_test.cc"], deps = [ ":cc_ops", @@ -322,6 +360,7 @@ tf_gen_op_wrappers_cc( "io_ops", "linalg_ops", "logging_ops", + "lookup_ops", "math_ops", "nn_ops", "no_op", @@ -343,6 +382,7 @@ tf_gen_op_wrappers_cc( tf_cc_test( name = "framework_cc_ops_test", + size = "small", srcs = ["framework/cc_ops_test.cc"], deps = [ ":cc_ops", @@ -376,6 +416,16 @@ tf_gen_op_wrappers_cc( visibility = ["//tensorflow:internal"], ) +tf_gen_op_wrappers_cc( + name = "functional_ops", + include_internal_ops = 1, + op_lib_names = [ + "functional_ops", + ], + pkg = "//tensorflow/core", + visibility = ["//tensorflow:internal"], +) + tf_gen_op_wrappers_cc( name = "resource_variable_ops", include_internal_ops = 1, @@ -410,6 +460,7 @@ cc_library_with_android_deps( "//tensorflow/core:lib", "//tensorflow/core:lib_internal", "//tensorflow/core:op_gen_lib", + "//tensorflow/core:op_gen_overrides_proto_cc", "//tensorflow/core:proto_text", "//tensorflow/core:protos_all_cc", ], @@ -432,7 +483,6 @@ cc_library( ":client_session", ":ops", ":scope", - "//tensorflow/core:all_kernels", "//tensorflow/core:core_cpu", "//tensorflow/core:lib_internal", "//tensorflow/core:tensorflow", @@ -451,13 +501,25 @@ cc_binary( name = "tutorials_example_trainer", srcs = ["tutorials/example_trainer.cc"], copts = tf_copts(), - linkopts = [ - "-lpthread", - "-lm", - ], + linkopts = select({ + "//tensorflow:windows": [], + "//tensorflow:windows_msvc": [], + "//tensorflow:darwin": [ + "-lm", + "-lpthread", + ], + "//tensorflow:ios": [ + "-lm", + "-lpthread", + ], + "//conditions:default": [ + "-lm", + "-lpthread", + "-lrt", + ], + }), deps = [ ":cc_ops", - "//tensorflow/core:all_kernels", "//tensorflow/core:core_cpu", "//tensorflow/core:framework", "//tensorflow/core:lib", @@ -485,11 +547,9 @@ cc_library( deps = [ ":coordinator", "//tensorflow/core:core_cpu", - "//tensorflow/core:framework", "//tensorflow/core:lib", "//tensorflow/core:lib_internal", "//tensorflow/core:protos_all_cc", - "//tensorflow/core:tensorflow", "//tensorflow/core/kernels:ops_util", ], ) @@ -519,8 +579,6 @@ cc_library( srcs = ["training/coordinator.cc"], hdrs = ["training/coordinator.h"], deps = [ - "//tensorflow/core:core_cpu", - "//tensorflow/core:framework", "//tensorflow/core:lib", "//tensorflow/core:lib_internal", "//tensorflow/core:protos_all_cc", @@ -530,6 +588,7 @@ cc_library( tf_cc_test( name = "coordinator_test", + size = "small", srcs = ["training/coordinator_test.cc"], deps = [ ":cc_ops", diff --git a/tensorflow/cc/client/client_session.cc b/tensorflow/cc/client/client_session.cc index 2732f3f5010d7522a1cf8631183e9b4df7ac86d8..ba056a8f3a84910aebf5079573cb64c19f41469d 100644 --- a/tensorflow/cc/client/client_session.cc +++ b/tensorflow/cc/client/client_session.cc @@ -16,6 +16,7 @@ limitations under the License. #include "tensorflow/cc/client/client_session.h" #include +#include #include #include "tensorflow/core/platform/env.h" @@ -31,7 +32,7 @@ class ClientSession::Impl { friend class ClientSession; Impl(Session* session, std::shared_ptr graph) - : session_(session), graph_(graph) {} + : session_(session), graph_(std::move(graph)) {} static SessionOptions MakeDefaultSessionOptions(const string& target); Status MaybeExtendGraph() const; @@ -112,10 +113,12 @@ Status ClientSession::Run(const RunOptions& run_options, const FeedType& inputs, feeds.emplace_back(feed.first.name(), feed.second.tensor); } std::vector output_tensor_names; + output_tensor_names.reserve(fetch_outputs.size()); for (auto const& output : fetch_outputs) { output_tensor_names.push_back(output.name()); } std::vector target_node_names; + target_node_names.reserve(run_outputs.size()); for (auto const& output : run_outputs) { target_node_names.push_back(output.node()->name()); } diff --git a/tensorflow/cc/client/client_session_test.cc b/tensorflow/cc/client/client_session_test.cc index 9c0f00f2b128c7d06bc7c0ca8579d4ca8e530fe8..dfbac9788e16e9c7c65abcd1ea213b51d5d5d060 100644 --- a/tensorflow/cc/client/client_session_test.cc +++ b/tensorflow/cc/client/client_session_test.cc @@ -49,7 +49,7 @@ TEST(ClientSessionTest, Feed) { TEST(ClientSessionTest, Extend) { Scope root = Scope::NewRootScope(); - auto a = Placeholder(root, DT_INT32); + auto a = Placeholder(root, DT_INT32, Placeholder::Shape({2})); auto c = Add(root, a, {2, 2}); ClientSession session(root); std::vector outputs; diff --git a/tensorflow/cc/framework/cc_op_gen.cc b/tensorflow/cc/framework/cc_op_gen.cc index 22cd7fb0d438db9d9f7f29f5386c7a9722afe43d..38a17598b8e4161f96ab8134823de033d3284440 100644 --- a/tensorflow/cc/framework/cc_op_gen.cc +++ b/tensorflow/cc/framework/cc_op_gen.cc @@ -18,8 +18,12 @@ limitations under the License. #include #include "tensorflow/cc/framework/cc_op_gen.h" +#include "tensorflow/core/framework/attr_value.pb.h" #include "tensorflow/core/framework/attr_value_util.h" #include "tensorflow/core/framework/op_gen_lib.h" +#include "tensorflow/core/framework/op_gen_overrides.pb.h" +#include "tensorflow/core/framework/tensor.pb.h" +#include "tensorflow/core/framework/tensor_shape.pb.h" #include "tensorflow/core/framework/types.pb_text.h" #include "tensorflow/core/lib/gtl/map_util.h" #include "tensorflow/core/lib/gtl/stl_util.h" @@ -126,7 +130,11 @@ string PrintString(const string& str) { return strings::StrCat("\"", str_util::CEscape(str), "\""); } -string PrintTensorShape(const TensorShape& shape) { +string PrintTensorShape(const TensorShapeProto& shape_proto) { + PartialTensorShape shape(shape_proto); + if (shape.IsIdenticalTo(PartialTensorShape())) { + return "::tensorflow::PartialTensorShape() /* unknown */"; + } string ret = "{"; for (int d = 0; d < shape.dims(); ++d) { if (d > 0) strings::StrAppend(&ret, ", "); @@ -188,7 +196,13 @@ string PrintTensor(const TensorProto& tensor_proto) { } } -string PrintAttrValue(string op, const AttrValue& attr_value) { +string PrintTensorProto(const TensorProto& proto) { + return strings::StrCat("Input::Initializer(", "{", PrintTensor(proto), "}, ", + PrintTensorShape(proto.tensor_shape()), + ").AsTensorProto()"); +} + +string PrintAttrValue(const string& op, const AttrValue& attr_value) { switch (attr_value.value_case()) { case AttrValue::kS: return PrintString(attr_value.s()); @@ -203,12 +217,9 @@ string PrintAttrValue(string op, const AttrValue& attr_value) { case AttrValue::kType: return EnumName_DataType(attr_value.type()); case AttrValue::kShape: - return PrintTensorShape(TensorShape(attr_value.shape())); + return PrintTensorShape(attr_value.shape()); case AttrValue::kTensor: - return strings::StrCat( - "Input::Initializer(", "{", PrintTensor(attr_value.tensor()), "}, ", - PrintTensorShape(TensorShape(attr_value.tensor().tensor_shape())), - ").AsTensorProto()"); + return PrintTensorProto(attr_value.tensor()); case AttrValue::kList: { string ret = "{"; if (attr_value.list().s_size() > 0) { @@ -241,8 +252,14 @@ string PrintAttrValue(string op, const AttrValue& attr_value) { } else if (attr_value.list().shape_size() > 0) { for (int i = 0; i < attr_value.list().shape_size(); ++i) { if (i > 0) strings::StrAppend(&ret, ", "); - strings::StrAppend( - &ret, PrintTensorShape(TensorShape(attr_value.list().shape(i)))); + strings::StrAppend(&ret, + PrintTensorShape(attr_value.list().shape(i))); + } + } else if (attr_value.list().tensor_size() > 0) { + for (int i = 0; i < attr_value.list().tensor_size(); ++i) { + if (i > 0) strings::StrAppend(&ret, ", "); + strings::StrAppend(&ret, + PrintTensorProto(attr_value.list().tensor(i))); } } strings::StrAppend(&ret, "}"); @@ -292,8 +309,8 @@ std::pair AttrTypeName(StringPiece attr_type) { {"list(bool)", {"gtl::ArraySlice", true}}, {"type", {"DataType", false}}, {"list(type)", {"DataTypeSlice", true}}, - {"shape", {"TensorShape", false}}, - {"list(shape)", {"gtl::ArraySlice", true}}, + {"shape", {"PartialTensorShape", false}}, + {"list(shape)", {"gtl::ArraySlice", true}}, {"tensor", {"TensorProto", true}}, {"list(tensor)", {"gtl::ArraySlice", true}}, {"func", {"NameAttrList", true}}, @@ -717,7 +734,7 @@ void OpInfo::GetOutput(string* out) const { // One output, no need for NameRangeMap if (is_list_output[0]) { strings::StrAppend(out, - " for (int64 i = 0; i < ret->num_outputs(); ++i)\n"); + " for (int32 i = 0; i < ret->num_outputs(); ++i)\n"); strings::StrAppend(out, " this->", output_names[0], ".push_back(Output(ret, i));\n"); } else { @@ -727,11 +744,10 @@ void OpInfo::GetOutput(string* out) const { return; } strings::StrAppend(out, " ::tensorflow::NameRangeMap _outputs_range;\n"); - strings::StrAppend( - out, - " ::tensorflow::Status _status_ = " - "::tensorflow::NameRangesForNode(ret->def(), ret->op_def(), " - "nullptr, &_outputs_range);\n"); + strings::StrAppend(out, + " ::tensorflow::Status _status_ = " + "::tensorflow::NameRangesForNode(*ret, ret->op_def(), " + "nullptr, &_outputs_range);\n"); strings::StrAppend(out, " if (!_status_.ok()) {\n", " ", scope_str, ".UpdateStatus(_status_);\n", " return;\n"); strings::StrAppend(out, " }\n\n"); @@ -740,7 +756,7 @@ void OpInfo::GetOutput(string* out) const { const string arg_range = strings::StrCat( "_outputs_range[\"", graph_op_def.output_arg(i).name(), "\"]"); if (is_list_output[i]) { - strings::StrAppend(out, " for (int64 i = ", arg_range, ".first; i < ", + strings::StrAppend(out, " for (int32 i = ", arg_range, ".first; i < ", arg_range, ".second; ++i)\n"); strings::StrAppend(out, " this->", output_names[i], ".push_back(Output(ret, i));\n"); @@ -796,12 +812,8 @@ string OpInfo::GetConstructorBody() const { strings::StrAppend(&body, " ", scope_str, ".UpdateStatus(builder.Finalize(", scope_str, ".graph(), &ret));\n"); strings::StrAppend(&body, " ", return_on_error, "\n"); - - // TODO(b/28152992): Enable this code-path once we have converted - // all python shape functions to call their C++ versions. - - // strings::StrAppend(&body, " ", scope_str, ".UpdateStatus(", scope_str, - // ".refiner()->AddNode(ret));\n"); + strings::StrAppend(&body, " ", scope_str, ".UpdateStatus(", scope_str, + ".DoShapeInference(ret));\n"); GetOutput(&body); return body; diff --git a/tensorflow/cc/framework/cc_ops_test.cc b/tensorflow/cc/framework/cc_ops_test.cc index 6dc0d84c16d5b534341575b384997cc398c80bec..5da23036eaadbef270ba839357dc4613bf3bf490 100644 --- a/tensorflow/cc/framework/cc_ops_test.cc +++ b/tensorflow/cc/framework/cc_ops_test.cc @@ -32,10 +32,11 @@ Output Linear(const Scope& scope, Input x, Input w, Input b) { return BiasAdd(cop_scopes.last, m, b); } -void GetColocationConstraints(Output tensor, std::vector* constraints) { +void GetColocationConstraints(const Output& tensor, + std::vector* constraints) { constraints->clear(); - TF_EXPECT_OK( - GetNodeAttr(tensor.op().node()->def(), kColocationAttrName, constraints)); + TF_EXPECT_OK(GetNodeAttr(tensor.op().node()->attrs(), kColocationAttrName, + constraints)); } } // namespace @@ -158,11 +159,11 @@ TEST(CCOpTest, KernelLabel) { Scope root = Scope::NewRootScope(); auto add = Add(root.WithKernelLabel("AddWithKernelLabel"), 1.0f, 2.0f); TF_EXPECT_OK(root.status()); - const auto& attrs = add.z.op().node()->def().attr(); - ASSERT_TRUE(attrs.find("_kernel") != attrs.end()); - auto kernel_attr = attrs.find("_kernel")->second; - TF_EXPECT_OK(AttrValueHasType(kernel_attr, "string")); - EXPECT_EQ(kernel_attr.s(), "AddWithKernelLabel"); + AttrSlice attrs = add.z.op().node()->attrs(); + const auto* kernel_attr = attrs.Find("_kernel"); + ASSERT_TRUE(kernel_attr); + TF_EXPECT_OK(AttrValueHasType(*kernel_attr, "string")); + EXPECT_EQ(kernel_attr->s(), "AddWithKernelLabel"); } TEST(CCOpTest, ColocateWith) { @@ -189,8 +190,7 @@ TEST(CCOpTest, ColocateWith) { Scope with_colocate = root.ColocateWith(c3).ColocateWith(c4); auto c6 = Const(with_colocate.WithOpName("c6").ClearColocation(), 7); - const auto& attrs = c6.op().node()->def().attr(); - EXPECT_TRUE(attrs.find("_class") == attrs.end()); + EXPECT_FALSE(c6.op().node()->attrs().Find("_class")); } TEST(CCOpTest, TemplatedConst) { diff --git a/tensorflow/cc/framework/grad_op_registry.cc b/tensorflow/cc/framework/grad_op_registry.cc index 0d6a377b507161c4420a6076b9ee71e799e0223b..254705736e7711e58aa87054f36c8a19eebd4f0d 100644 --- a/tensorflow/cc/framework/grad_op_registry.cc +++ b/tensorflow/cc/framework/grad_op_registry.cc @@ -32,7 +32,13 @@ bool GradOpRegistry::Register(const string& op, GradFunc func) { Status GradOpRegistry::Lookup(const string& op, GradFunc* func) const { auto iter = registry_.find(op); if (iter == registry_.end()) { - return errors::NotFound("No gradient defined for op: ", op); + const string error_msg = + "No gradient defined for op: " + op + + ". Please see " + "https://www.tensorflow.org/code/" + "tensorflow/cc/gradients/README.md" + " for instructions on how to add C++ gradients."; + return errors::NotFound(error_msg); } *func = iter->second; return Status::OK(); diff --git a/tensorflow/cc/framework/gradient_checker.cc b/tensorflow/cc/framework/gradient_checker.cc index 849a8eed6f23fb8dd1290d1bfa9db9c47d5d9f9d..f3a7c138c4e53c261175a501698246f98b1b5327 100644 --- a/tensorflow/cc/framework/gradient_checker.cc +++ b/tensorflow/cc/framework/gradient_checker.cc @@ -22,8 +22,6 @@ limitations under the License. #include "tensorflow/core/lib/core/errors.h" namespace tensorflow { -using namespace ops; // NOLINT(build/namespaces) - namespace { // TODO(andydavis) Support returning relative error (as opposed to max error) @@ -39,14 +37,16 @@ Status ComputeTheoreticalJacobianTranspose( const std::vector& x_shapes, const std::vector& x_datas, const OutputList& ys, const std::vector& y_shapes, - std::vector& jacobian_ts) { - int y_num = y_shapes.size(); - int x_num = x_shapes.size(); + std::vector* jacobian_ts) { + size_t y_num = y_shapes.size(); + size_t x_num = x_shapes.size(); // Call AddSymbolicGradients to get 'dxs' (we will feed 'dys'). OutputList dys; + dys.reserve(y_shapes.size()); for (const auto& y_shape : y_shapes) { // TODO(suharshs): This currently assumes that all x's are the same type. - dys.push_back(Cast(scope, Const(scope, 1.0, y_shape), xs[0].type())); + dys.push_back( + ops::Cast(scope, ops::Const(scope, 1.0, y_shape), xs[0].type())); } OutputList dxs; TF_RETURN_IF_ERROR(AddSymbolicGradients(scope, ys, xs, dys, &dxs)); @@ -84,7 +84,7 @@ Status ComputeTheoreticalJacobianTranspose( for (int x_idx = 0; x_idx < x_num; x_idx++) { const int64 x_size = x_shapes[x_idx].num_elements(); - auto jacobian = jacobian_ts[x_idx * y_num + y_idx].matrix(); + auto jacobian = (*jacobian_ts)[x_idx * y_num + y_idx].matrix(); auto dx_flat = dxout[x_idx].flat(); for (int r = 0; r < x_size; ++r) { jacobian(r, c) = dx_flat(r); @@ -97,20 +97,20 @@ Status ComputeTheoreticalJacobianTranspose( return Status::OK(); } -Status EvaluateGraph(ClientSession& session, const OutputList& xs, - const OutputList& ys, std::vector& x_datas, +Status EvaluateGraph(ClientSession* session, const OutputList& xs, + const OutputList& ys, std::vector* x_datas, std::vector* y_datas) { // Create the feed list. ClientSession::FeedType feed_list; - for (int i = 0; i < x_datas.size(); i++) { - feed_list.insert({xs[i], x_datas[i]}); + for (int i = 0; i < x_datas->size(); i++) { + feed_list.insert({xs[i], (*x_datas)[i]}); } - TF_RETURN_IF_ERROR(session.Run(feed_list, ys, y_datas)); + TF_RETURN_IF_ERROR(session->Run(feed_list, ys, y_datas)); for (int y_idx = 0; y_idx < y_datas->size(); y_idx++) { - for (int x_idx = 0; x_idx < x_datas.size(); x_idx++) { + for (int x_idx = 0; x_idx < x_datas->size(); x_idx++) { Tensor y_data = (*y_datas)[y_idx]; - if (y_data.SharesBufferWith(x_datas[x_idx])) { + if (y_data.SharesBufferWith((*x_datas)[x_idx])) { // Create copies of outputs that share a buffer with any inputs since // the underlying buffer of the input Tensors are not copied for some // operations (i.e. Identity), which can lead to incorrect results for @@ -128,14 +128,14 @@ Status ComputeNumericJacobianTranspose(const Scope& scope, const OutputList& xs, const OutputList& ys, const std::vector& y_shapes, const T delta, - std::vector& x_datas, - std::vector& jacobian_ts) { - int y_num = y_shapes.size(); - int x_num = x_shapes.size(); + std::vector* x_datas, + std::vector* jacobian_ts) { + size_t y_num = y_shapes.size(); + size_t x_num = x_shapes.size(); ClientSession session(scope); for (int x_idx = 0; x_idx < x_num; x_idx++) { - auto x_data_flat = x_datas[x_idx].flat(); + auto x_data_flat = (*x_datas)[x_idx].flat(); const int64 x_size = x_shapes[x_idx].num_elements(); // Compute the numeric Jacobian one column at a time by perturbing each @@ -147,11 +147,11 @@ Status ComputeNumericJacobianTranspose(const Scope& scope, const OutputList& xs, // Evaluate at positive delta. x_data_flat(r) = v + delta; std::vector y_pos; - TF_RETURN_IF_ERROR(EvaluateGraph(session, xs, ys, x_datas, &y_pos)); + TF_RETURN_IF_ERROR(EvaluateGraph(&session, xs, ys, x_datas, &y_pos)); // Evaluate at negative delta. x_data_flat(r) = v - delta; std::vector y_neg; - TF_RETURN_IF_ERROR(EvaluateGraph(session, xs, ys, x_datas, &y_neg)); + TF_RETURN_IF_ERROR(EvaluateGraph(&session, xs, ys, x_datas, &y_neg)); for (int y_idx = 0; y_idx < y_num; y_idx++) { // Compute element-wise centered difference and store in each Jacobian. @@ -159,7 +159,7 @@ Status ComputeNumericJacobianTranspose(const Scope& scope, const OutputList& xs, auto y_neg_flat = y_neg[y_idx].flat(); const int64 y_size = y_shapes[y_idx].num_elements(); const T scale = 2 * delta; - auto jacobian = jacobian_ts[x_idx * y_num + y_idx].matrix(); + auto jacobian = (*jacobian_ts)[x_idx * y_num + y_idx].matrix(); for (int c = 0; c < y_size; ++c) { jacobian(r, c) = (y_pos_flat(c) - y_neg_flat(c)) / scale; } @@ -175,11 +175,11 @@ template void InitJacobians(const OutputList& xs, const std::vector& x_shapes, const std::vector& y_shapes, - std::vector& jacobians) { - int y_num = y_shapes.size(); - int x_num = x_shapes.size(); + std::vector* jacobians) { + size_t y_num = y_shapes.size(); + size_t x_num = x_shapes.size(); - jacobians.resize(y_num * x_num); + jacobians->resize(y_num * x_num); for (int x_idx = 0; x_idx < x_num; x_idx++) { const int64 x_size = x_shapes[x_idx].num_elements(); for (int y_idx = 0; y_idx < y_num; y_idx++) { @@ -187,7 +187,7 @@ void InitJacobians(const OutputList& xs, Tensor jacobian_t(xs[x_idx].type(), {x_size, y_size}); auto jacobian_t_flat = jacobian_t.flat(); jacobian_t_flat.setZero(); - jacobians[x_idx * y_num + y_idx] = std::move(jacobian_t); + (*jacobians)[x_idx * y_num + y_idx] = std::move(jacobian_t); } } } @@ -197,23 +197,23 @@ Status ComputeGradientErrorInternal(const Scope& scope, const OutputList& xs, const std::vector& x_shapes, const OutputList& ys, const std::vector& y_shapes, - std::vector& x_datas, + std::vector* x_datas, T* max_error) { // Initialize theoretical Jacobians to zeros. std::vector jacobian_ts; - InitJacobians(xs, x_shapes, y_shapes, jacobian_ts); + InitJacobians(xs, x_shapes, y_shapes, &jacobian_ts); // Compute theoretical Jacobian. TF_RETURN_IF_ERROR(ComputeTheoreticalJacobianTranspose( - scope, xs, x_shapes, x_datas, ys, y_shapes, jacobian_ts)); + scope, xs, x_shapes, *x_datas, ys, y_shapes, &jacobian_ts)); // Initialize numeric Jacobian to zeros. std::vector jacobian_ns; - InitJacobians(xs, x_shapes, y_shapes, jacobian_ns); + InitJacobians(xs, x_shapes, y_shapes, &jacobian_ns); // Compute numeric Jacobian. TF_RETURN_IF_ERROR(ComputeNumericJacobianTranspose( - scope, xs, x_shapes, ys, y_shapes, 1e-3, x_datas, jacobian_ns)); + scope, xs, x_shapes, ys, y_shapes, 1e-3, x_datas, &jacobian_ns)); for (int i = 0; i < jacobian_ts.size(); i++) { // Compute the maximum error between theoretical and numeric Jacobians. @@ -256,7 +256,7 @@ Status ComputeGradientError(const Scope& scope, const OutputList& xs, } // Compute gradient error. return ComputeGradientErrorInternal(scope, xs, x_shapes, ys, y_shapes, - x_datas, max_error); + &x_datas, max_error); } template @@ -267,7 +267,7 @@ Status ComputeGradientError(const Scope& scope, const Output& x, std::vector x_datas(1, Tensor(x_init_value)); // Compute gradient error. return ComputeGradientErrorInternal(scope, {x}, {x_datas[0].shape()}, {y}, - {y_shape}, x_datas, max_error); + {y_shape}, &x_datas, max_error); } #define INSTANTIATE_GRAD_ERR_TYPE(T) \ diff --git a/tensorflow/cc/framework/gradients.cc b/tensorflow/cc/framework/gradients.cc index 2c60f947a55479e27937b98de91d80b559d32576..66a943410e2757ea5a5c55351c1fc20d5a5e3154 100644 --- a/tensorflow/cc/framework/gradients.cc +++ b/tensorflow/cc/framework/gradients.cc @@ -65,7 +65,7 @@ class SymbolicGradientBuilder { // gradients for the node associated with `src`. Status BackpropAlongEdge(const Output& dst_grad, const Output& src); - // Adds a node to the graph (returned in`grad`) that sums the in-bound + // Adds a node to the graph (returned in `grad`) that sums the in-bound // gradients to `src` (if there are more than one). Status SumGradients(const Output& src, Output* grad); @@ -152,12 +152,12 @@ Status SymbolicGradientBuilder::Initialize() { grad_outputs_->resize(inputs_.size()); // Populate `output_nodes_` from node ids in `outputs_`. output_nodes_.reserve(outputs_.size()); - for (int i = 0; i < outputs_.size(); ++i) { + for (size_t i = 0; i < outputs_.size(); ++i) { output_nodes_.insert(outputs_[i].node()->id()); } // Populate `input_nodes_` from Outputs in `inputs_`. input_nodes_.reserve(inputs_.size()); - for (int i = 0; i < inputs_.size(); ++i) { + for (size_t i = 0; i < inputs_.size(); ++i) { input_nodes_.insert({inputs_[i], i}); } @@ -210,8 +210,8 @@ Status SymbolicGradientBuilder::Initialize() { { // Initialize backprop with `grad_inputs_`. - const int num_dy = grad_inputs_.size(); - for (int i = 0; i < num_dy; ++i) { + const size_t num_dy = grad_inputs_.size(); + for (size_t i = 0; i < num_dy; ++i) { TF_RETURN_IF_ERROR(BackpropAlongEdge(grad_inputs_[i], outputs_[i])); } } @@ -308,7 +308,7 @@ Status SymbolicGradientBuilder::AddGradients() { continue; } - const int num_no_grad = no_grad_dy_indices.size(); + const size_t num_no_grad = no_grad_dy_indices.size(); if (IsPrimitiveOpWithNoGrad(n->type_string()) || num_no_grad == num_y) { // No grad defined for this op, or all outputs returned 'NoGradient': // Backprop 'NoGradient' along the in edges. @@ -341,7 +341,7 @@ Status SymbolicGradientBuilder::AddGradients() { // gradient function to the src node/output to which it should be // backproped. Maybe grad functions can return a vector of Output pairs to // make this association explicit. - int dx_index = 0; + size_t dx_index = 0; for (const Edge* e : n->in_edges()) { if (e->IsControlEdge()) continue; if (dx_index == dx.size()) { @@ -352,6 +352,23 @@ Status SymbolicGradientBuilder::AddGradients() { BackpropAlongEdge(dx[dx_index++], {e->src(), e->src_output()})); } } + + // Check if any input nodes still have pending gradients and have not been + // processed yet. This happens if not all outputs of a node are in 'inputs_'. + std::unordered_map requested_grads; + for (const Output& nout : inputs_) { + if (pending_[nout.node()->id()] > 0) { + DCHECK_GT(nout.node()->num_outputs(), 1); + int idx = input_nodes_[nout]; + DCHECK(((*grad_outputs_)[idx].node() == nullptr)); + TF_RETURN_IF_ERROR(SumGradients(nout, &(*grad_outputs_)[idx])); + ++requested_grads[nout.node()]; + } + } + for (const auto& p : requested_grads) { + int num_requested_inputs = p.first->num_outputs() - pending_[p.first->id()]; + CHECK_EQ(num_requested_inputs, p.second); + } return Status::OK(); } @@ -367,6 +384,19 @@ Status AddSymbolicGradients(const Scope& scope, return builder.AddGradients(); } +Status AddSymbolicGradients(const Scope& scope, + const std::vector& outputs, + const std::vector& inputs, + std::vector* grad_outputs) { + std::vector grad_inputs; + grad_inputs.reserve(outputs.size()); + for (const Output& output : outputs) { + grad_inputs.emplace_back(ops::OnesLike(scope, output)); + } + return AddSymbolicGradients(scope, outputs, inputs, grad_inputs, + grad_outputs); +} + Output NoGradient() { return SymbolicGradientBuilder::NoGradient(); } } // end namespace tensorflow diff --git a/tensorflow/cc/framework/gradients.h b/tensorflow/cc/framework/gradients.h index d076bc43b4fbb1c8911b52c5ab258b7e9837113b..717f6f0636d3dd1a546ef7477b100bbfc86ba13d 100644 --- a/tensorflow/cc/framework/gradients.h +++ b/tensorflow/cc/framework/gradients.h @@ -27,16 +27,19 @@ namespace tensorflow { /// derivatives of some loss function 'L' w.r.t 'outputs'), adds gradient nodes /// to the graph associated with 'scope', which compute (and return in /// 'grad_outputs') the symbolic partial derivatives of 'L' w.r.t 'inputs'. -/// - -// TODO(andydavis) Add overload of this function with no 'grad_inputs' arg. -// Implementation will fill in 'OnesLike' for all shapes in 'outputs'. Status AddSymbolicGradients(const Scope& scope, const std::vector& outputs, const std::vector& inputs, const std::vector& grad_inputs, std::vector* grad_outputs); +// Same as above, but uses 'OnesLike' for all shapes in +// 'outputs' as grad_inputs. +Status AddSymbolicGradients(const Scope& scope, + const std::vector& outputs, + const std::vector& inputs, + std::vector* grad_outputs); + /// Returns a sentinel Output that represents 'no gradient' (i.e. no gradient /// flows along some graph edge during backpropagation). /// Can be returned in 'grad_outputs' by an invocation of 'AddSymbolicGradients' diff --git a/tensorflow/cc/framework/gradients_test.cc b/tensorflow/cc/framework/gradients_test.cc index 6c2c2fcd1e2c5941dadebfbc78fb5bae9122e7c3..24af7d567b267332610eba2c8c8c57681fa0559b 100644 --- a/tensorflow/cc/framework/gradients_test.cc +++ b/tensorflow/cc/framework/gradients_test.cc @@ -17,6 +17,7 @@ limitations under the License. #include "tensorflow/cc/framework/grad_op_registry.h" #include "tensorflow/cc/framework/testutil.h" #include "tensorflow/cc/ops/standard_ops.h" +#include "tensorflow/core/framework/graph.pb.h" #include "tensorflow/core/framework/node_def_util.h" #include "tensorflow/core/framework/tensor_testutil.h" #include "tensorflow/core/lib/core/status_test_util.h" @@ -40,7 +41,7 @@ class GradientsTest : public ::testing::Test { TF_ASSERT_OK(scope_test_.ToGraphDef(&gdef_test)); GraphDef gdef_exp; TF_ASSERT_OK(scope_expected_.ToGraphDef(&gdef_exp)); - TF_EXPECT_GRAPH_EQ(gdef_test, gdef_exp); + TF_EXPECT_GRAPH_EQ(gdef_exp, gdef_test); } Scope scope_expected_; @@ -98,6 +99,32 @@ TEST_F(GradientsTest, OneMatMul) { CompareTestAndExpectedGraphs(); } +TEST_F(GradientsTest, OneMatMul_InferGradInputs) { + for (const bool expected : {false, true}) { + const Scope& scope = expected ? scope_expected_ : scope_test_; + // Construct forward graph. + auto x = Const(scope, {{1.0, 2.0}, {3.0, 4.0}}); + auto y = Const(scope, {{1.0, 0.0}, {0.0, 1.0}}); + auto z = MatMul(scope, x, y); + TF_ASSERT_OK(scope.status()); + CHECK_NOTNULL(z.node()); + + if (expected) { + // Construct backward graph. + // The gradients function adds a OnesLike to create a dz of ones with the + // shape of z. + auto dz = OnesLike(scope, z); + auto dx = MatMul(scope, dz, y, MatMul::TransposeB(true)); + auto dy = MatMul(scope, x, dz, MatMul::TransposeA(true)); + } else { + // Call AddSymbolicGradients. + std::vector grad_outputs; + TF_ASSERT_OK(AddSymbolicGradients(scope, {z}, {x, y}, &grad_outputs)); + } + } + CompareTestAndExpectedGraphs(); +} + TEST_F(GradientsTest, TwoMatMuls_Chained) { for (const bool expected : {false, true}) { const Scope& scope = expected ? scope_expected_ : scope_test_; @@ -233,8 +260,44 @@ TEST_F(GradientsTest, StackUnstack_StopBackprop) { CompareTestAndExpectedGraphs(); } +TEST_F(GradientsTest, StackUnstack_SubsetOfUnstackOutputs) { + // Constructs an unstack with three outputs, and takes the gradient with + // respect to only two of the outputs. Tests that the output gradients are + // computed. + for (const bool expected : {false, true}) { + const Scope& scope = expected ? scope_expected_ : scope_test_; + // Construct forward graph. + auto c = Const(scope, 1, {3, 4, 2}); + auto unpack = Unstack(scope, c, 3); + auto x = Identity(scope, unpack.output[0]); + auto y = Identity(scope, unpack.output[1]); + auto z = Identity(scope, unpack.output[2]); + TF_ASSERT_OK(scope.status()); + + // Construct grad inputs. + auto dy = Const(scope, 4, {4, 2}); + auto dz = Const(scope, 5, {4, 2}); + + if (expected) { + // Construct backward graph. + auto g1 = Identity(scope, dy); + auto g2 = Identity(scope, dz); + } else { + // Call AddSymbolicGradients. + std::vector grad_outputs; + TF_ASSERT_OK(AddSymbolicGradients(scope, {y, z}, + {unpack.output[1], unpack.output[2]}, + {dy, dz}, &grad_outputs)); + ASSERT_EQ(grad_outputs.size(), 2); + EXPECT_TRUE(grad_outputs[0].node() != nullptr); + EXPECT_TRUE(grad_outputs[1].node() != nullptr); + } + } + CompareTestAndExpectedGraphs(); +} + TEST_F(GradientsTest, DependentGradOutputs) { - // Tests that dependant gradients (in this case the gradients w.r.t to the + // Tests that dependent gradients (in this case the gradients w.r.t to the // output and one input of MatMul) are computed properly. // Create two chained MatMul ops. diff --git a/tensorflow/cc/framework/ops.cc b/tensorflow/cc/framework/ops.cc index 50df891a4c434ad58e962d7a31599df08cedaeb7..920a8e7955631ba0d33d2d36506703e107420a69 100644 --- a/tensorflow/cc/framework/ops.cc +++ b/tensorflow/cc/framework/ops.cc @@ -20,7 +20,7 @@ namespace tensorflow { Operation::Operation(Node* n) : inputs_(GetInputs(n)), node_(n) {} -Output Operation::input(int i) const { +Output Operation::input(int32 i) const { CHECK_NOTNULL(node_); CHECK_GE(i, 0); CHECK_LT(i, node_->num_inputs()); @@ -37,14 +37,14 @@ Output Operation::input(int i) const { return Output(inputs_[i].first, inputs_[i].second); } -Output Operation::output(int i) const { +Output Operation::output(int32 i) const { CHECK_NOTNULL(node_); CHECK_GE(i, 0); CHECK_LT(i, node_->num_outputs()); return Output(node_, i); } -uint64 Operation::hash(int64 index) const { +uint64 Operation::hash(int32 index) const { return ::tensorflow::Hash64(reinterpret_cast(&node_), sizeof(Node*), index); } diff --git a/tensorflow/cc/framework/ops.h b/tensorflow/cc/framework/ops.h index 889d5db31dd06fd25b7a72e209a8d7f37b8429ca..8d4154220c4b18f9286094b10c1b1e96eb4e31e7 100644 --- a/tensorflow/cc/framework/ops.h +++ b/tensorflow/cc/framework/ops.h @@ -39,22 +39,22 @@ class Operation { Operation() : node_(nullptr) {} explicit Operation(Node* n); - int num_inputs() const { return node_->num_inputs(); } - DataType input_type(int o) const { return node_->input_type(o); } - Output input(int i) const; + int32 num_inputs() const { return node_->num_inputs(); } + DataType input_type(int32 o) const { return node_->input_type(o); } + Output input(int32 i) const; - int num_outputs() const { return node_->num_outputs(); } - DataType output_type(int o) const { return node_->output_type(o); } - Output output(int i) const; + int32 num_outputs() const { return node_->num_outputs(); } + DataType output_type(int32 o) const { return node_->output_type(o); } + Output output(int32 i) const; Node* node() const { return node_; } - uint64 hash(int64 index) const; + uint64 hash(int32 index) const; bool operator==(const Operation& other) const { return node_ == other.node_; } private: - typedef std::vector> Inputs; + typedef std::vector> Inputs; static Inputs GetInputs(Node* node); Inputs inputs_; @@ -66,12 +66,12 @@ class Output { public: Output() = default; explicit Output(Node* n) : op_(n) {} - Output(Node* n, int64 index) : op_(n), index_(index) {} - Output(const Operation& op, int64 index) : op_(op), index_(index) {} + Output(Node* n, int32 index) : op_(n), index_(index) {} + Output(const Operation& op, int32 index) : op_(op), index_(index) {} Operation op() const { return op_; } Node* node() const { return op().node(); } - int64 index() const { return index_; } + int32 index() const { return index_; } DataType type() const { return op_.output_type(index_); } string name() const { return strings::StrCat(node()->name(), ":", index()); } bool operator==(const Output& other) const { @@ -82,14 +82,14 @@ class Output { private: Operation op_ = Operation(nullptr); - int64 index_ = 0; + int32 index_ = 0; }; /// Hash class that can be used for e.g. storing Outputs in an unordered_map struct OutputHash { std::size_t operator()(const Output& output) const { return Hash64Combine(std::hash()(output.node()), - std::hash()(output.index())); + std::hash()(output.index())); } }; @@ -230,12 +230,12 @@ class Input { /// Constructor specifying a node name, index and datatype. This should only /// be used for specifying a backward edge, needed by control flow. - Input(const string& name, int i, DataType dt) + Input(const string& name, int32 i, DataType dt) : node_name_(name), index_(i), data_type_(dt) {} Node* node() const { return output_.node(); } string node_name() const { return node_name_; } - int index() const { return node_name_.empty() ? output_.index() : index_; } + int32 index() const { return node_name_.empty() ? output_.index() : index_; } DataType data_type() const { return data_type_; } Status status() const { return status_; } const Tensor& tensor() const { return tensor_; } @@ -245,7 +245,7 @@ class Input { Output output_ = Output(Operation(nullptr), 0); Tensor tensor_; const string node_name_ = ""; - int index_ = 0; + int32 index_ = 0; DataType data_type_ = DT_INVALID; }; diff --git a/tensorflow/cc/framework/scope.cc b/tensorflow/cc/framework/scope.cc index 571c6e1e579f630db473ffc1312d1a1f3162f475..71642492627422e09c19b7bcb4dc522846cf08b1 100644 --- a/tensorflow/cc/framework/scope.cc +++ b/tensorflow/cc/framework/scope.cc @@ -16,7 +16,7 @@ limitations under the License. #include #include -#include "tensorflow/cc/framework/scope.h" +#include "tensorflow/cc/framework/scope_internal.h" #include "tensorflow/core/common_runtime/shape_refiner.h" #include "tensorflow/core/framework/node_def_util.h" #include "tensorflow/core/graph/graph_constructor.h" @@ -24,77 +24,6 @@ limitations under the License. namespace tensorflow { -class Scope::Impl { - private: - friend class Scope; - - // Tag types to choose the constructor to dispatch. - struct Tags { - enum class ScopeName; - enum class OpName; - enum class ControlDeps; - enum class Device; - enum class SingleUseScope; - enum class ExitOnError; - enum class KernelLabel; - enum class Colocate; - }; - - // A NameMap is used to keep track of suffixes for names used in a scope. A - // name that has not been used so far in a scope will get no suffix. Later - // uses of the same name will get suffixes _1, _2, _3, etc. Multiple scopes - // can share the same NameMap. For instance, a new scope created using - // WithControlDependencies() should would share the same NameMap with the - // parent. - typedef std::unordered_map NameMap; - - Impl(Graph* graph, Status* status, NameMap* name_map, ShapeRefiner* refiner); - Impl(const Scope& other, Tags::ScopeName, const string& name, - bool copy_names); - Impl(const Scope& other, Tags::OpName, const string& name, - const string& op_name); - Impl(const Scope& other, Tags::ControlDeps, - std::vector control_deps, bool clear_control_deps); - Impl(const Scope& other, Tags::Device, const string& device); - Impl(const Scope& other, Tags::SingleUseScope, const string& op_name); - Impl(const Scope& other, Tags::ExitOnError); - Impl(const Scope& other, Tags::KernelLabel, const string& kernel_label); - Impl(const Scope& other, Tags::Colocate, const Operation& colocate_with_op, - bool clear_colocations); - - std::unordered_set GetColocationConstraints( - const Operation& colocate_with_op) const; - - // Helper functions to get a unique names. - string GetUniqueName(const string& prefix, bool check_single_use) const; - string GetNameForOp(const string& default_name) const; - - bool single_use_scope() const { return scope_used_ != nullptr; } - - // The graph, status, and name maps are shared by all child scopes - // created from a single 'root' scope. A root scope is created by calling the - // Scope::NewRootScope function, which creates a new graph, a new status and - // the name maps. - std::shared_ptr graph_ = nullptr; - std::shared_ptr status_ = nullptr; - std::shared_ptr name_map_ = nullptr; - std::shared_ptr refiner_ = nullptr; - - // If scope_used_ is not nullptr, op_name_ should be empty and - // GetUniqueNameForOp can only be called once on this scope. More calls to - // GetUniqueNameForOp will cause an error status to be set on this scope. - std::shared_ptr scope_used_ = nullptr; - - const std::vector control_deps_; - - const string name_ = ""; - const string op_name_ = ""; - const bool exit_on_error_ = false; - const string kernel_label_ = ""; - const string device_ = ""; - const std::unordered_set colocation_constraints_; -}; - Scope::Scope(Impl* impl) : impl_(impl) {} Scope::Scope(const Scope& other) : impl_(new Impl(*other.impl())) {} @@ -108,19 +37,41 @@ Scope& Scope::operator=(const Scope& other) { } Scope::Impl::Impl(Graph* graph, Status* status, NameMap* name_map, - ShapeRefiner* refiner) + ShapeRefiner* refiner, bool disable_shape_inference) + : graph_(graph), + status_(status), + name_map_(name_map), + refiner_(refiner), + scope_used_(nullptr), + colocation_constraints_(), + disable_shape_inference_(disable_shape_inference) {} + +Scope::Impl::Impl(const std::shared_ptr& graph, + const std::shared_ptr& status, + const std::shared_ptr& name_map, + const std::shared_ptr& refiner) : graph_(graph), status_(status), name_map_(name_map), refiner_(refiner), scope_used_(nullptr), - colocation_constraints_() {} + colocation_constraints_(), + disable_shape_inference_(false) {} Scope Scope::NewRootScope() { Graph* graph = new Graph(OpRegistry::Global()); ShapeRefiner* refiner = - new ShapeRefiner(graph->versions().producer(), graph->op_registry()); - return Scope(new Impl(graph, new Status, new Impl::NameMap, refiner)); + new ShapeRefiner(graph->versions(), graph->op_registry()); + return Scope(new Impl(graph, new Status, new Impl::NameMap, refiner, + /* disable_shape_inference */ false)); +} + +Scope Scope::DisabledShapeInferenceScope() { + Graph* graph = new Graph(OpRegistry::Global()); + ShapeRefiner* refiner = + new ShapeRefiner(graph->versions(), graph->op_registry()); + return Scope(new Impl(graph, new Status, new Impl::NameMap, refiner, + /* disable_shape_inference */ true)); } Scope::Impl::Impl(const Scope& other, Tags::ScopeName, const string& name, @@ -137,7 +88,8 @@ Scope::Impl::Impl(const Scope& other, Tags::ScopeName, const string& name, exit_on_error_(other.impl()->exit_on_error_), kernel_label_(other.impl()->kernel_label_), device_(other.impl()->device_), - colocation_constraints_(other.impl()->colocation_constraints_) {} + colocation_constraints_(other.impl()->colocation_constraints_), + disable_shape_inference_(other.impl()->disable_shape_inference_) {} Scope::Impl::Impl(const Scope& other, Tags::OpName, const string& name, const string& op_name) @@ -152,7 +104,8 @@ Scope::Impl::Impl(const Scope& other, Tags::OpName, const string& name, exit_on_error_(other.impl()->exit_on_error_), kernel_label_(other.impl()->kernel_label_), device_(other.impl()->device_), - colocation_constraints_(other.impl()->colocation_constraints_) {} + colocation_constraints_(other.impl()->colocation_constraints_), + disable_shape_inference_(other.impl()->disable_shape_inference_) {} Scope::Impl::Impl(const Scope& other, Tags::ControlDeps, std::vector control_deps, bool clear_control_deps) @@ -173,7 +126,8 @@ Scope::Impl::Impl(const Scope& other, Tags::ControlDeps, exit_on_error_(other.impl()->exit_on_error_), kernel_label_(other.impl()->kernel_label_), device_(other.impl()->device_), - colocation_constraints_(other.impl()->colocation_constraints_) {} + colocation_constraints_(other.impl()->colocation_constraints_), + disable_shape_inference_(other.impl()->disable_shape_inference_) {} Scope::Impl::Impl(const Scope& other, Tags::Device, const string& device) : graph_(other.impl()->graph_), @@ -187,7 +141,8 @@ Scope::Impl::Impl(const Scope& other, Tags::Device, const string& device) exit_on_error_(other.impl()->exit_on_error_), kernel_label_(other.impl()->kernel_label_), device_(device), - colocation_constraints_(other.impl()->colocation_constraints_) {} + colocation_constraints_(other.impl()->colocation_constraints_), + disable_shape_inference_(other.impl()->disable_shape_inference_) {} Scope::Impl::Impl(const Scope& other, Tags::SingleUseScope, const string& op_name) @@ -202,7 +157,8 @@ Scope::Impl::Impl(const Scope& other, Tags::SingleUseScope, exit_on_error_(other.impl()->exit_on_error_), kernel_label_(other.impl()->kernel_label_), device_(other.impl()->device_), - colocation_constraints_(other.impl()->colocation_constraints_) {} + colocation_constraints_(other.impl()->colocation_constraints_), + disable_shape_inference_(other.impl()->disable_shape_inference_) {} Scope::Impl::Impl(const Scope& other, Tags::ExitOnError) : graph_(other.impl()->graph_), @@ -216,7 +172,8 @@ Scope::Impl::Impl(const Scope& other, Tags::ExitOnError) exit_on_error_(true), kernel_label_(other.impl()->kernel_label_), device_(other.impl()->device_), - colocation_constraints_(other.impl()->colocation_constraints_) {} + colocation_constraints_(other.impl()->colocation_constraints_), + disable_shape_inference_(other.impl()->disable_shape_inference_) {} Scope::Impl::Impl(const Scope& other, Tags::KernelLabel, const string& kernel_label) @@ -231,7 +188,8 @@ Scope::Impl::Impl(const Scope& other, Tags::KernelLabel, exit_on_error_(other.impl()->exit_on_error_), kernel_label_(kernel_label), device_(other.impl()->device_), - colocation_constraints_(other.impl()->colocation_constraints_) {} + colocation_constraints_(other.impl()->colocation_constraints_), + disable_shape_inference_(other.impl()->disable_shape_inference_) {} Scope::Impl::Impl(const Scope& other, Tags::Colocate, const Operation& colocate_with_op, bool clear_colocations) @@ -249,14 +207,15 @@ Scope::Impl::Impl(const Scope& other, Tags::Colocate, colocation_constraints_( clear_colocations ? std::unordered_set() - : other.impl()->GetColocationConstraints(colocate_with_op)) {} + : other.impl()->GetColocationConstraints(colocate_with_op)), + disable_shape_inference_(other.impl()->disable_shape_inference_) {} std::unordered_set Scope::Impl::GetColocationConstraints( const Operation& colocate_with_op) const { std::unordered_set current_constraints(colocation_constraints_); - const NodeDef& node_def = colocate_with_op.node()->def(); + const AttrSlice attrs = colocate_with_op.node()->attrs(); std::vector node_constraints; - if (GetNodeAttr(node_def, kColocationAttrName, &node_constraints).ok()) { + if (GetNodeAttr(attrs, kColocationAttrName, &node_constraints).ok()) { for (const string& entry : node_constraints) { StringPiece s(entry); if (s.Consume(kColocationGroupPrefix)) { @@ -277,7 +236,7 @@ std::shared_ptr Scope::graph_as_shared_ptr() const { return impl()->graph_; } -Status Scope::status() const { return *impl()->status_; }; +Status Scope::status() const { return *impl()->status_; } const std::vector& Scope::control_deps() const { return impl()->control_deps_; @@ -464,4 +423,31 @@ CompositeOpScopes Scope::GetCompositeOpScopes( } } +Status Scope::DoShapeInference(Node* node) const { + if (impl_->disable_shape_inference_) return Status::OK(); + return impl_->refiner_->AddNode(node); +} + +class InternalScope { + public: + // NewScope doesn't take ownership of the inputs. + static Scope NewScope(Graph* graph, Status* status, ShapeRefiner* refiner) { + Scope::Impl::NameMap* name_map = new Scope::Impl::NameMap; + for (const Node* node : graph->nodes()) { + (*name_map)[node->name()] = 0; + } + // We provide null destructors for these shared ptrs (except for name_map) + // since the caller owns them and doesn't want the scope to destroy them. + return Scope(new Scope::Impl( + std::shared_ptr(graph, [](Graph*) {}), + std::shared_ptr(status, [](Status*) {}), + std::shared_ptr(name_map), + std::shared_ptr(refiner, [](ShapeRefiner*) {}))); + } +}; + +Scope NewInternalScope(Graph* graph, Status* status, ShapeRefiner* refiner) { + return InternalScope::NewScope(graph, status, refiner); +} + } // namespace tensorflow diff --git a/tensorflow/cc/framework/scope.h b/tensorflow/cc/framework/scope.h index ce70da709630bd402be9c75b3f6a5d638cd4a588..0335f6357d0cb4bf0d586a17856bbf46f23d34d9 100644 --- a/tensorflow/cc/framework/scope.h +++ b/tensorflow/cc/framework/scope.h @@ -167,7 +167,8 @@ class Scope { // START_SKIP_DOXYGEN - /// Update the builder with properties accumulated in this scope. + /// Update the builder with properties accumulated in this scope. Does not set + /// status(). // TODO(skyewm): NodeBuilder is not part of public API void UpdateBuilder(NodeBuilder* builder) const; // END_SKIP_DOXYGEN @@ -199,15 +200,31 @@ class Scope { // edges from the source and to the sink node, resolves back edges // by name), and makes sure the resulting graph is valid. Status ToGraph(Graph* g) const; + + // Calls AddNode() using this scope's ShapeRefiner. This exists in the public + // API to prevent custom op wrappers from needing access to shape_refiner.h or + // scope_internal.h. + // TODO(skyewm): remove this from public API + Status DoShapeInference(Node* node) const; + + // Creates a new root scope that causes all DoShapeInference() calls to return + // Status::OK() (on the returned scope and any subscopes). Used for testing. + // TODO(skyewm): fix tests that still require this and eventually remove, or + // at least remove from public API + static Scope DisabledShapeInferenceScope(); // END_SKIP_DOXYGEN const std::vector& control_deps() const; - private: + // START_SKIP_DOXYGEN class Impl; - std::unique_ptr impl_; Impl* impl() { return impl_.get(); } const Impl* impl() const { return impl_.get(); } + // END_SKIP_DOXYGEN + + private: + 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 new file mode 100644 index 0000000000000000000000000000000000000000..968c366550ef6f46557cd9b5662d9d0719b31531 --- /dev/null +++ b/tensorflow/cc/framework/scope_internal.h @@ -0,0 +1,120 @@ +/* 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 THIRD_PARTY_TENSORFLOW_CC_FRAMEWORK_SCOPE_INTERNAL_H_ +#define THIRD_PARTY_TENSORFLOW_CC_FRAMEWORK_SCOPE_INTERNAL_H_ + +#include "tensorflow/cc/framework/scope.h" + +namespace tensorflow { + +class ShapeRefiner; + +// NewInternalScope returns a new scope which doesn't take ownership of +// graph, status, name_map, and refiner. +// This is intended to enable the C API (which are used by other language +// bindings) to create a Scope and access C++ functionality (i.e. gradients). +Scope NewInternalScope(Graph* graph, Status* status, ShapeRefiner* refiner); + +class Scope::Impl { + public: + // A NameMap is used to keep track of suffixes for names used in a scope. A + // name that has not been used so far in a scope will get no suffix. Later + // uses of the same name will get suffixes _1, _2, _3, etc. Multiple scopes + // can share the same NameMap. For instance, a new scope created using + // WithControlDependencies() should would share the same NameMap with the + // parent. + typedef std::unordered_map NameMap; + + Impl(const std::shared_ptr& graph, + const std::shared_ptr& status, + const std::shared_ptr& name_map, + const std::shared_ptr& refiner); + + const string& name() const { return name_; } + const std::vector& control_deps() const { return control_deps_; } + + private: + friend class Scope; + + // Tag types to choose the constructor to dispatch. + struct Tags { + enum class ScopeName; + enum class OpName; + enum class ControlDeps; + enum class Device; + enum class SingleUseScope; + enum class ExitOnError; + enum class KernelLabel; + enum class Colocate; + }; + + Impl(Graph* graph, Status* status, NameMap* name_map, ShapeRefiner* refiner, + bool disable_shape_inference); + Impl(const Scope& other, Tags::ScopeName, const string& name, + bool copy_names); + Impl(const Scope& other, Tags::OpName, const string& name, + const string& op_name); + Impl(const Scope& other, Tags::ControlDeps, + std::vector control_deps, bool clear_control_deps); + Impl(const Scope& other, Tags::Device, const string& device); + Impl(const Scope& other, Tags::SingleUseScope, const string& op_name); + Impl(const Scope& other, Tags::ExitOnError); + Impl(const Scope& other, Tags::KernelLabel, const string& kernel_label); + Impl(const Scope& other, Tags::Colocate, const Operation& colocate_with_op, + bool clear_colocations); + + std::unordered_set GetColocationConstraints( + const Operation& colocate_with_op) const; + + // Helper functions to get a unique names. + string GetUniqueName(const string& prefix, bool check_single_use) const; + string GetNameForOp(const string& default_name) const; + + bool single_use_scope() const { return scope_used_ != nullptr; } + + // The graph, status, and name maps are shared by all child scopes + // created from a single 'root' scope. A root scope is created by calling the + // Scope::NewRootScope function, which creates a new graph, a new status and + // the name maps. + std::shared_ptr graph_ = nullptr; + std::shared_ptr status_ = nullptr; + std::shared_ptr name_map_ = nullptr; + std::shared_ptr refiner_ = nullptr; + + // If scope_used_ is not nullptr, op_name_ should be empty and + // GetUniqueNameForOp can only be called once on this scope. More calls to + // GetUniqueNameForOp will cause an error status to be set on this scope. + std::shared_ptr scope_used_ = nullptr; + + const std::vector control_deps_; + + // The fully-qualified name of this scope (i.e. includes any parent scope + // names). + const string name_ = ""; + const string op_name_ = ""; + const bool exit_on_error_ = false; + const string kernel_label_ = ""; + const string device_ = ""; + const std::unordered_set colocation_constraints_; + + // If true, Scope::DoShapeInference() always returns Status:OK(). + // TODO(skyewm): remove this when possible + const bool disable_shape_inference_; +}; + +} // namespace tensorflow + +#endif // THIRD_PARTY_TENSORFLOW_CC_FRAMEWORK_SCOPE_INTERNAL_H_ diff --git a/tensorflow/cc/framework/test_op.cc b/tensorflow/cc/framework/test_op.cc index fe0d907df0b6d95416807fd6750ae183c2e59a90..b76842a9a0ca0783063c8a3a88ebef0dff7a51b7 100644 --- a/tensorflow/cc/framework/test_op.cc +++ b/tensorflow/cc/framework/test_op.cc @@ -13,6 +13,7 @@ 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 { @@ -24,6 +25,7 @@ REGISTER_OP("ThrowAway1") .Attr("scope: int") .Attr("builder: int = 1") .Attr("while: int") + .SetShapeFn(shape_inference::UnknownShape) .Doc(R"doc( Op to test keywords and reserved words in input and attr names. @@ -36,12 +38,20 @@ REGISTER_OP("ThrowAway2") .Attr("scope: int = 2") .Attr("throw_away2: int = 2") .Attr("attrs: int = 4") - .Attr("node: int = 4"); + .Attr("node: int = 4") + .SetShapeFn(shape_inference::UnknownShape); -REGISTER_OP("ThrowAway3").Output("node: int32"); +REGISTER_OP("ThrowAway3") + .Output("node: int32") + .SetShapeFn(shape_inference::UnknownShape); -REGISTER_OP("ThrowAway4").Input("node: int32"); +REGISTER_OP("ThrowAway4") + .Input("node: int32") + .SetShapeFn(shape_inference::UnknownShape); -REGISTER_OP("ThrowAway5").Output("foo: int32").Attr("node: int = 4"); +REGISTER_OP("ThrowAway5") + .Output("foo: int32") + .Attr("node: int = 4") + .SetShapeFn(shape_inference::UnknownShape); } // namespace tensorflow diff --git a/tensorflow/cc/framework/testutil.cc b/tensorflow/cc/framework/testutil.cc index b0746913a16b78fd3855572b042c47e8b7445e77..ca78f31db513f043d02594e100e549cb16e92795 100644 --- a/tensorflow/cc/framework/testutil.cc +++ b/tensorflow/cc/framework/testutil.cc @@ -15,6 +15,8 @@ limitations under the License. #include "tensorflow/cc/framework/testutil.h" +#include + #include "tensorflow/cc/client/client_session.h" #include "tensorflow/core/framework/tensor_testutil.h" #include "tensorflow/core/graph/default_device.h" @@ -30,7 +32,7 @@ void GetTensors(const Scope& scope, OutputList tensors, void GetTensor(const Scope& scope, Output tensor, Tensor* out) { std::vector outputs; - GetTensors(scope, {tensor}, &outputs); + GetTensors(scope, {std::move(tensor)}, &outputs); *out = outputs[0]; } diff --git a/tensorflow/cc/gradients/array_grad.cc b/tensorflow/cc/gradients/array_grad.cc index 26abd2438e652f29a1d25caf689ab0606a12b00a..6545e4ee3eb406436937a43ddac66d017af8e108 100644 --- a/tensorflow/cc/gradients/array_grad.cc +++ b/tensorflow/cc/gradients/array_grad.cc @@ -43,9 +43,9 @@ Status PackGrad(const Scope& scope, const Operation& op, const std::vector& grad_inputs, std::vector* grad_outputs) { int N; - TF_RETURN_IF_ERROR(GetNodeAttr(op.node()->def(), "N", &N)); + TF_RETURN_IF_ERROR(GetNodeAttr(op.node()->attrs(), "N", &N)); int axis; - TF_RETURN_IF_ERROR(GetNodeAttr(op.node()->def(), "axis", &axis)); + TF_RETURN_IF_ERROR(GetNodeAttr(op.node()->attrs(), "axis", &axis)); grad_outputs->reserve(N); auto grad_op = Unstack(scope, grad_inputs[0], N, Unstack::Axis(axis)); @@ -60,7 +60,7 @@ Status UnpackGrad(const Scope& scope, const Operation& op, const std::vector& grad_inputs, std::vector* grad_outputs) { int axis; - TF_RETURN_IF_ERROR(GetNodeAttr(op.node()->def(), "axis", &axis)); + TF_RETURN_IF_ERROR(GetNodeAttr(op.node()->attrs(), "axis", &axis)); grad_outputs->push_back(Stack(scope, grad_inputs, Stack::Axis(axis))); return scope.status(); } @@ -100,6 +100,17 @@ Status QuantizeAndDequantizeV2Grad(const Scope& scope, const Operation& op, } REGISTER_GRADIENT_OP("QuantizeAndDequantizeV2", QuantizeAndDequantizeV2Grad); +Status QuantizeAndDequantizeV3Grad(const Scope& scope, const Operation& op, + const std::vector& grad_inputs, + std::vector* grad_outputs) { + grad_outputs->push_back(Identity(scope, grad_inputs[0])); + grad_outputs->push_back(NoGradient()); + grad_outputs->push_back(NoGradient()); + grad_outputs->push_back(NoGradient()); + return scope.status(); +} +REGISTER_GRADIENT_OP("QuantizeAndDequantizeV3", QuantizeAndDequantizeV3Grad); + Status SplitGrad(const Scope& scope, const Operation& op, const std::vector& grad_inputs, std::vector* grad_outputs) { @@ -162,7 +173,7 @@ Status CheckNumericsGrad(const Scope& scope, const Operation& op, const std::vector& grad_inputs, std::vector* grad_outputs) { string message; - TF_RETURN_IF_ERROR(GetNodeAttr(op.node()->def(), "message", &message)); + TF_RETURN_IF_ERROR(GetNodeAttr(op.node()->attrs(), "message", &message)); string err_msg = strings::StrCat( "Not a number (NaN) or infinity (Inf) values detected in gradient. ", message); @@ -215,9 +226,9 @@ Status ReverseSequenceGrad(const Scope& scope, const Operation& op, std::vector* grad_outputs) { auto seq_lengths = op.input(1); int batch_dim; - TF_RETURN_IF_ERROR(GetNodeAttr(op.node()->def(), "batch_dim", &batch_dim)); + TF_RETURN_IF_ERROR(GetNodeAttr(op.node()->attrs(), "batch_dim", &batch_dim)); int seq_dim; - TF_RETURN_IF_ERROR(GetNodeAttr(op.node()->def(), "seq_dim", &seq_dim)); + TF_RETURN_IF_ERROR(GetNodeAttr(op.node()->attrs(), "seq_dim", &seq_dim)); grad_outputs->push_back( ReverseSequence(scope, grad_inputs[0], seq_lengths, seq_dim, ReverseSequence::BatchDim(batch_dim))); @@ -247,6 +258,18 @@ Status ScatterNdGrad(const Scope& scope, const Operation& op, } REGISTER_GRADIENT_OP("ScatterNd", ScatterNdGrad); +Status ScatterNdNonAliasingAddGrad(const Scope& scope, const Operation& op, + const std::vector& grad_inputs, + std::vector* grad_outputs) { + auto indices = op.input(1); + grad_outputs->push_back(Identity(scope, grad_inputs[0])); + grad_outputs->push_back(NoGradient()); + grad_outputs->push_back(GatherNd(scope, grad_inputs[0], indices)); + return scope.status(); +} +REGISTER_GRADIENT_OP("ScatterNdNonAliasingAdd", ScatterNdNonAliasingAddGrad); + +template Status PadGrad(const Scope& scope, const Operation& op, const std::vector& grad_inputs, std::vector* grad_outputs) { @@ -259,15 +282,21 @@ Status PadGrad(const Scope& scope, const Operation& op, auto begin = Reshape(scope, pad_before, {-1}); grad_outputs->push_back(Slice(scope, grad_inputs[0], begin, Shape(scope, x))); grad_outputs->push_back(NoGradient()); + // PadV2 adds a "constant_values" input. + if (IsPadV2) { + grad_outputs->push_back(NoGradient()); + } return scope.status(); } -REGISTER_GRADIENT_OP("Pad", PadGrad); +REGISTER_GRADIENT_OP("Pad", PadGrad); +REGISTER_GRADIENT_OP("PadV2", PadGrad); Status SpaceToBatchGrad(const Scope& scope, const Operation& op, const std::vector& grad_inputs, std::vector* grad_outputs) { int block_size; - TF_RETURN_IF_ERROR(GetNodeAttr(op.node()->def(), "block_size", &block_size)); + TF_RETURN_IF_ERROR( + GetNodeAttr(op.node()->attrs(), "block_size", &block_size)); grad_outputs->push_back( BatchToSpace(scope, grad_inputs[0], op.input(1), block_size)); grad_outputs->push_back(NoGradient()); @@ -290,7 +319,8 @@ Status BatchToSpaceGrad(const Scope& scope, const Operation& op, const std::vector& grad_inputs, std::vector* grad_outputs) { int block_size; - TF_RETURN_IF_ERROR(GetNodeAttr(op.node()->def(), "block_size", &block_size)); + TF_RETURN_IF_ERROR( + GetNodeAttr(op.node()->attrs(), "block_size", &block_size)); grad_outputs->push_back( SpaceToBatch(scope, grad_inputs[0], op.input(1), block_size)); grad_outputs->push_back(NoGradient()); @@ -313,7 +343,8 @@ Status SpaceToDepthGrad(const Scope& scope, const Operation& op, const std::vector& grad_inputs, std::vector* grad_outputs) { int block_size; - TF_RETURN_IF_ERROR(GetNodeAttr(op.node()->def(), "block_size", &block_size)); + TF_RETURN_IF_ERROR( + GetNodeAttr(op.node()->attrs(), "block_size", &block_size)); grad_outputs->push_back(DepthToSpace(scope, grad_inputs[0], block_size)); return scope.status(); } @@ -323,7 +354,8 @@ Status DepthToSpaceGrad(const Scope& scope, const Operation& op, const std::vector& grad_inputs, std::vector* grad_outputs) { int block_size; - TF_RETURN_IF_ERROR(GetNodeAttr(op.node()->def(), "block_size", &block_size)); + TF_RETURN_IF_ERROR( + GetNodeAttr(op.node()->attrs(), "block_size", &block_size)); grad_outputs->push_back(SpaceToDepth(scope, grad_inputs[0], block_size)); return scope.status(); } @@ -333,7 +365,7 @@ Status MirrorPadGrad(const Scope& scope, const Operation& op, const std::vector& grad_inputs, std::vector* grad_outputs) { string mode; - TF_RETURN_IF_ERROR(GetNodeAttr(op.node()->def(), "mode", &mode)); + TF_RETURN_IF_ERROR(GetNodeAttr(op.node()->attrs(), "mode", &mode)); grad_outputs->push_back(tensorflow::ops::internal::MirrorPadGrad( scope, grad_inputs[0], op.input(1), mode)); grad_outputs->push_back(NoGradient()); @@ -346,7 +378,7 @@ Status MirrorPadGradGrad(const Scope& scope, const Operation& op, const std::vector& grad_inputs, std::vector* grad_outputs) { string mode; - TF_RETURN_IF_ERROR(GetNodeAttr(op.node()->def(), "mode", &mode)); + TF_RETURN_IF_ERROR(GetNodeAttr(op.node()->attrs(), "mode", &mode)); grad_outputs->push_back(MirrorPad(scope, grad_inputs[0], op.input(1), mode)); grad_outputs->push_back(NoGradient()); return scope.status(); diff --git a/tensorflow/cc/gradients/array_grad_test.cc b/tensorflow/cc/gradients/array_grad_test.cc index 5798b5b509fc14e6c9d95d4fd42aca893254f775..1777e181451b267f52a418888912ed1393bdf8b1 100644 --- a/tensorflow/cc/gradients/array_grad_test.cc +++ b/tensorflow/cc/gradients/array_grad_test.cc @@ -233,6 +233,28 @@ TEST_F(ArrayGradTest, ScatterNdGrad_SliceIndexing) { RunTest(updates, updates_shape, y, y_shape); } +TEST_F(ArrayGradTest, ScatterNdNonAliasingAddGrad_SimpleIndexing) { + TensorShape updates_shape({4}); + TensorShape input_shape({8}); + auto input = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(input_shape)); + auto updates = + Placeholder(scope_, DT_FLOAT, Placeholder::Shape(updates_shape)); + auto indices = Const(scope_, {{4}, {3}, {1}, {7}}); + auto y = ScatterNdNonAliasingAdd(scope_, input, indices, updates); + RunTest({input, updates}, {input_shape, updates_shape}, {y}, {input_shape}); +} + +TEST_F(ArrayGradTest, ScatterNdNonAliasingAddGrad_SliceIndexing) { + TensorShape updates_shape({2, 4, 4}); + TensorShape input_shape({4, 4, 4}); + auto input = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(input_shape)); + auto updates = + Placeholder(scope_, DT_FLOAT, Placeholder::Shape(updates_shape)); + auto indices = Const(scope_, {{0}, {2}}); + auto y = ScatterNdNonAliasingAdd(scope_, input, indices, updates); + RunTest({input, updates}, {input_shape, updates_shape}, {y}, {input_shape}); +} + TEST_F(ArrayGradTest, PadGrad) { TensorShape x_shape({2, 3}); auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape)); diff --git a/tensorflow/cc/gradients/math_grad.cc b/tensorflow/cc/gradients/math_grad.cc index aff0653139538820a705371ee9446a3d38ca69b5..d90654f2e9a89da56ef45d82b875c123d80f4633 100644 --- a/tensorflow/cc/gradients/math_grad.cc +++ b/tensorflow/cc/gradients/math_grad.cc @@ -13,14 +13,40 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ +#include "tensorflow/cc/ops/array_ops_internal.h" +#include "tensorflow/cc/ops/math_ops_internal.h" #include "tensorflow/cc/ops/standard_ops.h" #include "tensorflow/cc/framework/grad_op_registry.h" +#include "tensorflow/cc/framework/gradients.h" namespace tensorflow { namespace ops { namespace { +// Logical operations have no gradients. +REGISTER_NO_GRADIENT_OP("Less"); +REGISTER_NO_GRADIENT_OP("LessEqual"); +REGISTER_NO_GRADIENT_OP("Greater"); +REGISTER_NO_GRADIENT_OP("GreaterEqual"); +REGISTER_NO_GRADIENT_OP("Equal"); +REGISTER_NO_GRADIENT_OP("ApproximateEqual"); +REGISTER_NO_GRADIENT_OP("NotEqual"); +REGISTER_NO_GRADIENT_OP("LogicalAnd"); +REGISTER_NO_GRADIENT_OP("LogicalOr"); +REGISTER_NO_GRADIENT_OP("LogicalNot"); + +// Conjugate helper function returns the conjugate of an Output if it +// is complex valued. +Output ConjugateHelper(const Scope& scope, const Output& out) { + DataType dtype = out.type(); + if (dtype == DT_COMPLEX64 || dtype == DT_COMPLEX128) { + return Conj(scope, out); + } else { + return out; + } +} + // TODO(andydavis) Add control dependencies to gradient functions (as needed). Status AbsGrad(const Scope& scope, const Operation& op, @@ -44,9 +70,9 @@ REGISTER_GRADIENT_OP("Neg", NegGrad); Status InvGrad(const Scope& scope, const Operation& op, const std::vector& grad_inputs, std::vector* grad_outputs) { - // dx = dy * (-1 * (y * y)) + // Use the built-in operator. grad_outputs->push_back( - Mul(scope, grad_inputs[0], Neg(scope, Square(scope, op.output(0))))); + internal::ReciprocalGrad(scope, op.output(0), grad_inputs[0])); return scope.status(); } REGISTER_GRADIENT_OP("Inv", InvGrad); @@ -55,10 +81,12 @@ REGISTER_GRADIENT_OP("Reciprocal", InvGrad); Status SquareGrad(const Scope& scope, const Operation& op, const std::vector& grad_inputs, std::vector* grad_outputs) { - // dx = dy * (2 * x) + // dy/dx = (2 * x) auto two = Cast(scope, Const(scope, 2), op.input(0).type()); + auto dydx = Mul(scope, two, op.input(0)); + // grad(x) = grad(y) * conj(dy/dx) grad_outputs->push_back( - Mul(scope, grad_inputs[0], Mul(scope, two, op.input(0)))); + Mul(scope, grad_inputs[0], ConjugateHelper(scope, dydx))); return scope.status(); } REGISTER_GRADIENT_OP("Square", SquareGrad); @@ -66,13 +94,9 @@ REGISTER_GRADIENT_OP("Square", SquareGrad); Status SqrtGrad(const Scope& scope, const Operation& op, const std::vector& grad_inputs, std::vector* grad_outputs) { - // y = sqrt(x) - // dy/dx = 0.5 * (1 / sqrt(x)) = 0.5 * (1 / y) - // dx = dy * (0.5 * (1 / y)) - auto y_inv = Reciprocal(scope, op.output(0)); - auto half = Cast(scope, Const(scope, 0.5), op.input(0).type()); - auto dx = Mul(scope, grad_inputs[0], Mul(scope, half, y_inv)); - grad_outputs->push_back(dx); + // Use the built-in operator. + grad_outputs->push_back( + internal::SqrtGrad(scope, op.output(0), grad_inputs[0])); return scope.status(); } REGISTER_GRADIENT_OP("Sqrt", SqrtGrad); @@ -80,16 +104,9 @@ REGISTER_GRADIENT_OP("Sqrt", SqrtGrad); Status RsqrtGrad(const Scope& scope, const Operation& op, const std::vector& grad_inputs, std::vector* grad_outputs) { - // y = 1/x^1/2 = x^-1/2 - // dy/dx = -1/2 * x^-3/2 = -1/2 * x^-1/2 * x^-1 = -1/2 * y * x^-1 - // dx = dy * (-1/2 * y * x^-1) - auto x_inv = Reciprocal(scope, op.input(0)); - auto y = op.output(0); - auto neghalf = Cast(scope, Const(scope, -0.5), op.input(0).type()); - auto a = Mul(scope, neghalf, x_inv); - auto b = Mul(scope, a, y); - auto dx = Mul(scope, grad_inputs[0], b); - grad_outputs->push_back(dx); + // Use the built-in operator. + grad_outputs->push_back( + internal::RsqrtGrad(scope, op.output(0), grad_inputs[0])); return scope.status(); } REGISTER_GRADIENT_OP("Rsqrt", RsqrtGrad); @@ -97,10 +114,11 @@ REGISTER_GRADIENT_OP("Rsqrt", RsqrtGrad); Status ExpGrad(const Scope& scope, const Operation& op, const std::vector& grad_inputs, std::vector* grad_outputs) { - // y = exp(x) - // dy/dx = exp(x) - // dx = dy * y - grad_outputs->push_back(Mul(scope, grad_inputs[0], op.output(0))); + // dy/dx = exp(x) = y + // grad(x) = grad(y) * conj(dy/dx) + // = grad(y) * conj(y) + grad_outputs->push_back( + Mul(scope, grad_inputs[0], ConjugateHelper(scope, op.output(0)))); return scope.status(); } REGISTER_GRADIENT_OP("Exp", ExpGrad); @@ -108,10 +126,12 @@ REGISTER_GRADIENT_OP("Exp", ExpGrad); Status Expm1Grad(const Scope& scope, const Operation& op, const std::vector& grad_inputs, std::vector* grad_outputs) { - // f(x) = expm1(x) - // df/dx = exp(x) - // dx = dy * exp(x) - grad_outputs->push_back(Mul(scope, grad_inputs[0], Exp(scope, op.input(0)))); + // y = expm1(x) + // dy/dx = exp(x) + auto dydx = Exp(scope, op.input(0)); + // grad(x) = grad(y) * conj(dy/dx) + grad_outputs->push_back( + Mul(scope, grad_inputs[0], ConjugateHelper(scope, dydx))); return scope.status(); } REGISTER_GRADIENT_OP("Expm1", Expm1Grad); @@ -119,11 +139,12 @@ REGISTER_GRADIENT_OP("Expm1", Expm1Grad); Status LogGrad(const Scope& scope, const Operation& op, const std::vector& grad_inputs, std::vector* grad_outputs) { - // f(x) = log(x) = y - // df/dx = 1 / x - // dx = dy * (1 / x) + // y = log(x) + // dy/dx = 1 / x + auto dydx = Reciprocal(scope, op.input(0)); + // grad(x) = grad(y) * conj(dy/dx) grad_outputs->push_back( - Mul(scope, grad_inputs[0], Reciprocal(scope, op.input(0)))); + Mul(scope, grad_inputs[0], ConjugateHelper(scope, dydx))); return scope.status(); } REGISTER_GRADIENT_OP("Log", LogGrad); @@ -131,40 +152,111 @@ REGISTER_GRADIENT_OP("Log", LogGrad); Status Log1pGrad(const Scope& scope, const Operation& op, const std::vector& grad_inputs, std::vector* grad_outputs) { - // f(x) = log1p(x) = y - // df/dx = 1 / (1 + x) - // dx = dy * (1 / (1 + x)) + // y = log1p(x) + // dy/dx = 1 / (1 + x) auto one = Cast(scope, Const(scope, 1.0), op.input(0).type()); + auto dydx = Reciprocal(scope, Add(scope, one, op.input(0))); + // grad(x) = grad(y) * conj(dy/dx) grad_outputs->push_back( - Div(scope, grad_inputs[0], Add(scope, one, op.input(0)))); + Mul(scope, grad_inputs[0], ConjugateHelper(scope, dydx))); return scope.status(); } REGISTER_GRADIENT_OP("Log1p", Log1pGrad); +Status SinhGrad(const Scope& scope, const Operation& op, + const std::vector& grad_inputs, + std::vector* grad_outputs) { + // y = sinh(x) + // dy/dx = cosh(x) + auto dydx = Cosh(scope, op.input(0)); + // grad(x) = grad(y) * conj(dy/dx) + grad_outputs->push_back( + Mul(scope, grad_inputs[0], ConjugateHelper(scope, dydx))); + return scope.status(); +} +REGISTER_GRADIENT_OP("Sinh", SinhGrad); + +Status CoshGrad(const Scope& scope, const Operation& op, + const std::vector& grad_inputs, + std::vector* grad_outputs) { + // y = cosh(x) + // dy/dx = sinh(x) + auto dydx = Sinh(scope, op.input(0)); + // grad(x) = grad(y) * conj(dy/dx) + grad_outputs->push_back( + Mul(scope, grad_inputs[0], ConjugateHelper(scope, dydx))); + return scope.status(); +} +REGISTER_GRADIENT_OP("Cosh", CoshGrad); + Status TanhGrad(const Scope& scope, const Operation& op, const std::vector& grad_inputs, std::vector* grad_outputs) { - // y = tanh(x) - // dy/dx = 1 - (tanh(x))^2 = 1 - y^2 - // dx = dy * (1 - y^2) - auto y2 = Square(scope, op.output(0)); - auto one = Cast(scope, Const(scope, 1.0), op.input(0).type()); - auto dx = Mul(scope, grad_inputs[0], Sub(scope, one, y2)); - grad_outputs->push_back(dx); + // Use the built-in operator. + // Note that the built-in operator does not return the conjugate of + // the gradient. + auto grad = grad_inputs[0]; + // Optimization to avoid calculating conj(y) until the gradient is + // evaluated. + Scope grad_scope = scope.WithControlDependencies(grad); + auto y = ConjugateHelper(grad_scope, op.output(0)); + grad_outputs->push_back(internal::TanhGrad(scope, y, grad)); return scope.status(); } REGISTER_GRADIENT_OP("Tanh", TanhGrad); +Status AsinhGrad(const Scope& scope, const Operation& op, + const std::vector& grad_inputs, + std::vector* grad_outputs) { + // y = asinh(x) + // dy/dx = 1 / cosh(y) + auto dydx = Reciprocal(scope, Cosh(scope, op.output(0))); + // grad(x) = grad(y) * conj(dy/dx) + grad_outputs->push_back( + Mul(scope, grad_inputs[0], ConjugateHelper(scope, dydx))); + return scope.status(); +} +REGISTER_GRADIENT_OP("Asinh", AsinhGrad); + +Status AcoshGrad(const Scope& scope, const Operation& op, + const std::vector& grad_inputs, + std::vector* grad_outputs) { + // y = acosh(x) + // dy/dx = 1 / sinh(y) + auto dydx = Reciprocal(scope, Sinh(scope, op.output(0))); + // grad(x) = grad(y) * conj(dy/dx) + grad_outputs->push_back( + Mul(scope, grad_inputs[0], ConjugateHelper(scope, dydx))); + return scope.status(); +} +REGISTER_GRADIENT_OP("Acosh", AcoshGrad); + +Status AtanhGrad(const Scope& scope, const Operation& op, + const std::vector& grad_inputs, + std::vector* grad_outputs) { + // y = atanh(x) + // dy/dx = 1 / (1 - x^2) + auto one = Cast(scope, Const(scope, 1.0), op.input(0).type()); + auto dydx = Reciprocal(scope, Sub(scope, one, Square(scope, op.input(0)))); + // grad(x) = grad(y) * conj(dy/dx) + grad_outputs->push_back( + Mul(scope, grad_inputs[0], ConjugateHelper(scope, dydx))); + return scope.status(); +} +REGISTER_GRADIENT_OP("Atanh", AtanhGrad); + Status SigmoidGrad(const Scope& scope, const Operation& op, const std::vector& grad_inputs, std::vector* grad_outputs) { - // y = 1 / (1 + exp(-x)) - // dy/dx = y * (1 - y) - // dx = dy * y * (1 - y) - auto y = op.output(0); - auto one = Cast(scope, Const(scope, 1.0), op.input(0).type()); - auto dx = Mul(scope, grad_inputs[0], Mul(scope, y, Sub(scope, one, y))); - grad_outputs->push_back(dx); + // Use the built-in operator. + // Note that the built-in operator does not return the conjugate of + // the gradient. + auto grad = grad_inputs[0]; + // Optimization to avoid calculating conj(y) until the gradient is + // evaluated. + Scope grad_scope = scope.WithControlDependencies(grad); + auto y = ConjugateHelper(grad_scope, op.output(0)); + grad_outputs->push_back(internal::SigmoidGrad(scope, y, grad)); return scope.status(); } REGISTER_GRADIENT_OP("Sigmoid", SigmoidGrad); @@ -185,9 +277,10 @@ Status SinGrad(const Scope& scope, const Operation& op, std::vector* grad_outputs) { // y = sin(x) // dy/dx = cos(x) - // dx = dy * cos(x) - auto dx = Mul(scope, grad_inputs[0], Cos(scope, op.input(0))); - grad_outputs->push_back(dx); + auto dydx = Cos(scope, op.input(0)); + // grad(x) = grad(y) * conj(dy/dx) + grad_outputs->push_back( + Mul(scope, grad_inputs[0], ConjugateHelper(scope, dydx))); return scope.status(); } REGISTER_GRADIENT_OP("Sin", SinGrad); @@ -197,9 +290,10 @@ Status CosGrad(const Scope& scope, const Operation& op, std::vector* grad_outputs) { // y = cos(x) // dy/dx = -sin(x) - // dx = dy * -sin(x) - auto dx = Mul(scope, grad_inputs[0], Neg(scope, Sin(scope, op.input(0)))); - grad_outputs->push_back(dx); + auto dydx = Neg(scope, Sin(scope, op.input(0))); + // grad(x) = grad(y) * conj(dy/dx) + grad_outputs->push_back( + Mul(scope, grad_inputs[0], ConjugateHelper(scope, dydx))); return scope.status(); } REGISTER_GRADIENT_OP("Cos", CosGrad); @@ -208,12 +302,12 @@ Status AsinGrad(const Scope& scope, const Operation& op, const std::vector& grad_inputs, std::vector* grad_outputs) { // y = asin(x) - // dy/dx = 1 / (1 - x * x)^1/2 - // dx = dy * (1 / (1 - x * x)^1/2) + // dy/dx = 1 / sqrt(1 - x^2) auto x2 = Square(scope, op.input(0)); auto one = Cast(scope, Const(scope, 1.0), op.input(0).type()); auto dydx = Reciprocal(scope, Sqrt(scope, Sub(scope, one, x2))); - auto dx = Mul(scope, grad_inputs[0], dydx); + // grad(x) = grad(y) * conj(dy/dx) + auto dx = Mul(scope, grad_inputs[0], ConjugateHelper(scope, dydx)); grad_outputs->push_back(dx); return scope.status(); } @@ -239,9 +333,9 @@ Status TanGrad(const Scope& scope, const Operation& op, std::vector* grad_outputs) { // y = tan(x) // dy/dx = sec(x)^2 = 1 / cos(x)^2 - // dx = dy * (1 / cos(x)^2) auto dydx = Square(scope, Reciprocal(scope, Cos(scope, op.input(0)))); - auto dx = Mul(scope, grad_inputs[0], dydx); + // grad(x) = grad(y) * conj(dy/dx) + auto dx = Mul(scope, grad_inputs[0], ConjugateHelper(scope, dydx)); grad_outputs->push_back(dx); return scope.status(); } @@ -261,6 +355,157 @@ Status AtanGrad(const Scope& scope, const Operation& op, } REGISTER_GRADIENT_OP("Atan", AtanGrad); +// BinaryGradCommon handles the setup for binary ops that broadcast +// their inputs. +Status BinaryGradCommon(const Scope& scope, const Operation& op, + std::vector* grad_outputs, const Output& gx_1, + const Output& gx_2) { + auto sx_1 = Shape(scope, op.input(0)); + auto sx_2 = Shape(scope, op.input(1)); + auto rx = internal::BroadcastGradientArgs(scope, sx_1, sx_2); + auto dx_1 = Reshape(scope, Sum(scope, gx_1, rx.r0), sx_1); + auto dx_2 = Reshape(scope, Sum(scope, gx_2, rx.r1), sx_2); + grad_outputs->push_back(dx_1); + grad_outputs->push_back(dx_2); + return scope.status(); +} + +Status AddGrad(const Scope& scope, const Operation& op, + const std::vector& grad_inputs, + std::vector* grad_outputs) { + // y = x_1 + x_2 + // dy/dx_1 = dy/dx_2 = 1 + auto gx_1 = Identity(scope, grad_inputs[0]); + auto gx_2 = Identity(scope, grad_inputs[0]); + return BinaryGradCommon(scope, op, grad_outputs, gx_1, gx_2); +} +REGISTER_GRADIENT_OP("Add", AddGrad); + +Status SubGrad(const Scope& scope, const Operation& op, + const std::vector& grad_inputs, + std::vector* grad_outputs) { + // y = x_1 - x_2 + // dy/dx_1 = 1 + // dy/dx_2 = -1 + auto gx_1 = Identity(scope, grad_inputs[0]); + auto gx_2 = Neg(scope, grad_inputs[0]); + return BinaryGradCommon(scope, op, grad_outputs, gx_1, gx_2); +} +REGISTER_GRADIENT_OP("Sub", SubGrad); + +Status MulGrad(const Scope& scope, const Operation& op, + const std::vector& grad_inputs, + std::vector* grad_outputs) { + auto x_1 = ConjugateHelper(scope, op.input(0)); + auto x_2 = ConjugateHelper(scope, op.input(1)); + // y = x_1 * x_2 + // dy/dx_1 = x_2 + // dy/dx_2 = x_1 + auto gx_1 = Mul(scope, grad_inputs[0], x_2); + auto gx_2 = Mul(scope, grad_inputs[0], x_1); + return BinaryGradCommon(scope, op, grad_outputs, gx_1, gx_2); +} +REGISTER_GRADIENT_OP("Mul", MulGrad); + +Status DivGrad(const Scope& scope, const Operation& op, + const std::vector& grad_inputs, + std::vector* grad_outputs) { + auto x_1 = ConjugateHelper(scope, op.input(0)); + auto x_2 = ConjugateHelper(scope, op.input(1)); + // y = x_1 / x_2 + // dy/dx_1 = 1/x_2 + // dy/dx_2 = -x_1/x_2^2 + auto gx_1 = Div(scope, grad_inputs[0], x_2); + auto gx_2 = Mul(scope, grad_inputs[0], + Div(scope, Div(scope, Neg(scope, x_1), x_2), x_2)); + return BinaryGradCommon(scope, op, grad_outputs, gx_1, gx_2); +} +REGISTER_GRADIENT_OP("Div", DivGrad); + +Status RealDivGrad(const Scope& scope, const Operation& op, + const std::vector& grad_inputs, + std::vector* grad_outputs) { + auto x_1 = ConjugateHelper(scope, op.input(0)); + auto x_2 = ConjugateHelper(scope, op.input(1)); + // y = x_1 / x_2 + // dy/dx_1 = 1/x_2 + // dy/dx_2 = -x_1/x_2^2 + auto gx_1 = RealDiv(scope, grad_inputs[0], x_2); + auto gx_2 = Mul(scope, grad_inputs[0], + RealDiv(scope, RealDiv(scope, Neg(scope, x_1), x_2), x_2)); + return BinaryGradCommon(scope, op, grad_outputs, gx_1, gx_2); +} +REGISTER_GRADIENT_OP("RealDiv", RealDivGrad); + +Status SquaredDifferenceGrad(const Scope& scope, const Operation& op, + const std::vector& grad_inputs, + std::vector* grad_outputs) { + auto x_1 = ConjugateHelper(scope, op.input(0)); + auto x_2 = ConjugateHelper(scope, op.input(1)); + // y = (x_1 - x_2)^2 + // dy/dx_1 = 2 * (x_1 - x_2) + // dy/dx_2 = -2 * (x_1 - x_2) + auto two = Cast(scope, Const(scope, 2), grad_inputs[0].type()); + auto gx_1 = Mul(scope, grad_inputs[0], Mul(scope, two, Sub(scope, x_1, x_2))); + auto gx_2 = Neg(scope, gx_1); + return BinaryGradCommon(scope, op, grad_outputs, gx_1, gx_2); +} +REGISTER_GRADIENT_OP("SquaredDifference", SquaredDifferenceGrad); + +Status AddNGrad(const Scope& scope, const Operation& op, + const std::vector& grad_inputs, + std::vector* grad_outputs) { + // AddN doesn't support broadcasting, so all the inputs must be the + // same shape. + // Note: + // dy/dx_k = d(x_1 + x_2 + ... + x_n)/dx_k = 1 for all x_k + // hence dx_k = dy for all x_k + // So the gradient for AddN just transfers the incoming gradient to + // all outgoing gradients. + auto incoming = Identity(scope, grad_inputs[0]); + for (int32 i = 0; i < op.num_inputs(); ++i) { + grad_outputs->push_back(incoming); + } + return scope.status(); +} +REGISTER_GRADIENT_OP("AddN", AddNGrad); + +// MaximumMinimumGradCommon adds shared ops to calculate gradients for +// the binary Maximum and Minimum ops. +Status MaximumMinimumGradCommon(const Scope& scope, const Operation& op, + const std::vector& grad_inputs, + std::vector* grad_outputs, + const Output& comparator) { + // comparator is a boolean tensor, with + // y = x_1 at points where comparator is true, and x_2 otherwise + // Therefore + // dy/dx_1 = 1 where comparator is true, and 0 otherwise. + // dy/dx_2 = 0 where comparator is true, and 1 otherwise. + auto grad = grad_inputs[0]; + auto zeros = ZerosLike(scope, grad); + auto gx_1 = Where3(scope, comparator, grad, zeros); + auto gx_2 = Where3(scope, LogicalNot(scope, comparator), grad, zeros); + return BinaryGradCommon(scope, op, grad_outputs, gx_1, gx_2); +} + +Status MaximumGrad(const Scope& scope, const Operation& op, + const std::vector& grad_inputs, + std::vector* grad_outputs) { + auto comparator = GreaterEqual(scope, op.input(0), op.input(1)); + return MaximumMinimumGradCommon(scope, op, grad_inputs, grad_outputs, + comparator); +} +REGISTER_GRADIENT_OP("Maximum", MaximumGrad); + +Status MinimumGrad(const Scope& scope, const Operation& op, + const std::vector& grad_inputs, + std::vector* grad_outputs) { + auto comparator = LessEqual(scope, op.input(0), op.input(1)); + return MaximumMinimumGradCommon(scope, op, grad_inputs, grad_outputs, + comparator); +} +REGISTER_GRADIENT_OP("Minimum", MinimumGrad); + Status RealGrad(const Scope& scope, const Operation& op, const std::vector& grad_inputs, std::vector* grad_outputs) { @@ -281,6 +526,22 @@ Status ImagGrad(const Scope& scope, const Operation& op, } REGISTER_GRADIENT_OP("Imag", ImagGrad); +Status AngleGrad(const Scope& scope, const Operation& op, + const std::vector& grad_inputs, + std::vector* grad_outputs) { + // y = Angle(x) + // dx = -dy / (Im(x) + iRe(x)) = -dy * z + auto re = Real(scope, op.input(0)); + auto im = Imag(scope, op.input(0)); + auto z_inv = Reciprocal(scope, Complex(scope, im, re)); + auto zero = Cast(scope, Const(scope, 0), grad_inputs[0].type()); + auto grad = Complex(scope, grad_inputs[0], zero); + auto dx = Neg(scope, Mul(scope, grad, z_inv)); + grad_outputs->push_back(dx); + return scope.status(); +} +REGISTER_GRADIENT_OP("Angle", AngleGrad); + Status ConjGrad(const Scope& scope, const Operation& op, const std::vector& grad_inputs, std::vector* grad_outputs) { @@ -289,6 +550,209 @@ Status ConjGrad(const Scope& scope, const Operation& op, } REGISTER_GRADIENT_OP("Conj", ConjGrad); +// Integer division x / y, assuming x and y >=0, but treats x/0 = x +Output SafeDivHelper(const Scope& scope, const Output& x, const Output& y) { + return Div(scope, x, Maximum(scope, y, Const(scope, 1))); +} + +// Helper function for reduction ops. +// +// input_shape: 1-D Tensor, the shape of the Tensor being reduced. +// axes: 1-D Tensor, the reduction axes. +// Note that the reduction indices are in the range +// -rank(input_shape), rank(input_shape) +// returns a 1-D Tensor, the output shape as if keep_dims were set to True. +Output ReducedShapeHelper(const Scope& scope, const Output& input_shape, + const Output& reduction_axes) { + auto zero = Const(scope, 0); + auto one = Const(scope, 1); + + // Running example in comments + // input_shape = [2, 3, 5, 7] + // axes = [1, 2] + // The result (a shape after a reduction with keep_dims=True) + // [2, 1, 1, 7] + // + // We can treat each entry in axes as an index into input_shape that + // should be replaced by 1. + // We use DynamicStitch to do this. + + // input_rank = 4 + auto input_rank = Size(scope, input_shape); + + // Normalize any negative indices in the reduction_axes to positive + // values. + auto axes = Mod(scope, Add(scope, reduction_axes, input_rank), input_rank); + + // This [0..input_rank) range of integers is used in DynamicStitch to + // first copy input_shape to the result. + // input_rank_range = [0, 1, 2, 3] + auto input_rank_range = Range(scope, zero, input_rank, one); + + // A 1-filled tensor with the same shape as axes. DynamicStitch will + // merge these 1s (using axes for indices) to the correct + // position in the result. + // axes_ones = [1, 1] + auto axes_ones = OnesLike(scope, axes); + + // using DynamicStitch: + // indices = { input_rank_range, axes } + // = { [0, 1, 2, 3], [1, 2] } + // data = { input_shape, axes_ones } + // = { [2, 3, 5, 7], [1, 1] } + // The input_rank_range entry in indices first replicates the + // input_shape to the result. + // The axes entry in indices then moves a 1 to each of its entries, + // resulting in + // [2, 1, 1, 7] + std::vector indices = {input_rank_range, axes}; + std::vector data = {input_shape, axes_ones}; + return DynamicStitch(scope, indices, data); +} + +// SumGradHelper returns the gradient for the Sum operator, and is used +// by SumGrad and MeanGrad. +Output SumGradHelper(const Scope& scope, const Operation& op, + const std::vector& grad_inputs) { + // The partial derivative for any input along a "reduced" dimension + // is just 1, so we only need replicate the output gradient on such a + // dimension to its "expanded" shape. + // Running example: + // input is + // [[a, b, c], + // [d, e, f]] + // reduction_indices = [1] + // Sum = [a + b + c, d + e + f] + // if the gradient is [g1, g2] + // We want the propagated gradient to be + // [[g1, g1, g1], + // [g2, g2, g2]] + + // input_shape = [2, 3] + auto input_shape = Shape(scope, op.input(0)); + + // output_shape_kept_dims = [2, 1] + auto output_shape_kept_dims = + ReducedShapeHelper(scope, input_shape, op.input(1)); + + // This step "flips" any 1s with values from the input_shape, and + // replaces remaining entries with 1. This creates a shape that + // shows how much each dimension in the incoming gradient should be + // replicated. + // tile_scaling = [1, 3] + auto tile_scaling = SafeDivHelper(scope, input_shape, output_shape_kept_dims); + + // grad = [[g1], [g2]] + auto grad = Reshape(scope, grad_inputs[0], output_shape_kept_dims); + + // tile(grad, tile_scaling) = [[g1, g1, g1], [g2, g2, g2]] + return Tile(scope, grad, tile_scaling); +} + +Status SumGrad(const Scope& scope, const Operation& op, + const std::vector& grad_inputs, + std::vector* grad_outputs) { + grad_outputs->push_back(SumGradHelper(scope, op, grad_inputs)); + + // Stop propagation along reduction_indices + grad_outputs->push_back(NoGradient()); + return scope.status(); +} +REGISTER_GRADIENT_OP("Sum", SumGrad); + +Status MeanGrad(const Scope& scope, const Operation& op, + const std::vector& grad_inputs, + std::vector* grad_outputs) { + // The Mean gradient is just like the Sum gradient, except that + // all gradients are also divided by the size of reduced groups. + auto sum_grad = SumGradHelper(scope, op, grad_inputs); + + // The product of all entries in a tensor's shape is the total + // number of entries in the tensor. This step calculates + // n_input_entries/n_output_entries + // = group_size + auto input_shape = Shape(scope, op.input(0)); + auto output_shape = Shape(scope, op.output(0)); + auto zero = Const(scope, 0); + auto group_size = SafeDivHelper(scope, Prod(scope, input_shape, zero), + Prod(scope, output_shape, zero)); + + // propagate sum_grad/group_size + grad_outputs->push_back( + Div(scope, sum_grad, Cast(scope, group_size, sum_grad.type()))); + + // Stop propagation along reduction_indices + grad_outputs->push_back(NoGradient()); + return scope.status(); +} +REGISTER_GRADIENT_OP("Mean", MeanGrad); + +Status MinOrMaxGrad(const Scope& scope, const Operation& op, + const std::vector& grad_inputs, + std::vector* grad_outputs) { + // The partial derivative for any input along a "reduced" dimension + // is 1 when it is the min (or max) and 0 everywhere else. So the + // gradient calculation is identical for both operators. + // + // There's a special case for propagating gradients when there are + // multiple minima (or maxima) - we choose to divide the gradient + // equally among all matching inputs. + // + // Please note this comment + // https://github.com/tensorflow/tensorflow/issues/4886#issuecomment-256836063 + // for details. + + // Running example: + // input: [[5, 5, 5], + // [1, 2, -3]] + // reduction_indices: [1] + auto input = op.input(0); + auto reduction_indices = op.input(1); + + // [2, 3] + auto input_shape = Shape(scope, input); + + // [2, 1] + auto output_shape_kept_dims = + ReducedShapeHelper(scope, input_shape, reduction_indices); + + // for op=min (say) + // output = [5, -3] + // y = [[5], + // [-3]] + auto y = Reshape(scope, op.output(0), output_shape_kept_dims); + + // reshape([g1, g2], [2, 1]) = [[g1], + // [g2]] + auto grad = Reshape(scope, grad_inputs[0], output_shape_kept_dims); + + // indicators = equal(y, input) + // = equal([[5], [[5, 5, 5], + // [-3]], [1, 2, -3]]) + // = [[1, 1, 1], + // [0, 0, 1]] + auto indicators = Cast(scope, Equal(scope, y, input), grad_inputs[0].type()); + + // [[3], + // [1]] + auto num_selected = Reshape(scope, Sum(scope, indicators, reduction_indices), + output_shape_kept_dims); + + // [[1/3, 1/3, 1/3], + // [0, 0, 1]] + auto scale = Div(scope, indicators, num_selected); + + // [[g1/3, g1/3, g1/3], + // [0, 0, g2]] + grad_outputs->push_back(Mul(scope, scale, grad)); + + // Stop propagation along reduction_indices + grad_outputs->push_back(NoGradient()); + return scope.status(); +} +REGISTER_GRADIENT_OP("Min", MinOrMaxGrad); +REGISTER_GRADIENT_OP("Max", MinOrMaxGrad); + // MatMulGrad helper function used to compute two MatMul operations // based on input matrix transposition combinations. Status MatMulGradHelper(const Scope& scope, const bool is_batch, @@ -324,7 +788,7 @@ Status MatMulGradCommon(const Scope& scope, const Operation& op, const string& attr_adj_x, const string& attr_adj_y, std::vector* grad_outputs) { DataType dtype; - TF_RETURN_IF_ERROR(GetNodeAttr(op.output(0).node()->def(), "T", &dtype)); + TF_RETURN_IF_ERROR(GetNodeAttr(op.output(0).node()->attrs(), "T", &dtype)); if (dtype == DT_COMPLEX64 || dtype == DT_COMPLEX128) { return errors::Unimplemented( "MatMul gradient for complex data type is not supported yet."); @@ -332,8 +796,10 @@ Status MatMulGradCommon(const Scope& scope, const Operation& op, bool ta; bool tb; - TF_RETURN_IF_ERROR(GetNodeAttr(op.output(0).node()->def(), attr_adj_x, &ta)); - TF_RETURN_IF_ERROR(GetNodeAttr(op.output(0).node()->def(), attr_adj_y, &tb)); + TF_RETURN_IF_ERROR( + GetNodeAttr(op.output(0).node()->attrs(), attr_adj_x, &ta)); + TF_RETURN_IF_ERROR( + GetNodeAttr(op.output(0).node()->attrs(), attr_adj_y, &tb)); if (!ta && !tb) { return MatMulGradHelper(scope, is_batch, grad_inputs[0], false, op.input(1), diff --git a/tensorflow/cc/gradients/math_grad_test.cc b/tensorflow/cc/gradients/math_grad_test.cc index d7278929d4651f17d25670934b15e6da33d6a960..5b1558dd820862b18e486b347254c3a249bd016c 100644 --- a/tensorflow/cc/gradients/math_grad_test.cc +++ b/tensorflow/cc/gradients/math_grad_test.cc @@ -14,6 +14,7 @@ limitations under the License. ==============================================================================*/ #include "tensorflow/cc/framework/grad_op_registry.h" +#include "tensorflow/cc/framework/gradient_checker.h" #include "tensorflow/cc/framework/testutil.h" #include "tensorflow/cc/gradients/grad_testutil.h" #include "tensorflow/cc/ops/standard_ops.h" @@ -45,7 +46,12 @@ class CWiseUnaryGradTest : public ::testing::Test { EXPM1, LOG, LOG1P, + SINH, + COSH, TANH, + ASINH, + ACOSH, + ATANH, SIGMOID, SIGN, SIN, @@ -56,23 +62,25 @@ class CWiseUnaryGradTest : public ::testing::Test { ATAN }; - void TestCWiseGrad(UnaryOpType op_type, std::function x_fn, - std::function dy_fn, - std::function dx_fn) { - Tensor x(DT_FLOAT, {2, 3, 2}); - auto x_flat = x.flat(); + template + void TestCWiseGrad(UnaryOpType op_type, const std::function& x_fn, + const std::function& dy_fn, + const std::function& dx_fn) { + DataType dtype = DataTypeToEnum::v(); + Tensor x(dtype, {2, 3, 2}); + auto x_flat = x.flat(); for (int i = 0; i < x_flat.size(); ++i) { x_flat(i) = x_fn(i); } - Tensor dy(DT_FLOAT, {2, 3, 2}); - auto dy_flat = dy.flat(); + Tensor dy(dtype, {2, 3, 2}); + auto dy_flat = dy.flat(); for (int i = 0; i < dy_flat.size(); ++i) { dy_flat(i) = dy_fn(x_flat(i)); } - Tensor dx(DT_FLOAT, {2, 3, 2}); - auto dx_flat = dx.flat(); + Tensor dx(dtype, {2, 3, 2}); + auto dx_flat = dx.flat(); for (int i = 0; i < dx_flat.size(); ++i) { dx_flat(i) = dx_fn(x_flat(i), dy_flat(i)); } @@ -109,9 +117,24 @@ class CWiseUnaryGradTest : public ::testing::Test { case LOG1P: y = Log1p(scope_, x); break; + case SINH: + y = Sinh(scope_, x); + break; + case COSH: + y = Cosh(scope_, x); + break; case TANH: y = Tanh(scope_, x); break; + case ASINH: + y = Asinh(scope_, x); + break; + case ACOSH: + y = Acosh(scope_, x); + break; + case ATANH: + y = Atanh(scope_, x); + break; case SIGMOID: y = Sigmoid(scope_, x); break; @@ -146,7 +169,19 @@ class CWiseUnaryGradTest : public ::testing::Test { test::ExpectClose(output, dx); } - float RV(std::vector v) { return v[random::New64() % v.size()]; } + float RV(const std::vector& v) { + return v[random::New64() % v.size()]; + } + + complex64 CRV(const std::vector& v) { + return v[random::New64() % v.size()]; + } + + complex64 conjugate(const complex64& val) { + return complex64(val.real(), -val.imag()); + } + + const complex64 one_{1.0, 0}; Scope scope_; }; @@ -155,14 +190,14 @@ TEST_F(CWiseUnaryGradTest, Abs) { auto x_fn = [this](const int i) { return RV({-1, 0, 1}); }; auto dy_fn = [this](const float x) { return x + RV({-2, 2, -3, 3, -4, 4}); }; auto dx_fn = [this](const float x, const float dy) { return x * dy; }; - TestCWiseGrad(ABS, x_fn, dy_fn, dx_fn); + TestCWiseGrad(ABS, x_fn, dy_fn, dx_fn); } TEST_F(CWiseUnaryGradTest, Neg) { auto x_fn = [this](const int i) { return RV({-1, 0, 1}); }; auto dy_fn = [this](const float x) { return x + RV({-2, 2, -3, 3, -4, 4}); }; auto dx_fn = [this](const float x, const float dy) { return -dy; }; - TestCWiseGrad(NEG, x_fn, dy_fn, dx_fn); + TestCWiseGrad(NEG, x_fn, dy_fn, dx_fn); } TEST_F(CWiseUnaryGradTest, Reciprocal) { @@ -171,14 +206,36 @@ TEST_F(CWiseUnaryGradTest, Reciprocal) { auto dx_fn = [this](const float x, const float dy) { return -(1 / (x * x)) * dy; }; - TestCWiseGrad(INV, x_fn, dy_fn, dx_fn); + TestCWiseGrad(INV, x_fn, dy_fn, dx_fn); +} + +TEST_F(CWiseUnaryGradTest, Reciprocal_Complex) { + auto x_fn = [this](const int i) { return CRV({{-1, 0}, {1, 0}, {2, -1}}); }; + auto dy_fn = [this](const complex64 x) { + return x + CRV({{-2, 2}, {-3, 3}, {1, -4}}); + }; + auto dx_fn = [this](const complex64 x, const complex64 dy) { + return -conjugate(one_ / (x * x)) * dy; + }; + TestCWiseGrad(INV, x_fn, dy_fn, dx_fn); } TEST_F(CWiseUnaryGradTest, Square) { auto x_fn = [this](const int i) { return RV({0, -1, 1, -2, 2, -3, 3}); }; auto dy_fn = [this](const float x) { return RV({0, -7, 7, -8, 8, -9, 9}); }; auto dx_fn = [this](const float x, const float dy) { return 2 * x * dy; }; - TestCWiseGrad(SQUARE, x_fn, dy_fn, dx_fn); + TestCWiseGrad(SQUARE, x_fn, dy_fn, dx_fn); +} + +TEST_F(CWiseUnaryGradTest, Square_Complex) { + auto x_fn = [this](const int i) { return CRV({{-1, 0}, {1, 0}, {2, -1}}); }; + auto dy_fn = [this](const complex64& x) { + return x + CRV({{-2, 2}, {-3, 3}, {1, -4}}); + }; + auto dx_fn = [this](const complex64& x, const complex64& dy) { + return conjugate(complex64(2, 0) * x) * dy; + }; + TestCWiseGrad(SQUARE, x_fn, dy_fn, dx_fn); } TEST_F(CWiseUnaryGradTest, Sqrt) { @@ -187,7 +244,18 @@ TEST_F(CWiseUnaryGradTest, Sqrt) { auto dx_fn = [this](const float x, const float dy) { return dy * 0.5 * (1.0 / std::sqrt(x)); }; - TestCWiseGrad(SQRT, x_fn, dy_fn, dx_fn); + TestCWiseGrad(SQRT, x_fn, dy_fn, dx_fn); +} + +TEST_F(CWiseUnaryGradTest, Sqrt_Complex) { + auto x_fn = [this](const int i) { return CRV({{-1, 0}, {1, 0}, {2, -1}}); }; + auto dy_fn = [this](const complex64& x) { + return x + CRV({{-2, 2}, {-3, 3}, {1, -4}}); + }; + auto dx_fn = [this](const complex64& x, const complex64& dy) { + return conjugate(complex64(0.5, 0) / std::sqrt(x)) * dy; + }; + TestCWiseGrad(SQRT, x_fn, dy_fn, dx_fn); } TEST_F(CWiseUnaryGradTest, Rsqrt) { @@ -196,7 +264,18 @@ TEST_F(CWiseUnaryGradTest, Rsqrt) { auto dx_fn = [this](const float x, const float dy) { return dy * -0.5 * (1 / std::sqrt(x)) * (1 / x); }; - TestCWiseGrad(RSQRT, x_fn, dy_fn, dx_fn); + TestCWiseGrad(RSQRT, x_fn, dy_fn, dx_fn); +} + +TEST_F(CWiseUnaryGradTest, Rsqrt_Complex) { + auto x_fn = [this](const int i) { return CRV({{-1, 0}, {1, 0}, {2, -1}}); }; + auto dy_fn = [this](const complex64& x) { + return x + CRV({{-2, 2}, {-3, 3}, {1, -4}}); + }; + auto dx_fn = [this](const complex64& x, const complex64& dy) { + return conjugate(complex64(-0.5, 0) / std::sqrt(x) / x) * dy; + }; + TestCWiseGrad(RSQRT, x_fn, dy_fn, dx_fn); } TEST_F(CWiseUnaryGradTest, Exp) { @@ -205,7 +284,18 @@ TEST_F(CWiseUnaryGradTest, Exp) { auto dx_fn = [this](const float x, const float dy) { return dy * std::exp(x); }; - TestCWiseGrad(EXP, x_fn, dy_fn, dx_fn); + TestCWiseGrad(EXP, x_fn, dy_fn, dx_fn); +} + +TEST_F(CWiseUnaryGradTest, Exp_Complex) { + auto x_fn = [this](const int i) { return CRV({{-1, 0}, {1, 0}, {2, -1}}); }; + auto dy_fn = [this](const complex64& x) { + return x + CRV({{-2, 2}, {-3, 3}, {1, -4}}); + }; + auto dx_fn = [this](const complex64& x, const complex64& dy) { + return dy * conjugate(std::exp(x)); + }; + TestCWiseGrad(EXP, x_fn, dy_fn, dx_fn); } TEST_F(CWiseUnaryGradTest, Expm1) { @@ -214,14 +304,36 @@ TEST_F(CWiseUnaryGradTest, Expm1) { auto dx_fn = [this](const float x, const float dy) { return dy * std::exp(x); }; - TestCWiseGrad(EXPM1, x_fn, dy_fn, dx_fn); + TestCWiseGrad(EXPM1, x_fn, dy_fn, dx_fn); +} + +TEST_F(CWiseUnaryGradTest, Expm1_Complex) { + auto x_fn = [this](const int i) { return CRV({{-1, 0}, {1, 0}, {2, -1}}); }; + auto dy_fn = [this](const complex64& x) { + return x + CRV({{-2, 2}, {-3, 3}, {1, -4}}); + }; + auto dx_fn = [this](const complex64& x, const complex64& dy) { + return dy * conjugate(std::exp(x)); + }; + TestCWiseGrad(EXPM1, x_fn, dy_fn, dx_fn); } TEST_F(CWiseUnaryGradTest, Log) { auto x_fn = [this](const int i) { return RV({-1, 1, -2, 2, -3, 3, -4, 4}); }; auto dy_fn = [this](const float x) { return x + RV({-2, 2, -3, 3, -4, 4}); }; auto dx_fn = [this](const float x, const float dy) { return dy * (1.0 / x); }; - TestCWiseGrad(LOG, x_fn, dy_fn, dx_fn); + TestCWiseGrad(LOG, x_fn, dy_fn, dx_fn); +} + +TEST_F(CWiseUnaryGradTest, Log_Complex) { + auto x_fn = [this](const int i) { return CRV({{-1, 0}, {1, 0}, {2, -1}}); }; + auto dy_fn = [this](const complex64& x) { + return x + CRV({{-2, 2}, {-3, 3}, {1, -4}}); + }; + auto dx_fn = [this](const complex64& x, const complex64& dy) { + return dy * conjugate(one_ / x); + }; + TestCWiseGrad(LOG, x_fn, dy_fn, dx_fn); } TEST_F(CWiseUnaryGradTest, Log1p) { @@ -230,7 +342,64 @@ TEST_F(CWiseUnaryGradTest, Log1p) { auto dx_fn = [this](const float x, const float dy) { return dy * (1.0 / (1.0 + x)); }; - TestCWiseGrad(LOG1P, x_fn, dy_fn, dx_fn); + TestCWiseGrad(LOG1P, x_fn, dy_fn, dx_fn); +} + +TEST_F(CWiseUnaryGradTest, Log1p_Complex) { + auto x_fn = [this](const int i) { + return CRV({{0, 0}, {1e-6, 0}, {2, -1}, {1, 2}, {3, 4}}); + }; + auto dy_fn = [this](const complex64& x) { + return x + CRV({{-2, 2}, {-3, 3}, {1, -4}}); + }; + auto dx_fn = [this](const complex64& x, const complex64& dy) { + return dy / (one_ + conjugate(x)); + }; + TestCWiseGrad(LOG1P, x_fn, dy_fn, dx_fn); +} + +TEST_F(CWiseUnaryGradTest, Sinh) { + auto x_fn = [this](const int i) { return RV({0, -1, 1, -2, 2, -3, 3}); }; + auto dy_fn = [this](const float x) { return x + RV({-2, 2, -3, 3, -4, 4}); }; + auto dx_fn = [this](const float x, const float dy) { + return dy * std::cosh(x); + }; + TestCWiseGrad(SINH, x_fn, dy_fn, dx_fn); +} + +TEST_F(CWiseUnaryGradTest, Sinh_Complex) { + auto x_fn = [this](const int i) { + return CRV({{1, 0}, {0, 1}, {2, -1}, {1, 2}, {3, 4}}); + }; + auto dy_fn = [this](const complex64& x) { + return x + CRV({{-2, 2}, {-3, 3}, {1, -4}}); + }; + auto dx_fn = [this](const complex64& x, const complex64& dy) { + return dy * conjugate(std::cosh(x)); + }; + TestCWiseGrad(SINH, x_fn, dy_fn, dx_fn); +} + +TEST_F(CWiseUnaryGradTest, Cosh) { + auto x_fn = [this](const int i) { return RV({0, -1, 1, -2, 2, -3, 3}); }; + auto dy_fn = [this](const float x) { return x + RV({-2, 2, -3, 3, -4, 4}); }; + auto dx_fn = [this](const float x, const float dy) { + return dy * std::sinh(x); + }; + TestCWiseGrad(COSH, x_fn, dy_fn, dx_fn); +} + +TEST_F(CWiseUnaryGradTest, Cosh_Complex) { + auto x_fn = [this](const int i) { + return CRV({{1, 0}, {0, 1}, {2, -1}, {1, 2}, {3, 4}}); + }; + auto dy_fn = [this](const complex64& x) { + return x + CRV({{-2, 2}, {-3, 3}, {1, -4}}); + }; + auto dx_fn = [this](const complex64& x, const complex64& dy) { + return dy * conjugate(std::sinh(x)); + }; + TestCWiseGrad(COSH, x_fn, dy_fn, dx_fn); } TEST_F(CWiseUnaryGradTest, Tanh) { @@ -240,7 +409,93 @@ TEST_F(CWiseUnaryGradTest, Tanh) { const float y = std::tanh(x); return dy * (1.0 - y * y); }; - TestCWiseGrad(TANH, x_fn, dy_fn, dx_fn); + TestCWiseGrad(TANH, x_fn, dy_fn, dx_fn); +} + +TEST_F(CWiseUnaryGradTest, Tanh_Complex) { + auto x_fn = [this](const int i) { + return CRV({{1, 0}, {0, 1}, {2, -1}, {1, 2}, {3, 4}}); + }; + auto dy_fn = [this](const complex64& x) { + return x + CRV({{-2, 2}, {-3, 3}, {1, -4}}); + }; + auto dx_fn = [this](const complex64& x, const complex64& dy) { + const complex64 y = std::tanh(x); + return dy * conjugate((one_ - y * y)); + }; + TestCWiseGrad(TANH, x_fn, dy_fn, dx_fn); +} + +TEST_F(CWiseUnaryGradTest, Asinh) { + auto x_fn = [this](const int i) { return RV({0, -1, 1, -2, 2, -3, 3}); }; + auto dy_fn = [this](const float x) { return x + RV({-2, 2, -3, 3, -4, 4}); }; + auto dx_fn = [this](const float x, const float dy) { + auto y = std::asinh(x); + return dy / std::cosh(y); + }; + TestCWiseGrad(ASINH, x_fn, dy_fn, dx_fn); +} + +TEST_F(CWiseUnaryGradTest, Asinh_Complex) { + auto x_fn = [this](const int i) { + return CRV({{1, 0}, {0, 1}, {2, -1}, {1, 2}, {3, 4}}); + }; + auto dy_fn = [this](const complex64& x) { + return x + CRV({{-2, 2}, {-3, 3}, {1, -4}}); + }; + auto dx_fn = [this](const complex64& x, const complex64& dy) { + auto y = std::asinh(x); + return dy / conjugate(std::cosh(y)); + }; + TestCWiseGrad(ASINH, x_fn, dy_fn, dx_fn); +} + +TEST_F(CWiseUnaryGradTest, Acosh) { + auto x_fn = [this](const int i) { return RV({1, 2, 3, 4, 5, 6, 7}); }; + auto dy_fn = [this](const float x) { + return x + RV({8, 9, 10, 11, 12, 13, 14}); + }; + auto dx_fn = [this](const float x, const float dy) { + auto y = std::acosh(x); + return dy / std::sinh(y); + }; + TestCWiseGrad(ACOSH, x_fn, dy_fn, dx_fn); +} + +TEST_F(CWiseUnaryGradTest, Acosh_Complex) { + auto x_fn = [this](const int i) { + return CRV({{1, 1}, {2, 1}, {1, 4}, {1, 2}, {3, 4}}); + }; + auto dy_fn = [this](const complex64& x) { + return x + CRV({{2, 2}, {3, 3}, {1, 4}}); + }; + auto dx_fn = [this](const complex64& x, const complex64& dy) { + auto y = std::acosh(x); + return dy / conjugate(std::sinh(y)); + }; + TestCWiseGrad(ACOSH, x_fn, dy_fn, dx_fn); +} + +TEST_F(CWiseUnaryGradTest, Atanh) { + auto x_fn = [this](const int i) { return RV({0, -0.5, 0.5, -0.1, 0.1}); }; + auto dy_fn = [this](const float x) { return x + RV({-2, 2, -3, 3, -4, 4}); }; + auto dx_fn = [this](const float x, const float dy) { + return dy * (1. / (1. - x * x)); + }; + TestCWiseGrad(ATANH, x_fn, dy_fn, dx_fn); +} + +TEST_F(CWiseUnaryGradTest, Atanh_Complex) { + auto x_fn = [this](const int i) { + return CRV({{0.1, 0}, {0, 0.1}, {0.2, -0.1}, {0.1, 0.2}, {0.3, 0.4}}); + }; + auto dy_fn = [this](const complex64& x) { + return x + CRV({{-2, 2}, {-3, 3}, {1, -4}}); + }; + auto dx_fn = [this](const complex64& x, const complex64& dy) { + return dy / conjugate(one_ - x * x); + }; + TestCWiseGrad(ATANH, x_fn, dy_fn, dx_fn); } TEST_F(CWiseUnaryGradTest, Sigmoid) { @@ -250,14 +505,28 @@ TEST_F(CWiseUnaryGradTest, Sigmoid) { const float y = 1.0 / (1.0 + std::exp(-x)); return dy * y * (1.0 - y); }; - TestCWiseGrad(SIGMOID, x_fn, dy_fn, dx_fn); + TestCWiseGrad(SIGMOID, x_fn, dy_fn, dx_fn); +} + +TEST_F(CWiseUnaryGradTest, Sigmoid_Complex) { + auto x_fn = [this](const int i) { + return CRV({{1, 0}, {0, 0}, {2, -1}, {1, 2}, {3, 4}}); + }; + auto dy_fn = [this](const complex64& x) { + return x + CRV({{-2, 2}, {-3, 3}, {1, -4}}); + }; + auto dx_fn = [this](const complex64& x, const complex64& dy) { + const complex64 y = one_ / (one_ + std::exp(-x)); + return dy * conjugate(y * (one_ - y)); + }; + TestCWiseGrad(SIGMOID, x_fn, dy_fn, dx_fn); } TEST_F(CWiseUnaryGradTest, Sign) { auto x_fn = [this](const int i) { return RV({0, -1, 1, -2, 2, -3, 3}); }; auto dy_fn = [this](const float x) { return x + RV({-2, 2, -3, 3, -4, 4}); }; auto dx_fn = [this](const float x, const float dy) { return 0.0; }; - TestCWiseGrad(SIGN, x_fn, dy_fn, dx_fn); + TestCWiseGrad(SIGN, x_fn, dy_fn, dx_fn); } TEST_F(CWiseUnaryGradTest, Sin) { @@ -266,7 +535,20 @@ TEST_F(CWiseUnaryGradTest, Sin) { auto dx_fn = [this](const float x, const float dy) { return dy * std::cos(x); }; - TestCWiseGrad(SIN, x_fn, dy_fn, dx_fn); + TestCWiseGrad(SIN, x_fn, dy_fn, dx_fn); +} + +TEST_F(CWiseUnaryGradTest, Sin_Complex) { + auto x_fn = [this](const int i) { + return CRV({{1, 0}, {0, 1}, {2, -1}, {1, 2}, {3, 4}}); + }; + auto dy_fn = [this](const complex64& x) { + return x + CRV({{-2, 2}, {-3, 3}, {1, -4}}); + }; + auto dx_fn = [this](const complex64& x, const complex64& dy) { + return dy * conjugate(std::cos(x)); + }; + TestCWiseGrad(SIN, x_fn, dy_fn, dx_fn); } TEST_F(CWiseUnaryGradTest, Cos) { @@ -275,7 +557,20 @@ TEST_F(CWiseUnaryGradTest, Cos) { auto dx_fn = [this](const float x, const float dy) { return dy * -1.0 * std::sin(x); }; - TestCWiseGrad(COS, x_fn, dy_fn, dx_fn); + TestCWiseGrad(COS, x_fn, dy_fn, dx_fn); +} + +TEST_F(CWiseUnaryGradTest, Cos_Complex) { + auto x_fn = [this](const int i) { + return CRV({{1, 0}, {0, 1}, {2, -1}, {1, 2}, {3, 4}}); + }; + auto dy_fn = [this](const complex64& x) { + return x + CRV({{-2, 2}, {-3, 3}, {1, -4}}); + }; + auto dx_fn = [this](const complex64& x, const complex64& dy) { + return dy * conjugate(-std::sin(x)); + }; + TestCWiseGrad(COS, x_fn, dy_fn, dx_fn); } TEST_F(CWiseUnaryGradTest, Asin) { @@ -284,7 +579,24 @@ TEST_F(CWiseUnaryGradTest, Asin) { auto dx_fn = [this](const float x, const float dy) { return dy * (1.0 / std::sqrt(1.0 - x * x)); }; - TestCWiseGrad(ASIN, x_fn, dy_fn, dx_fn); + TestCWiseGrad(ASIN, x_fn, dy_fn, dx_fn); +} + +TEST_F(CWiseUnaryGradTest, Asin_Complex) { + auto x_fn = [this](const int i) { + return CRV({{1, 0}, {0, 1}, {2, -1}, {1, 2}, {3, 4}}); + }; + auto dy_fn = [this](const complex64& x) { + return x + CRV({{-2, 2}, {-3, 3}, {1, -4}}); + }; + auto dx_fn = [this](const complex64& x, const complex64& dy) { + return dy / conjugate(std::sqrt(one_ - x * x)); + }; + // TODO(kbsriram) + // Enable test when the asin kernel supports complex numbers + if (false) { + TestCWiseGrad(ASIN, x_fn, dy_fn, dx_fn); + } } TEST_F(CWiseUnaryGradTest, Acos) { @@ -293,7 +605,24 @@ TEST_F(CWiseUnaryGradTest, Acos) { auto dx_fn = [this](const float x, const float dy) { return dy * (-1.0 / std::sqrt(1.0 - x * x)); }; - TestCWiseGrad(ACOS, x_fn, dy_fn, dx_fn); + TestCWiseGrad(ACOS, x_fn, dy_fn, dx_fn); +} + +TEST_F(CWiseUnaryGradTest, Acos_Complex) { + auto x_fn = [this](const int i) { + return CRV({{1, 0}, {0, 1}, {2, -1}, {1, 2}, {3, 4}}); + }; + auto dy_fn = [this](const complex64& x) { + return x + CRV({{-2, 2}, {-3, 3}, {1, -4}}); + }; + auto dx_fn = [this](const complex64& x, const complex64& dy) { + return dy / -conjugate(std::sqrt(one_ - x * x)); + }; + // TODO(kbsriram) + // Add test when the acos kernel supports complex numbers + if (false) { + TestCWiseGrad(ACOS, x_fn, dy_fn, dx_fn); + } } TEST_F(CWiseUnaryGradTest, Tan) { @@ -303,7 +632,25 @@ TEST_F(CWiseUnaryGradTest, Tan) { const float cosx = std::cos(x); return dy * (1 / (cosx * cosx)); }; - TestCWiseGrad(TAN, x_fn, dy_fn, dx_fn); + TestCWiseGrad(TAN, x_fn, dy_fn, dx_fn); +} + +TEST_F(CWiseUnaryGradTest, Tan_Complex) { + auto x_fn = [this](const int i) { + return CRV({{1, 0}, {0, 1}, {2, -1}, {1, 2}, {3, 4}}); + }; + auto dy_fn = [this](const complex64& x) { + return x + CRV({{-2, 2}, {-3, 3}, {1, -4}}); + }; + auto dx_fn = [this](const complex64& x, const complex64& dy) { + const complex64 cosx = std::cos(x); + return dy / conjugate(cosx * cosx); + }; + // TODO(kbsriram) + // Enable when tan kernel supports complex inputs + if (false) { + TestCWiseGrad(TAN, x_fn, dy_fn, dx_fn); + } } TEST_F(CWiseUnaryGradTest, Atan) { @@ -312,7 +659,24 @@ TEST_F(CWiseUnaryGradTest, Atan) { auto dx_fn = [this](const float x, const float dy) { return dy * (1 / (1 + x * x)); }; - TestCWiseGrad(ATAN, x_fn, dy_fn, dx_fn); + TestCWiseGrad(ATAN, x_fn, dy_fn, dx_fn); +} + +TEST_F(CWiseUnaryGradTest, Atan_Complex) { + auto x_fn = [this](const int i) { + return CRV({{1, 0}, {0, 1}, {2, -1}, {1, 2}, {3, 4}}); + }; + auto dy_fn = [this](const complex64& x) { + return x + CRV({{-2, 2}, {-3, 3}, {1, -4}}); + }; + auto dx_fn = [this](const complex64& x, const complex64& dy) { + return dy / (one_ + x * x); + }; + // TODO(kbsriram) + // Add test when the atan kernel supports complex numbers + if (false) { + TestCWiseGrad(ATAN, x_fn, dy_fn, dx_fn); + } } class CWiseUnaryComplexGradTest : public ::testing::Test { @@ -320,7 +684,7 @@ class CWiseUnaryComplexGradTest : public ::testing::Test { CWiseUnaryComplexGradTest() : scope_(Scope::NewRootScope().WithDevice("/cpu:0")) {} - enum UnaryOpType { REAL, IMAG, CONJ }; + enum UnaryOpType { REAL, IMAG, ANGLE, CONJ }; void TestCWiseGradComplex(UnaryOpType op_type, const Tensor& x, const Tensor& dy, const Tensor& dx_expected) { @@ -332,6 +696,9 @@ class CWiseUnaryComplexGradTest : public ::testing::Test { case IMAG: y = Imag(scope_, x); break; + case ANGLE: + y = Angle(scope_, x); + break; case CONJ: y = Conj(scope_, x); break; @@ -366,6 +733,17 @@ TEST_F(CWiseUnaryComplexGradTest, Imag) { TestCWiseGradComplex(IMAG, x, dy, dx_expected); } +TEST_F(CWiseUnaryComplexGradTest, Angle) { + Tensor x = test::AsTensor( + {{1, -1}, {-2, 2}, {3, -3}, {-4, 4}, {8, -8}, {-9, 9}}, {2, 3}); + Tensor dy = test::AsTensor({11, -12, 13, -14, 15, -16}, {2, 3}); + Tensor dx_expected = test::AsTensor( + {{5.5, 5.5}, {3, 3}, + {2.1666666666666665, 2.1666666666666665}, {1.75, 1.75}, + {0.9375, 0.9375}, {0.8888888888888888, 0.8888888888888888}}, {2, 3}); + TestCWiseGradComplex(ANGLE, x, dy, dx_expected); +} + TEST_F(CWiseUnaryComplexGradTest, Conj) { Tensor x = test::AsTensor( {{1, -1}, {-2, 2}, {3, -3}, {-4, 4}, {8, -8}, {-9, 9}}, {2, 3}); @@ -534,5 +912,184 @@ TEST_F(MathGradTest, BatchMatMulGrad_TransposeX_TransposeY) { TestMatMulGrad(true, true, true); } +class NaryGradTest : public ::testing::Test { + protected: + NaryGradTest() : scope_(Scope::NewRootScope().WithDevice("/cpu:0")) {} + + void RunTest(const OutputList& xs, const std::vector& x_shapes, + const OutputList& ys, const std::vector& y_shapes) { + TF_ASSERT_OK(scope_.status()); + float max_error; + TF_ASSERT_OK( + ComputeGradientError(scope_, xs, x_shapes, ys, y_shapes, &max_error)); + EXPECT_LT(max_error, 1e-3); + } + + void RunTest(const Output& x, const Tensor& x_init_value, const Output& y, + const TensorShape& y_shape) { + TF_ASSERT_OK(scope_.status()); + float max_error; + TF_ASSERT_OK( + ComputeGradientError(scope_, x, x_init_value, y, y_shape, &max_error)); + EXPECT_LT(max_error, 1e-3); + } + + Scope scope_; +}; + +TEST_F(NaryGradTest, Sum) { + TensorShape x_shape({2, 3, 5, 7}); + auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape)); + auto y = Sum(scope_, x, {1, -1}); + // y's shape is the result of reducing x along axes 1 and -1 (= 3) + TensorShape y_shape({2, 5}); + RunTest({x}, {x_shape}, {y}, {y_shape}); +} + +TEST_F(NaryGradTest, Mean) { + TensorShape x_shape({2, 3, 5, 7}); + auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape)); + auto y = Mean(scope_, x, {1, -1}); + // y's shape is the result of reducing x along axes 1 and -1 (= 3) + TensorShape y_shape({2, 5}); + RunTest({x}, {x_shape}, {y}, {y_shape}); +} + +TEST_F(NaryGradTest, Min) { + TensorShape x_shape({2, 3}); + auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape)); + auto y = Min(scope_, x, {-1}); + // y's shape is the result of reducing x along axes -1 (= 1) + TensorShape y_shape({2}); + Tensor x_init_value = + test::AsTensor({0.5f, 0.7f, 0.2f, 1.0f, 1.5f, -2.8f}, x_shape); + RunTest(x, x_init_value, y, y_shape); +} + +TEST_F(NaryGradTest, Max) { + TensorShape x_shape({2, 3}); + auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape)); + auto y = Max(scope_, x, {-1}); + // y's shape is the result of reducing x along axes -1 (= 1) + TensorShape y_shape({2}); + Tensor x_init_value = + test::AsTensor({0.5f, 0.7f, 0.2f, 1.0f, 1.5f, -2.8f}, x_shape); + RunTest(x, x_init_value, y, y_shape); +} + +TEST_F(NaryGradTest, MinMulti) { + // Test gradient when there are multiple minima. + // Note that we cannot directly use a test Tensor with multiple + // minima, as the numeric estimator will calculate incorrect + // gradients when perturbing each entry in the Tensor (which then + // changes how many minima exist.) + // Instead, we use a single input that broadcast-multiplies a larger + // tensor with equal values, and apply reduce_min to the multiplied + // result. + TensorShape x_shape({1}); + auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape)); + auto all_same = Mul(scope_, Const(scope_, {1.f, 1.f, 1.f}), x); + auto y = Min(scope_, all_same, {0}); + // y is a [3] shaped tensor reduced along dimension 0, so it is [1] shaped + TensorShape y_shape({1}); + RunTest({x}, {x_shape}, {y}, {y_shape}); +} + +TEST_F(NaryGradTest, MaxMulti) { + TensorShape x_shape({1}); + auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape)); + auto all_same = Mul(scope_, Const(scope_, {1.f, 1.f, 1.f}), x); + auto y = Max(scope_, all_same, {0}); + TensorShape y_shape({1}); + RunTest({x}, {x_shape}, {y}, {y_shape}); +} + +TEST_F(NaryGradTest, AddN) { + TensorShape shape({3, 2, 5}); + std::vector xs; + xs.push_back(Placeholder(scope_, DT_FLOAT, Placeholder::Shape(shape))); + xs.push_back(Placeholder(scope_, DT_FLOAT, Placeholder::Shape(shape))); + xs.push_back(Placeholder(scope_, DT_FLOAT, Placeholder::Shape(shape))); + auto y = AddN(scope_, xs); + RunTest(xs, {shape, shape, shape}, {y}, {shape}); +} + +TEST_F(NaryGradTest, Add) { + TensorShape x1_shape({3, 2, 5}); + TensorShape x2_shape({2, 5}); + auto x1 = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x1_shape)); + auto x2 = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x2_shape)); + auto y = Add(scope_, x1, x2); + RunTest({x1, x2}, {x1_shape, x2_shape}, {y}, {x1_shape}); +} + +TEST_F(NaryGradTest, Sub) { + TensorShape x1_shape({3, 2, 5}); + TensorShape x2_shape({2, 5}); + auto x1 = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x1_shape)); + auto x2 = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x2_shape)); + auto y = Sub(scope_, x1, x2); + RunTest({x1, x2}, {x1_shape, x2_shape}, {y}, {x1_shape}); +} + +TEST_F(NaryGradTest, Mul) { + TensorShape x1_shape({3, 2, 5}); + TensorShape x2_shape({2, 5}); + auto x1 = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x1_shape)); + auto x2 = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x2_shape)); + auto y = Mul(scope_, x1, x2); + RunTest({x1, x2}, {x1_shape, x2_shape}, {y}, {x1_shape}); +} + +TEST_F(NaryGradTest, Div) { + TensorShape x_shape({3, 2, 5}); + auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape)); + // Test x / (1 + |x|) rather than x_1 / x_2 to avoid triggering large + // division errors in the numeric estimator used by the gradient checker. + auto y = Div(scope_, x, Add(scope_, Const(scope_, 1), Abs(scope_, x))); + RunTest({x}, {x_shape}, {y}, {x_shape}); +} + +TEST_F(NaryGradTest, RealDiv) { + TensorShape x_shape({3, 2, 5}); + auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape)); + // Test x / (1 + |x|) rather than x_1 / x_2 to avoid triggering large + // division errors in the numeric estimator used by the gradient checker. + auto y = + RealDiv(scope_, x, Add(scope_, Const(scope_, 1), Abs(scope_, x))); + RunTest({x}, {x_shape}, {y}, {x_shape}); +} + +TEST_F(NaryGradTest, SquaredDifference) { + TensorShape x1_shape({3, 2, 5}); + TensorShape x2_shape({2, 5}); + auto x1 = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x1_shape)); + auto x2 = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x2_shape)); + auto y = SquaredDifference(scope_, x1, x2); + RunTest({x1, x2}, {x1_shape, x2_shape}, {y}, {x1_shape}); +} + +TEST_F(NaryGradTest, Maximum) { + TensorShape shape({3, 2}); + auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(shape)); + auto y = Maximum(scope_, x, Const(scope_, 1.0f)); + // Select values away from 1.0f to avoid instability when computing + // finite differences. + Tensor x_init_value = + test::AsTensor({0.5f, 1.5f, -1.2f, 3.0f, 0.1f, 2.8f}, {3, 2}); + RunTest(x, x_init_value, y, shape); +} + +TEST_F(NaryGradTest, Minimum) { + TensorShape shape({3, 2}); + auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(shape)); + auto y = Minimum(scope_, x, Const(scope_, 1.0f)); + // Select values away from 1.0f to avoid instability when computing + // finite differences. + Tensor x_init_value = + test::AsTensor({0.5f, 1.5f, -1.2f, 3.0f, 0.1f, 2.8f}, {3, 2}); + RunTest(x, x_init_value, y, shape); +} + } // namespace } // namespace tensorflow diff --git a/tensorflow/cc/gradients/nn_grad.cc b/tensorflow/cc/gradients/nn_grad.cc index 5e5203d09055d65cb1dcc16e091f6e5028ee7ae1..ccb58e7f915e0f32fba57df78d045dcb0b55cebf 100644 --- a/tensorflow/cc/gradients/nn_grad.cc +++ b/tensorflow/cc/gradients/nn_grad.cc @@ -46,6 +46,19 @@ Status SoftmaxGrad(const Scope& scope, const Operation& op, } REGISTER_GRADIENT_OP("Softmax", SoftmaxGrad); +Status LogSoftmaxGrad(const Scope& scope, const Operation& op, + const std::vector& grad_inputs, + std::vector* grad_outputs) { + + auto softmax = Exp(scope, op.output(0)); + auto sum = Sum(scope, grad_inputs[0], {1}, Sum::KeepDims(true)); + auto mul = Mul(scope, sum, softmax); + auto dx = Sub(scope, grad_inputs[0], mul); + grad_outputs->push_back(dx); + return scope.status(); +} +REGISTER_GRADIENT_OP("LogSoftmax", LogSoftmaxGrad); + Status ReluGradHelper(const Scope& scope, const Operation& op, const std::vector& grad_inputs, std::vector* grad_outputs) { @@ -73,6 +86,38 @@ Status EluGradHelper(const Scope& scope, const Operation& op, } REGISTER_GRADIENT_OP("Elu", EluGradHelper); +Status SeluGradHelper(const Scope& scope, const Operation& op, + const std::vector& grad_inputs, + std::vector* grad_outputs) { + auto dx = internal::SeluGrad(scope, grad_inputs[0], op.output(0)); + grad_outputs->push_back(dx); + return scope.status(); +} +REGISTER_GRADIENT_OP("Selu", SeluGradHelper); + +Status L2LossGrad(const Scope& scope, const Operation& op, + const std::vector& grad_inputs, + std::vector* grad_outputs) { + grad_outputs->push_back(Mul(scope, op.input(0), grad_inputs[0])); + return scope.status(); +} +REGISTER_GRADIENT_OP("L2Loss", L2LossGrad); + +Status BiasAddGradHelper(const Scope& scope, const Operation& op, + const std::vector& grad_inputs, + std::vector* grad_outputs) { + string data_format; + BiasAddGrad::Attrs input_attrs; + TF_RETURN_IF_ERROR( + GetNodeAttr(op.output(0).node()->attrs(), "data_format", &data_format)); + input_attrs.DataFormat(data_format); + auto dx_1 = BiasAddGrad(scope, grad_inputs[0], input_attrs); + grad_outputs->push_back(Identity(scope, grad_inputs[0])); + grad_outputs->push_back(dx_1); + return scope.status(); +} +REGISTER_GRADIENT_OP("BiasAdd", BiasAddGradHelper); + } // anonymous namespace } // namespace ops } // namespace tensorflow diff --git a/tensorflow/cc/gradients/nn_grad_test.cc b/tensorflow/cc/gradients/nn_grad_test.cc index 70c9bd4e08b2b46866a44becc8fe1305fec48ea9..affc1e1dbe6526bd468e07bc6803cbf9b7b54db2 100644 --- a/tensorflow/cc/gradients/nn_grad_test.cc +++ b/tensorflow/cc/gradients/nn_grad_test.cc @@ -36,7 +36,7 @@ class NNGradTest : public ::testing::Test { float max_error; TF_ASSERT_OK(ComputeGradientError(scope_, {x}, {x_shape}, {y}, {y_shape}, &max_error)); - EXPECT_LT(max_error, 1e-4); + EXPECT_LT(max_error, 2e-4); } void RunTest(const Output& x, const Tensor& x_init_value, const Output& y, @@ -44,7 +44,16 @@ class NNGradTest : public ::testing::Test { float max_error; TF_ASSERT_OK( ComputeGradientError(scope_, x, x_init_value, y, y_shape, &max_error)); - EXPECT_LT(max_error, 1e-4); + EXPECT_LT(max_error, 2e-4); + } + + void RunTest(const OutputList& xs, const std::vector& x_shapes, + const OutputList& ys, const std::vector& y_shapes) { + TF_ASSERT_OK(scope_.status()); + float max_error; + TF_ASSERT_OK( + ComputeGradientError(scope_, xs, x_shapes, ys, y_shapes, &max_error)); + EXPECT_LT(max_error, 2e-4); } Scope scope_; @@ -57,6 +66,19 @@ TEST_F(NNGradTest, SoftmaxGrad) { RunTest(x, shape, y, shape); } +TEST_F(NNGradTest, LogSoftmaxGrad) { + TensorShape shape({5, 3}); + auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(shape)); + auto y = LogSoftmax(scope_, x); + // Avoid numerical instability when computing finite differences. + Tensor x_init_value = test::AsTensor( + {-0.9f, -0.7f, -0.5f, -0.3f, -0.1f, + 0.1f, 0.3f, 0.5f, 0.7f, 0.8f, + -0.1f, 0.1f, 0.1f, 0.1f, 1.2f}, + {5, 3}); + RunTest(x, x_init_value, y, shape); +} + TEST_F(NNGradTest, ReluGrad) { TensorShape shape({5, 2}); auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(shape)); @@ -90,5 +112,32 @@ TEST_F(NNGradTest, EluGrad) { RunTest(x, x_init_value, y, shape); } +TEST_F(NNGradTest, SeluGrad) { + TensorShape shape({5, 2}); + auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(shape)); + auto y = Selu(scope_, x); + Tensor x_init_value = test::AsTensor( + {-0.9f, -0.7f, -0.5f, -0.3f, -0.1f, 0.1f, 0.3f, 0.5f, 0.7f, 0.9f}, + {5, 2}); + RunTest(x, x_init_value, y, shape); +} + +TEST_F(NNGradTest, L2LossGrad) { + TensorShape x_shape({5, 2}); + TensorShape y_shape({1}); + auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape)); + auto y = L2Loss(scope_, x); + RunTest(x, x_shape, y, y_shape); +} + +TEST_F(NNGradTest, BiasAddGradHelper) { + TensorShape shape({4, 5}); + TensorShape bias_shape({5}); + auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(shape)); + auto bias = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(bias_shape)); + auto y = BiasAdd(scope_, x, bias); + RunTest({x,bias}, {shape, bias_shape}, {y}, {shape}); +} + } // namespace } // namespace tensorflow diff --git a/tensorflow/cc/ops/const_op.cc b/tensorflow/cc/ops/const_op.cc index b37b8b67d755ec0e642bfdc126f5f12720c6097a..0030c2b2a7b69afe2151e88ef5d6d0755f72bfa7 100644 --- a/tensorflow/cc/ops/const_op.cc +++ b/tensorflow/cc/ops/const_op.cc @@ -34,7 +34,9 @@ Output Const(const Scope& scope, const Input::Initializer& val) { .Attr("dtype", val.tensor.dtype()); scope.UpdateBuilder(&builder); scope.UpdateStatus(builder.Finalize(graph, &ret)); + if (!scope.ok()) return Output(); + scope.UpdateStatus(scope.DoShapeInference(ret)); if (!scope.ok()) return Output(); return Output(ret); diff --git a/tensorflow/cc/ops/const_op.h b/tensorflow/cc/ops/const_op.h index e8cb6cf1dd1b1881733048ff4a1abbd05826f139..516800920f282be0590ef72b26a7fdd8b92a38f9 100644 --- a/tensorflow/cc/ops/const_op.h +++ b/tensorflow/cc/ops/const_op.h @@ -56,6 +56,8 @@ Output Const(const Scope& scope, const Input::Initializer& val) { scope.UpdateBuilder(&cast_builder); Node* ret; scope.UpdateStatus(cast_builder.Finalize(scope.graph(), &ret)); + if (!scope.ok()) return Output(); + scope.UpdateStatus(scope.DoShapeInference(ret)); return Output(ret, 0); } diff --git a/tensorflow/cc/ops/const_op_test.cc b/tensorflow/cc/ops/const_op_test.cc index 5a4770f879ff9a1422a63a88bd2b67ba201a0567..3184edeb3307cafcbfbc41c6477fd092ab613b46 100644 --- a/tensorflow/cc/ops/const_op_test.cc +++ b/tensorflow/cc/ops/const_op_test.cc @@ -28,9 +28,9 @@ void ExpectNodeEqual(const Node* n, gtl::ArraySlice values, TensorShape shape) { EXPECT_TRUE(n->IsConstant()); Tensor tensor; - TF_EXPECT_OK(GetNodeAttr(n->def(), "value", &tensor)); + TF_EXPECT_OK(GetNodeAttr(n->attrs(), "value", &tensor)); DataType dtype; - TF_EXPECT_OK(GetNodeAttr(n->def(), "dtype", &dtype)); + TF_EXPECT_OK(GetNodeAttr(n->attrs(), "dtype", &dtype)); EXPECT_EQ(tensor.dtype(), dtype); test::ExpectTensorEqual(tensor, test::AsTensor(values, shape)); } @@ -39,9 +39,9 @@ void ExpectTypeAndShape(const Node* n, DataType expected_dtype, TensorShape expected_shape) { EXPECT_TRUE(n->IsConstant()); Tensor tensor; - TF_EXPECT_OK(GetNodeAttr(n->def(), "value", &tensor)); + TF_EXPECT_OK(GetNodeAttr(n->attrs(), "value", &tensor)); DataType dtype; - TF_EXPECT_OK(GetNodeAttr(n->def(), "dtype", &dtype)); + TF_EXPECT_OK(GetNodeAttr(n->attrs(), "dtype", &dtype)); EXPECT_EQ(dtype, expected_dtype); EXPECT_EQ(expected_shape, TensorShape(tensor.shape())); } diff --git a/tensorflow/cc/ops/op_gen_overrides.pbtxt b/tensorflow/cc/ops/op_gen_overrides.pbtxt index cd94ddf4a1b67d3b98da7769db95bbda294e76db..777e54d342ff83baac41c78dc66dc27d5c4f0f5c 100644 --- a/tensorflow/cc/ops/op_gen_overrides.pbtxt +++ b/tensorflow/cc/ops/op_gen_overrides.pbtxt @@ -22,7 +22,7 @@ op { name: "Where" input_rename: { from: "input" to: "condition" } } op { name: "ThreadUnsafeUnigramCandidateSampler", skip: true } # control_flow_ops -# TODO(josh11b): Hide Switch and Merge once we write and migrate users to +# TODO(joshl): Hide Switch and Merge once we write and migrate users to # a Cond() API. #op { name: "Switch" hide: true } #op { name: "Merge" hide: true } @@ -100,6 +100,10 @@ op { name: "Stack" skip: true } op { name: "StackClose" skip: true } op { name: "StackPop" skip: true } op { name: "StackPush" skip: true } +op { name: "StackV2" skip: true } +op { name: "StackCloseV2" skip: true } +op { name: "StackPopV2" skip: true } +op { name: "StackPushV2" skip: true } op { name: "TensorArrayCloseV2" skip: true } op { name: "TensorArrayCloseV3" rename_to: "TensorArrayClose" } @@ -139,6 +143,31 @@ op { name: "SelfAdjointEigV2" rename_to: "SelfAdjointEig" } # logging_ops op { name: "AudioSummaryV2" rename_to: "AudioSummary" } + +# lookup_ops +op { name: "LookupTableFind" skip: true } +op { name: "LookupTableFindV2" rename_to: "LookupTableFind" } +op { name: "LookupTableInsert" skip: true } +op { name: "LookupTableInsertV2" rename_to: "LookupTableInsert" } +op { name: "LookupTableSize" skip: true } +op { name: "LookupTableSizeV2" rename_to: "LookupTableSize" } +op { name: "LookupTableExport" skip: true } +op { name: "LookupTableExportV2" rename_to: "LookupTableExport" } +op { name: "LookupTableImport" skip: true } +op { name: "LookupTableImportV2" rename_to: "LookupTableImport" } +op { name: "HashTable" skip: true } +op { name: "HashTableV2" rename_to: "HashTable" } +op { name: "MutableHashTable" skip: true } +op { name: "MutableHashTableV2" rename_to: "MutableHashTable" } +op { name: "MutableHashTableOfTensors" skip: true } +op { name: "MutableHashTableOfTensorsV2" rename_to: "MutableHashTableOfTensors" } +op { name: "MutableDenseHashTable" skip: true } +op { name: "MutableDenseHashTableV2" rename_to: "MutableDenseHashTable" } +op { name: "InitializeTable" skip: true } +op { name: "InitializeTableV2" rename_to: "InitializeTable" } +op { name: "InitializeTableFromTextFile" skip: true } +op { name: "InitializeTableFromTextFileV2" rename_to: "InitializeTableFromTextFile" } + # math_ops op { name: "All" alias: "ReduceAll" input_rename: { from: "reduction_indices" to: "axis" } } op { name: "Any" alias: "ReduceAny" input_rename: { from: "reduction_indices" to: "axis" } } @@ -173,6 +202,7 @@ op { name: "MaxPoolGradWithArgmax" hide: true } op { name: "ReluGrad" hide: true } op { name: "Relu6Grad" hide: true } op { name: "EluGrad" hide: true } +op { name: "SeluGrad" hide: true } op { name: "SoftplusGrad" hide: true } op { name: "SoftsignGrad" hide: true } op { name: "FractionalAvgPoolGrad" hide: true } diff --git a/tensorflow/cc/ops/standard_ops.h b/tensorflow/cc/ops/standard_ops.h index e117ddd04274958aa1285c3c1dc2b9fe5b9d64d1..0c021f0b3ac02c596e0511e650a3caa0002c25d1 100644 --- a/tensorflow/cc/ops/standard_ops.h +++ b/tensorflow/cc/ops/standard_ops.h @@ -25,6 +25,7 @@ limitations under the License. #include "tensorflow/cc/ops/io_ops.h" #include "tensorflow/cc/ops/linalg_ops.h" #include "tensorflow/cc/ops/logging_ops.h" +#include "tensorflow/cc/ops/lookup_ops.h" #include "tensorflow/cc/ops/math_ops.h" #include "tensorflow/cc/ops/nn_ops.h" #include "tensorflow/cc/ops/no_op.h" diff --git a/tensorflow/cc/ops/while_loop.cc b/tensorflow/cc/ops/while_loop.cc new file mode 100644 index 0000000000000000000000000000000000000000..27da77bbe068fd4be0eec40590a204fe6dedd235 --- /dev/null +++ b/tensorflow/cc/ops/while_loop.cc @@ -0,0 +1,223 @@ +/* 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/cc/ops/while_loop.h" + +#include "tensorflow/cc/framework/scope_internal.h" +#include "tensorflow/cc/ops/control_flow_ops_internal.h" +#include "tensorflow/cc/ops/standard_ops.h" +#include "tensorflow/core/common_runtime/shape_refiner.h" +#include "tensorflow/core/graph/node_builder.h" + +namespace tensorflow { +namespace ops { + +namespace { + +// Utility function for converting to internal C++ datatypes. +OutputTensor ToOutputTensor(const Output& output) { + return OutputTensor(output.node(), output.index()); +} + +// Utility function for converting to internal C++ datatypes. +std::vector ToOutputTensors(const std::vector& outputs) { + std::vector result(outputs.size()); + for (int i = 0; i < outputs.size(); ++i) { + result[i] = ToOutputTensor(outputs[i]); + } + return result; +} + +// Utility function for converting to internal C++ datatypes. +std::vector ToNodes(const std::vector& outputs) { + std::vector result(outputs.size()); + for (int i = 0; i < outputs.size(); ++i) { + result[i] = outputs[i].node(); + } + return result; +} + +// Manually generates the name of the `loop_var_idx`-th NextIteration node of a +// loop being constructed with `scope`. This is used to define the backedge +// before the NextIteration node is created. +string NextIterationName(const Scope& scope, int loop_var_idx) { + string result; + const string& prefix = scope.impl()->name(); + if (!prefix.empty()) strings::StrAppend(&result, prefix, "/"); + strings::StrAppend(&result, "NextIteration"); + if (loop_var_idx > 0) strings::StrAppend(&result, "_", loop_var_idx); + return result; +} + +// Creates the `loop_var_idx`-th Merge node of a loop being constructed with +// `scope`. `enter_output` is the `loop_var_idx`-th Enter node's output. +Status CreateMerge(const Scope& scope, int loop_var_idx, + const Output& enter_output, Output* merge_output) { + // The merge nodes accept the while loop's back edges as an input (i.e. the + // not-yet-created next iteration nodes). Use the underlying NodeBuilder API + // directly to create the back edge. + NodeBuilder::NodeOut enter_input(enter_output.node(), enter_output.index()); + + const int next_output_index = 0; + DataType dtype = enter_output.node()->output_type(0); + NodeBuilder::NodeOut next_input(NextIterationName(scope, loop_var_idx), + next_output_index, dtype); + + std::vector input_list({enter_input, next_input}); + const string unique_name = scope.GetUniqueNameForOp("Merge"); + NodeBuilder builder = NodeBuilder(unique_name, "Merge").Input(input_list); + scope.UpdateBuilder(&builder); + + Node* merge_node; + TF_RETURN_IF_ERROR(builder.Finalize(scope.graph(), &merge_node)); + TF_RETURN_IF_ERROR(scope.DoShapeInference(merge_node)); + *merge_output = Output(merge_node, 0); + return Status::OK(); +} + +// Creates the condition subgraph defined by `cond`. +Status CreateCond(const Scope& scope, const CondGraphBuilderFn& cond, + const std::vector& inputs, Output* output) { + // The control dependency is for constants in the cond graph, and other ops + // that do not depend on the loop variables. This ensures that these ops are + // in the while loop frame (since they will indirectly depend on an Enter node + // defining the frame) and that they are executed once per loop iteration. + // + // TODO(skyewm): the control dep will be added to all nodes in the cond graph. + // This is at best unnecessary, and at worst may prevent different parts of + // different loop iterations from executing in parallel. + Scope cond_scope = + scope.NewSubScope("cond").WithControlDependencies(inputs[0]); + Output raw_cond_out; + TF_RETURN_IF_ERROR(cond(cond_scope, inputs, &raw_cond_out)); + if (raw_cond_out.type() != DT_BOOL) { + return errors::InvalidArgument( + "BuildWhileLoop: 'cond' argument must return a boolean output, got ", + DataTypeString(raw_cond_out.type())); + } + *output = LoopCond(scope, raw_cond_out).output; + return Status::OK(); +} + +// Create the bdoy subgraph defined by `body`. `outputs` must be non-null and +// empty. +Status CreateBody(const Scope& scope, const BodyGraphBuilderFn& body, + const std::vector& inputs, + std::vector* outputs) { + DCHECK(outputs != nullptr); + DCHECK(outputs->empty()); + + // The control dependency is analogous to that in CreateCond(). + Scope body_scope = + scope.NewSubScope("body").WithControlDependencies(inputs[0]); + TF_RETURN_IF_ERROR(body(body_scope, inputs, outputs)); + const size_t num_loop_vars = inputs.size(); + if (outputs->size() != num_loop_vars) { + return errors::InvalidArgument( + "BuildWhileLoop: 'body' argument expected to return ", num_loop_vars, + "outputs, got ", outputs->size()); + } + // TODO(skyewm): check output types/shapes + return Status::OK(); +} + +} // namespace + +// A while loop with a single loop variable looks like this: +// +// (output) +// ^ +---------------+ +// | | body subgraph +-------------+ +// Exit +---------------+ | +// ^ ^ | +// | | | +// Switch<--------+ v +// ^ | NextIteration +// | +------+--------+ | +// +---->| cond subgraph | | +// | +---------------+ | +// Merge<---------------------------+ +// ^ +// | +// Enter +// ^ +// | +// (input) +// +// If there are multiple loop variables, each of the control flow ops is +// duplicated for each loop variable. +// TODO(skyewm): link to public version of design doc +Status BuildWhileLoop(const Scope& scope, const std::vector& inputs, + const CondGraphBuilderFn& cond, + const BodyGraphBuilderFn& body, const string& frame_name, + OutputList* outputs) { + DCHECK(!inputs.empty()); + DCHECK(outputs != nullptr); + DCHECK(outputs->empty()); + + TF_RETURN_IF_ERROR(scope.status()); + const size_t num_loop_vars = inputs.size(); + + std::vector enter_outputs(num_loop_vars); + for (int i = 0; i < num_loop_vars; ++i) { + enter_outputs[i] = internal::Enter(scope, inputs[i], frame_name); + } + TF_RETURN_IF_ERROR(scope.status()); + + std::vector merge_outputs(num_loop_vars); + for (int i = 0; i < num_loop_vars; ++i) { + TF_RETURN_IF_ERROR( + CreateMerge(scope, i, enter_outputs[i], &merge_outputs[i])); + } + + Output cond_out; + TF_RETURN_IF_ERROR(CreateCond(scope, cond, merge_outputs, &cond_out)); + + std::vector switch_trues(num_loop_vars); + std::vector switch_falses(num_loop_vars); + for (int i = 0; i < num_loop_vars; ++i) { + auto switch_i = Switch(scope, merge_outputs[i], cond_out); + switch_trues[i] = switch_i.output_true; + switch_falses[i] = switch_i.output_false; + } + TF_RETURN_IF_ERROR(scope.status()); + + std::vector body_outputs; + TF_RETURN_IF_ERROR(CreateBody(scope, body, switch_trues, &body_outputs)); + + std::vector next_outputs(num_loop_vars); + for (int i = 0; i < num_loop_vars; ++i) { + next_outputs[i] = NextIteration(scope, body_outputs[i]); + DCHECK_EQ(next_outputs[i].node()->name(), NextIterationName(scope, i)); + } + TF_RETURN_IF_ERROR(scope.status()); + + // Create the backedges from the NextIteration nodes to the Merge nodes. + for (int i = 0; i < num_loop_vars; ++i) { + const int merge_backedge_output_index = 1; + scope.graph()->AddEdge(next_outputs[i].node(), next_outputs[i].index(), + merge_outputs[i].node(), + merge_backedge_output_index); + } + + outputs->resize(num_loop_vars); + for (int i = 0; i < num_loop_vars; ++i) { + (*outputs)[i] = internal::Exit(scope, switch_falses[i]); + } + return scope.status(); +} + +} // namespace ops +} // namespace tensorflow diff --git a/tensorflow/cc/ops/while_loop.h b/tensorflow/cc/ops/while_loop.h new file mode 100644 index 0000000000000000000000000000000000000000..253d5d8935cf1632f06c1f3ce728a68fc85391bf --- /dev/null +++ b/tensorflow/cc/ops/while_loop.h @@ -0,0 +1,64 @@ +/* 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 THIRD_PARTY_TENSORFLOW_CC_OPS_WHILE_LOOP_H_ +#define THIRD_PARTY_TENSORFLOW_CC_OPS_WHILE_LOOP_H_ + +#include "tensorflow/cc/framework/ops.h" +#include "tensorflow/cc/framework/scope.h" + +namespace tensorflow { +namespace ops { + +// Function that takes cond graph inputs and returns cond graph boolean output. +// 'output' need not be set if an error is returned. +typedef std::function& inputs, + Output* output)> + CondGraphBuilderFn; + +// Function that takes body graph inputs and returns body graph outputs. +// 'outputs' need not be populated if an error is returned. +typedef std::function& inputs, + std::vector* outputs)> + BodyGraphBuilderFn; + +// Constructs a while loop. +// +// Arguments: +// * scope: used to construct the while loop. +// * inputs: the initial values of the loop variables. Must be non-empty. +// * cond: a function that builds the condition graph of the loop. Takes the +// current loop variables as inputs and returns a scalar boolean Output +// indicating whether the loop should continue. +// * body: a function that builds the body graph of the loop. Takes the current +// loop variables as inputs and returns the updated loop variables. +// * frame_name: the frame name to use for this while loop. This should be a +// unique name. This will be used as a prefix for created operations. +// * outputs: output param that returns final loop variable outputs in non-error +// case. Must be non-null and empty. +// +// Returns an error if the while loop could not be fully constructed. +// +// TODO(skyewm): clean up partially-constructed loop in error case +// TODO(skyewm): create public interface to this method +Status BuildWhileLoop(const Scope& scope, const std::vector& inputs, + const CondGraphBuilderFn& cond, + const BodyGraphBuilderFn& body, const string& frame_name, + OutputList* outputs); + +} // namespace ops +} // namespace tensorflow + +#endif // THIRD_PARTY_TENSORFLOW_CC_OPS_WHILE_LOOP_H_ diff --git a/tensorflow/cc/saved_model/constants.h b/tensorflow/cc/saved_model/constants.h index 94a3b3cf465a279e3bb44344739499ad670119c3..c940df8a8761d97a859be3af30980ff79ca3577a 100644 --- a/tensorflow/cc/saved_model/constants.h +++ b/tensorflow/cc/saved_model/constants.h @@ -21,6 +21,9 @@ namespace tensorflow { /// SavedModel assets directory. constexpr char kSavedModelAssetsDirectory[] = "assets"; +/// SavedModel assets.extra directory. +constexpr char kSavedModelAssetsExtraDirectory[] = "assets.extra"; + /// SavedModel assets key for graph collection-def. constexpr char kSavedModelAssetsKey[] = "saved_model_assets"; diff --git a/tensorflow/cc/saved_model/loader.cc b/tensorflow/cc/saved_model/loader.cc index b144bfc33e46c3db192cfb1e3ef8a0633e9fa519..f98abc8a817eca7bc129bb03a2ad31b97d957065 100644 --- a/tensorflow/cc/saved_model/loader.cc +++ b/tensorflow/cc/saved_model/loader.cc @@ -20,9 +20,11 @@ limitations under the License. #include "tensorflow/cc/saved_model/constants.h" #include "tensorflow/core/lib/io/path.h" #include "tensorflow/core/lib/monitoring/counter.h" +#include "tensorflow/core/lib/strings/strcat.h" #include "tensorflow/core/platform/env.h" #include "tensorflow/core/platform/protobuf_internal.h" #include "tensorflow/core/protobuf/saved_model.pb.h" +#include "tensorflow/core/protobuf/saver.pb.h" #include "tensorflow/core/public/session.h" #include "tensorflow/core/public/session_options.h" #include "tensorflow/core/util/tensor_bundle/naming.h" @@ -36,7 +38,7 @@ auto* load_attempt_count = monitoring::Counter<2>::New( "status"); auto* load_latency = monitoring::Counter<1>::New( "/tensorflow/cc/saved_model/load_latency", - "Latency in microseconds for SavedModels that were succesfully loaded.", + "Latency in microseconds for SavedModels that were successfully loaded.", "model_path"); constexpr char kLoadAttemptFail[] = "fail"; constexpr char kLoadAttemptSuccess[] = "success"; @@ -75,8 +77,16 @@ Status FindMetaGraphDefToLoad(const SavedModel& saved_model_proto, return Status::OK(); } } + string tags_as_string = "{ "; + for (const string& tag : tags) { + tags_as_string = strings::StrCat(tags_as_string, tag, " "); + } + tags_as_string = strings::StrCat(tags_as_string, "}"); return Status(error::Code::NOT_FOUND, - "Could not find meta graph def matching supplied tags."); + "Could not find meta graph def matching supplied tags: " + + tags_as_string + + ". To inspect available tag-sets in the SavedModel, please " + "use the SavedModel CLI: `saved_model_cli`"); } Status LoadMetaGraphIntoSession(const MetaGraphDef& meta_graph_def, diff --git a/tensorflow/cc/saved_model/loader_test.cc b/tensorflow/cc/saved_model/loader_test.cc index cef29e7b071e538a60193fd998acc0fb29c2cea3..0ad6b33bba5fcceaca68e2f179cef2232c689a80 100644 --- a/tensorflow/cc/saved_model/loader_test.cc +++ b/tensorflow/cc/saved_model/loader_test.cc @@ -133,9 +133,9 @@ TEST_F(LoaderTest, NoTagMatch) { Status st = LoadSavedModel(session_options, run_options, export_dir, {"missing-tag"}, &bundle); EXPECT_FALSE(st.ok()); - EXPECT_TRUE( - StringPiece(st.error_message()) - .contains("Could not find meta graph def matching supplied tags.")) + EXPECT_TRUE(StringPiece(st.error_message()) + .contains("Could not find meta graph def matching supplied " + "tags: { missing-tag }")) << st.error_message(); } @@ -151,7 +151,7 @@ TEST_F(LoaderTest, NoTagMatchMultiple) { EXPECT_FALSE(st.ok()); EXPECT_TRUE( StringPiece(st.error_message()) - .contains("Could not find meta graph def matching supplied tags.")) + .contains("Could not find meta graph def matching supplied tags: ")) << st.error_message(); } diff --git a/tensorflow/cc/saved_model/tag_constants.h b/tensorflow/cc/saved_model/tag_constants.h index 48ab1158e462af25c27a728e404a041516e82057..2b0b2d5c7fb33768494c1781669c1adcb875a579 100644 --- a/tensorflow/cc/saved_model/tag_constants.h +++ b/tensorflow/cc/saved_model/tag_constants.h @@ -18,10 +18,13 @@ limitations under the License. namespace tensorflow { +/// Tag for the `gpu` graph. +constexpr char kSavedModelTagGpu[] = "gpu"; + /// Tag for the `serving` graph. constexpr char kSavedModelTagServe[] = "serve"; -/// Tag for the `training` graph.` +/// Tag for the `training` graph. constexpr char kSavedModelTagTrain[] = "train"; } // namespace tensorflow diff --git a/tensorflow/cc/saved_model/testdata/half_plus_two/00000123/saved_model.pb b/tensorflow/cc/saved_model/testdata/half_plus_two/00000123/saved_model.pb index d4ca4403d7576cfa8c2b5efaad08ae58785e6ef7..4a4fd025d9de0da7bd8e70eb0f017b851f686f28 100755 Binary files a/tensorflow/cc/saved_model/testdata/half_plus_two/00000123/saved_model.pb and b/tensorflow/cc/saved_model/testdata/half_plus_two/00000123/saved_model.pb differ diff --git a/tensorflow/cc/saved_model/testdata/half_plus_two_main_op/00000123/saved_model.pb b/tensorflow/cc/saved_model/testdata/half_plus_two_main_op/00000123/saved_model.pb index cf6234821a9516de962e00f25456fcb3ead7f681..daa272aead04c9a215334cfc0b6e10819833b3cd 100644 Binary files a/tensorflow/cc/saved_model/testdata/half_plus_two_main_op/00000123/saved_model.pb and b/tensorflow/cc/saved_model/testdata/half_plus_two_main_op/00000123/saved_model.pb differ diff --git a/tensorflow/cc/saved_model/testdata/half_plus_two_pbtxt/00000123/saved_model.pbtxt b/tensorflow/cc/saved_model/testdata/half_plus_two_pbtxt/00000123/saved_model.pbtxt index b24ebc7ddaa70350f204dadaf1f2bd34ef237c28..9d7813a0a162b4fad3b55d9ec252178b5b224783 100755 --- a/tensorflow/cc/saved_model/testdata/half_plus_two_pbtxt/00000123/saved_model.pbtxt +++ b/tensorflow/cc/saved_model/testdata/half_plus_two_pbtxt/00000123/saved_model.pbtxt @@ -284,6 +284,7 @@ meta_graphs { type: "shape" default_value { shape { + unknown_rank: true } } } @@ -447,7 +448,7 @@ meta_graphs { } } tags: "serve" - tensorflow_version: "1.0.0" + tensorflow_version: "1.1.0-rc2" tensorflow_git_version: "unknown" } graph_def { @@ -885,6 +886,7 @@ meta_graphs { key: "shape" value { shape { + unknown_rank: true } } } @@ -1714,7 +1716,7 @@ meta_graphs { dtype: DT_STRING tensor_shape { } - string_val: "_temp_d286b725003942fd8bac94b6c67e7c0c/part" + string_val: "_temp_80e928f1e0c844239d136d1ea966099d/part" } } } @@ -2444,7 +2446,7 @@ meta_graphs { input: "^save/restore_shard" } versions { - producer: 21 + producer: 23 } } saver_def { diff --git a/tensorflow/cc/training/coordinator.cc b/tensorflow/cc/training/coordinator.cc index 4618c932c310eefe775ccf9d8c38fbe1eea702ca..4511d0432068e5463121de3deb9d6bc6177b5370 100644 --- a/tensorflow/cc/training/coordinator.cc +++ b/tensorflow/cc/training/coordinator.cc @@ -116,18 +116,14 @@ void Coordinator::WaitForStop() { } Status Coordinator::ExportCostGraph(CostGraphDef* cost_graph) const { - RunMetadata tmp_metadata; - { - mutex_lock l(runners_lock_); - for (auto& t : runners_) { - Status s = t->ExportRunMetadata(&tmp_metadata); - if (!s.ok()) { - return s; - } + mutex_lock l(runners_lock_); + for (auto& t : runners_) { + Status s = t->ExportCostGraph(cost_graph); + if (!s.ok()) { + return s; } } - cost_graph->MergeFrom(tmp_metadata.cost_graph()); return Status::OK(); } -} // namespace +} // namespace tensorflow diff --git a/tensorflow/cc/training/coordinator.h b/tensorflow/cc/training/coordinator.h index 632418c5ca5f523defe781a780ca0987202f59e4..0e01b19cd98bc797b7bb25da55c05d96f3eb93c7 100644 --- a/tensorflow/cc/training/coordinator.h +++ b/tensorflow/cc/training/coordinator.h @@ -36,8 +36,8 @@ class RunnerInterface { public: virtual ~RunnerInterface() {} virtual Status Join() = 0; - virtual Status ExportRunMetadata(RunMetadata* metadata) const { - return Status(error::INVALID_ARGUMENT, "No RunMetadata to export."); + virtual Status ExportCostGraph(CostGraphDef* cost_graph) const { + return Status(error::INVALID_ARGUMENT, "No cost model to export."); } /// Returns true iff the runner is running, i.e. if it is trying to populate /// its queue. diff --git a/tensorflow/cc/training/coordinator_test.cc b/tensorflow/cc/training/coordinator_test.cc index a87913deafe01bba9d54c9ad59580b04e1945866..48874033841008ada0a5ee207a74331f24247ed4 100644 --- a/tensorflow/cc/training/coordinator_test.cc +++ b/tensorflow/cc/training/coordinator_test.cc @@ -55,7 +55,7 @@ TEST(CoordinatorTest, TestStopAndWaitOnStop) { class MockQueueRunner : public RunnerInterface { public: - MockQueueRunner(Coordinator* coord) { + explicit MockQueueRunner(Coordinator* coord) { coord_ = coord; join_counter_ = nullptr; thread_pool_.reset(new thread::ThreadPool(Env::Default(), "test-pool", 10)); @@ -79,7 +79,7 @@ class MockQueueRunner : public RunnerInterface { status, counter, start)); } - Status Join() { + Status Join() override { if (join_counter_ != nullptr) { (*join_counter_)++; } diff --git a/tensorflow/cc/training/queue_runner.cc b/tensorflow/cc/training/queue_runner.cc index 6b615916813519d7eaa94e69e846dcbfb87623bc..b6f240350b8954ea75f2d08374c4436508c717c2 100644 --- a/tensorflow/cc/training/queue_runner.cc +++ b/tensorflow/cc/training/queue_runner.cc @@ -49,7 +49,12 @@ Status QueueRunner::Init(const QueueRunnerDef& queue_runner_def) { enqueue_op_names_.insert(enqueue_op_names_.end(), queue_runner_def.enqueue_op_name().begin(), queue_runner_def.enqueue_op_name().end()); - runs_ = enqueue_op_names_.size(); + size_t op_names_size = enqueue_op_names_.size(); + if (op_names_size > kint32max) { + return Status(error::INVALID_ARGUMENT, + "Enqueue ops to run cannot exceed kint32max"); + } + runs_ = static_cast(op_names_size); if (runs_ == 0) { return Status(error::INVALID_ARGUMENT, "Empty enqueue ops to run."); } @@ -82,9 +87,9 @@ QueueRunner::~QueueRunner() { Status QueueRunner::Start(Session* sess) { return Start(sess, 0); } -Status QueueRunner::StartAndCollectRunMetadata(Session* sess, - const RunOptions* run_options) { - SetRunArgumentsAndRunMetadata(run_options); +Status QueueRunner::StartAndCollectCostGraph(Session* sess, + const RunOptions& run_options) { + SetRunArgumentsAndCostGraph(run_options); return Start(sess, 0); } @@ -115,10 +120,9 @@ Status QueueRunner::Start(Session* sess, int wait_for) { return Status::OK(); } -Status QueueRunner::StartAndCollectRunMetadata(Session* session, - int wait_for_ms, - const RunOptions* run_options) { - SetRunArgumentsAndRunMetadata(run_options); +Status QueueRunner::StartAndCollectCostGraph(Session* session, int wait_for_ms, + const RunOptions& run_options) { + SetRunArgumentsAndCostGraph(run_options); return Start(session, wait_for_ms); } @@ -127,7 +131,7 @@ void QueueRunner::Stop(Session* sess) { coord_->WaitForStop(); } if (!cancel_op_name_.empty()) { - UpdateStatus(RealRun(sess, cancel_op_name_)); + UpdateStatus(RealRun(sess, cancel_op_name_, false)); } stopped_ = true; } @@ -162,7 +166,7 @@ void QueueRunner::Run(Session* sess, const string& enqueue_op) { if (coord_ && coord_->ShouldStop()) { break; } - status = RealRun(sess, enqueue_op); + status = RealRun(sess, enqueue_op, true); if (first_iteration) { if (!status.ok()) { mutex_lock l(mu_); @@ -183,9 +187,11 @@ void QueueRunner::Run(Session* sess, const string& enqueue_op) { // will be run anway in this case. if (IsQueueClosed(status) && (!coord_ || !coord_->ShouldStop())) { if (last_run && !close_op_name_.empty()) { - UpdateStatus(RealRun(sess, close_op_name_)); + UpdateStatus(RealRun(sess, close_op_name_, false)); } } else if (!status.ok()) { + LOG(ERROR) << "Queue runner thread got a failure status: " + << status.ToString(); UpdateStatus(status); if (coord_) { coord_->RequestStop().IgnoreError(); @@ -198,34 +204,33 @@ Status QueueRunner::GetStatus() { return status_; } -Status QueueRunner::ExportRunMetadata(RunMetadata* metadata) const { - if (!rm_mu_) { +Status QueueRunner::ExportCostGraph(CostGraphDef* cost_graph) const { + if (!cg_mu_) { return Status(error::FAILED_PRECONDITION, - "This QueueRunner doesn't collect and store RunMetadata."); + "This QueueRunner doesn't collect a cost graph."); } - mutex_lock l(*rm_mu_); - metadata->MergeFrom(*run_metadata_); + mutex_lock l(*cg_mu_); + cost_graph->MergeFrom(*cost_graph_); return Status::OK(); } -void QueueRunner::SetRunArgumentsAndRunMetadata(const RunOptions* run_options) { - rm_mu_.reset(new mutex()); +void QueueRunner::SetRunArgumentsAndCostGraph(const RunOptions& run_options) { + cg_mu_.reset(new mutex()); { - mutex_lock l(*rm_mu_); - run_metadata_.reset(new RunMetadata()); - } - if (run_options) { - run_options_ = *run_options; + mutex_lock l(*cg_mu_); + cost_graph_.reset(new CostGraphDef()); } + run_options_ = run_options; } -Status QueueRunner::RealRun(Session* sess, const string& op) { +Status QueueRunner::RealRun(Session* sess, const string& op, + bool update_costs) { Status s; - if (rm_mu_) { + if (update_costs && cg_mu_) { RunMetadata metadata; s = sess->Run(run_options_, {}, {}, {op}, nullptr, &metadata); - mutex_lock l(*rm_mu_); - run_metadata_->MergeFrom(metadata); + mutex_lock l(*cg_mu_); + cost_graph_->Swap(metadata.mutable_cost_graph()); } else { s = sess->Run({}, {}, {op}, nullptr); } diff --git a/tensorflow/cc/training/queue_runner.h b/tensorflow/cc/training/queue_runner.h index c69f28793a95990901961e835e004b019b98dbdc..2d3450032388bfee96055f23cf621af0fa4731ae 100644 --- a/tensorflow/cc/training/queue_runner.h +++ b/tensorflow/cc/training/queue_runner.h @@ -60,15 +60,15 @@ class QueueRunner : public RunnerInterface { Status Start(Session* sess); /// Starts the queue runner with the given session and sets the run arguments - /// for sess->Run. It also collects and stores the run metedata. - Status StartAndCollectRunMetadata(Session* sess, - const RunOptions* run_options = nullptr); + /// for sess->Run. It also collects and stores the cost model. + Status StartAndCollectCostGraph(Session* sess, + const RunOptions& run_options = RunOptions()); /// Starts the queue runner with the given session, and wait for up to the /// specified time (in milliseconds) for the queues to start to fill up. Status Start(Session* sess, int wait_for_ms); - Status StartAndCollectRunMetadata(Session* session, int wait_for_ms, - const RunOptions* run_options = nullptr); + Status StartAndCollectCostGraph(Session* session, int wait_for_ms, + const RunOptions& run_options = RunOptions()); /// Requests to stop and runs the cancel op. It would be called in a separate /// thread when coordinator is set. If there is no coordinator it should be @@ -82,11 +82,11 @@ class QueueRunner : public RunnerInterface { /// Returns the latest status. Status GetStatus(); - // Returns the stored run metadata. - Status ExportRunMetadata(RunMetadata* metadata) const override; + // Returns the stored cost model. + Status ExportCostGraph(CostGraphDef* cost_graph) const override; private: - QueueRunner() : coord_(nullptr), stopped_(false), rm_mu_(nullptr) {} + QueueRunner() : coord_(nullptr), stopped_(false), cg_mu_(nullptr) {} // Initializes the instance with the QueueRunnerDef proto. Status Init(const QueueRunnerDef& queue_runner_def); @@ -105,9 +105,9 @@ class QueueRunner : public RunnerInterface { bool IsRunning() const override { return !stopped_; } - void SetRunArgumentsAndRunMetadata(const RunOptions* run_options); + void SetRunArgumentsAndCostGraph(const RunOptions& run_options); - Status RealRun(Session* sess, const string& op); + Status RealRun(Session* sess, const string& op, bool update_costs); string queue_name_; std::vector enqueue_op_names_; @@ -130,8 +130,8 @@ class QueueRunner : public RunnerInterface { mutex cb_mu_; std::vector> callbacks_; - mutable std::unique_ptr rm_mu_; - std::unique_ptr run_metadata_ GUARDED_BY(rm_mu_); + mutable std::unique_ptr cg_mu_; + std::unique_ptr cost_graph_ GUARDED_BY(cg_mu_); RunOptions run_options_; }; diff --git a/tensorflow/cc/training/queue_runner_test.cc b/tensorflow/cc/training/queue_runner_test.cc index c37a69a7f76b6d83634d0b01e2038c4e6b4fa22e..075bc8a55afb5425b4bee8bca7960b39a30f10ff 100644 --- a/tensorflow/cc/training/queue_runner_test.cc +++ b/tensorflow/cc/training/queue_runner_test.cc @@ -44,6 +44,7 @@ using ops::FIFOQueue; using ops::QueueClose; using ops::QueueDequeue; using ops::QueueEnqueue; +using ops::RandomNormal; using ops::Square; using ops::Variable; @@ -84,7 +85,7 @@ QueueRunnerDef BuildQueueRunnerDef( const std::string& close_op, const std::string& cancel_op, const std::vector& queue_closed_error_codes) { QueueRunnerDef queue_runner_def; - *queue_runner_def.mutable_queue_name() = kQueueName; + *queue_runner_def.mutable_queue_name() = queue_name; for (const std::string& enqueue_op : enqueue_ops) { *queue_runner_def.mutable_enqueue_op_name()->Add() = enqueue_op; } @@ -345,37 +346,54 @@ TEST(QueueRunnerTest, CallbackCalledOnError) { } TEST(QueueRunnerTest, RunMetaDataTest) { + Scope root = Scope::NewRootScope(); + auto q0 = FIFOQueue(root.WithOpName(kQueueName), {DataType::DT_FLOAT}); + Output rnd = RandomNormal(root.WithOpName("rnd"), {1, 1}, DataType::DT_FLOAT); + Output square = Square(root.WithOpName(kSquareOpName), rnd); + auto enqueue0 = QueueEnqueue(root.WithOpName(kEnqueueOp0), q0, {square}); + auto close0 = QueueClose(root.WithOpName(kCloseOp0), q0); + auto cancel0 = QueueClose(root.WithOpName(kCancelOp0), q0, + QueueClose::CancelPendingEnqueues(true)); + auto dequeue0 = + QueueDequeue(root.WithOpName(kDequeueOp0), q0, {DataType::DT_FLOAT}); + + GraphDef graph_def; + TF_EXPECT_OK(root.ToGraphDef(&graph_def)); + for (auto& node : *graph_def.mutable_node()) { + node.set_device("/cpu:0"); + } SessionOptions sess_options; sess_options.config.mutable_graph_options()->set_build_cost_model(1); std::unique_ptr session(NewSession(sess_options)); - GraphDef graph_def = BuildSimpleGraph(); TF_CHECK_OK(session->Create(graph_def)); - TF_CHECK_OK(session->Run({}, {}, {kAssignOpName}, nullptr)); - RunOptions run_options; - run_options.set_trace_level(RunOptions::HARDWARE_TRACE); - - QueueRunnerDef queue_runner_def = BuildQueueRunnerDef( - kQueueName, {kCountUpToOpName}, kSquareOpName, "", {}); + QueueRunnerDef queue_runner_def = + BuildQueueRunnerDef(kQueueName, {kEnqueueOp0}, kCloseOp0, kCancelOp0, {}); std::unique_ptr qr; TF_EXPECT_OK(QueueRunner::New(queue_runner_def, &qr)); - TF_CHECK_OK(qr->StartAndCollectRunMetadata(session.get(), &run_options)); + RunOptions run_options; + TF_CHECK_OK(qr->StartAndCollectCostGraph(session.get(), run_options)); - TF_EXPECT_OK(qr->Join()); - RunMetadata run_metadata; - TF_CHECK_OK(qr->ExportRunMetadata(&run_metadata)); + // Make sure there was at least one element enqueued in q0: this prevents a + // race condition where we close the queue before it was populated. + std::vector dq0; + TF_EXPECT_OK(session->Run({}, {kDequeueOp0}, {}, &dq0)); + // Second call to run dequeue op is to make sure the cost graph has been + // stored. + TF_EXPECT_OK(session->Run({}, {kDequeueOp0}, {}, &dq0)); + + CostGraphDef cost_graph; + TF_CHECK_OK(qr->ExportCostGraph(&cost_graph)); + EXPECT_TRUE(cost_graph.node_size() > 0); - EXPECT_TRUE(run_metadata.has_cost_graph()); + qr->Stop(session.get()); } TEST(QueueRunnerTest, NoRunMetaDataTest) { GraphDef graph_def = BuildSimpleGraph(); auto session = BuildSessionAndInitVariable(graph_def); - RunOptions run_options; - run_options.set_trace_level(RunOptions::HARDWARE_TRACE); - QueueRunnerDef queue_runner_def = BuildQueueRunnerDef( kQueueName, {kCountUpToOpName}, kSquareOpName, "", {}); std::unique_ptr qr; @@ -383,8 +401,8 @@ TEST(QueueRunnerTest, NoRunMetaDataTest) { TF_CHECK_OK(qr->Start(session.get())); TF_EXPECT_OK(qr->Join()); - RunMetadata run_metadata; - EXPECT_EQ(qr->ExportRunMetadata(&run_metadata).code(), + CostGraphDef cost_graph; + EXPECT_EQ(qr->ExportCostGraph(&cost_graph).code(), error::FAILED_PRECONDITION); } diff --git a/tensorflow/cc/tutorials/example_trainer.cc b/tensorflow/cc/tutorials/example_trainer.cc index f2ecd2eddc28da94ac1c2404c02324e7782831c3..3675d72ee354533a7d84b5e8783cde452d8d60c9 100644 --- a/tensorflow/cc/tutorials/example_trainer.cc +++ b/tensorflow/cc/tutorials/example_trainer.cc @@ -101,7 +101,7 @@ void ConcurrentSteps(const Options* opts, int session_index) { std::unique_ptr session(NewSession(options)); GraphDef def = CreateGraphDef(); if (options.target.empty()) { - graph::SetDefaultDevice(opts->use_gpu ? "/gpu:0" : "/cpu:0", &def); + graph::SetDefaultDevice(opts->use_gpu ? "/device:GPU:0" : "/cpu:0", &def); } TF_CHECK_OK(session->Create(def)); @@ -227,7 +227,7 @@ int main(int argc, char* argv[]) { argv[dst++] = f; } argv[dst++] = nullptr; - argc = unknown_flags.size() + 1; + argc = static_cast(unknown_flags.size() + 1); tensorflow::port::InitMain(argv[0], &argc, &argv); tensorflow::example::ConcurrentSessions(opts); } diff --git a/tensorflow/compiler/aot/BUILD b/tensorflow/compiler/aot/BUILD index c52a56b6428fb8a8415ed53477ba3e81c57b0ded..f956602ba221bbbb3c2fc9c7df7d452da833c002 100644 --- a/tensorflow/compiler/aot/BUILD +++ b/tensorflow/compiler/aot/BUILD @@ -20,6 +20,7 @@ cc_library( cc_test( name = "runtime_test", + size = "small", srcs = ["runtime_test.cc"], deps = [ ":runtime", @@ -73,7 +74,7 @@ cc_library( "//tensorflow/compiler/xla:util", "//tensorflow/compiler/xla:xla_data_proto", "//tensorflow/compiler/xla/client:client_library", - "//tensorflow/compiler/xla/client:local_client", + "//tensorflow/compiler/xla/client:compile_only_client", "//tensorflow/compiler/xla/service:compiler", "//tensorflow/compiler/xla/service/cpu:cpu_compiler", "//tensorflow/core:core_cpu", @@ -88,6 +89,7 @@ cc_library( cc_test( name = "codegen_test", + size = "small", srcs = ["codegen_test.cc"], data = ["codegen_test_h.golden"], deps = [ @@ -101,10 +103,12 @@ cc_test( cc_test( name = "tfcompile_util_test", + size = "small", srcs = ["tfcompile_util_test.cc"], deps = [ ":tfcompile_lib", "//tensorflow/core:lib", + "//tensorflow/core:protos_all_cc", "//tensorflow/core:test", "//tensorflow/core:test_main", ], @@ -123,9 +127,7 @@ cc_library( deps = [ ":tfcompile_lib", ":tfcompile_proto", - "//tensorflow/compiler/xla/legacy_flags:compiler_functor_flags", - "//tensorflow/compiler/xla/legacy_flags:cpu_compiler_flags", - "//tensorflow/compiler/xla/legacy_flags:cpu_runtime_flags", + "//tensorflow/compiler/xla/legacy_flags:debug_options_flags", "//tensorflow/compiler/xla/service:compiler", "//tensorflow/core:core_cpu", "//tensorflow/core:core_cpu_internal", @@ -151,6 +153,40 @@ tf_library( tags = ["manual"], ) +# A test of tf_library that includes a graph with an unknown op, but where +# the compilation works because the unknown op is not needed for the fetches. +tf_library( + name = "test_graph_tfunknownop", + testonly = 1, + config = "test_graph_tfunknownop.config.pbtxt", + cpp_class = "UnknownOpAddComp", + graph = "test_graph_tfunknownop.pbtxt", + tags = ["manual"], +) + +# A test of tf_library that includes a graph with an unknown op, but where +# the compilation works because the op between the unknown op and the +# fetches is a feed. +tf_library( + name = "test_graph_tfunknownop2", + testonly = 1, + config = "test_graph_tfunknownop2.config.pbtxt", + cpp_class = "UnknownOpAddComp", + graph = "test_graph_tfunknownop.pbtxt", + tags = ["manual"], +) + +# A test of tf_library that includes a graph with an unknown op, but where +# the compilation works because the unknown op is fed. +tf_library( + name = "test_graph_tfunknownop3", + testonly = 1, + config = "test_graph_tfunknownop3.config.pbtxt", + cpp_class = "UnknownOpAddComp", + graph = "test_graph_tfunknownop.pbtxt", + tags = ["manual"], +) + # Utility library for benchmark binaries, used by the *_benchmark rules that are # added by the tfcompile bazel macro. cc_library( @@ -194,6 +230,7 @@ test_suite( tests = [ ":benchmark_test", ":test_graph_tfadd_test", + ":test_graph_tfunknownop_test", "//tensorflow/compiler/aot/tests:all_tests", ], ) diff --git a/tensorflow/compiler/aot/benchmark.cc b/tensorflow/compiler/aot/benchmark.cc index 0c5e2c103ea9476ac19c6119dd44a3229ede7a12..ff72038281227ce8adb20ddc25e678c49a2448db 100644 --- a/tensorflow/compiler/aot/benchmark.cc +++ b/tensorflow/compiler/aot/benchmark.cc @@ -40,7 +40,7 @@ namespace benchmark { // the implementation without pulling in all of the Env dependencies. static double NowMicros() { struct timeval tv; - gettimeofday(&tv, NULL); + gettimeofday(&tv, nullptr); return static_cast(tv.tv_sec) * 1000000 + tv.tv_usec; } diff --git a/tensorflow/compiler/aot/codegen.cc b/tensorflow/compiler/aot/codegen.cc index 042a72745a78c4a11b22c85e3a094d78c4ab2ed5..bbdb342a623f5d4435e437fbb94e282b685751c9 100644 --- a/tensorflow/compiler/aot/codegen.cc +++ b/tensorflow/compiler/aot/codegen.cc @@ -152,8 +152,7 @@ Status AddRewritesForShape(int i, const xla::Shape& shape, string RewriteWithName(const string& name, string code, const std::vector>& rewrites) { str_util::ReplaceAllPairs(&code, rewrites); - str_util::ReplaceAll(&code, "{{NAME}}", name); - return code; + return str_util::StringReplace(code, "{{NAME}}", name, /*replace_all=*/true); } // Generate methods for args (inputs). @@ -366,7 +365,7 @@ Status GenerateHeader(const HeaderOpts& opts, const Config& config, #include "tensorflow/core/platform/macros.h" #include "tensorflow/core/platform/types.h" -namespace Eigen { class ThreadPoolDevice; } +namespace Eigen { struct ThreadPoolDevice; } // (Implementation detail) Entry point to the function in the object file. extern "C" void {{ENTRY}}( diff --git a/tensorflow/compiler/aot/codegen_test_h.golden b/tensorflow/compiler/aot/codegen_test_h.golden index 46d7c03006a1344df17fc99c8b837f31ee86feb9..01963c6df4682ec8c23a93201d7fbbab63558060 100644 --- a/tensorflow/compiler/aot/codegen_test_h.golden +++ b/tensorflow/compiler/aot/codegen_test_h.golden @@ -15,7 +15,7 @@ #include "tensorflow/core/platform/macros.h" #include "tensorflow/core/platform/types.h" -namespace Eigen { class ThreadPoolDevice; } +namespace Eigen { struct ThreadPoolDevice; } // (Implementation detail) Entry point to the function in the object file. extern "C" void entry_point( diff --git a/tensorflow/compiler/aot/compile.cc b/tensorflow/compiler/aot/compile.cc index 1284155c07b1a253d42e7641354626eb153f0c35..a485d2e555ac8176c8b8950ef6d58757a2c71aaa 100644 --- a/tensorflow/compiler/aot/compile.cc +++ b/tensorflow/compiler/aot/compile.cc @@ -24,10 +24,11 @@ limitations under the License. #include "tensorflow/compiler/aot/flags.h" #include "tensorflow/compiler/aot/tfcompile_util.h" +#include "tensorflow/compiler/tf2xla/shape_util.h" #include "tensorflow/compiler/tf2xla/xla_compiler.h" #include "tensorflow/compiler/tf2xla/xla_op_registry.h" #include "tensorflow/compiler/xla/client/client_library.h" -#include "tensorflow/compiler/xla/client/local_client.h" +#include "tensorflow/compiler/xla/client/compile_only_client.h" #include "tensorflow/compiler/xla/service/compiler.h" #include "tensorflow/compiler/xla/service/cpu/cpu_compiler.h" #include "tensorflow/compiler/xla/shape_util.h" @@ -40,6 +41,7 @@ limitations under the License. #include "tensorflow/core/framework/graph_def_util.h" #include "tensorflow/core/framework/op.h" #include "tensorflow/core/framework/tensor_shape.h" +#include "tensorflow/core/framework/versions.pb.h" #include "tensorflow/core/graph/algorithm.h" #include "tensorflow/core/graph/graph.h" #include "tensorflow/core/graph/graph_constructor.h" @@ -77,66 +79,60 @@ Status DumpGraph(const MainFlags& flags, const string& name, return WriteTextProto(Env::Default(), file, graph_def); } -string TensorIdToString(const TensorId& id) { - return strings::StrCat(id.node_name(), ":", id.output_index()); -} - typedef std::unordered_map NodeMap; // Each feed id identifies the positional output of some node, which may consist -// of multiple edges. For each feed node, replaces all matching edges so that -// they point from a new _Arg node instead. +// of multiple edges. AddPlaceholdersForFeeds has already replaced each fed +// tensor with a placeholder. For each feed tensor, replaces all edges so they +// point from a new _Arg node instead. Status AddArgNodes(Graph* graph, const NodeMap& node_map, - const protobuf::RepeatedPtrField& feeds) { + const protobuf::RepeatedPtrField& feeds, + const std::unordered_map& feed_remapping) { for (int arg_index = 0; arg_index < feeds.size(); ++arg_index) { const Feed& feed = feeds[arg_index]; - const TensorId& id = feed.id(); - auto it = node_map.find(id.node_name()); - if (it == node_map.end()) { - return errors::NotFound("Can't find feed id: ", TensorIdToString(id)); - } - const Node* feed_node = it->second; - if (id.output_index() >= feed_node->num_outputs()) { - return errors::InvalidArgument("Invalid feed id: ", TensorIdToString(id), - ", output index should be < ", - feed_node->num_outputs()); + // All feeds have been replaced by placeholders. + const int output_index = 0; + + const string key = TensorIdToString(feed.id()); + const auto remap_it = feed_remapping.find(key); + auto node_it = node_map.find(remap_it->second); + if (node_it == node_map.end()) { + // Strip off the aot_feed_#/ prefix. + StringPiece name(remap_it->second); + const auto index = name.find('/'); + if (index > 0) name.remove_prefix(index + 1); + return errors::InvalidArgument( + "Node is fed but not needed for fetching: ", name); } - // TODO(toddw): Invoke shape inference on the graph and add a "_shape" attr - // if we can determine it. That way the graph will be initialized with - // whatever shapes we can infer, while the user can still explicitly specify - // or override them. + const Node* feed_node = node_it->second; + + // TODO(toddw): Invoke shape inference in AddPlaceholdersForFeeds and add a + // "_shape" attr if we can determine it. That way the graph will be + // initialized with whatever shapes we can infer, while the user can still + // explicitly specify or override them. Node* arg_node = nullptr; TF_RETURN_IF_ERROR( NodeBuilder(strings::StrCat("_arg_", arg_index), kArgOp) - .Attr("T", BaseType(feed_node->output_type(id.output_index()))) + .Attr("T", BaseType(feed_node->output_type(output_index))) .Attr("index", arg_index) - .Attr(kFeedIdAttr, TensorIdToString(id)) + .Attr(kFeedIdAttr, TensorIdToString(feed.id())) .Attr(kShapeAttr, TensorShape(feed.shape())) .Attr(kDebugNameAttr, feed.name()) .Finalize(graph, &arg_node)); + // Collects out-edges from the feed node that have a matching edge index; - // these will be replaced with edges from the arg node instead. Also - // replaces all control edges from Placeholder feed nodes; similar code - // exists in subgraph::RewriteGraphForExecution. - // TODO(toddw): Why only replace control edges from Placeholder? + // these will be replaced with edges from the arg node instead. // // We must collect the edges first and process them in a second pass, since // removing the edge from the graph invalidates feed_node->out_edges. std::vector feed_edges; for (const Edge* edge : feed_node->out_edges()) { - if (edge->src_output() == id.output_index() || - (edge->src_output() == Graph::kControlSlot && - feed_node->type_string() == "Placeholder")) { + if (edge->src_output() == output_index) { feed_edges.push_back(edge); } } for (const Edge* edge : feed_edges) { - if (edge->src_output() == id.output_index()) { - graph->AddEdge(arg_node, 0, edge->dst(), edge->dst_input()); - } else { - CHECK_EQ(edge->src_output(), Graph::kControlSlot); - graph->AddControlEdge(arg_node, edge->dst()); - } + graph->AddEdge(arg_node, 0, edge->dst(), edge->dst_input()); graph->RemoveEdge(edge); } } @@ -178,13 +174,16 @@ Status AddRetvalNodes(Graph* graph, const NodeMap& node_map, // fetch ids respectively), and rewrites the edges so that inputs flow from _Arg // nodes, and outputs flow to _Retval nodes. This allows the symbolic graph // execution to know the input and output args for the generated function. -Status RewriteAndPruneGraph(Graph* graph, const Config& config, - const MainFlags& flags) { +Status RewriteAndPruneGraph( + Graph* graph, const Config& config, + const std::unordered_map& feed_remapping, + const MainFlags& flags) { NodeMap node_map; for (Node* n : graph->nodes()) { node_map[n->name()] = n; } - TF_RETURN_IF_ERROR(AddArgNodes(graph, node_map, config.feed())); + TF_RETURN_IF_ERROR( + AddArgNodes(graph, node_map, config.feed(), feed_remapping)); std::unordered_set retval_nodes; TF_RETURN_IF_ERROR( AddRetvalNodes(graph, node_map, config.fetch(), &retval_nodes)); @@ -200,17 +199,17 @@ Status RewriteAndPruneGraph(Graph* graph, const Config& config, for (const Fetch& fetch : config.fetch()) { missing_fetches.insert(TensorIdToString(fetch.id())); } - for (const Node* n : graph->nodes()) { + for (const Node* n : graph->op_nodes()) { if (n->type_string() == kArgOp) { string feed_id; - TF_RETURN_IF_ERROR(GetNodeAttr(n->def(), kFeedIdAttr, &feed_id)); + TF_RETURN_IF_ERROR(GetNodeAttr(n->attrs(), kFeedIdAttr, &feed_id)); if (missing_feeds.erase(feed_id) == 0) { return errors::Aborted(kArgOp, " node found with unknown feed id: ", feed_id); } } else if (n->type_string() == kRetvalOp) { string fetch_id; - TF_RETURN_IF_ERROR(GetNodeAttr(n->def(), kFetchIdAttr, &fetch_id)); + TF_RETURN_IF_ERROR(GetNodeAttr(n->attrs(), kFetchIdAttr, &fetch_id)); if (missing_fetches.erase(fetch_id) == 0) { return errors::Aborted(kRetvalOp, " node found with unknown fetch id: ", fetch_id); @@ -234,7 +233,7 @@ Status CollectArgNodes(const Graph& graph, std::vector* arg_nodes) { for (Node* n : graph.nodes()) { if (n->type_string() == kArgOp) { int index; - TF_RETURN_IF_ERROR(GetNodeAttr(n->def(), "index", &index)); + TF_RETURN_IF_ERROR(GetNodeAttr(n->attrs(), "index", &index)); auto insert_result = indexed_arg_nodes.insert({index, n}); if (!insert_result.second) { const Node* dup = insert_result.first->second; @@ -264,9 +263,11 @@ Status CreateXlaArgs(const Graph& graph, for (const Node* node : arg_nodes) { XlaCompiler::Argument arg; arg.kind = XlaCompiler::Argument::kParameter; - TF_RETURN_IF_ERROR(GetNodeAttr(node->def(), "T", &arg.type)); - TF_RETURN_IF_ERROR(GetNodeAttr(node->def(), kShapeAttr, &arg.shape)); - TF_RETURN_IF_ERROR(GetNodeAttr(node->def(), kDebugNameAttr, &arg.name)); + TF_RETURN_IF_ERROR(GetNodeAttr(node->attrs(), "T", &arg.type)); + TensorShape shape; + TF_RETURN_IF_ERROR(GetNodeAttr(node->attrs(), kShapeAttr, &shape)); + TF_RETURN_IF_ERROR(TensorShapeToXLAShape(arg.type, shape, &arg.shape)); + TF_RETURN_IF_ERROR(GetNodeAttr(node->attrs(), kDebugNameAttr, &arg.name)); xla_args->push_back(arg); } return Status::OK(); @@ -274,8 +275,8 @@ Status CreateXlaArgs(const Graph& graph, // Converts the TensorFlow graph into an XLA computation, by executing the // graph symbolically, with each op building up the XLA HLO. -Status ConvertGraphToXla(xla::LocalClient* client, std::unique_ptr graph, - const FunctionLibraryDefinition* flib_def, +Status ConvertGraphToXla(xla::CompileOnlyClient* client, + std::unique_ptr graph, xla::Computation* computation, bool* has_context_arg) { // Create a device and context to convert the graph into an XLA computation. XlaOpRegistry::RegisterCompilationKernels(); @@ -289,18 +290,19 @@ Status ConvertGraphToXla(xla::LocalClient* client, std::unique_ptr graph, // Compile the graph into an XLA computation. XlaCompiler::Options compiler_options; compiler_options.client = client; - compiler_options.device_type = DeviceType(DEVICE_CPU_XLA_JIT); + DeviceType device_type(DEVICE_CPU_XLA_JIT); + compiler_options.device_type = &device_type; + compiler_options.flib_def = &graph->flib_def(); + compiler_options.graph_def_version = graph->versions().producer(); compiler_options.allow_cpu_custom_calls = true; XlaCompiler compiler(compiler_options); - std::unique_ptr flib_run(NewFunctionLibraryRuntime( - compiler.device_mgr(), Env::Default(), compiler.device(), - graph->versions().producer(), flib_def, OptimizerOptions())); XlaCompiler::CompilationResult result; - TF_RETURN_IF_ERROR(compiler.CompileGraph("tfcompile", std::move(graph), - flib_run.get(), xla_args, &result)); + TF_RETURN_IF_ERROR(compiler.CompileGraph(XlaCompiler::CompileOptions(), + "tfcompile", std::move(graph), + xla_args, &result)); *has_context_arg = result.requires_runtime_context; - *computation = std::move(result.computation); + *computation = std::move(*result.computation); int num_const_results = 0; for (int i = 0; i < result.outputs.size(); ++i) { @@ -334,7 +336,8 @@ Status ConvertGraphToXla(xla::LocalClient* client, std::unique_ptr graph, } // Compiles the XLA computation into executable code. -Status CompileXla(xla::LocalClient* client, const xla::Computation& computation, +Status CompileXla(xla::CompileOnlyClient* client, + const xla::Computation& computation, const xla::cpu::CpuAotCompilationOptions& aot_opts, CompileResult* compile_result) { // Retrieves arg and result layouts from the computation. @@ -348,10 +351,11 @@ Status CompileXla(xla::LocalClient* client, const xla::Computation& computation, compile_result->program_shape = *pshape_or.ValueOrDie(); xla::ProgramShape* pshape = &compile_result->program_shape; std::vector arg_layouts; + arg_layouts.reserve(pshape->parameters_size()); for (int i = 0; i < pshape->parameters_size(); ++i) { arg_layouts.push_back(pshape->mutable_parameters(i)); } - xla::LocalClient::AheadOfTimeComputationInstance instance; + xla::CompileOnlyClient::AotComputationInstance instance; instance.computation = &computation; instance.argument_layouts = std::move(arg_layouts); instance.result_layout = &pshape->result(); @@ -366,29 +370,46 @@ Status CompileXla(xla::LocalClient* client, const xla::Computation& computation, std::move(aot_or.ValueOrDie().back())); compile_result->entry_point = aot_opts.entry_point_name(); compile_result->pointer_size = - xla::LocalClient::PointerSizeForTriple(aot_opts.triple()); + xla::CompileOnlyClient::PointerSizeForTriple(aot_opts.triple()); return Status::OK(); } } // namespace Status InitGraph(const GraphDef& graph_def, const Config& config, - const MainFlags& flags, const FunctionLibraryDefinition* flib, - std::unique_ptr* graph) { + const MainFlags& flags, std::unique_ptr* graph) { TF_RETURN_IF_ERROR(ValidateConfig(config)); - std::unique_ptr g(new Graph(flib)); - GraphDef copy_def(graph_def); - TF_RETURN_IF_ERROR(AddDefaultAttrsToGraphDef(©_def, *g->op_registry(), - 0 /*node_offset*/)); + + FunctionLibraryDefinition flib_def(OpRegistry::Global(), graph_def.library()); + std::unique_ptr g(new Graph(flib_def)); + + // Replace references to fed tensors with references to newly added + // placeholders. + GraphDef first_copy_def = graph_def; + + // Maps from name:port of a feed to the name:port of the placeholder to use. + std::unordered_map feed_remapping; + TF_RETURN_IF_ERROR(AddPlaceholdersForFeeds(config, g->op_registry(), + &feed_remapping, &first_copy_def)); + + // Prune the GraphDef first so that unknown ops that we aren't compiling get + // filtered out. + GraphDef second_copy_def; + TF_RETURN_IF_ERROR( + PruneGraphDefInto(config, first_copy_def, &second_copy_def)); + + TF_RETURN_IF_ERROR(AddDefaultAttrsToGraphDef( + &second_copy_def, *g->op_registry(), 0 /*node_offset*/)); + + TF_RETURN_IF_ERROR(ConvertGraphDefToGraph(GraphConstructorOptions(), + second_copy_def, g.get())); TF_RETURN_IF_ERROR( - ConvertGraphDefToGraph(GraphConstructorOptions(), copy_def, g.get())); - TF_RETURN_IF_ERROR(RewriteAndPruneGraph(g.get(), config, flags)); + RewriteAndPruneGraph(g.get(), config, feed_remapping, flags)); *graph = std::move(g); return Status::OK(); } Status CompileGraph(std::unique_ptr graph, const MainFlags& flags, - const FunctionLibraryDefinition* flib, CompileResult* compile_result) { // Converts the graph into an XLA computation, and compiles the // computation. @@ -396,11 +417,11 @@ Status CompileGraph(std::unique_ptr graph, const MainFlags& flags, namespace gpu = perftools::gputools; gpu::Platform* cpu_platform = gpu::MultiPlatformManager::PlatformWithName("Host").ValueOrDie(); - xla::LocalClient* client = - xla::ClientLibrary::GetOrCreateLocalClient(cpu_platform).ValueOrDie(); + xla::CompileOnlyClient* client = + xla::ClientLibrary::GetOrCreateCompileOnlyClient(cpu_platform) + .ValueOrDie(); xla::Computation computation; - TF_RETURN_IF_ERROR(ConvertGraphToXla(client, std::move(graph), flib, - &computation, + TF_RETURN_IF_ERROR(ConvertGraphToXla(client, std::move(graph), &computation, &compile_result->has_context_arg)); if (!flags.debug_dir.empty()) { TF_ASSIGN_OR_RETURN(std::unique_ptr module, diff --git a/tensorflow/compiler/aot/compile.h b/tensorflow/compiler/aot/compile.h index 8e9c64820baf0cb672cead59954098f10a9c9a32..e929272b2e4760e39cddba7e585cb12a7d2d7e98 100644 --- a/tensorflow/compiler/aot/compile.h +++ b/tensorflow/compiler/aot/compile.h @@ -56,8 +56,7 @@ extern const char* const kDebugNameAttr; // compute the outputs. If dump_graphs is true, graph rewrites will be dumped // for debugging. Status InitGraph(const GraphDef& graph_def, const Config& config, - const MainFlags& flags, const FunctionLibraryDefinition* flib, - std::unique_ptr* graph); + const MainFlags& flags, std::unique_ptr* graph); // CompileResult describes the output of CompileGraph, where the object file // data and meta-information is available in aot. @@ -83,7 +82,6 @@ struct CompileResult { // // The XLA compilation options are specified in the flags. Status CompileGraph(std::unique_ptr graph, const MainFlags& flags, - const FunctionLibraryDefinition* flib, CompileResult* result); } // namespace tfcompile diff --git a/tensorflow/compiler/aot/runtime.cc b/tensorflow/compiler/aot/runtime.cc index 208de5498dbee6773683ac1aa2b33400a8a21f35..5772776666129ed55a479c8917e69df3f3ce2fc0 100644 --- a/tensorflow/compiler/aot/runtime.cc +++ b/tensorflow/compiler/aot/runtime.cc @@ -31,6 +31,8 @@ namespace { inline void* aligned_malloc(size_t size, int minimum_alignment) { #if defined(__ANDROID__) || defined(OS_ANDROID) || defined(OS_CYGWIN) return memalign(minimum_alignment, size); +#elif defined(COMPILER_MSVC) + return _aligned_malloc(size, minimum_alignment); #else // !__ANDROID__ && !OS_ANDROID && !OS_CYGWIN void* ptr = nullptr; // posix_memalign requires that the requested alignment be at least @@ -45,7 +47,13 @@ inline void* aligned_malloc(size_t size, int minimum_alignment) { #endif } -inline void aligned_free(void* aligned_memory) { free(aligned_memory); } +inline void aligned_free(void* aligned_memory) { +#if defined(COMPILER_MSVC) + _aligned_free(aligned_memory); +#else + free(aligned_memory); +#endif +} size_t align_to(size_t n, size_t align) { return (((n - 1) / align) + 1) * align; diff --git a/tensorflow/compiler/aot/test_graph_tfadd.config.pbtxt b/tensorflow/compiler/aot/test_graph_tfadd.config.pbtxt index 5625c0ab03893c997245a6449d145b9149b48627..f2d9c34b2d1d68aa80245a6f3379b3759bb9f4b9 100644 --- a/tensorflow/compiler/aot/test_graph_tfadd.config.pbtxt +++ b/tensorflow/compiler/aot/test_graph_tfadd.config.pbtxt @@ -6,7 +6,7 @@ feed { } } feed { - id { node_name: "y_const" } + id { node_name: "y_reshape" } shape { dim { size: 1 } } diff --git a/tensorflow/compiler/aot/test_graph_tfadd.pbtxt b/tensorflow/compiler/aot/test_graph_tfadd.pbtxt index 91c900e06d7547fe9a377a427b6ca56b9e46942d..665c9fe28721b25c544c30ecd1b4dfc399934314 100644 --- a/tensorflow/compiler/aot/test_graph_tfadd.pbtxt +++ b/tensorflow/compiler/aot/test_graph_tfadd.pbtxt @@ -4,15 +4,7 @@ node { attr { key: "value" value { - tensor { - dtype: DT_INT32 - tensor_shape { - dim { - size: 1 - } - } - int_val: 1 - } + tensor { dtype: DT_INT32 tensor_shape { dim { size: 1 } } int_val: 1 } } } attr { @@ -28,15 +20,7 @@ node { attr { key: "value" value { - tensor { - dtype: DT_INT32 - tensor_shape { - dim { - size: 1 - } - } - int_val: 2 - } + tensor { dtype: DT_INT32 tensor_shape { dim { size: 1 } } int_val: 2 } } } attr { @@ -46,11 +30,20 @@ node { } } } +node { + name : "y_reshape" + op : "Reshape" + input : "y_const" + input : "y_shape" + attr { key: "T" value { type: DT_INT32 } } + # Attribute TShape not specified; needs to be set to its default + # by tfcompile. +} node { name : "x_y_sum" op : "Add" input : "x_const" - input : "y_const" + input : "y_reshape" attr { key : "T" value { diff --git a/tensorflow/compiler/aot/test_graph_tfunknownop.config.pbtxt b/tensorflow/compiler/aot/test_graph_tfunknownop.config.pbtxt new file mode 100644 index 0000000000000000000000000000000000000000..5625c0ab03893c997245a6449d145b9149b48627 --- /dev/null +++ b/tensorflow/compiler/aot/test_graph_tfunknownop.config.pbtxt @@ -0,0 +1,16 @@ +# Text form of tensorflow.tfcompile.Config proto. +feed { + id { node_name: "x_const" } + shape { + dim { size: 1 } + } +} +feed { + id { node_name: "y_const" } + shape { + dim { size: 1 } + } +} +fetch { + id { node_name: "x_y_sum" } +} diff --git a/tensorflow/compiler/aot/test_graph_tfunknownop.pbtxt b/tensorflow/compiler/aot/test_graph_tfunknownop.pbtxt new file mode 100644 index 0000000000000000000000000000000000000000..48b881bb9462dc30944a1377d4d2a2c58b9dfe43 --- /dev/null +++ b/tensorflow/compiler/aot/test_graph_tfunknownop.pbtxt @@ -0,0 +1,58 @@ +node { + name : "x_const" + op : "Const" + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { dim { size: 1 } } + int_val: 1 + } + } + } + attr { key : "dtype" value { type: DT_INT32 } } +} +node { + name : "y_const" + op : "Const" + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { dim { size: 1 } } + int_val: 2 + } + } + } + attr { key: "dtype" value { type: DT_INT32 } } +} +node { + name : "x_y_sum" + op : "Add" + input : "x_const" + input : "y_const" + attr { key : "T" value { type: DT_INT32 } } +} +node { + name : "z" + op : "SomeUnknownOp" + input : "x_const" +} +node { + name : "z_identity" + op : "Identity" + input : "z:1" + attr { key : "T" value { type: DT_INT32 } } +} +node { + name : "x_z_sum" + op : "Add" + input : "x_const" + input : "z_identity" + attr { key : "T" value { type: DT_INT32 } } +} +versions { + producer: 15 +} diff --git a/tensorflow/compiler/aot/test_graph_tfunknownop2.config.pbtxt b/tensorflow/compiler/aot/test_graph_tfunknownop2.config.pbtxt new file mode 100644 index 0000000000000000000000000000000000000000..7370ed370d314052ed23d4ceca22cab7def65485 --- /dev/null +++ b/tensorflow/compiler/aot/test_graph_tfunknownop2.config.pbtxt @@ -0,0 +1,25 @@ +# Text form of tensorflow.tfcompile.Config proto. +feed { + id { node_name: "x_const" } + shape { + dim { size: 1 } + } +} +feed { + id { node_name: "y_const" } + shape { + dim { size: 1 } + } +} +feed { + id { node_name: "z_identity"} + shape { + dim { size: 1 } + } +} +fetch { + id { node_name: "x_y_sum" } +} +fetch { + id { node_name: "x_z_sum" } +} diff --git a/tensorflow/compiler/aot/test_graph_tfunknownop3.config.pbtxt b/tensorflow/compiler/aot/test_graph_tfunknownop3.config.pbtxt new file mode 100644 index 0000000000000000000000000000000000000000..b2d7d5457427775fe2f00e079ced6b23c3308230 --- /dev/null +++ b/tensorflow/compiler/aot/test_graph_tfunknownop3.config.pbtxt @@ -0,0 +1,26 @@ +# Text form of tensorflow.tfcompile.Config proto. +feed { + id { node_name: "x_const" } + shape { + dim { size: 1 } + } +} +feed { + id { node_name: "y_const" } + shape { + dim { size: 1 } + } +} +feed { + id { node_name: "z" output_index: 1} + shape { + dim { size: 1 } + } + type: DT_INT32 +} +fetch { + id { node_name: "x_y_sum" } +} +fetch { + id { node_name: "x_z_sum" } +} diff --git a/tensorflow/compiler/aot/tests/BUILD b/tensorflow/compiler/aot/tests/BUILD index 59d13e5393445330ba5f1c5a54b73de6b3b4c0d8..05d338e4c53e139e0d1ad0d7f838d2dcecd3a335 100644 --- a/tensorflow/compiler/aot/tests/BUILD +++ b/tensorflow/compiler/aot/tests/BUILD @@ -33,6 +33,7 @@ py_binary( "//tensorflow/python:framework_for_generated_wrappers", "//tensorflow/python:math_ops", "//tensorflow/python:platform", + "//tensorflow/python:session", "//tensorflow/python:training", "//tensorflow/python:variables", ], @@ -43,15 +44,16 @@ genrule( testonly = 1, outs = [ "test_graph_tfadd.pb", - "test_graph_tfadd_with_ckpt.pb", "test_graph_tfadd_with_ckpt.ckpt", - "test_graph_tfadd_with_ckpt_saver.pb", + "test_graph_tfadd_with_ckpt.pb", "test_graph_tfadd_with_ckpt_saver.ckpt", + "test_graph_tfadd_with_ckpt_saver.pb", "test_graph_tfadd_with_ckpt_saver.saver", + "test_graph_tffunction.pb", "test_graph_tfgather.pb", "test_graph_tfmatmul.pb", "test_graph_tfmatmulandadd.pb", - "test_graph_tffunction.pb", + "test_graph_tfsplits.pb", ], cmd = "$(location :make_test_graphs) --out_dir $(@D)", tags = ["manual"], @@ -124,6 +126,15 @@ tf_library( tags = ["manual"], ) +tf_library( + name = "test_graph_tfsplits", + testonly = 1, + config = "test_graph_tfsplits.config.pbtxt", + cpp_class = "SplitsComp", + graph = "test_graph_tfsplits.pb", + tags = ["manual"], +) + cc_test( name = "tfcompile_test", srcs = ["tfcompile_test.cc"], @@ -136,6 +147,7 @@ cc_test( ":test_graph_tfgather", ":test_graph_tfmatmul", ":test_graph_tfmatmulandadd", + ":test_graph_tfsplits", "//tensorflow/core:test", "//tensorflow/core:test_main", "//third_party/eigen3", diff --git a/tensorflow/compiler/aot/tests/make_test_graphs.py b/tensorflow/compiler/aot/tests/make_test_graphs.py index 98c13958d3729bc6c7f554630e236892be130a4a..a898eab1d1ab0eb5d55983bf366753c968887296 100644 --- a/tensorflow/compiler/aot/tests/make_test_graphs.py +++ b/tensorflow/compiler/aot/tests/make_test_graphs.py @@ -107,6 +107,23 @@ def tffunction(_): test_func(x, y, name='func_call') # pylint: disable=unexpected-keyword-arg +def tfsplits(_): + """A more complex graph, including splits.""" + x = array_ops.placeholder(dtypes.float32, shape=[2, 2], name='x') + y = array_ops.placeholder(dtypes.float32, shape=[2, 2], name='y') + for _ in range(3): + x0, x1 = array_ops.split(x, 2, 0) + y0, y1 = array_ops.split(y, 2, 0) + x0 += 1 + y0 += 1 + z = math_ops.matmul(x, y, name='x_y_prod') + a = array_ops.concat([x0, y1], axis=0, name='concat_x0_y1') + b = array_ops.concat([y0, x1], axis=0, name='concat_y0_x1') + x = math_ops.matmul(a, b, name='a_b') + y = math_ops.add(x, z) + array_ops.identity(y, name='result') + + def write_graph(build_graph, out_dir): """Build a graph using build_graph and write it out.""" g = ops.Graph() @@ -125,6 +142,7 @@ def main(_): write_graph(tfmatmul, FLAGS.out_dir) write_graph(tfmatmulandadd, FLAGS.out_dir) write_graph(tffunction, FLAGS.out_dir) + write_graph(tfsplits, FLAGS.out_dir) if __name__ == '__main__': diff --git a/tensorflow/compiler/aot/tests/test_graph_tfsplits.config.pbtxt b/tensorflow/compiler/aot/tests/test_graph_tfsplits.config.pbtxt new file mode 100644 index 0000000000000000000000000000000000000000..85fc7da44285df043e841aa3c9af48b47bb68903 --- /dev/null +++ b/tensorflow/compiler/aot/tests/test_graph_tfsplits.config.pbtxt @@ -0,0 +1,18 @@ +# Text form of tensorflow.tfcompile.Config proto. +feed { + id { node_name: "x" } + shape { + dim { size: 2 } + dim { size: 2 } + } +} +feed { + id { node_name: "y" } + shape { + dim { size: 2 } + dim { size: 2 } + } +} +fetch { + id { node_name: "result" } +} diff --git a/tensorflow/compiler/aot/tests/tfcompile_test.cc b/tensorflow/compiler/aot/tests/tfcompile_test.cc index 76343b9752199fc4d26e4988452cd3c055bb5d96..07562e59c8dac942f41af69c289c9f29a9767a6a 100644 --- a/tensorflow/compiler/aot/tests/tfcompile_test.cc +++ b/tensorflow/compiler/aot/tests/tfcompile_test.cc @@ -24,6 +24,7 @@ limitations under the License. #include "tensorflow/compiler/aot/tests/test_graph_tfgather.h" #include "tensorflow/compiler/aot/tests/test_graph_tfmatmul.h" #include "tensorflow/compiler/aot/tests/test_graph_tfmatmulandadd.h" +#include "tensorflow/compiler/aot/tests/test_graph_tfsplits.h" #include "tensorflow/core/platform/test.h" namespace tensorflow { @@ -392,6 +393,34 @@ TEST(TFCompileTest, Function) { EXPECT_EQ(add_fn.result0_data(), add_fn.results()[0]); } +TEST(TFCompileTest, Splits) { + Eigen::ThreadPool tp(1); + Eigen::ThreadPoolDevice device(&tp, tp.NumThreads()); + + SplitsComp fn; + + fn.set_thread_pool(&device); + // x = [[1, 2], [3, 4]] + fn.arg0(0, 0) = 1; + fn.arg0(0, 1) = 2; + fn.arg0(1, 0) = 3; + fn.arg0(1, 1) = 4; + + // y = [[10, 20], [30, 40]] + fn.arg1(0, 0) = 10; + fn.arg1(0, 1) = 20; + fn.arg1(1, 0) = 30; + fn.arg1(1, 1) = 40; + EXPECT_TRUE(fn.Run()); + EXPECT_EQ(fn.error_msg(), ""); + const float expected[] = {7.86375557e+10, 1.34274679e+11, 1.92741717e+12, + 3.29964742e+12}; + EXPECT_NEAR(expected[0], fn.result0(0, 0), 1e4); + EXPECT_NEAR(expected[1], fn.result0(0, 1), 1e4); + EXPECT_NEAR(expected[2], fn.result0(1, 0), 1e4); + EXPECT_NEAR(expected[3], fn.result0(1, 1), 1e4); +} + } // namespace } // namespace tfcompile } // namespace tensorflow diff --git a/tensorflow/compiler/aot/tfcompile.bzl b/tensorflow/compiler/aot/tfcompile.bzl index 64e5bfd602cb2898dcbe57bfa0949c954f17acc1..f9896988dc2db5a6d7ca343e1ad3ae7d2aaad5a0 100644 --- a/tensorflow/compiler/aot/tfcompile.bzl +++ b/tensorflow/compiler/aot/tfcompile.bzl @@ -27,6 +27,15 @@ def tf_library(name, graph, config, deps=None, tags=None): """Runs tfcompile to compile a TensorFlow graph into executable code. + Given an invocation of tf_library(name="foo", ...), generates the following + build targets: + foo: A cc_library containing the generated header and computation. + foo_test: A cc_test with simple tests and benchmarks. Only created if + gen_test=True. + foo_benchmark: A cc_binary that runs a minimal-dependency benchmark, useful + for mobile devices or other platforms that can't compile the + full test libraries. Only created if gen_benchmark=True. + Args: name: The name of the build rule. graph: The TensorFlow GraphDef to compile. If the file ends in '.pbtxt' it @@ -169,6 +178,9 @@ def tf_library(name, graph, config, "//tensorflow/compiler/tf2xla/kernels:index_ops_kernel_argmax_float_2d", "//tensorflow/compiler/aot:runtime", "//tensorflow/compiler/tf2xla:xla_local_runtime_context", + "//tensorflow/compiler/xla/service/cpu:cpu_runtime_avx", + "//tensorflow/compiler/xla/service/cpu:cpu_runtime_neon", + "//tensorflow/compiler/xla/service/cpu:cpu_runtime_sse4_1", "//tensorflow/compiler/xla/service/cpu:runtime_conv2d", "//tensorflow/compiler/xla/service/cpu:runtime_matmul", "//tensorflow/compiler/xla/service/cpu:runtime_single_threaded_conv2d", @@ -279,8 +291,11 @@ def target_llvm_triple(): # TODO(toddw): Add target_triple for other targets. For details see: # http://llvm.org/docs/doxygen/html/Triple_8h_source.html return select({ + "//tensorflow:android_armeabi": "armv5-none-android", "//tensorflow:android_arm": "armv7-none-android", "//tensorflow:android_arm64": "aarch64-none-android", "//tensorflow:android_x86": "i686-none-android", + "//tensorflow:linux_ppc64le": "ppc64le-ibm-linux-gnu", + "//tensorflow:darwin": "x86_64-none-darwin", "//conditions:default": "x86_64-pc-linux", }) diff --git a/tensorflow/compiler/aot/tfcompile.proto b/tensorflow/compiler/aot/tfcompile.proto index be3f5043501c71c844a00b5a5b23fa4285c00ec6..cd83840d894f2a28ca70c54f3320a6287b4a0a20 100644 --- a/tensorflow/compiler/aot/tfcompile.proto +++ b/tensorflow/compiler/aot/tfcompile.proto @@ -7,6 +7,7 @@ option java_multiple_files = true; option java_package = "org.tensorflow.tfcompile"; import "tensorflow/core/framework/tensor_shape.proto"; +import "tensorflow/core/framework/types.proto"; // TensorId identifies a tensor in a TensorFlow graph, by specifying the output // index of a particular node in the graph. If the output of the named node @@ -23,6 +24,12 @@ message Feed { TensorId id = 1; TensorShapeProto shape = 2; string name = 3; // Optional name for generated code. + + // Optional data type. This is not normally required, as the graph itself + // contains this information. However, if the node being fed is an op that + // is not linked into the tfcompile binary, then the type cannot be inferred + // from the node; in this case, the type should be set here. + DataType type = 4; }; // Fetch represents a single fetch tensor in the graph, which corresponds to an diff --git a/tensorflow/compiler/aot/tfcompile_main.cc b/tensorflow/compiler/aot/tfcompile_main.cc index 85ef9560bbf1a7130dd6b140d552d96c2a0e21d6..be2cfe4734e0493ba41a1bda23606a65d2cb4af4 100644 --- a/tensorflow/compiler/aot/tfcompile_main.cc +++ b/tensorflow/compiler/aot/tfcompile_main.cc @@ -23,9 +23,7 @@ limitations under the License. #include "tensorflow/compiler/aot/flags.h" #include "tensorflow/compiler/aot/tfcompile.pb.h" #include "tensorflow/compiler/aot/tfcompile_util.h" -#include "tensorflow/compiler/xla/legacy_flags/compiler_functor_flags.h" -#include "tensorflow/compiler/xla/legacy_flags/cpu_compiler_flags.h" -#include "tensorflow/compiler/xla/legacy_flags/cpu_runtime_flags.h" +#include "tensorflow/compiler/xla/legacy_flags/debug_options_flags.h" #include "tensorflow/compiler/xla/service/compiler.h" #include "tensorflow/core/framework/function.h" #include "tensorflow/core/framework/graph.pb.h" @@ -52,7 +50,8 @@ const char kUsageHeader[] = "header file that gives access to the functionality in the object file.\n" "A typical invocation looks like this:\n" "\n" - " $ tfcompile --graph=mygraph.pb --config=myfile.pbtxt\n" + " $ tfcompile --graph=mygraph.pb --config=myfile.pbtxt " + "--cpp_class=\"mynamespace::MyComputation\"\n" "\n"; Status ReadProtoFile(const string& kind, const string& fname, @@ -73,6 +72,9 @@ void ParseTensorId(const string& name, TensorId* id) { Status Main(const MainFlags& flags) { // Process config. Config config; + if (flags.config.empty()) { + return errors::InvalidArgument("Must specify --config"); + } TF_RETURN_IF_ERROR(ReadProtoFile("config", flags.config, &config)); TF_RETURN_IF_ERROR(ValidateConfig(config)); if (flags.dump_fetch_nodes) { @@ -85,15 +87,16 @@ Status Main(const MainFlags& flags) { } // Read and initialize the graph. + if (flags.graph.empty()) { + return errors::InvalidArgument("Must specify --graph"); + } GraphDef graph_def; TF_RETURN_IF_ERROR(ReadProtoFile("graph", flags.graph, &graph_def)); std::unique_ptr graph; - FunctionLibraryDefinition flib(OpRegistry::Global(), graph_def.library()); - TF_RETURN_IF_ERROR(InitGraph(graph_def, config, flags, &flib, &graph)); + TF_RETURN_IF_ERROR(InitGraph(graph_def, config, flags, &graph)); CompileResult compile_result; - TF_RETURN_IF_ERROR( - CompileGraph(std::move(graph), flags, &flib, &compile_result)); + TF_RETURN_IF_ERROR(CompileGraph(std::move(graph), flags, &compile_result)); // Write output files. Env* env = Env::Default(); @@ -101,6 +104,9 @@ Status Main(const MainFlags& flags) { TF_RETURN_IF_ERROR(WriteStringToFile(env, flags.out_object, StringPiece(obj.data(), obj.size()))); HeaderOpts header_opts; + if (flags.cpp_class.empty()) { + return errors::InvalidArgument("Must specify --cpp_class"); + } TF_RETURN_IF_ERROR(ParseCppClass(flags.cpp_class, &header_opts.class_name, &header_opts.namespaces)); string header; @@ -118,12 +124,11 @@ int main(int argc, char** argv) { flags.target_triple = "x86_64-pc-linux"; flags.out_object = "out.o"; flags.out_header = "out.h"; + flags.entry_point = "entry"; std::vector flag_list; AppendMainFlags(&flag_list, &flags); - xla::legacy_flags::AppendCompilerFunctorFlags(&flag_list); - xla::legacy_flags::AppendCpuCompilerFlags(&flag_list); - xla::legacy_flags::AppendCpuRuntimeFlags(&flag_list); + xla::legacy_flags::AppendDebugOptionsFlags(&flag_list); tensorflow::string usage = tensorflow::tfcompile::kUsageHeader; usage += tensorflow::Flags::Usage(argv[0], flag_list); @@ -131,12 +136,16 @@ int main(int argc, char** argv) { QCHECK(parsed_flags_ok) << "\n" << usage; tensorflow::port::InitMain(usage.c_str(), &argc, &argv); - QCHECK(argc == 1 && !flags.config.empty() && - (flags.dump_fetch_nodes || - (!flags.graph.empty() && !flags.entry_point.empty()))) - << "\n" - << usage; - - TF_QCHECK_OK(tensorflow::tfcompile::Main(flags)); + QCHECK(argc == 1) << "\nERROR: This command does not take any arguments " + "other than flags\n\n" + << usage; + tensorflow::Status status = tensorflow::tfcompile::Main(flags); + if (status.code() == tensorflow::error::INVALID_ARGUMENT) { + std::cerr << "INVALID ARGUMENTS: " << status.error_message() << "\n\n" + << usage; + return 1; + } else { + TF_QCHECK_OK(status); + } return 0; } diff --git a/tensorflow/compiler/aot/tfcompile_util.cc b/tensorflow/compiler/aot/tfcompile_util.cc index fd073a2e2623b4b24ddc58360525886f3fc1b3ac..629187d621c2b98e764b66d59c8a78182726ee19 100644 --- a/tensorflow/compiler/aot/tfcompile_util.cc +++ b/tensorflow/compiler/aot/tfcompile_util.cc @@ -15,13 +15,22 @@ limitations under the License. #include "tensorflow/compiler/aot/tfcompile_util.h" +#include #include +#include #include "tensorflow/compiler/aot/tfcompile.pb.h" +#include "tensorflow/core/framework/graph.pb.h" +#include "tensorflow/core/framework/graph_def_util.h" +#include "tensorflow/core/framework/node_def.pb.h" #include "tensorflow/core/framework/tensor_shape.h" +#include "tensorflow/core/framework/tensor_shape.pb.h" +#include "tensorflow/core/framework/versions.pb.h" +#include "tensorflow/core/graph/tensor_id.h" #include "tensorflow/core/lib/core/errors.h" #include "tensorflow/core/lib/core/status.h" #include "tensorflow/core/lib/core/stringpiece.h" +#include "tensorflow/core/lib/strings/strcat.h" namespace tensorflow { namespace tfcompile { @@ -115,5 +124,164 @@ Status ValidateConfig(const Config& config) { return Status::OK(); } +Status AddPlaceholdersForFeeds( + const Config& config, const OpRegistryInterface* op_registry, + std::unordered_map* feed_remapping, GraphDef* graph_def) { + struct PlaceholderInfo { + const Feed* feed = nullptr; // point to Feed in . + string placeholder_name; + DataType data_type = DT_INVALID; + }; + + // Put each fed tensor into a map by name:port. A map is used for determinism + // when creating placeholders (genrules want deterministic output). + std::map placeholder_info; + for (int i = 0; i < config.feed_size(); ++i) { + const Feed* feed = &config.feed(i); + const string name_port = TensorIdToString(feed->id()); + auto& info = placeholder_info[name_port]; + info.feed = feed; + info.placeholder_name = strings::StrCat( + "aot_feed_", feed->id().output_index(), "/", feed->id().node_name()); + (*feed_remapping)[name_port] = info.placeholder_name; + } + + // Verify node exists and determine data type. + std::unordered_map name_to_node; + for (int i = 0; i < graph_def->node_size(); ++i) { + name_to_node[graph_def->node(i).name()] = &graph_def->node(i); + } + for (auto it = placeholder_info.begin(); it != placeholder_info.end(); ++it) { + PlaceholderInfo& info = it->second; + const TensorId& feed_id = info.feed->id(); + + // Find the existing node and determine data type. + auto node_it = name_to_node.find(feed_id.node_name()); + if (node_it == name_to_node.end()) { + return errors::NotFound("Can't find feed node: ", + TensorIdToString(feed_id)); + } + const NodeDef* existing = node_it->second; + + if (info.feed->type() != DT_INVALID) { + info.data_type = info.feed->type(); + } else { + // Build the node in order to infer its type. + + // Must first add default attrs as well, so do this in a copied GraphDef. + GraphDef gd; + *gd.mutable_versions() = graph_def->versions(); + *gd.add_node() = *existing; + TF_RETURN_IF_ERROR( + AddDefaultAttrsToGraphDef(&gd, *op_registry, 0 /*node_offset*/)); + + // Now build the node from the copied node def. + Graph g(op_registry); + g.set_versions(graph_def->versions()); + Status status; + Node* feed_node = g.AddNode(gd.node(0), &status); + TF_RETURN_IF_ERROR(status); + info.data_type = + BaseType(feed_node->output_type(info.feed->id().output_index())); + } + } + + // Create placeholders. Note that we could avoid creating a placeholder for + // feeds which are already placeholders, but we omit that to avoid more cases + // in this code. + for (auto it = placeholder_info.begin(); it != placeholder_info.end(); ++it) { + const PlaceholderInfo& info = it->second; + NodeDef* d = graph_def->add_node(); + d->set_name(info.placeholder_name); + d->set_op("PlaceholderV2"); + auto& attr_map = *d->mutable_attr(); + attr_map["dtype"].set_type(info.data_type); + *attr_map["shape"].mutable_shape() = info.feed->shape(); + } + + // Rewrite references to the fed tensors to refer to the placeholder. + for (int i = 0; i < graph_def->node_size(); ++i) { + NodeDef* node_def = graph_def->mutable_node(i); + for (int j = 0; j < node_def->input_size(); ++j) { + auto id = ParseTensorName(node_def->input(j)); + auto it = placeholder_info.find(id.ToString()); + if (it != placeholder_info.end()) { + node_def->set_input(j, it->second.placeholder_name); + } + } + } + + return Status::OK(); +} + +Status PruneGraphDefInto(const Config& config, const GraphDef& in, + GraphDef* out) { + *out = in; + out->clear_node(); + + // Tensors needed for feeding. + std::set> feed_tensors; + for (const auto& feed_config : config.feed()) { + feed_tensors.insert(std::make_pair(feed_config.id().node_name(), + feed_config.id().output_index())); + } + + // Maps node name to reachability. + std::unordered_map> node_by_name; + for (const NodeDef& node : in.node()) { + node_by_name[node.name()] = std::pair(false, &node); + } + + // Traverse. + std::queue name_queue; + for (int i = 0; i < config.fetch_size(); ++i) { + name_queue.push(config.fetch(i).id().node_name()); + } + while (!name_queue.empty()) { + const string name = name_queue.front(); + name_queue.pop(); + + auto find_it = node_by_name.find(name); + if (find_it == node_by_name.end()) { + return errors::InvalidArgument("While pruning graph, node ", name, + " needed but not found in the graph."); + } + auto& map_entry = find_it->second; + if (map_entry.first) { + continue; + } + map_entry.first = true; + + // Push input nodes of the currently visited node to name_queue. + for (const string& in_edge : map_entry.second->input()) { + auto id = ParseTensorName(in_edge); + const string node_name = id.first.ToString(); + if (feed_tensors.find(std::make_pair(node_name, id.second)) == + feed_tensors.end()) { + name_queue.push(node_name); + } else { + // The input tensor is from an edge that is being fed. Therefore, + // we skip recursing down that edge, to avoid requiring nodes that + // may not be needed (note that the input node may still be added + // to name_queue later if one of its output edges is not being fed). + } + } + } + + // Copy over, preserving order of original and only nodes that are reachable + // from the fetches. + out->mutable_node()->Reserve(in.node_size()); + for (const NodeDef& node : in.node()) { + if (node_by_name[node.name()].first) { + *out->add_node() = node; + } + } + return Status::OK(); +} + +string TensorIdToString(const TensorId& id) { + return strings::StrCat(id.node_name(), ":", id.output_index()); +} + } // namespace tfcompile } // namespace tensorflow diff --git a/tensorflow/compiler/aot/tfcompile_util.h b/tensorflow/compiler/aot/tfcompile_util.h index 651d75d0d02bdac110159996498778d2c57ddf78..365f7b0e7b19a495ade13a7cff4140cdae68cad2 100644 --- a/tensorflow/compiler/aot/tfcompile_util.h +++ b/tensorflow/compiler/aot/tfcompile_util.h @@ -16,7 +16,11 @@ limitations under the License. #ifndef TENSORFLOW_COMPILER_AOT_TFCOMPILE_UTIL_H_ #define TENSORFLOW_COMPILER_AOT_TFCOMPILE_UTIL_H_ +#include + #include "tensorflow/compiler/aot/tfcompile.pb.h" +#include "tensorflow/core/framework/graph.pb.h" +#include "tensorflow/core/framework/op.h" #include "tensorflow/core/lib/core/status.h" #include "tensorflow/core/lib/core/stringpiece.h" @@ -30,6 +34,23 @@ Status ValidateCppIdent(StringPiece ident, StringPiece msg); // ValidateConfig returns OK iff config is valid. Status ValidateConfig(const Config& config); +// Modifies to include placeholders for each fed tensor, and +// update references to the fed tensors to refer to the placeholders. +// The existing nodes referenced by the feeds are not removed or modified +// (except where their input edges are modified by the replacement of other +// feeds). +Status AddPlaceholdersForFeeds( + const Config& config, const OpRegistryInterface* op_registry, + std::unordered_map* feed_remapping, GraphDef* graph_def); + +// Returns in a copy of , pruned to only include fetches from +// . +Status PruneGraphDefInto(const Config& config, const GraphDef& in, + GraphDef* out); + +// Returns node:port for the given . +string TensorIdToString(const TensorId& id); + } // namespace tfcompile } // namespace tensorflow diff --git a/tensorflow/compiler/aot/tfcompile_util_test.cc b/tensorflow/compiler/aot/tfcompile_util_test.cc index 108ab1eab7bf3b087e8049c5b24d652d871789c8..5a92851ceb972ca63a8a3845eb4730fe198762dd 100644 --- a/tensorflow/compiler/aot/tfcompile_util_test.cc +++ b/tensorflow/compiler/aot/tfcompile_util_test.cc @@ -15,16 +15,18 @@ limitations under the License. #include "tensorflow/compiler/aot/tfcompile_util.h" +#include "tensorflow/core/framework/node_def.pb.h" #include "tensorflow/core/lib/core/status.h" #include "tensorflow/core/lib/core/status_test_util.h" #include "tensorflow/core/lib/core/stringpiece.h" +#include "tensorflow/core/lib/strings/strcat.h" #include "tensorflow/core/platform/test.h" namespace tensorflow { namespace tfcompile { namespace { -void ExpectErrorContains(Status status, StringPiece str) { +void ExpectErrorContains(const Status& status, StringPiece str) { EXPECT_NE(Status::OK(), status); EXPECT_TRUE(StringPiece(status.error_message()).contains(str)) << "expected error: " << status.error_message() << " to contain: " << str; @@ -180,6 +182,65 @@ TEST(ValidateConfig, ConflictingFetchName) { ExpectErrorContains(ValidateConfig(config), "conflicting fetch name"); } +static Config FetchesConfig(std::vector fetches) { + Config config; + for (const auto& fetch_node_name : fetches) { + auto* fetch = config.add_fetch(); + fetch->set_name(strings::StrCat("fetch_", fetch_node_name)); + fetch->mutable_id()->set_node_name(fetch_node_name); + } + return config; +} + +TEST(PruneGraphDefInto, Basic) { + GraphDef def; + auto* n = def.add_node(); + n->set_name("a"); + n->add_input("b:0"); + n->add_input("^c"); + + GraphDef copy; + ExpectErrorContains(PruneGraphDefInto(FetchesConfig({"missing"}), def, ©), + "node missing needed"); + ExpectErrorContains(PruneGraphDefInto(FetchesConfig({"a"}), def, ©), + "node b needed"); + + n = def.add_node(); + n->set_name("b"); + ExpectErrorContains(PruneGraphDefInto(FetchesConfig({"a"}), def, ©), + "node c needed"); + n->add_input("d:1"); + + n = def.add_node(); + n->set_name("c"); + n->add_input("d:1"); + + n = def.add_node(); + n->set_name("d"); + + // Graph is full, no pruning done. + // Graph right now has diamond from d: + // d --> b --> a + // d --> c --> a + TF_EXPECT_OK(PruneGraphDefInto(FetchesConfig({"a"}), def, ©)); + EXPECT_EQ(def.DebugString(), copy.DebugString()); + GraphDef pruned_a = copy; + + // Add some unrelated fields that use b and c, but are not needed for a. + n = def.add_node(); + n->set_name("e"); + n->add_input("^d"); + n->add_input("b:2"); + copy.Clear(); + TF_EXPECT_OK(PruneGraphDefInto(FetchesConfig({"a"}), def, ©)); + EXPECT_EQ(pruned_a.DebugString(), copy.DebugString()); + + // Fetch "a" and "e" to get the original graph. + copy.Clear(); + TF_EXPECT_OK(PruneGraphDefInto(FetchesConfig({"a", "e"}), def, ©)); + EXPECT_EQ(def.DebugString(), copy.DebugString()); +} + } // namespace } // namespace tfcompile } // namespace tensorflow diff --git a/tensorflow/compiler/jit/BUILD b/tensorflow/compiler/jit/BUILD index c16fe56122fca8cf8a88d6098b2374285f33e9f2..02e7ca64e5f08b153417a52c1aba734694b83f3b 100644 --- a/tensorflow/compiler/jit/BUILD +++ b/tensorflow/compiler/jit/BUILD @@ -15,10 +15,16 @@ package_group( ) package( - default_visibility = [":internal"], + default_visibility = [ + ":internal", + "//tensorflow/compiler/plugin/executor:__pkg__", + ], ) +load("//tensorflow:tensorflow.bzl", "cc_header_only_library") load("//tensorflow:tensorflow.bzl", "tf_kernel_library") +load("@local_config_cuda//cuda:build_defs.bzl", "if_cuda") +load("@local_config_cuda//cuda:build_defs.bzl", "if_cuda_is_configured") # Target that bundles up the XLA CPU and GPU JIT devices. cc_library( @@ -27,15 +33,17 @@ cc_library( deps = [ ":xla_cpu_device", ":xla_cpu_jit", + "//tensorflow/compiler/plugin", + ] + if_cuda_is_configured([ ":xla_gpu_device", ":xla_gpu_jit", - ], + ]), alwayslink = 1, ) cc_library( name = "xla_cpu_jit", - visibility = [":friends"], + visibility = ["//visibility:public"], deps = [ ":jit_compilation_passes", "//tensorflow/compiler/jit/kernels:xla_local_launch_op", @@ -47,13 +55,13 @@ cc_library( cc_library( name = "xla_gpu_jit", - visibility = [":friends"], - deps = [ + visibility = ["//visibility:public"], + deps = if_cuda([ ":jit_compilation_passes", "//tensorflow/compiler/jit/kernels:xla_local_launch_op", "//tensorflow/compiler/tf2xla/kernels:xla_ops", "//tensorflow/compiler/xla/service:gpu_plugin", - ], + ]), alwayslink = 1, ) @@ -67,7 +75,7 @@ cc_library( "//tensorflow/compiler/jit/kernels:xla_device_launch_op", "//tensorflow/compiler/tf2xla:xla_compiler", "//tensorflow/compiler/tf2xla/kernels:xla_ops", - "//tensorflow/compiler/xla/service:cpu_plugin", + "//tensorflow/compiler/xla/service:cpu_plugin", # buildcleaner: keep "//tensorflow/core:core_cpu_internal", "//tensorflow/core:lib", ], @@ -84,7 +92,7 @@ cc_library( "//tensorflow/compiler/jit/kernels:xla_device_launch_op", "//tensorflow/compiler/tf2xla:xla_compiler", "//tensorflow/compiler/tf2xla/kernels:xla_ops", - "//tensorflow/compiler/xla/service:gpu_plugin", + "//tensorflow/compiler/xla/service:gpu_plugin", # buildcleaner: keep "//tensorflow/core:core_cpu_internal", "//tensorflow/core:lib", ], @@ -125,7 +133,6 @@ cc_library( "//tensorflow/compiler/tf2xla:xla_compiler", "//tensorflow/compiler/tf2xla/kernels:xla_ops", "//tensorflow/compiler/xla:literal_util", - "//tensorflow/compiler/xla:statusor", "//tensorflow/compiler/xla/client:client_library", "//tensorflow/compiler/xla/client:global_data", "//tensorflow/compiler/xla/client:local_client", @@ -133,9 +140,10 @@ cc_library( "//tensorflow/core:core_cpu_internal", "//tensorflow/core:framework", "//tensorflow/core:lib", + "//tensorflow/core:lib_internal", + "//tensorflow/core:protos_all_cc", "//tensorflow/core:stream_executor_no_cuda", "//tensorflow/core:tensorflow_opensource", - "//tensorflow/core/kernels:assign_op", "//tensorflow/core/kernels:constant_op", "//tensorflow/core/kernels:control_flow_ops", "//tensorflow/core/kernels:identity_op", @@ -155,7 +163,6 @@ cc_library( "//tensorflow/compiler/tf2xla:xla_compiler", "//tensorflow/compiler/xla/client:client_library", "//tensorflow/compiler/xla/client:local_client", - "//tensorflow/compiler/xla/service:cpu_plugin", "//tensorflow/core:core_cpu", "//tensorflow/core:core_cpu_internal", "//tensorflow/core:framework", @@ -176,22 +183,34 @@ cc_library( alwayslink = 1, ) +cc_library( + name = "graph_to_functiondef", + srcs = ["graph_to_functiondef.cc"], + hdrs = ["graph_to_functiondef.h"], + deps = [ + "//tensorflow/core:core_cpu", + "//tensorflow/core:framework", + "//tensorflow/core:lib", + "//tensorflow/core:protos_all_cc", + ], +) + cc_library( name = "compilation_passes", srcs = [ "build_xla_launch_ops_pass.cc", "encapsulate_subgraphs_pass.cc", - "graph_to_functiondef.cc", "mark_for_compilation_pass.cc", ], hdrs = [ "build_xla_launch_ops_pass.h", "encapsulate_subgraphs_pass.h", - "graph_to_functiondef.h", "mark_for_compilation_pass.h", ], deps = [ ":common", + ":graph_to_functiondef", + ":union_find", "//tensorflow/compiler/jit/graphcycles", "//tensorflow/compiler/jit/kernels:parallel_check_op", "//tensorflow/compiler/jit/kernels:xla_local_launch_op", @@ -211,12 +230,38 @@ cc_library( ], ) +cc_library( + name = "union_find", + hdrs = ["union_find.h"], +) + +cc_test( + name = "graph_to_functiondef_test", + size = "small", + srcs = [ + "graph_to_functiondef_test.cc", + ], + deps = [ + ":graph_to_functiondef", + "//tensorflow/cc:cc_ops", + "//tensorflow/cc:cc_ops_internal", + "//tensorflow/cc:function_ops", + "//tensorflow/cc:ops", + "//tensorflow/compiler/tf2xla:xla_compiler", + "//tensorflow/compiler/tf2xla/kernels:xla_ops", + "//tensorflow/core:core_cpu", + "//tensorflow/core:framework_internal", + "//tensorflow/core:test", + "//tensorflow/core:test_main", + "//tensorflow/core:testlib", + ], +) + cc_test( name = "compilation_passes_test", size = "small", srcs = [ "encapsulate_subgraphs_pass_test.cc", - "graph_to_functiondef_test.cc", "mark_for_compilation_pass_test.cc", ], deps = [ @@ -250,3 +295,15 @@ filegroup( ), visibility = ["//tensorflow:__subpackages__"], ) + +# 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. +cc_header_only_library( + name = "xla_jit_headers_lib", + visibility = ["//visibility:public"], + deps = [ + ":xla_cpu_device", + ":xla_cpu_jit", + ":xla_gpu_device", + ":xla_gpu_jit", + ], +) diff --git a/tensorflow/compiler/jit/build_xla_launch_ops_pass.cc b/tensorflow/compiler/jit/build_xla_launch_ops_pass.cc index abb68f73d7e3870f733c350be0dc99ab21a6b083..ef56ccf8e8f4257b72120472d0b7a56946bd92e0 100644 --- a/tensorflow/compiler/jit/build_xla_launch_ops_pass.cc +++ b/tensorflow/compiler/jit/build_xla_launch_ops_pass.cc @@ -23,7 +23,6 @@ limitations under the License. #include "tensorflow/compiler/tf2xla/xla_op_registry.h" #include "tensorflow/core/common_runtime/function.h" #include "tensorflow/core/common_runtime/optimization_registry.h" -#include "tensorflow/core/framework/graph.pb.h" #include "tensorflow/core/framework/graph_def_util.h" #include "tensorflow/core/framework/node_def_builder.h" #include "tensorflow/core/framework/node_def_util.h" @@ -32,7 +31,6 @@ limitations under the License. #include "tensorflow/core/graph/graph_constructor.h" #include "tensorflow/core/lib/core/status.h" #include "tensorflow/core/lib/hash/hash.h" -#include "tensorflow/core/protobuf/config.pb.h" #include "tensorflow/core/public/version.h" namespace tensorflow { @@ -66,9 +64,9 @@ static Status ReplaceNodeWithXlaLaunch(Graph* graph, Node* node) { int num_constant_args, num_resource_args; TF_RETURN_IF_ERROR( - GetNodeAttr(node->def(), kXlaNumConstantArgsAttr, &num_constant_args)); + GetNodeAttr(node->attrs(), kXlaNumConstantArgsAttr, &num_constant_args)); TF_RETURN_IF_ERROR( - GetNodeAttr(node->def(), kXlaNumResourceArgsAttr, &num_resource_args)); + GetNodeAttr(node->attrs(), kXlaNumResourceArgsAttr, &num_resource_args)); if (num_constant_args < 0 || num_resource_args < 0 || num_constant_args + num_resource_args > node->num_inputs()) { @@ -88,7 +86,7 @@ static Status ReplaceNodeWithXlaLaunch(Graph* graph, Node* node) { Node* launch_node; TF_RETURN_IF_ERROR(BuildLaunchNode( graph->NewName(node->name()), node->type_string(), node->def().attr(), - node->def().device(), const_dtypes, num_resource_args, arg_dtypes, + node->requested_device(), const_dtypes, num_resource_args, arg_dtypes, node->output_types(), graph, &launch_node)); launch_node->set_assigned_device_name(node->assigned_device_name()); @@ -125,9 +123,9 @@ static Status ReplaceNodeWithXlaLaunch(Graph* graph, Node* node) { Status BuildXlaLaunchOpsPass::Run(const GraphOptimizationPassOptions& options) { Graph* graph = options.graph->get(); - for (Node* n : graph->nodes()) { + for (Node* n : graph->op_nodes()) { // In all cases, only try to compile computational nodes. - if (!n->IsOp() || n->IsSend() || n->IsRecv() || n->IsControlFlow()) { + if (n->IsSend() || n->IsRecv() || n->IsControlFlow()) { continue; } @@ -173,7 +171,8 @@ Status CreateXlaLaunchOp(FunctionLibraryRuntime* flr, const NodeDef& ndef, FunctionLibraryRuntime::Handle handle; // If ndef is not instantiable, e.g., the function does not exist, // simply bail out. - TF_RETURN_IF_ERROR(flr->Instantiate(ndef.op(), ndef.attr(), &handle)); + TF_RETURN_IF_ERROR( + flr->Instantiate(ndef.op(), AttrSlice(&ndef.attr()), &handle)); const FunctionBody* fbody = flr->GetFunctionBody(handle); CHECK(fbody); // Can't be nullptr since we just instantiated it. std::vector const_args(fbody->arg_types.size()); diff --git a/tensorflow/compiler/jit/encapsulate_subgraphs_pass.cc b/tensorflow/compiler/jit/encapsulate_subgraphs_pass.cc index 1d2793d3c55f4436a07e4f632887561202d0498e..22899ebeebc929055518893b358f7950d380d6f6 100644 --- a/tensorflow/compiler/jit/encapsulate_subgraphs_pass.cc +++ b/tensorflow/compiler/jit/encapsulate_subgraphs_pass.cc @@ -25,9 +25,7 @@ limitations under the License. #include "tensorflow/compiler/tf2xla/dump_graph.h" #include "tensorflow/core/common_runtime/function.h" #include "tensorflow/core/common_runtime/optimization_registry.h" -#include "tensorflow/core/framework/attr_value.pb.h" #include "tensorflow/core/framework/function.h" -#include "tensorflow/core/framework/graph.pb.h" #include "tensorflow/core/framework/graph_def_util.h" #include "tensorflow/core/framework/node_def_builder.h" #include "tensorflow/core/framework/node_def_util.h" @@ -88,9 +86,12 @@ class Encapsulator { // Build a FunctionDef for each subgraph, and add it 'library'. The values of // the 'group_attribute' annotations become the function names. + // If 'reuse_existing_functions' is set, use an existing function with the + // same name, if any. // If 'rewrite_subgraph_fn' is set, it is applied to each subgraph before // function conversion. Status BuildFunctionDefs(const RewriteSubgraphFn& rewrite_subgraph_fn, + bool reuse_existing_functions, FunctionLibraryDefinition* library); // Write a copy of the input graph to 'graph_out', where the subgraphs are @@ -110,8 +111,8 @@ class Encapsulator { // returned by _Retval nodes. std::unique_ptr graph; - // Which device are these nodes on? Used both to check that all nodes - // are assigned to the same device, and to assign a device to the call node. + // Which device are these nodes on? Used to assign a device to the call + // node. string device; // NodeDef for the function call node. @@ -162,7 +163,7 @@ static const char* const kRetValOp = "_Retval"; // none. string Encapsulator::GetFunctionNameAttr(Node const* node) const { string attr; - if (!GetNodeAttr(node->def(), group_attribute_, &attr).ok()) { + if (!GetNodeAttr(node->attrs(), group_attribute_, &attr).ok()) { attr.clear(); } return attr; @@ -175,8 +176,7 @@ Status Encapsulator::SplitIntoSubgraphs() { std::unordered_map node_images; // Copy all marked nodes to a subgraph. Do nothing for unmarked nodes. - for (Node* node : graph_in_->nodes()) { - if (node->IsSource() || node->IsSink()) continue; + for (Node* node : graph_in_->op_nodes()) { string func_id = GetFunctionNameAttr(node); if (func_id.empty()) continue; @@ -190,16 +190,10 @@ Status Encapsulator::SplitIntoSubgraphs() { image->ClearAttr(group_attribute_); node_images[node] = image; - // Check the device matches any existing device. - string device = node->assigned_device_name().empty() - ? node->def().device() - : node->assigned_device_name(); - if (subgraph.device.empty()) { - subgraph.device = device; - } else if (subgraph.device != device) { - s.Update(errors::InvalidArgument( - "Mismatched devices for nodes to be grouped by Encapsulator")); + subgraph.device = node->assigned_device_name().empty() + ? node->requested_device() + : node->assigned_device_name(); } } @@ -236,9 +230,16 @@ Status Encapsulator::SplitIntoSubgraphs() { // Create a new _Retval node DataType dtype = edge->src()->output_type(edge->src_output()); + if (IsRefType(dtype)) { + return errors::InvalidArgument( + "Ref Tensors (e.g., Variables) are not supported: tensor ", + edge->src()->name(), ":", edge->src_output()); + } + NodeDef ret_def; ret_def.set_op(kRetValOp); - ret_def.set_name(src_subgraph.graph->NewName("output")); + ret_def.set_name(strings::StrCat(edge->src()->name(), "_", + edge->src_output(), "_retval")); AddNodeAttr("T", dtype, &ret_def); AddNodeAttr("index", ret_index, &ret_def); Node* ret = src_subgraph.graph->AddNode(ret_def, &s); @@ -263,8 +264,16 @@ Status Encapsulator::SplitIntoSubgraphs() { // This is the first time we have seen this tensor. Create an _Arg node. DataType dtype = edge->dst()->input_type(edge->dst_input()); + if (IsRefType(dtype)) { + return errors::InvalidArgument( + "Ref Tensors (e.g., Variables) are not supported: tensor ", + edge->src()->name(), ":", edge->src_output()); + } + NodeDef arg_def; - NodeDefBuilder builder(dst_subgraph.graph->NewName("input"), kArgOp); + NodeDefBuilder builder(strings::StrCat(edge->src()->name(), "_", + edge->src_output(), "_arg"), + kArgOp); builder.Attr("T", dtype); builder.Attr("index", arg_index); s = builder.Finalize(&arg_def); @@ -291,11 +300,11 @@ Status Encapsulator::SplitIntoSubgraphs() { } Status Encapsulator::BuildFunctionDefs( - const RewriteSubgraphFn& rewrite_subgraph_fn, + const RewriteSubgraphFn& rewrite_subgraph_fn, bool reuse_existing_functions, FunctionLibraryDefinition* library) { // For each subgraph, build a FunctionDef. for (auto& subgraph_entry : subgraphs_) { - const string& name = subgraph_entry.first; + string name = subgraph_entry.first; Subgraph& subgraph = subgraph_entry.second; subgraph.call_node_def.set_op(name); @@ -332,6 +341,8 @@ Status Encapsulator::BuildFunctionDefs( for (auto& result : subgraph.results) { result.second = output_permutation[result.second]; } + + name = subgraph.call_node_def.op(); } FunctionDef fdef; @@ -346,7 +357,9 @@ Status Encapsulator::BuildFunctionDefs( strings::StrCat("encapsulate_fdef_", name), fdef); } - TF_RETURN_IF_ERROR(library->AddFunctionDef(fdef)); + if (!reuse_existing_functions || library->Find(name) == nullptr) { + TF_RETURN_IF_ERROR(library->AddFunctionDef(fdef)); + } } return Status::OK(); } @@ -423,8 +436,7 @@ Status Encapsulator::BuildOutputGraph(bool parallel_checking, std::unordered_map node_images; // Copy all unmarked nodes to the output graph. - for (Node* node : graph_in_->nodes()) { - if (node->IsSource() || node->IsSink()) continue; + for (Node* node : graph_in_->op_nodes()) { string func_id = GetFunctionNameAttr(node); // Don't copy nodes that going to be encapsulated, unless parallel checking @@ -545,14 +557,16 @@ Status Encapsulator::BuildOutputGraph(bool parallel_checking, Status EncapsulateSubgraphsInFunctions( string group_attribute, const Graph& graph_in, const RewriteSubgraphFn& rewrite_subgraph_fn, bool parallel_checking, - std::unique_ptr* graph_out, FunctionLibraryDefinition* library) { + bool reuse_existing_functions, std::unique_ptr* graph_out, + FunctionLibraryDefinition* library) { Status s; Encapsulator encapsulator(std::move(group_attribute), &graph_in); s = encapsulator.SplitIntoSubgraphs(); if (!s.ok()) return s; - s = encapsulator.BuildFunctionDefs(rewrite_subgraph_fn, library); + s = encapsulator.BuildFunctionDefs(rewrite_subgraph_fn, + reuse_existing_functions, library); if (!s.ok()) return s; std::unique_ptr out(new Graph(library)); @@ -566,10 +580,10 @@ Status EncapsulateSubgraphsInFunctions( // Finds the types of the _Arg nodes, indexed by position. static Status GetArgTypes(const Graph& graph, DataTypeVector* types) { - for (Node* n : graph.nodes()) { + for (Node* n : graph.op_nodes()) { if (n->type_string() == kArgOp) { int index; - TF_RETURN_IF_ERROR(GetNodeAttr(n->def(), "index", &index)); + TF_RETURN_IF_ERROR(GetNodeAttr(n->attrs(), "index", &index)); if (index < 0 || index >= types->size()) { return errors::InvalidArgument("Invalid argument number"); } @@ -583,10 +597,10 @@ static Status GetArgTypes(const Graph& graph, DataTypeVector* types) { // 'permutation' that maps old indices to new indices. static Status RenumberArguments(Graph* graph, const std::vector& permutation) { - for (Node* n : graph->nodes()) { + for (Node* n : graph->op_nodes()) { if (n->type_string() == kArgOp) { int index; - TF_RETURN_IF_ERROR(GetNodeAttr(n->def(), "index", &index)); + TF_RETURN_IF_ERROR(GetNodeAttr(n->attrs(), "index", &index)); if (index < 0 || index >= permutation.size()) { return errors::InvalidArgument("Invalid argument number"); } @@ -610,15 +624,18 @@ Status EncapsulateSubgraphsPass::Run( FunctionLibraryDefinition* const library = options.flib_def; OptimizerOptions opts; - std::unique_ptr flr( - NewFunctionLibraryRuntime(nullptr, options.session_options->env, nullptr, - TF_GRAPH_DEF_VERSION, library, opts)); - - auto rewrite_subgraph = [&flr]( - std::unique_ptr* subgraph, std::vector* input_permutation, - std::vector* output_permutation, NodeDef* node) { + std::unique_ptr pflr( + new ProcessFunctionLibraryRuntime(nullptr, options.session_options->env, + TF_GRAPH_DEF_VERSION, library, opts)); + FunctionLibraryRuntime* flr = + pflr->GetFLR(ProcessFunctionLibraryRuntime::kDefaultFLRDevice); + + auto rewrite_subgraph = [flr](std::unique_ptr* subgraph, + std::vector* input_permutation, + std::vector* output_permutation, + NodeDef* node) { // Optimize the subgraph. - OptimizeGraph(flr.get(), subgraph); + OptimizeGraph(flr, subgraph); const int num_args = input_permutation->size(); std::vector const_args(num_args); @@ -674,7 +691,8 @@ Status EncapsulateSubgraphsPass::Run( TF_RETURN_IF_ERROR(EncapsulateSubgraphsInFunctions( kXlaClusterAttr, **options.graph, rewrite_subgraph, - flags->tf_xla_parallel_checking, &graph_out, library)); + flags->tf_xla_parallel_checking, /*reuse_existing_functions=*/false, + &graph_out, library)); if (VLOG_IS_ON(1)) { dump_graph::DumpGraphToFile("after_encapsulate_subgraphs", *graph_out, @@ -688,7 +706,7 @@ Status EncapsulateSubgraphsPass::Run( bool IsXlaCompiledKernel(const Node& node) { bool is_compiled = false; bool has_compilation_attr = - GetNodeAttr(node.def(), kXlaCompiledKernelAttr, &is_compiled).ok() && + GetNodeAttr(node.attrs(), kXlaCompiledKernelAttr, &is_compiled).ok() && is_compiled; return has_compilation_attr ? is_compiled : false; } diff --git a/tensorflow/compiler/jit/encapsulate_subgraphs_pass.h b/tensorflow/compiler/jit/encapsulate_subgraphs_pass.h index 3ca7dfbf6a0ec29d9517139ffb952298d503cabc..b0987f76c91ed48df52fab303ea6052ebd8fd336 100644 --- a/tensorflow/compiler/jit/encapsulate_subgraphs_pass.h +++ b/tensorflow/compiler/jit/encapsulate_subgraphs_pass.h @@ -34,6 +34,8 @@ namespace tensorflow { // 'input_permutation' and 'output_permutation' are initialized to the identity // permutation. 'nodedef' is the NodeDef for the call to the function under // construction, provided to allow additional attributes to be set. +// The rewrite may also change the NodeDef's operator name, and that +// name will be used as the name of the generated function. typedef std::function* graph, std::vector* input_permutation, std::vector* output_permutation, NodeDef* node_def)> @@ -53,6 +55,9 @@ typedef std::function* graph_out, FunctionLibraryDefinition* library); + bool reuse_existing_functions, std::unique_ptr* graph_out, + FunctionLibraryDefinition* library); // The attribute that marks function calls produced by the encapsulate // subgraphs pass and that should in turn be compiled via _XlaLaunch operators. diff --git a/tensorflow/compiler/jit/encapsulate_subgraphs_pass_test.cc b/tensorflow/compiler/jit/encapsulate_subgraphs_pass_test.cc index faab7bd3d25d2491cf74faeb3b06acf4c2d6a054..4a1dbaf05dc7824835f3567c6abcf48222720230 100644 --- a/tensorflow/compiler/jit/encapsulate_subgraphs_pass_test.cc +++ b/tensorflow/compiler/jit/encapsulate_subgraphs_pass_test.cc @@ -13,6 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ +#include + #include "tensorflow/compiler/jit/encapsulate_subgraphs_pass.h" #include "tensorflow/cc/framework/ops.h" @@ -76,7 +78,7 @@ bool EqualFunctionDefLibrary(const FunctionDefLibrary& expected, #define TF_EXPECT_FUNCTIONDEFLIBRARY_EQ(expected, actual) \ do { \ string diff; \ - EXPECT_TRUE(EqualFunctionDefLibrary(actual, expected, &diff)) \ + EXPECT_TRUE(EqualFunctionDefLibrary(expected, actual, &diff)) \ << diff << "\nActual: " << actual.DebugString(); \ } while (false) @@ -101,15 +103,15 @@ Node* Input(const GraphDefBuilder::Options& opts) { } Node* Unary(ops::NodeOut a, const GraphDefBuilder::Options& opts) { - return ops::UnaryOp("UnaryTest", a, opts); + return ops::UnaryOp("UnaryTest", std::move(a), opts); } Node* Binary(ops::NodeOut a, ops::NodeOut b, const GraphDefBuilder::Options& opts) { - return ops::BinaryOp("BinaryTest", a, b, opts); + return ops::BinaryOp("BinaryTest", std::move(a), std::move(b), opts); } -Node* AddNLike(std::vector inputs, +Node* AddNLike(const std::vector& inputs, const GraphDefBuilder::Options& opts) { if (opts.HaveError()) return nullptr; NodeBuilder node_builder(opts.GetNameForOp("AddN"), "AddNLikeTest", @@ -127,7 +129,7 @@ Node* RetOp(int index, ops::NodeOut a, const GraphDefBuilder::Options& opts) { if (opts.HaveError()) return nullptr; NodeBuilder node_builder(opts.GetNameForOp("Retval"), "_Retval", opts.op_registry()); - node_builder.Input(a).Attr("index", index); + node_builder.Input(std::move(a)).Attr("index", index); return opts.FinalizeBuilder(&node_builder); } @@ -144,8 +146,9 @@ Status Encapsulate(GraphDef* graphdef, FunctionDefLibrary* library) { std::unique_ptr graph_out; s = EncapsulateSubgraphsInFunctions("_encapsulate", *graph, - /* rewrite_subgraph_fn= */ {}, - /* parallel_checking= */ false, + /*rewrite_subgraph_fn=*/{}, + /*parallel_checking=*/false, + /*reuse_existing_functions=*/false, &graph_out, lib_def.get()); if (!s.ok()) return s; @@ -205,12 +208,12 @@ TEST(EncapsulateSubgraphsTest, OneFunction) { *library_expected.add_function() = test::function::XTimesTwo(); *library_expected.add_function() = FunctionDefHelper::Create( - "F1", {"input__0:float", "input__1:float"}, {"output__2:float"}, {}, + "F1", {"a_0_arg:float", "b_0_arg:float"}, {"c_0_retval:float"}, {}, { - {{"C"}, "UnaryTest", {"input__0"}}, - {{"c"}, "BinaryTest", {"input__1", "C:o:0"}, {}, {"C"}}, + {{"C"}, "UnaryTest", {"a_0_arg"}}, + {{"c"}, "BinaryTest", {"b_0_arg", "C:o:0"}, {}, {"C"}}, }, - {{"output__2", "c:o:0"}}); + {{"c_0_retval", "c:o:0"}}); { std::unique_ptr lib_def( @@ -261,17 +264,17 @@ TEST(EncapsulateSubgraphsTest, TwoFunctions) { *library_expected.add_function() = test::function::XTimesTwo(); *library_expected.add_function() = FunctionDefHelper::Create( - "F1", {"input__0:float"}, {"output__1:float"}, {}, + "F1", {"a_0_arg:float"}, {"c_0_retval:float"}, {}, { - {{"C"}, "UnaryTest", {"input__0"}}, + {{"C"}, "UnaryTest", {"a_0_arg"}}, }, - {{"output__1", "C:o:0"}}); + {{"c_0_retval", "C:o:0"}}); *library_expected.add_function() = FunctionDefHelper::Create( - "F2", {"input__0:float", "input__1:float"}, {"output__2:float"}, {}, + "F2", {"b_0_arg:float", "c_0_arg:float"}, {"d_0_retval:float"}, {}, { - {{"D"}, "BinaryTest", {"input__0", "input__1"}}, + {{"D"}, "BinaryTest", {"b_0_arg", "c_0_arg"}}, }, - {{"output__2", "D:o:0"}}); + {{"d_0_retval", "D:o:0"}}); { std::unique_ptr lib_def( @@ -340,7 +343,8 @@ TEST(EncapsulateSubgraphsTest, InputDeduplication) { std::unique_ptr graph; TF_ASSERT_OK(EncapsulateSubgraphsInFunctions( "_cluster", graph_before_encapsulation, /*rewrite_subgraph_fn=*/{}, - /*parallel_checking=*/false, &graph, &library)); + /*parallel_checking=*/false, /*reuse_existing_functions=*/false, &graph, + &library)); std::vector expected_nodes = {"cluster1", "cluster2", "mul", "x"}; EXPECT_EQ(expected_nodes, GraphNodes(*graph)); @@ -371,7 +375,8 @@ TEST(EncapsulateSubgraphsTest, ParallelChecking) { std::unique_ptr graph; TF_ASSERT_OK(EncapsulateSubgraphsInFunctions( "_cluster", graph_before_encapsulation, /*rewrite_subgraph_fn=*/{}, - /*parallel_checking=*/true, &graph, &library)); + /*parallel_checking=*/true, /*reuse_existing_functions=*/false, &graph, + &library)); std::vector expected_nodes = { "add1", "add2", "cluster1", "cluster1_parallel_check/_0", diff --git a/tensorflow/compiler/jit/graph_to_functiondef.cc b/tensorflow/compiler/jit/graph_to_functiondef.cc index ce943471fb07fe02f18596247ccfddb94bd35158..6fa21fa6204dcc9446081d07e2a59ccace216713 100644 --- a/tensorflow/compiler/jit/graph_to_functiondef.cc +++ b/tensorflow/compiler/jit/graph_to_functiondef.cc @@ -120,14 +120,12 @@ Status GraphToFunctionDef(const Graph& graph, const string& name, std::unordered_map return_values; NodeNameMapping node_names; - for (Node const* node : graph.nodes()) { - if (!node->IsOp()) continue; - + for (Node const* node : graph.op_nodes()) { if (node->type_string() == kArgOp) { int index; DataType type; - TF_RETURN_IF_ERROR(GetNodeAttr(node->def(), "T", &type)); - TF_RETURN_IF_ERROR(GetNodeAttr(node->def(), "index", &index)); + TF_RETURN_IF_ERROR(GetNodeAttr(node->attrs(), "T", &type)); + TF_RETURN_IF_ERROR(GetNodeAttr(node->attrs(), "index", &index)); while (fdef->signature().input_arg_size() <= index) { fdef->mutable_signature()->add_input_arg(); } @@ -143,8 +141,8 @@ Status GraphToFunctionDef(const Graph& graph, const string& name, if (node->type_string() == kRetValOp) { int index; DataType type; - TF_RETURN_IF_ERROR(GetNodeAttr(node->def(), "T", &type)); - TF_RETURN_IF_ERROR(GetNodeAttr(node->def(), "index", &index)); + TF_RETURN_IF_ERROR(GetNodeAttr(node->attrs(), "T", &type)); + TF_RETURN_IF_ERROR(GetNodeAttr(node->attrs(), "index", &index)); while (fdef->signature().output_arg_size() <= index) { fdef->mutable_signature()->add_output_arg(); } @@ -153,17 +151,19 @@ Status GraphToFunctionDef(const Graph& graph, const string& name, argdef->set_type(type); const string normalized = node_names.Normalize(node->name()); argdef->set_name(normalized); - CHECK_EQ(node->in_edges().size(), 1) << node->DebugString(); - Edge const* edge = *node->in_edges().begin(); + Edge const* edge; + TF_CHECK_OK(node->input_edge(0, &edge)); return_values[normalized] = strings::StrCat(edge->src()->name(), ":", edge->src_output()); continue; } NodeDef* node_def = fdef->add_node_def(); - node_def->CopyFrom(node->def()); + *node_def = node->def(); + if (!node->assigned_device_name().empty()) { + node_def->set_device(node->assigned_device_name()); + } node_def->set_name(node_names.Uniquify(node->name())); - node_def->clear_device(); // Reset input names based on graph rather than the NodeDef. node_def->clear_input(); @@ -204,8 +204,8 @@ Status GraphToFunctionDef(const Graph& graph, const string& name, // Populate tensor_renaming. NameRangeMap output_ranges; - TF_RETURN_IF_ERROR(NameRangesForNode(node->def(), node->op_def(), nullptr, - &output_ranges)); + TF_RETURN_IF_ERROR( + NameRangesForNode(*node, node->op_def(), nullptr, &output_ranges)); for (const auto& output : output_ranges) { for (int i = output.second.first; i < output.second.second; ++i) { const string tensor_name = strings::StrCat( diff --git a/tensorflow/compiler/jit/graph_to_functiondef_test.cc b/tensorflow/compiler/jit/graph_to_functiondef_test.cc index 5c09e96a4c2817e5a871a91ca6c68de87dc3b762..676db7c4dd2fd7047e8ae9bb190daf18af6ac7cf 100644 --- a/tensorflow/compiler/jit/graph_to_functiondef_test.cc +++ b/tensorflow/compiler/jit/graph_to_functiondef_test.cc @@ -82,5 +82,38 @@ TEST(GraphToFunctionDefTest, Basics) { EXPECT_TRUE(fdefs_equal) << diff; } +// Regression test for a crash if there was a control edge to a _Retval node. +TEST(GraphToFunctionDefTest, ControlDependencies) { + Scope root = Scope::NewRootScope().ExitOnError(); + auto a = ops::_Arg(root.WithOpName("a"), DT_FLOAT, 0); + auto b = ops::Neg(root.WithOpName("b").WithControlDependencies(a), a); + auto c = ops::_Retval(root.WithOpName("c").WithControlDependencies(b), b, 0); + + GraphDef graph_def; + TF_EXPECT_OK(root.ToGraphDef(&graph_def)); + + std::unique_ptr graph(new Graph(OpRegistry::Global())); + GraphConstructorOptions options; + TF_EXPECT_OK(ConvertGraphDefToGraph(options, graph_def, graph.get())); + + FunctionDef fdef; + TF_EXPECT_OK(GraphToFunctionDef(*graph, "test_fn", &fdef)); + + FunctionDef fdef_expected = FunctionDefHelper::Create( + "test_fn", // function name + {"a: float"}, // inputs + {"c: float"}, // outputs + {}, // attrs + { + // nodes in the function body + {{"b"}, "Neg", {"a", "^a"}, {{"T", DT_FLOAT}}}, + }, + {{"c", "b:y:0"}}); // return values + + string diff; + bool fdefs_equal = EqualFunctionDef(fdef_expected, fdef, &diff); + EXPECT_TRUE(fdefs_equal) << diff; +} + } // namespace } // namespace tensorflow diff --git a/tensorflow/compiler/jit/graphcycles/graphcycles.cc b/tensorflow/compiler/jit/graphcycles/graphcycles.cc index 2139ffed4b9bfc5c0dd437faf2beae6d57471725..bc68afb322b5cfc814ce0537254ba14053ae4550 100644 --- a/tensorflow/compiler/jit/graphcycles/graphcycles.cc +++ b/tensorflow/compiler/jit/graphcycles/graphcycles.cc @@ -108,7 +108,7 @@ int32 GraphCycles::NewNode() { if (rep_->free_nodes_.empty()) { Node* n = new Node; n->visited = false; - n->data = NULL; + n->data = nullptr; n->rank = rep_->nodes_.size(); rep_->nodes_.push_back(n); return n->rank; @@ -116,7 +116,7 @@ int32 GraphCycles::NewNode() { // Preserve preceding rank since the set of ranks in use must be // a permutation of [0,rep_->nodes_.size()-1]. int32 r = rep_->free_nodes_.back(); - rep_->nodes_[r]->data = NULL; + rep_->nodes_[r]->data = nullptr; rep_->free_nodes_.pop_back(); return r; } @@ -332,7 +332,7 @@ int GraphCycles::FindPath(int32 x, int32 y, int max_path_len, } bool GraphCycles::IsReachable(int32 x, int32 y) const { - return FindPath(x, y, 0, NULL) > 0; + return FindPath(x, y, 0, nullptr) > 0; } bool GraphCycles::IsReachableNonConst(int32 x, int32 y) { diff --git a/tensorflow/compiler/jit/graphcycles/graphcycles_test.cc b/tensorflow/compiler/jit/graphcycles/graphcycles_test.cc index f27a616ac9ddf5cc79f262b19e533a84a4515b44..e47b782207e9122740fe9d5daf1fa0dbaeb47754 100644 --- a/tensorflow/compiler/jit/graphcycles/graphcycles_test.cc +++ b/tensorflow/compiler/jit/graphcycles/graphcycles_test.cc @@ -230,7 +230,7 @@ TEST(GraphCycles, RandomizedTest) { int new_node = graph_cycles.NewNode(); ASSERT_NE(-1, new_node); VLOG(1) << "adding node " << new_node; - ASSERT_EQ(0, graph_cycles.GetNodeData(new_node)); + ASSERT_EQ(nullptr, graph_cycles.GetNodeData(new_node)); graph_cycles.SetNodeData( new_node, reinterpret_cast( static_cast(new_node + kDataOffset))); @@ -243,7 +243,7 @@ TEST(GraphCycles, RandomizedTest) { break; case 1: // Remove a node - if (nodes.size() > 0) { + if (!nodes.empty()) { int node_index = RandomNode(&rnd, &nodes); int node = nodes[node_index]; nodes[node_index] = nodes.back(); @@ -263,7 +263,7 @@ TEST(GraphCycles, RandomizedTest) { break; case 2: // Add an edge - if (nodes.size() > 0) { + if (!nodes.empty()) { int from = RandomNode(&rnd, &nodes); int to = RandomNode(&rnd, &nodes); if (EdgeIndex(&edges, nodes[from], nodes[to]) == -1) { @@ -282,7 +282,7 @@ TEST(GraphCycles, RandomizedTest) { break; case 3: // Remove an edge - if (edges.size() > 0) { + if (!edges.empty()) { int i = RandomEdge(&rnd, &edges); int from = edges[i].from; int to = edges[i].to; @@ -296,7 +296,7 @@ TEST(GraphCycles, RandomizedTest) { break; case 4: // Check a path - if (nodes.size() > 0) { + if (!nodes.empty()) { int from = RandomNode(&rnd, &nodes); int to = RandomNode(&rnd, &nodes); int32 path[2 * kMaxNodes]; @@ -343,7 +343,7 @@ TEST(GraphCycles, RandomizedTest) { ASSERT_NE(-1, new_node); VLOG(1) << "adding node " << new_node; ASSERT_GE(new_node, 0); - ASSERT_EQ(0, graph_cycles.GetNodeData(new_node)); + ASSERT_EQ(nullptr, graph_cycles.GetNodeData(new_node)); graph_cycles.SetNodeData( new_node, reinterpret_cast( static_cast(new_node + kDataOffset))); diff --git a/tensorflow/compiler/jit/kernels/BUILD b/tensorflow/compiler/jit/kernels/BUILD index c4116cb8b52adc191e9f695bc9a6e0cf413b4b5c..354c0fabfc78bcb9f5d63e84edc224fc33650ea9 100644 --- a/tensorflow/compiler/jit/kernels/BUILD +++ b/tensorflow/compiler/jit/kernels/BUILD @@ -2,6 +2,7 @@ licenses(["notice"]) # Apache 2.0 package( default_visibility = [ + "//tensorflow/compiler/plugin/executor:__pkg__", "//tensorflow/compiler/tf2xla:internal", ], ) @@ -35,9 +36,11 @@ cc_library( "//tensorflow/compiler/jit:common", "//tensorflow/compiler/jit:xla_compilation_cache", "//tensorflow/compiler/jit:xla_device", + "//tensorflow/compiler/tf2xla:common", "//tensorflow/compiler/tf2xla:xla_compiler", "//tensorflow/compiler/xla:statusor", "//tensorflow/compiler/xla/client:local_client", + "//tensorflow/compiler/xla/legacy_flags:debug_options_flags", "//tensorflow/core:core_cpu_internal", "//tensorflow/core:framework", "//tensorflow/core:lib", diff --git a/tensorflow/compiler/jit/kernels/parallel_check_op.cc b/tensorflow/compiler/jit/kernels/parallel_check_op.cc index c86e03118b53ddf4865b7995b1d197c3ef07ba29..bd4eefbc0bb960f8ddc1d238057e73a29a098f26 100644 --- a/tensorflow/compiler/jit/kernels/parallel_check_op.cc +++ b/tensorflow/compiler/jit/kernels/parallel_check_op.cc @@ -64,7 +64,7 @@ class ParallelCheckOp : public OpKernel { ok = (diff <= tolerance); } if (ok) continue; - LOG(ERROR) << "Op " << def().name() << " fails equality at output " + LOG(ERROR) << "Op " << name() << " fails equality at output " << input_idx << " type " << DataTypeString(dtype) << " element " << i << ": std_val=" << p0[i] << " test_val=" << p1[i] << " diff=" << (p0[i] - p1[i]); @@ -75,7 +75,7 @@ class ParallelCheckOp : public OpKernel { } void Compute(OpKernelContext* ctx) override { - VLOG(1) << "Compute " << def().name(); + VLOG(1) << "Compute " << name(); const int num_pairs = ctx->num_inputs() / 2; for (int i = 0; i < num_pairs; ++i) { CHECK_EQ(ctx->input_dtype(i), ctx->input_dtype(i + num_pairs)); @@ -113,7 +113,7 @@ class ParallelCheckOp : public OpKernel { LOG(FATAL) << "unimpl: " << ctx->input_dtype(i); } if (failed > 0) { - LOG(ERROR) << "check failed for " << def().name() << " output " << i + LOG(ERROR) << "check failed for " << name() << " output " << i << " num_elts: " << num_elts; legacy_flags::ParallelCheckOpFlags* flags = legacy_flags::GetParallelCheckOpFlags(); @@ -121,7 +121,7 @@ class ParallelCheckOp : public OpKernel { LOG(QFATAL) << "failfast on first parallel-check failure"; } } else { - VLOG(1) << "check passed for " << def().name() << " output " << i + VLOG(1) << "check passed for " << name() << " output " << i << " num_elts: " << num_elts; } diff --git a/tensorflow/compiler/jit/kernels/xla_device_launch_op.cc b/tensorflow/compiler/jit/kernels/xla_device_launch_op.cc index c741ccfb31efa8794ae745e2e52e3c91b20cfcfc..2b77e5aaf4e0983354c14a4e20656af0e0e4f84b 100644 --- a/tensorflow/compiler/jit/kernels/xla_device_launch_op.cc +++ b/tensorflow/compiler/jit/kernels/xla_device_launch_op.cc @@ -18,7 +18,9 @@ limitations under the License. #include "tensorflow/compiler/jit/defs.h" #include "tensorflow/compiler/jit/xla_device.h" #include "tensorflow/compiler/jit/xla_device_context.h" +#include "tensorflow/compiler/tf2xla/shape_util.h" #include "tensorflow/compiler/xla/client/local_client.h" +#include "tensorflow/compiler/xla/legacy_flags/debug_options_flags.h" #include "tensorflow/compiler/xla/statusor.h" #include "tensorflow/core/common_runtime/dma_helper.h" #include "tensorflow/core/framework/allocator.h" @@ -34,7 +36,7 @@ namespace tensorflow { namespace { -Status BuildCompilationCache(ResourceMgr* rm, XlaCompilationCache** compiler) { +Status BuildCompilationCache(ResourceMgr* rm, XlaCompilationCache** cache) { XlaDevice::Metadata* metadata; Status s = rm->Lookup(rm->default_container(), "xla_metadata", &metadata); @@ -42,12 +44,8 @@ Status BuildCompilationCache(ResourceMgr* rm, XlaCompilationCache** compiler) { return s; } core::ScopedUnref metadata_ref(metadata); - XlaCompiler::Options options; - options.device_type = metadata->jit_device_type(); - options.client = metadata->client(); - options.allow_cpu_custom_calls = false; - options.local_executable_has_hybrid_result = false; - *compiler = new XlaCompilationCache(options); + *cache = + new XlaCompilationCache(metadata->client(), metadata->jit_device_type()); return Status::OK(); } @@ -59,7 +57,7 @@ XlaDeviceLaunchOp::XlaDeviceLaunchOp(OpKernelConstruction* ctx) OP_REQUIRES_OK(ctx, ctx->GetAttr("function", &func)); function_ = *func; VLOG(1) << "XlaDeviceLaunch created function=" - << Canonicalize(function_.name(), function_.attr()); + << Canonicalize(function_.name(), AttrSlice(&function_.attr())); DataTypeVector constant_types; OP_REQUIRES_OK(ctx, ctx->GetAttr("Tconstants", &constant_types)); num_constant_args_ = constant_types.size(); @@ -85,29 +83,37 @@ std::vector SnapshotResourceVariables(OpKernelContext* ctx, void XlaDeviceLaunchOp::Compute(OpKernelContext* ctx) { VLOG(1) << "XlaDeviceLaunch::Compute " - << Canonicalize(function_.name(), function_.attr()); + << Canonicalize(function_.name(), AttrSlice(&function_.attr())); // We store information about the JIT-compiled XLA computation // in the ResourceMgr. ResourceMgr* rm = ctx->resource_manager(); OP_REQUIRES(ctx, rm, errors::Internal("No resource manager.")); - XlaCompilationCache* compiler; + XlaCompilationCache* cache; OP_REQUIRES_OK(ctx, rm->LookupOrCreate( - rm->default_container(), "xla_compiler", &compiler, - [rm](XlaCompilationCache** compiler) { - return BuildCompilationCache(rm, compiler); + rm->default_container(), "xla_compiler", &cache, + [rm](XlaCompilationCache** cache) { + return BuildCompilationCache(rm, cache); })); // Holds the reference to the JIT during evaluation. (We could probably // free it sooner because the ResourceMgr will retain a reference, but // this is more obviously correct.) - core::ScopedUnref compiler_ref(compiler); + core::ScopedUnref cache_ref(cache); std::vector variables = SnapshotResourceVariables(ctx, num_resource_args_); + XlaCompiler::Options options; + options.client = cache->client(); + options.device_type = &cache->device_type(); + options.flib_def = ctx->function_library()->GetFunctionLibraryDefinition(); + options.graph_def_version = ctx->function_library()->graph_def_version(); + options.allow_cpu_custom_calls = false; + options.local_executable_has_hybrid_result = false; + const XlaCompiler::CompilationResult* kernel; - OP_REQUIRES_OK(ctx, compiler->Compile(function_, num_constant_args_, - variables, ctx, &kernel, nullptr)); + OP_REQUIRES_OK(ctx, cache->Compile(options, function_, num_constant_args_, + variables, ctx, &kernel, nullptr)); VLOG(1) << "XLA compilation complete..."; @@ -117,7 +123,7 @@ void XlaDeviceLaunchOp::Compute(OpKernelContext* ctx) { // Runs the computation, if any. There might not be a computation if all // outputs were compile-time constants. std::vector> outputs; - if (!kernel->computation.IsNull()) { + if (!kernel->computation->IsNull()) { auto opaque_shape = xla::ShapeUtil::MakeOpaqueShape(); // Builds the inputs to the computation. @@ -145,11 +151,13 @@ void XlaDeviceLaunchOp::Compute(OpKernelContext* ctx) { xla::ExecutionOptions execution_options; *execution_options.mutable_shape_with_output_layout() = kernel->xla_output_shape; + *execution_options.mutable_debug_options() = + xla::legacy_flags::GetDebugOptionsFromFlags(); Env* env = Env::Default(); auto start_time = env->NowMicros(); VLOG(1) << "Executing XLA Computation..."; - auto result = compiler->client()->Execute(kernel->computation, arg_ptrs, - &execution_options, &profile); + auto result = cache->client()->Execute(*kernel->computation, arg_ptrs, + &execution_options, &profile); auto elapsed = env->NowMicros() - start_time; OP_REQUIRES(ctx, result.ok(), result.status()); @@ -158,7 +166,7 @@ void XlaDeviceLaunchOp::Compute(OpKernelContext* ctx) { if (xla::ShapeUtil::IsTuple(kernel->xla_output_shape)) { auto outputs_or_error = - compiler->client()->DeconstructTuple(*result.ValueOrDie()); + cache->client()->DeconstructTuple(*result.ValueOrDie()); OP_REQUIRES(ctx, outputs_or_error.ok(), outputs_or_error.status()); outputs = outputs_or_error.ConsumeValueOrDie(); } else { @@ -198,11 +206,14 @@ void XlaDeviceLaunchOp::Compute(OpKernelContext* ctx) { // Apply variable updates, if any. VLOG(2) << "Applying variable updates"; - for (int i = 0; i < kernel->variable_updates.size(); ++i) { - const XlaCompiler::VariableUpdate& write = kernel->variable_updates[i]; + for (int i = 0; i < kernel->resource_updates.size(); ++i) { + const XlaCompiler::ResourceUpdate& write = kernel->resource_updates[i]; OP_REQUIRES(ctx, write.input_index >= 0 && write.input_index < ctx->num_inputs(), errors::Internal("Invalid input index for variable write.")); + TensorShape write_shape; + OP_REQUIRES_OK(ctx, XLAShapeToTensorShape(write.shape, &write_shape)); + // This code is very close to being a clone of AssignVariableOp, but the // key difference is that the contents of an XLA device tensor cannot be // copied safely; instead we must use @@ -210,26 +221,27 @@ void XlaDeviceLaunchOp::Compute(OpKernelContext* ctx) { Var* variable = nullptr; // TODO(b/35625933): tensorflow::Var should contain a PersistentTensor, not // a Tensor. - OP_REQUIRES_OK(ctx, LookupOrCreateResource( - ctx, HandleFromInput(ctx, write.input_index), - &variable, [this, ctx, &write](Var** ptr) { - *ptr = new Var(write.type); - PersistentTensor unused; - Tensor* tmp; - TF_RETURN_IF_ERROR(ctx->allocate_persistent( - write.type, write.shape, &unused, &tmp)); - *(*ptr)->tensor() = *tmp; - return Status::OK(); - })); + OP_REQUIRES_OK(ctx, + LookupOrCreateResource( + ctx, HandleFromInput(ctx, write.input_index), &variable, + [this, ctx, &write, &write_shape](Var** ptr) { + *ptr = new Var(write.type); + PersistentTensor unused; + Tensor* tmp; + TF_RETURN_IF_ERROR(ctx->allocate_persistent( + write.type, write_shape, &unused, &tmp)); + *(*ptr)->tensor() = *tmp; + return Status::OK(); + })); core::ScopedUnref s(variable); mutex_lock ml(*variable->mu()); OP_REQUIRES(ctx, variable->tensor()->dtype() == write.type, errors::Internal("Mismatched type in variable write")); - if (!variable->tensor()->shape().IsSameSize(write.shape)) { + if (!variable->tensor()->shape().IsSameSize(write_shape)) { PersistentTensor unused; Tensor* tmp; - OP_REQUIRES_OK(ctx, ctx->allocate_persistent(write.type, write.shape, + OP_REQUIRES_OK(ctx, ctx->allocate_persistent(write.type, write_shape, &unused, &tmp)); *variable->tensor() = *tmp; } diff --git a/tensorflow/compiler/jit/kernels/xla_local_launch_op.cc b/tensorflow/compiler/jit/kernels/xla_local_launch_op.cc index 8b43c7c1564a340b70e8cfa271a3ef50379b46bc..40acc0d81d08230b373823e333cd5e3e407b9c4f 100644 --- a/tensorflow/compiler/jit/kernels/xla_local_launch_op.cc +++ b/tensorflow/compiler/jit/kernels/xla_local_launch_op.cc @@ -148,24 +148,28 @@ XlaLocalLaunchOp::XlaLocalLaunchOp(OpKernelConstruction* ctx) OP_REQUIRES(ctx, num_resource_args == 0, errors::Unimplemented( "XlaLocalLaunchOp does not support resource variables")); -} - -Status XlaLocalLaunchOp::BuildCompilationCache(XlaCompilationCache** compiler) { - gpu::Platform::Id platform_id; if (device_type_ == DeviceType(DEVICE_CPU)) { - platform_id = gpu::host::kHostPlatformId; + platform_id_ = gpu::host::kHostPlatformId; } else if (device_type_ == DeviceType(DEVICE_GPU)) { - platform_id = gpu::cuda::kCudaPlatformId; + platform_id_ = gpu::cuda::kCudaPlatformId; } else { - return errors::InvalidArgument("Unknown device type for local _XlaLaunch"); + ctx->SetStatus( + errors::InvalidArgument("Unknown device type for local _XlaLaunch")); + return; } +} - auto platform = gpu::MultiPlatformManager::PlatformWithId(platform_id); +Status XlaLocalLaunchOp::BuildCompilationCache(OpKernelContext* ctx, + XlaCompilationCache** cache) { + auto platform = gpu::MultiPlatformManager::PlatformWithId(platform_id_); if (!platform.ok()) { return StreamExecutorUtil::ConvertStatus(platform.status()); } - auto client = - xla::ClientLibrary::GetOrCreateLocalClient(platform.ValueOrDie()); + xla::LocalClientOptions client_options; + client_options.set_platform(platform.ValueOrDie()); + client_options.set_intra_op_parallelism_threads( + ctx->device()->tensorflow_cpu_worker_threads()->num_threads); + auto client = xla::ClientLibrary::GetOrCreateLocalClient(client_options); if (!client.ok()) { return client.status(); } @@ -175,18 +179,14 @@ Status XlaLocalLaunchOp::BuildCompilationCache(XlaCompilationCache** compiler) { return errors::InvalidArgument("No JIT device registered for ", device_type_.type()); } - XlaCompiler::Options options; - options.device_type = DeviceType(registration->compilation_device_name); - options.client = client.ValueOrDie(); - options.allow_cpu_custom_calls = (platform_id == gpu::host::kHostPlatformId); - options.local_executable_has_hybrid_result = true; - *compiler = new XlaCompilationCache(options); + *cache = new XlaCompilationCache( + client.ValueOrDie(), DeviceType(registration->compilation_device_name)); return Status::OK(); } void XlaLocalLaunchOp::Compute(OpKernelContext* ctx) { VLOG(1) << "XlaLocalLaunchOp::Compute " - << Canonicalize(function_.name(), function_.attr()); + << Canonicalize(function_.name(), AttrSlice(&function_.attr())); // We store information about the JIT-compiled XLA computation // in the ResourceMgr. ResourceMgr* rm = ctx->resource_manager(); @@ -195,23 +195,31 @@ void XlaLocalLaunchOp::Compute(OpKernelContext* ctx) { gpu::Stream* stream = ctx->op_device_context() ? ctx->op_device_context()->stream() : nullptr; - XlaCompilationCache* compiler; + XlaCompilationCache* cache; OP_REQUIRES_OK(ctx, rm->LookupOrCreate( - rm->default_container(), "xla_compiler", &compiler, - [this](XlaCompilationCache** compiler) { - return BuildCompilationCache(compiler); + rm->default_container(), "xla_cache", &cache, + [this, ctx](XlaCompilationCache** cache) { + return BuildCompilationCache(ctx, cache); })); // Hold the reference to the JIT during evaluation. (We could probably // free it sooner because the ResourceMgr will retain a reference, but // this is more obviously correct.) - core::ScopedUnref compiler_ref(compiler); + core::ScopedUnref cache_ref(cache); + + xla::LocalClient* client = static_cast(cache->client()); - xla::LocalClient* client = static_cast(compiler->client()); + XlaCompiler::Options options; + options.client = client; + options.device_type = &cache->device_type(); + options.flib_def = ctx->function_library()->GetFunctionLibraryDefinition(); + options.graph_def_version = ctx->function_library()->graph_def_version(); + options.allow_cpu_custom_calls = (platform_id_ == gpu::host::kHostPlatformId); + options.local_executable_has_hybrid_result = true; const XlaCompiler::CompilationResult* kernel; xla::LocalExecutable* executable; - OP_REQUIRES_OK(ctx, compiler->Compile(function_, num_constant_args_, {}, ctx, - &kernel, &executable)); + OP_REQUIRES_OK(ctx, cache->Compile(options, function_, num_constant_args_, {}, + ctx, &kernel, &executable)); VLOG(1) << "Executing XLA Computation..."; @@ -221,7 +229,7 @@ void XlaLocalLaunchOp::Compute(OpKernelContext* ctx) { std::unique_ptr output; bool output_is_tuple; - if (!kernel->computation.IsNull()) { + if (!kernel->computation->IsNull()) { // Build xla::ShapedBuffers that point directly to the Tensor buffers. std::vector> arg_buffers; arg_buffers.reserve(kernel->xla_input_shapes.size() + 1); @@ -260,8 +268,6 @@ void XlaLocalLaunchOp::Compute(OpKernelContext* ctx) { xla::ExecutableRunOptions run_options; run_options.set_stream(stream); run_options.set_allocator(&xla_allocator); - run_options.set_inter_op_thread_pool( - ctx->device()->tensorflow_cpu_worker_threads()->workers); run_options.set_intra_op_thread_pool(&ctx->eigen_cpu_device()); Env* env = Env::Default(); auto start_time = env->NowMicros(); diff --git a/tensorflow/compiler/jit/kernels/xla_local_launch_op.h b/tensorflow/compiler/jit/kernels/xla_local_launch_op.h index 8023206762951a4dafba900dd291f2ee9bdbbdf3..5e4d3336a91001fac1d222709f64300e777247c7 100644 --- a/tensorflow/compiler/jit/kernels/xla_local_launch_op.h +++ b/tensorflow/compiler/jit/kernels/xla_local_launch_op.h @@ -22,6 +22,7 @@ limitations under the License. #include "tensorflow/core/framework/op_kernel.h" #include "tensorflow/core/framework/tensor.h" #include "tensorflow/core/platform/macros.h" +#include "tensorflow/core/util/stream_executor_util.h" namespace tensorflow { @@ -43,11 +44,15 @@ class XlaLocalLaunchOp : public OpKernel { private: // Builds a XlaCompilationCache class suitable for the current device. - Status BuildCompilationCache(XlaCompilationCache** compiler); + Status BuildCompilationCache(OpKernelContext* ctx, + XlaCompilationCache** compiler); DeviceType device_type_; NameAttrList function_; int num_constant_args_; + + perftools::gputools::Platform::Id platform_id_; + TF_DISALLOW_COPY_AND_ASSIGN(XlaLocalLaunchOp); }; diff --git a/tensorflow/compiler/jit/mark_for_compilation_pass.cc b/tensorflow/compiler/jit/mark_for_compilation_pass.cc index b27c07d0d987aafef1943fd795293bd066ad36f6..2fe190e605f4e31f6190aee952b7d139c1f772c6 100644 --- a/tensorflow/compiler/jit/mark_for_compilation_pass.cc +++ b/tensorflow/compiler/jit/mark_for_compilation_pass.cc @@ -24,11 +24,13 @@ limitations under the License. #include "tensorflow/compiler/jit/defs.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/tf2xla/dump_graph.h" #include "tensorflow/compiler/tf2xla/xla_op_registry.h" #include "tensorflow/core/common_runtime/function.h" #include "tensorflow/core/framework/graph_def_util.h" #include "tensorflow/core/framework/memory_types.h" +#include "tensorflow/core/framework/node_def.pb.h" #include "tensorflow/core/framework/op_kernel.h" #include "tensorflow/core/framework/types.h" #include "tensorflow/core/graph/algorithm.h" @@ -52,20 +54,22 @@ bool HasXLAKernel(const Node& node, const DeviceType& jit_device_type) { // Make sure we don't recurse infinitely on recursive functions. const int kMaxRecursionDepth = 10; -bool IsCompilableCall(const NodeDef& call_def, DeviceType jit_device_type, - int depth, FunctionLibraryRuntime* lib_runtime); +bool IsCompilableCall(const NodeDef& call_def, + const DeviceType& jit_device_type, int depth, + FunctionLibraryRuntime* lib_runtime); -// Tests whether 'while_def' is a completely compilable loop. +// Tests whether 'while_node' is a completely compilable loop. // Every operator in the condition and body functions must be compilable for a // while loop to be compilable. -bool IsCompilableWhile(const NodeDef& while_def, DeviceType jit_device_type, - int depth, FunctionLibraryRuntime* lib_runtime) { - VLOG(2) << "Loop marking: " << while_def.op(); +bool IsCompilableWhile(const Node& while_node, + const DeviceType& jit_device_type, int depth, + FunctionLibraryRuntime* lib_runtime) { + VLOG(2) << "Loop marking: " << while_node.type_string(); const NameAttrList* name_attr; NodeDef call; Status status; - status = GetNodeAttr(while_def, "cond", &name_attr); + status = GetNodeAttr(while_node.attrs(), "cond", &name_attr); if (!status.ok()) { VLOG(2) << "Missing 'cond' attribute on While node."; return false; @@ -78,7 +82,7 @@ bool IsCompilableWhile(const NodeDef& while_def, DeviceType jit_device_type, VLOG(2) << "Can't compile loop condition: " << cond_func; return false; } - status = GetNodeAttr(while_def, "body", &name_attr); + status = GetNodeAttr(while_node.attrs(), "body", &name_attr); if (!status.ok()) { VLOG(2) << "Missing 'body' attribute on While node."; return false; @@ -98,8 +102,9 @@ bool IsCompilableWhile(const NodeDef& while_def, DeviceType jit_device_type, // Tests whether 'call_def' is a call to a completely compilable function. // Every operator in the function must be compilable for a function to be // compilable. -bool IsCompilableCall(const NodeDef& call_def, DeviceType jit_device_type, - int depth, FunctionLibraryRuntime* lib_runtime) { +bool IsCompilableCall(const NodeDef& call_def, + const DeviceType& jit_device_type, int depth, + FunctionLibraryRuntime* lib_runtime) { VLOG(2) << "Function marking: " << call_def.op(); if (depth > kMaxRecursionDepth) { @@ -109,7 +114,7 @@ bool IsCompilableCall(const NodeDef& call_def, DeviceType jit_device_type, FunctionLibraryRuntime::Handle handle; Status status = - lib_runtime->Instantiate(call_def.op(), call_def.attr(), &handle); + lib_runtime->Instantiate(call_def.op(), AttrSlice(call_def), &handle); if (!status.ok()) { VLOG(2) << "Could not instantiate " << call_def.op() << ": " << status; return false; @@ -129,13 +134,12 @@ bool IsCompilableCall(const NodeDef& call_def, DeviceType jit_device_type, return false; } - for (Node* node : fbody->graph->nodes()) { - if (node->IsSource() || node->IsSink()) continue; - if (node->def().op() == "_Arg" || node->def().op() == "_Retval") continue; - if (node->def().op() == "While") { + for (Node* node : fbody->graph->op_nodes()) { + if (node->type_string() == "_Arg" || node->type_string() == "_Retval") + continue; + if (node->type_string() == "While") { // Handle functional While loop (not in open source build). - return IsCompilableWhile(node->def(), jit_device_type, depth + 1, - lib_runtime); + return IsCompilableWhile(*node, jit_device_type, depth + 1, lib_runtime); } if (!HasXLAKernel(*node, jit_device_type) && !IsCompilableCall(node->def(), jit_device_type, depth + 1, @@ -159,10 +163,12 @@ Status DeviceTypeOfDevice(const string& device, DeviceType* device_type) { return Status::OK(); } -// Does `node` have a DT_RESOURCE typed argument? -bool HasResourceArgument(const Node& node) { +// Tests whether `node` has a DT_RESOURCE typed input or output. +bool HasResourceInputOrOutput(const Node& node) { return std::find(node.input_types().begin(), node.input_types().end(), - DT_RESOURCE) != node.input_types().end(); + DT_RESOURCE) != node.input_types().end() || + std::find(node.output_types().begin(), node.output_types().end(), + DT_RESOURCE) != node.output_types().end(); } Status FindCompilationCandidates( @@ -170,12 +176,13 @@ Status FindCompilationCandidates( const std::function& is_compilable_fn, std::unordered_set* candidates) { OptimizerOptions opts; - std::unique_ptr lib_runtime(NewFunctionLibraryRuntime( - nullptr, env, nullptr, TF_GRAPH_DEF_VERSION, flib_def, opts)); - - for (Node* node : graph.nodes()) { - if (node->IsSource() || node->IsSink()) continue; + std::unique_ptr pflr( + new ProcessFunctionLibraryRuntime(nullptr, env, TF_GRAPH_DEF_VERSION, + flib_def, opts)); + FunctionLibraryRuntime* lib_runtime = + pflr->GetFLR(ProcessFunctionLibraryRuntime::kDefaultFLRDevice); + for (Node* node : graph.op_nodes()) { DeviceType device_type(""); TF_RETURN_IF_ERROR( DeviceTypeOfDevice(node->assigned_device_name(), &device_type)); @@ -187,19 +194,19 @@ Status FindCompilationCandidates( XlaOpRegistry::GetCompilationDevice(device_type.type(), ®istration)); DeviceType jit_device_type(registration->compilation_device_name); if (!HasXLAKernel(*node, jit_device_type) && - !IsCompilableCall(node->def(), jit_device_type, 0, lib_runtime.get())) { + !IsCompilableCall(node->def(), jit_device_type, 0, lib_runtime)) { VLOG(2) << "Compilation rejected node: unsupported op " << node->name() - << ": " << node->def().op(); + << ": " << node->type_string(); continue; } - if (!registration->compile_resource_ops && HasResourceArgument(*node)) { - VLOG(2) << "Compilation rejected node: resource argument " << node->name() - << ": " << node->def().op(); + if (!registration->compile_resource_ops && + HasResourceInputOrOutput(*node)) { + VLOG(2) << "Compilation rejected node: resource input/output " + << node->name() << ": " << node->type_string(); continue; } - if (node->def().op() == "While" && - !IsCompilableWhile(node->def(), jit_device_type, 0, - lib_runtime.get())) { + if (node->type_string() == "While" && + !IsCompilableWhile(*node, jit_device_type, 0, lib_runtime)) { continue; } candidates->insert(node); @@ -207,70 +214,12 @@ Status FindCompilationCandidates( return Status::OK(); } -// Union-Find data structure used to compute clusters. We use our own -// implementation because we want one key feature: when merging clusters, we -// need to know which value becomes the representative of the merged clusters. -// We use the representatives to name nodes in a cycle detection graph, and we -// need to control which node is named. -// TODO(phawkins): consider merging this code with union-find implementations -// in Tensorflow, e.g., in SimplePlacer. -class Cluster { - public: - Cluster(); - - int Size() { return FindRoot()->size_; } - - // Merges this cluster with 'other'. This cluster's representative becomes - // the representative of the merged cluster; the representative of 'other' - // is ignored. - void Merge(Cluster* other); - - // Each cluster has an associated integer 'representative', initialized to -1 - // by default. - int GetRepresentative() { return FindRoot()->representative_; } - void SetRepresentative(int representative) { - FindRoot()->representative_ = representative; - } - - private: - // Finds the root element of the cluster. Performs path compression. - Cluster* FindRoot(); - - int representative_; - int rank_; - int size_; // Size of the cluster. - Cluster* parent_; +struct Cluster { + // Identifies the node that represents this cluster in the cycle detection + // graph. + int representative = -1; }; -Cluster::Cluster() - : representative_(-1), rank_(0), size_(1), parent_(nullptr) {} - -void Cluster::Merge(Cluster* other) { - Cluster* a = FindRoot(); - Cluster* b = other->FindRoot(); - if (a == b) return; - if (a->rank_ > b->rank_) { - b->parent_ = a; - a->size_ += b->size_; - return; - } - - a->parent_ = b; - if (a->rank_ == b->rank_) { - b->rank_++; - } - b->representative_ = a->representative_; - b->size_ += a->size_; -} - -Cluster* Cluster::FindRoot() { - if (!parent_) return this; - // Path compression: update intermediate nodes to point to the root of the - // equivalence class. - parent_ = parent_->FindRoot(); - return parent_; -} - } // anonymous namespace bool IsCompilable(FunctionLibraryRuntime* flr, const NodeDef& ndef) { @@ -311,15 +260,20 @@ Status MarkForCompilationPass::Run( ®istration)) { return false; } + + // Don't compile control trigger nodes. We won't preserve their deadness + // semantics correctly, so it's safest not to compile them. + if (node->IsControlTrigger()) 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->def(), kXlaCompileAttr, &compile); + Status status = GetNodeAttr(node->attrs(), kXlaCompileAttr, &compile); if (status.ok()) return compile; - status = fld->GetAttr(node->def(), kXlaCompileAttr, &compile); + status = fld->GetAttr(*node, kXlaCompileAttr, &compile); if (status.ok()) return compile; // Otherwise use the value of global_jit_level. @@ -433,10 +387,11 @@ Status MarkForCompilationPass::RunImpl( // Each compilation candidate belongs to a cluster. The cluster's // representative // names the node in the 'cycles' graph that represents the cluster. - std::vector clusters(graph->num_node_ids()); - std::deque worklist; + std::vector> clusters(graph->num_node_ids()); + std::deque*> worklist; for (Node* node : compilation_candidates) { - clusters[node->id()].SetRepresentative(node->id()); + Cluster& cluster = clusters[node->id()].Get(); + cluster.representative = node->id(); worklist.push_back(&clusters[node->id()]); } @@ -446,7 +401,7 @@ Status MarkForCompilationPass::RunImpl( // Repeatedly contract edges between clusters that are on the same device, // provided the contraction would not create a cycle. while (!worklist.empty()) { - int from = worklist.front()->GetRepresentative(); + int from = worklist.front()->Get().representative; worklist.pop_front(); Node* node_from = graph->FindNodeId(from); @@ -482,8 +437,8 @@ Status MarkForCompilationPass::RunImpl( // all nodes marked with _XlaCompile=true to also have a // _XlaScope property set (and raise an error otherwise); but // for now we don't do this. - if (GetNodeAttr(node_from->def(), kXlaScopeAttr, &from_scope).ok() && - GetNodeAttr(node_to->def(), kXlaScopeAttr, &to_scope).ok() && + if (GetNodeAttr(node_from->attrs(), kXlaScopeAttr, &from_scope).ok() && + GetNodeAttr(node_to->attrs(), kXlaScopeAttr, &to_scope).ok() && from_scope != to_scope) { continue; } @@ -519,7 +474,7 @@ Status MarkForCompilationPass::RunImpl( // Count the number of elements in each cluster. std::vector cluster_sizes(graph->num_node_ids()); for (const Node* n : compilation_candidates) { - int cluster = clusters[n->id()].GetRepresentative(); + int cluster = clusters[n->id()].Get().representative; cluster_sizes[cluster]++; } @@ -533,15 +488,14 @@ Status MarkForCompilationPass::RunImpl( // if compilation is enabled, otherwise there will be no such candidates). const int min_cluster_size = flags->tf_xla_min_cluster_size; for (Node* n : compilation_candidates) { - int cluster = clusters[n->id()].GetRepresentative(); + int cluster = clusters[n->id()].Get().representative; // Compile if the user marked this node _XlaCompile=true bool compile_attr = false; bool marked_for_compilation = false; - if (GetNodeAttr(n->def(), kXlaCompileAttr, &compile_attr).ok()) { + if (GetNodeAttr(n->attrs(), kXlaCompileAttr, &compile_attr).ok()) { marked_for_compilation = compile_attr; - } else if (options.flib_def - ->GetAttr(n->def(), kXlaCompileAttr, &compile_attr) + } else if (options.flib_def->GetAttr(*n, kXlaCompileAttr, &compile_attr) .ok()) { marked_for_compilation = compile_attr; } diff --git a/tensorflow/compiler/jit/mark_for_compilation_pass_test.cc b/tensorflow/compiler/jit/mark_for_compilation_pass_test.cc index 91e4a2b41c7026b6ca028ed6a7e61588d57e9e50..4b88da27a188ed4fa6125b3e7a84034efb1a0ec1 100644 --- a/tensorflow/compiler/jit/mark_for_compilation_pass_test.cc +++ b/tensorflow/compiler/jit/mark_for_compilation_pass_test.cc @@ -14,11 +14,13 @@ limitations under the License. ==============================================================================*/ #include "tensorflow/compiler/jit/mark_for_compilation_pass.h" -#include "tensorflow/compiler/jit/defs.h" #include "tensorflow/cc/framework/ops.h" #include "tensorflow/cc/ops/control_flow_ops_internal.h" #include "tensorflow/cc/ops/standard_ops.h" +#include "tensorflow/compiler/jit/defs.h" +#include "tensorflow/compiler/tf2xla/xla_op_kernel.h" +#include "tensorflow/compiler/tf2xla/xla_op_registry.h" #include "tensorflow/core/framework/node_def_util.h" #include "tensorflow/core/framework/op.h" #include "tensorflow/core/graph/graph_constructor.h" @@ -57,7 +59,7 @@ std::unordered_map GetClusters(const Graph& graph) { std::unordered_map ids; for (Node* node : graph.nodes()) { string cluster; - if (GetNodeAttr(node->def(), kXlaClusterAttr, &cluster).ok()) { + if (GetNodeAttr(node->attrs(), kXlaClusterAttr, &cluster).ok()) { CHECK(!cluster.empty()); ids[node->name()] = cluster; } @@ -455,5 +457,39 @@ TEST(XlaCompilationTest, CyclesWithDifferentScopesAndBridge) { EXPECT_EQ(clusters["B"], clusters["C"]); } +REGISTER_OP("ResourceInput").Input("a: resource").Output("o: float"); +REGISTER_OP("ResourceOutput").Input("a: float").Output("o: resource"); + +namespace { + +class DummyOp : public XlaOpKernel { + using XlaOpKernel::XlaOpKernel; + void Compile(XlaOpKernelContext* ctx) override {} +}; + +REGISTER_XLA_OP(Name("ResourceInput"), DummyOp); +REGISTER_XLA_OP(Name("ResourceOutput"), DummyOp); + +} // namespace + +TEST(XlaCompilationTest, Resources) { + std::unique_ptr graph(new Graph(OpRegistry::Global())); + GraphDef graphdef; + { + GraphDefBuilder builder(GraphDefBuilder::kFailImmediately); + Node* a = + ops::SourceOp("UncompilableNullary", builder.opts().WithName("A")); + Node* b = ops::UnaryOp("Relu", a, builder.opts().WithName("B")); + // We should not form clusters with resource ops by default. + Node* c = ops::UnaryOp("ResourceOutput", b, builder.opts().WithName("C")); + Node* d = ops::UnaryOp("ResourceInput", c, builder.opts().WithName("D")); + ops::UnaryOp("Relu", d, builder.opts().WithName("E")); + TF_EXPECT_OK(builder.ToGraph(graph.get())); + } + MarkForCompilation(&graph); + auto clusters = GetClusters(*graph); + EXPECT_EQ(0, clusters.size()); // Nothing should be compiled. +} + } // namespace } // namespace tensorflow diff --git a/tensorflow/compiler/jit/ops/BUILD b/tensorflow/compiler/jit/ops/BUILD index 8d1fa03cc0d74f3a61b3e2e1d6f2af07c0bcd23f..e5787ca4c8cff436e4404b8488970248b24a5eda 100644 --- a/tensorflow/compiler/jit/ops/BUILD +++ b/tensorflow/compiler/jit/ops/BUILD @@ -1,32 +1,20 @@ licenses(["notice"]) # Apache 2.0 package( - default_visibility = [ - "//tensorflow/compiler/tf2xla:internal", - ], + default_visibility = ["//tensorflow/compiler/tf2xla:internal"], ) cc_library( name = "xla_ops", - srcs = [ - "xla_ops.cc", - ], - deps = [ - "//tensorflow/core:framework", - "//tensorflow/core:lib", - "//tensorflow/core:protos_all_cc", - ], + srcs = ["xla_ops.cc"], + deps = ["//tensorflow/core:framework"], alwayslink = 1, ) cc_library( name = "parallel_check_op", srcs = ["parallel_check_op.cc"], - deps = [ - "//tensorflow/core:framework", - "//tensorflow/core:lib", - "//tensorflow/core:protos_all_cc", - ], + deps = ["//tensorflow/core:framework"], alwayslink = 1, ) diff --git a/tensorflow/compiler/jit/union_find.h b/tensorflow/compiler/jit/union_find.h new file mode 100644 index 0000000000000000000000000000000000000000..a1a7a6a4d0d4c8a8b75c464b46dc4cf344218125 --- /dev/null +++ b/tensorflow/compiler/jit/union_find.h @@ -0,0 +1,81 @@ +/* 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_UNION_FIND_H_ +#define TENSORFLOW_COMPILER_JIT_UNION_FIND_H_ + +namespace tensorflow { + +// Union-Find data structure. +// Each cluster has an associated value; when merging clusters we can control +// which value becomes the representative of the merged clusters. Values must be +// copyable. +template +class UnionFind { + public: + UnionFind() : rank_(0), size_(1), parent_(nullptr) {} + + // Returns the number of elements in a cluster. + int Size() { return FindRoot()->size_; } + + // Merges this cluster with 'other'. This cluster's value becomes + // the value of the merged cluster; the value of 'other' is ignored. + void Merge(UnionFind* other); + + // Each cluster has an associated value. Retrieves the value associated + // with this cluster. + T& Get() { return FindRoot()->value_; } + + private: + // Finds the root element of the cluster. Performs path compression. + UnionFind* FindRoot(); + + int rank_; + int size_; // Size of the cluster. + UnionFind* parent_; + T value_; +}; + +template +void UnionFind::Merge(UnionFind* other) { + UnionFind* a = FindRoot(); + UnionFind* b = other->FindRoot(); + if (a == b) return; + if (a->rank_ > b->rank_) { + b->parent_ = a; + a->size_ += b->size_; + return; + } + + a->parent_ = b; + if (a->rank_ == b->rank_) { + b->rank_++; + } + b->value_ = a->value_; + b->size_ += a->size_; +} + +template +UnionFind* UnionFind::FindRoot() { + if (!parent_) return this; + // Path compression: update intermediate nodes to point to the root of the + // equivalence class. + parent_ = parent_->FindRoot(); + return parent_; +} + +} // namespace tensorflow + +#endif // TENSORFLOW_COMPILER_JIT_UNION_FIND_H_ diff --git a/tensorflow/compiler/jit/xla_compilation_cache.cc b/tensorflow/compiler/jit/xla_compilation_cache.cc index 41abea02eb2d17423744dfb719ee9a3f6b8f1198..3c52316ccef0023472b2e888e0c31b07fc00e694 100644 --- a/tensorflow/compiler/jit/xla_compilation_cache.cc +++ b/tensorflow/compiler/jit/xla_compilation_cache.cc @@ -37,9 +37,9 @@ limitations under the License. namespace tensorflow { -XlaCompilationCache::XlaCompilationCache(const XlaCompiler::Options& options) - : compiler_(options) {} - +XlaCompilationCache::XlaCompilationCache(xla::Client* client, + DeviceType device_type) + : client_(client), device_type_(std::move(device_type)) {} XlaCompilationCache::~XlaCompilationCache() = default; string XlaCompilationCache::DebugString() { @@ -95,7 +95,7 @@ Status XlaCompilationCache::BuildSignature( const NameAttrList& function, int num_constant_args, const std::vector& variable_args, OpKernelContext* ctx, Signature* signature) { - signature->name = Canonicalize(function.name(), function.attr()); + signature->name = Canonicalize(function.name(), AttrSlice(&function.attr())); signature->arg_values.resize(num_constant_args); signature->arg_types.reserve(ctx->num_inputs() - num_constant_args); @@ -148,7 +148,8 @@ Status BuildArguments(int num_constant_args, XlaCompiler::Argument& arg = (*args)[input_num]; arg.kind = XlaCompiler::Argument::kConstant; arg.type = input.dtype(); - arg.shape = input.shape(); + TF_RETURN_IF_ERROR( + TensorShapeToXLAShape(input.dtype(), input.shape(), &arg.shape)); arg.constant_value = input; ++input_num; } @@ -169,7 +170,8 @@ Status BuildArguments(int num_constant_args, arg.constant_value = input; } arg.type = input.dtype(); - arg.shape = input.shape(); + TF_RETURN_IF_ERROR( + TensorShapeToXLAShape(input.dtype(), input.shape(), &arg.shape)); ++input_num; } @@ -182,19 +184,21 @@ Status BuildArguments(int num_constant_args, XlaCompiler::Argument& arg = (*args)[input_num]; arg.name = variable_args[variable_id].name; + arg.kind = XlaCompiler::Argument::kVariable; if (variable_args[variable_id].present) { const Tensor& value = variable_args[variable_id].value; - arg.kind = XlaCompiler::Argument::kVariable; arg.type = value.dtype(); - arg.shape = value.shape(); + TF_RETURN_IF_ERROR( + TensorShapeToXLAShape(value.dtype(), value.shape(), &arg.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.kind = XlaCompiler::Argument::kUninitializedVariable; + arg.initialized = false; arg.type = DT_INVALID; - arg.shape = TensorShape(); + arg.shape = xla::Shape(); } ++input_num; } @@ -205,8 +209,9 @@ Status BuildArguments(int num_constant_args, } // namespace Status XlaCompilationCache::Compile( - const NameAttrList& function, int num_constant_args, - const std::vector& variable_args, OpKernelContext* ctx, + const XlaCompiler::Options& options, const NameAttrList& function, + int num_constant_args, const std::vector& variable_args, + OpKernelContext* ctx, const XlaCompiler::CompilationResult** compilation_result, xla::LocalExecutable** executable) { VLOG(1) << "XlaCompilationCache::Compile " << DebugString(); @@ -263,21 +268,18 @@ Status XlaCompilationCache::Compile( TF_RETURN_IF_ERROR( BuildArguments(num_constant_args, variable_args, ctx, &args)); - std::unique_ptr flr(NewFunctionLibraryRuntime( - compiler_.device_mgr(), ctx->env(), compiler_.device(), - TF_GRAPH_DEF_VERSION, - ctx->function_library()->GetFunctionLibraryDefinition(), - OptimizerOptions(), nullptr /* custom_kernel_creator */)); - + XlaCompiler compiler(options); entry->compiled = true; - entry->compilation_status = compiler_.CompileFunction( - flr.get(), function, args, &entry->compilation_result); + entry->compilation_status = + compiler.CompileFunction(XlaCompiler::CompileOptions(), function, args, + &entry->compilation_result); } *compilation_result = &entry->compilation_result; if (entry->compilation_status.ok() && executable) { if (entry->executable == nullptr && - !entry->compilation_result.computation.IsNull()) { - entry->compilation_status = compiler_.BuildExecutable( + !entry->compilation_result.computation->IsNull()) { + XlaCompiler compiler(options); + entry->compilation_status = compiler.BuildExecutable( entry->compilation_result, &entry->executable); } *executable = entry->executable.get(); diff --git a/tensorflow/compiler/jit/xla_compilation_cache.h b/tensorflow/compiler/jit/xla_compilation_cache.h index ff67e48d1a9a9f16881c2e141b23ce8c479aef50..4ffcb68a3220b2354a3542e4c2a4d3e000969e0b 100644 --- a/tensorflow/compiler/jit/xla_compilation_cache.h +++ b/tensorflow/compiler/jit/xla_compilation_cache.h @@ -46,7 +46,7 @@ struct OptionalTensor { // bound. class XlaCompilationCache : public ResourceBase { public: - explicit XlaCompilationCache(const XlaCompiler::Options& options); + XlaCompilationCache(xla::Client* client, DeviceType device_type); ~XlaCompilationCache() override; // Compiles a function into a XlaCompiler::CompilationResult that can be used @@ -61,19 +61,21 @@ class XlaCompilationCache : public ResourceBase { // xla::LocalExecutable and sets `executable to point to it. The resulting // executable pointer may be null if the computation has no non-constant // outputs. - Status Compile(const NameAttrList& function, int num_constant_args, + Status Compile(const XlaCompiler::Options& options, + const NameAttrList& function, int num_constant_args, const std::vector& variable_args, OpKernelContext* ctx, const XlaCompiler::CompilationResult** compilation_result, xla::LocalExecutable** executable); - xla::Client* client() const { return compiler_.client(); } + xla::Client* client() const { return client_; } + const DeviceType& device_type() const { return device_type_; } string DebugString() override; private: - XlaCompiler compiler_; - std::unique_ptr function_library_runtime_; + xla::Client* const client_; + const DeviceType device_type_; // Describes the types, shapes and any compile-time constant arguments // to a kernel. Key that uniquely identifies a compilation output. diff --git a/tensorflow/compiler/jit/xla_cpu_device.cc b/tensorflow/compiler/jit/xla_cpu_device.cc index e5153c2fe543810c2240b1abed3e5027d4412142..e8b1f542ecfde1f3e671cbbd9806e2068200c7c3 100644 --- a/tensorflow/compiler/jit/xla_cpu_device.cc +++ b/tensorflow/compiler/jit/xla_cpu_device.cc @@ -25,8 +25,6 @@ limitations under the License. namespace tensorflow { -const char* const DEVICE_XLA_CPU = "XLA_CPU"; - class XlaCpuDeviceFactory : public DeviceFactory { public: Status CreateDevices(const SessionOptions& options, const string& name_prefix, diff --git a/tensorflow/compiler/jit/xla_device.cc b/tensorflow/compiler/jit/xla_device.cc index 3c6793b89420ed61259070f7bf637d6f4aa097d0..615e2230f42f63f893ad645e1ab9513d6c30abf5 100644 --- a/tensorflow/compiler/jit/xla_device.cc +++ b/tensorflow/compiler/jit/xla_device.cc @@ -31,15 +31,18 @@ limitations under the License. #include "tensorflow/core/framework/allocator.h" #include "tensorflow/core/framework/device_base.h" #include "tensorflow/core/framework/function.h" +#include "tensorflow/core/framework/kernel_def.pb.h" #include "tensorflow/core/framework/node_def_builder.h" #include "tensorflow/core/framework/op_kernel.h" #include "tensorflow/core/framework/tensor.h" +#include "tensorflow/core/framework/tensor.pb.h" #include "tensorflow/core/framework/types.h" #include "tensorflow/core/graph/graph_constructor.h" #include "tensorflow/core/lib/core/notification.h" #include "tensorflow/core/lib/core/status.h" #include "tensorflow/core/platform/logging.h" #include "tensorflow/core/platform/stream_executor_no_cuda.h" +#include "tensorflow/core/platform/tracing.h" #include "tensorflow/core/public/session_options.h" #include "tensorflow/core/public/version.h" #include "tensorflow/core/util/device_name_utils.h" @@ -108,12 +111,23 @@ const DeviceType& XlaDevice::Metadata::jit_device_type() const { string XlaDevice::Metadata::DebugString() { return "XLA device metadata"; } +/* static */ Status XlaDevice::GetMetadata(OpKernelContext* ctx, + Metadata** metadata) { + ResourceMgr* rm = ctx->resource_manager(); + if (rm == nullptr) { + return errors::Internal("No resource manager."); + } + TF_RETURN_IF_ERROR( + rm->Lookup(rm->default_container(), "xla_metadata", metadata)); + return Status::OK(); +} + XlaDevice::XlaDevice(const SessionOptions& options, const DeviceAttributes& attrs, int device_ordinal, const DeviceType& jit_device_name, perftools::gputools::Platform* platform, Allocator* xla_allocator) - : LocalDevice(options, attrs, xla_allocator), + : LocalDevice(options, attrs), device_ordinal_(device_ordinal), jit_device_name_(jit_device_name), xla_allocator_(xla_allocator), @@ -163,6 +177,10 @@ Status XlaDevice::FillContextMap(const Graph* graph, void XlaDevice::Compute(OpKernel* op_kernel, OpKernelContext* context) { VLOG(1) << "XlaDevice::Compute " << op_kernel->name() << ":" << op_kernel->type_string(); + // When TraceMe profiling is off (which is the default), the + // following TraceMe constructor is simply a conditional test of + // false value. Measurements show that its overhead is negligible. + port::Tracing::TraceMe trace_me(op_kernel->name(), op_kernel->type_string()); op_kernel->Compute(context); } @@ -170,6 +188,7 @@ void XlaDevice::ComputeAsync(AsyncOpKernel* op_kernel, OpKernelContext* context, AsyncOpKernel::DoneCallback done) { VLOG(1) << "XlaDevice::ComputeAsync " << op_kernel->name() << ":" << op_kernel->type_string(); + port::Tracing::TraceMe trace_me(op_kernel->name(), op_kernel->type_string()); op_kernel->ComputeAsync(context, done); } diff --git a/tensorflow/compiler/jit/xla_device.h b/tensorflow/compiler/jit/xla_device.h index 3de14f306168937bb0483e0c442984a02e2b1442..0badb390c6b7785b36f58c786e1d32a8d10d7c29 100644 --- a/tensorflow/compiler/jit/xla_device.h +++ b/tensorflow/compiler/jit/xla_device.h @@ -67,6 +67,10 @@ class XlaDevice : public LocalDevice { perftools::gputools::Platform* platform_; // Not owned. }; + // Sets `*metadata` to the XlaDevice Metadata in the resource manager of + // `ctx`. + static Status GetMetadata(OpKernelContext* ctx, Metadata** metadata); + // Factory function. 'platform_name' is the name of the XLA platform. // 'device_name' is the name of the Tensorflow device to create. // 'jit_device_name' is the name of the corresponding JIT device. diff --git a/tensorflow/compiler/jit/xla_device_context.cc b/tensorflow/compiler/jit/xla_device_context.cc index f329e83e14dfce68eff3feb720c1603bd36fa7d6..0ab81ebd5ffec0b3dd6aee509a6d4d2b41d156db 100644 --- a/tensorflow/compiler/jit/xla_device_context.cc +++ b/tensorflow/compiler/jit/xla_device_context.cc @@ -137,7 +137,7 @@ void XlaTransferManager::CopyDeviceTensorToCPU(const Tensor* device_tensor, done(result.status()); return; } - const void* src_ptr = xla::LiteralUtil::InternalData(*result.ValueOrDie()); + const void* src_ptr = result.ValueOrDie()->InternalData(); void* dst_ptr = DMAHelper::base(cpu_tensor); size_t total_bytes = cpu_tensor->TotalBytes(); memcpy(dst_ptr, src_ptr, total_bytes); diff --git a/tensorflow/compiler/jit/xla_device_ops.h b/tensorflow/compiler/jit/xla_device_ops.h index a52239df252b2b556987fa9701f43047765c60de..8699006ebc5aacafd46046a7c3f093356f687280 100644 --- a/tensorflow/compiler/jit/xla_device_ops.h +++ b/tensorflow/compiler/jit/xla_device_ops.h @@ -63,30 +63,10 @@ class XlaDeviceDummyOp : public OpKernel { REGISTER_KERNEL_BUILDER(Name("PlaceholderV2").Device(DEVICE), \ PlaceholderOp); \ \ - REGISTER_KERNEL_BUILDER(Name("ControlTrigger").Device(DEVICE), \ - ControlTriggerOp); \ - REGISTER_KERNEL_BUILDER(Name("Enter").Device(DEVICE), EnterOp); \ - REGISTER_KERNEL_BUILDER(Name("Exit").Device(DEVICE), ExitOp); \ - REGISTER_KERNEL_BUILDER(Name("NextIteration").Device(DEVICE), \ - NextIterationOp); \ - REGISTER_KERNEL_BUILDER(Name("Switch").Device(DEVICE).HostMemory("pred"), \ - SwitchOp); \ - REGISTER_KERNEL_BUILDER( \ - Name("Merge").Device(DEVICE).HostMemory("value_index"), MergeOp); \ - REGISTER_KERNEL_BUILDER(Name("LoopCond") \ - .Device(DEVICE) \ - .HostMemory("input") \ - .HostMemory("output"), \ - IdentityOp); \ - \ REGISTER_KERNEL_BUILDER( \ Name("VarHandleOp").Device(DEVICE).HostMemory("resource"), \ ResourceHandleOp); -// TODO(b/32507444): the registrations for the control flow operators are -// temporary and exist primarily to work around a bug in the graph partitioning -// code. - } // namespace tensorflow #endif // TENSORFLOW_COMPILER_JIT_XLA_DEVICE_OPS_H_ diff --git a/tensorflow/compiler/jit/xla_gpu_device.cc b/tensorflow/compiler/jit/xla_gpu_device.cc index d7da30718d177205c7f5f1c0f2a5f0604a5fdb1b..872588a24e0b401007439100e40185f96dc20fb8 100644 --- a/tensorflow/compiler/jit/xla_gpu_device.cc +++ b/tensorflow/compiler/jit/xla_gpu_device.cc @@ -25,8 +25,6 @@ limitations under the License. namespace tensorflow { -const char* const DEVICE_XLA_GPU = "XLA_GPU"; - class XlaGpuDeviceFactory : public DeviceFactory { public: Status CreateDevices(const SessionOptions& options, const string& name_prefix, diff --git a/tensorflow/compiler/plugin/BUILD b/tensorflow/compiler/plugin/BUILD new file mode 100644 index 0000000000000000000000000000000000000000..8c2e9a7c818c735370921b52a7da6db34186b1b9 --- /dev/null +++ b/tensorflow/compiler/plugin/BUILD @@ -0,0 +1,38 @@ +# 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. +# ============================================================================== + +"""Configuration file for an XLA plugin. +- please don't check in changes to this file +- to prevent changes appearing in git status, use: + git update-index --assume-unchanged tensorflow/compiler/plugin/BUILD + +To add additional devices to the XLA subsystem, add targets to the +dependency list in the 'plugin' target. For instance: + + deps = ["//tensorflow/compiler/plugin/example:plugin_lib"], +""" + +licenses(["notice"]) + +package( + default_visibility = ["//visibility:public"], +) + +cc_library( + name = "plugin", + deps = [ + "//tensorflow/compiler/plugin/executor:plugin_lib", + ], +) diff --git a/tensorflow/compiler/plugin/executor/BUILD b/tensorflow/compiler/plugin/executor/BUILD new file mode 100644 index 0000000000000000000000000000000000000000..ffecd68d921b287f55bdb6ba8eac865e7a9e9fd8 --- /dev/null +++ b/tensorflow/compiler/plugin/executor/BUILD @@ -0,0 +1,37 @@ +licenses(["restricted"]) + +package(default_visibility = ["//visibility:public"]) + +cc_library( + name = "plugin_lib", + srcs = glob([ + "*.cc", + ]), + hdrs = glob([ + "*.h", + ]), + deps = [ + "//tensorflow/compiler/jit:xla_device", + "//tensorflow/compiler/jit:xla_jit_headers_lib", + "//tensorflow/compiler/tf2xla:xla_compiler", + "//tensorflow/compiler/xla:xla_headers_lib", + "//tensorflow/compiler/xla/service", + "//tensorflow/compiler/xla/service:computation_placer", + "//tensorflow/compiler/xla/service:layout_assignment", + "//third_party/eigen3", + "@local_config_cuda//cuda:cuda_headers", + "@protobuf_archive//:protobuf_headers", + ], + alwayslink = 1, +) + +filegroup( + name = "all_files", + srcs = glob( + ["**/*"], + exclude = [ + "**/METADATA", + "**/OWNERS", + ], + ), +) diff --git a/tensorflow/compiler/plugin/executor/compiler.cc b/tensorflow/compiler/plugin/executor/compiler.cc new file mode 100644 index 0000000000000000000000000000000000000000..77193f06c4b6557498aee6a957db56dbd1e3659d --- /dev/null +++ b/tensorflow/compiler/plugin/executor/compiler.cc @@ -0,0 +1,127 @@ +/* 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/compiler/plugin/executor/compiler.h" +#include "tensorflow/compiler/plugin/executor/executable.h" +#include "tensorflow/compiler/xla/service/algebraic_simplifier.h" +#include "tensorflow/compiler/xla/service/computation_placer.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" +#include "tensorflow/compiler/xla/service/hlo_dce.h" +#include "tensorflow/compiler/xla/service/hlo_pass_fix.h" +#include "tensorflow/compiler/xla/service/hlo_pass_pipeline.h" +#include "tensorflow/compiler/xla/service/hlo_subcomputation_unification.h" +#include "tensorflow/compiler/xla/service/inliner.h" +#include "tensorflow/compiler/xla/service/layout_assignment.h" +#include "tensorflow/compiler/xla/service/reshape_mover.h" +#include "tensorflow/compiler/xla/status_macros.h" +#include "tensorflow/core/lib/core/errors.h" +#include "tensorflow/stream_executor/lib/initialize.h" +#include "tensorflow/stream_executor/lib/strcat.h" + +namespace xla { +namespace executorplugin { + +namespace se = ::perftools::gputools; +namespace sep = ::perftools::gputools::executorplugin; + +/* + * Run optimization passes on the module. The graph is transformed by + * each pass in the optimization pipeline. The service subdirectory + * contains useful optimization passes. + */ +Status ExecutorCompiler::RunHloOptimization(HloModule* hlo_module) { + HloPassPipeline pipeline("Executor"); + pipeline.AddPass(); + pipeline.AddPass(); + pipeline.AddPass(false); + + pipeline.AddPass>( + false, [](const Shape&, const Shape&) { return false; }); + pipeline.AddPass(); + pipeline.AddPass(); + pipeline.AddPass(true); + pipeline.AddPass( + hlo_module->mutable_entry_computation_layout()); + + pipeline.AddPass(); + pipeline.AddPass(); + return pipeline.Run(hlo_module).status(); +} + +StatusOr> ExecutorCompiler::Compile( + std::unique_ptr hlo_module, + se::StreamExecutor* stream_exec) { + TF_RET_CHECK(stream_exec != nullptr); + + VLOG(1) << "Generate graph " << hlo_module->name(); + + TF_RETURN_IF_ERROR(RunHloOptimization(hlo_module.get())); + + // Typically you would visit the HLO graph, building up a compiled equivalent + // In this case we are using an Hlo evaluator at execution time, so we don't + // need to compile anything + + // Create executable from only the Hlo module + std::unique_ptr executable; + executable.reset(new ExecutorExecutable(std::move(hlo_module))); + + return std::move(executable); +} + +StatusOr>> ExecutorCompiler::Compile( + std::vector> hlo_modules, + std::vector stream_execs) { + + return tensorflow::errors::Unimplemented( + "Compilation of multiple HLO modules is not supported on Executor."); +} + +StatusOr>> +ExecutorCompiler::CompileAheadOfTime( + std::vector> hlo_modules, + const AotCompilationOptions& aot_options) { + + return tensorflow::errors::InvalidArgument( + "AOT compilation not supported on Executor"); +} + +se::Platform::Id ExecutorCompiler::PlatformId() const { + return sep::kExecutorPlatformId; +} + +HloCostAnalysis::ShapeSizeFunction +ExecutorCompiler::ShapeSizeBytesFunction() const { + return ExecutorExecutable::ShapeSizeBytes; +} + +static std::unique_ptr CreateComputationPlacer() { + return xla::MakeUnique(); +} + +REGISTER_MODULE_INITIALIZER(executor_compiler, { + xla::Compiler::RegisterCompilerFactory(sep::kExecutorPlatformId, []() { + return xla::MakeUnique(); + }); + xla::ComputationPlacer::RegisterComputationPlacer(sep::kExecutorPlatformId, + &CreateComputationPlacer); +}); + +} // namespace executorplugin +} // namespace xla diff --git a/tensorflow/compiler/plugin/executor/compiler.h b/tensorflow/compiler/plugin/executor/compiler.h new file mode 100644 index 0000000000000000000000000000000000000000..d318eefc49f0f1983cf58802d56e71b799944b11 --- /dev/null +++ b/tensorflow/compiler/plugin/executor/compiler.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_EXECUTOR_COMPILER_H_ +#define TENSORFLOW_COMPILER_EXECUTOR_COMPILER_H_ + +#include + +#include "tensorflow/compiler/xla/service/compiler.h" +#include "tensorflow/compiler/xla/service/executable.h" +#include "tensorflow/compiler/xla/service/hlo_module.h" +#include "tensorflow/compiler/xla/service/hlo_module_config.h" + +#include "tensorflow/compiler/plugin/executor/platform_id.h" + +namespace xla { +namespace executorplugin { + +class ExecutorCompiler : public Compiler { + public: + ExecutorCompiler() {} + ~ExecutorCompiler() override {} + + StatusOr> Compile( + std::unique_ptr hlo_module, + perftools::gputools::StreamExecutor* stream_exec) override; + + StatusOr>> Compile( + std::vector> hlo_module, + std::vector stream_exec) override; + + StatusOr>> + CompileAheadOfTime( + std::vector> module, + const AotCompilationOptions& options) override; + + HloCostAnalysis::ShapeSizeFunction ShapeSizeBytesFunction() const override; + + perftools::gputools::Platform::Id PlatformId() const override; + + private: + Status RunHloOptimization(HloModule* hlo_module); + + TF_DISALLOW_COPY_AND_ASSIGN(ExecutorCompiler); +}; + +} // namespace executorplugin +} // namespace xla + +#endif // TENSORFLOW_COMPILER_EXECUTOR_COMPILER_H_ diff --git a/tensorflow/compiler/plugin/executor/device.cc b/tensorflow/compiler/plugin/executor/device.cc new file mode 100644 index 0000000000000000000000000000000000000000..d902f9df6a50161dacf12a5b234c1304ead353d5 --- /dev/null +++ b/tensorflow/compiler/plugin/executor/device.cc @@ -0,0 +1,65 @@ +/* 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/jit/kernels/xla_device_launch_op.h" +#include "tensorflow/compiler/jit/xla_device.h" +#include "tensorflow/compiler/jit/xla_device_ops.h" +#include "tensorflow/compiler/tf2xla/xla_op_registry.h" + +namespace tensorflow { + +const char* const DEVICE_XLA_EXEC = "XLA_EXEC"; +const char* const DEVICE_EXEC_XLA_JIT = "XLA_EXEC_JIT"; + +constexpr std::array kExecAllTypes = { + {DT_INT32, DT_FLOAT, DT_BOOL, DT_DOUBLE, DT_INT64}}; + +class XlaExaDeviceFactory : public DeviceFactory { + public: + Status CreateDevices(const SessionOptions& options, const string& name_prefix, + std::vector* devices) override; +}; + +Status XlaExaDeviceFactory::CreateDevices(const SessionOptions& options, + const string& name_prefix, + std::vector* devices) { + static XlaDeviceOpRegistrations* registrations = + RegisterXlaDeviceKernels(DEVICE_XLA_EXEC, DEVICE_EXEC_XLA_JIT); + (void)registrations; + + std::unique_ptr device; + TF_RETURN_IF_ERROR(XlaDevice::Create("Executor", DEVICE_XLA_EXEC, 0, + DEVICE_EXEC_XLA_JIT, options, + name_prefix, &device)); + devices->push_back(device.release()); + return Status::OK(); +} + +// Set priority to be below the default priority (50), so that Executor is not +// selected as a high priority device over other default devices. +// See constructor comments for Registrar in +// tensorflow/core/common_runtime/device_factory.h for a list of priority for +// devices. +REGISTER_LOCAL_DEVICE_FACTORY(DEVICE_XLA_EXEC, XlaExaDeviceFactory, 40); + +// Kernel registrations + +static bool OpFilter(KernelDef* kdef) { return true; } + +REGISTER_XLA_LAUNCH_KERNEL(DEVICE_XLA_EXEC, XlaDeviceLaunchOp, kExecAllTypes); +REGISTER_XLA_DEVICE_KERNELS(DEVICE_XLA_EXEC, kExecAllTypes); +REGISTER_XLA_BACKEND(DEVICE_EXEC_XLA_JIT, kExecAllTypes, OpFilter); + +} // namespace tensorflow diff --git a/tensorflow/compiler/plugin/executor/executable.cc b/tensorflow/compiler/plugin/executor/executable.cc new file mode 100644 index 0000000000000000000000000000000000000000..e866dfef059b5c1e8af0ebc7457aad19231cfe45 --- /dev/null +++ b/tensorflow/compiler/plugin/executor/executable.cc @@ -0,0 +1,149 @@ +/* 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/plugin/executor/executable.h" +#include "tensorflow/compiler/plugin/executor/executor.h" +#include "tensorflow/compiler/xla/literal_util.h" +#include "tensorflow/compiler/xla/service/hlo_evaluator.h" +#include "tensorflow/compiler/xla/shape_util.h" + +namespace xla { +namespace executorplugin { + +namespace se = ::perftools::gputools; +namespace sep = ::perftools::gputools::executorplugin; + +ExecutorExecutable::ExecutorExecutable(std::unique_ptr hlo_module) + : Executable(std::move(hlo_module)) {} + +ExecutorExecutable::~ExecutorExecutable() {} + +static se::DeviceMemoryBase AllocateSingleOutput( + sep::ExecutorExecutor* executor, const Literal& literal) { + int64 size(xla::ShapeUtil::ByteSizeOf(literal.shape())); + void* buf = executor->Allocate(size); + const void* src = literal.InternalData(); + memcpy(buf, src, size); + return se::DeviceMemoryBase(buf, size); +} + +static se::DeviceMemoryBase AllocateOutputBuffer( + sep::ExecutorExecutor* executor, const Literal& literal) { + const Shape& shape = literal.shape(); + if (shape.element_type() != xla::TUPLE) { + return AllocateSingleOutput(executor, literal); + } else { + int64 size(xla::ShapeUtil::ByteSizeOf(shape, sizeof(void*))); + void** buf = reinterpret_cast(executor->Allocate(size)); + void** buf_rc = buf; + for (int64 n = 0; n < xla::ShapeUtil::TupleElementCount(shape); n++) { + se::DeviceMemoryBase out = + AllocateSingleOutput(executor, literal.tuple_literals(n)); + *buf++ = out.opaque(); + } + + return se::DeviceMemoryBase(buf_rc, size); + } +} + +StatusOr ExecutorExecutable::ExecuteOnStream( + const ServiceExecutableRunOptions* run_options, + tensorflow::gtl::ArraySlice arguments, + HloExecutionProfile* hlo_execution_profile) { + se::Stream* stream = run_options->stream(); + + VLOG(1) << "Execute " << module().name(); + if (VLOG_IS_ON(2)) { + for (const auto& a : arguments) { + VLOG(2) << "-- argument " << a.opaque(); + } + } + + uint64 start_micros = tensorflow::Env::Default()->NowMicros(); + + HloComputation* computation = module().entry_computation(); + if (computation->num_parameters() != arguments.size()) { + return tensorflow::errors::Internal( + "Mismatch between argument count and graph parameter count."); + } + + // Create the arguments as an vector of XLA literals + std::vector> arg_literals; + std::vector arg_literals_ptrs; + for (int64 p = 0; p < computation->num_parameters(); p++) { + // Create the input literal for the parameter + HloInstruction* param = computation->parameter_instruction(p); + arg_literals.emplace_back(Literal::CreateFromShape(param->shape())); + arg_literals_ptrs.push_back(arg_literals.back().get()); + + // Copy in the data from the stream_executor buffers + void* buffer = arg_literals.back()->MutableInternalData(); + memcpy(buffer, arguments[p].opaque(), + ShapeUtil::ByteSizeOf(param->shape())); + } + + // Execute the graph using the evaluator + HloEvaluator evaluator; + TF_ASSIGN_OR_RETURN(std::unique_ptr output, + evaluator.Evaluate(*computation, arg_literals_ptrs)); + + // Copy the result into the return buffer + perftools::gputools::StreamExecutor* executor(stream->parent()); + sep::ExecutorExecutor* executorExecutor( + static_cast(executor->implementation())); + + se::DeviceMemoryBase ret = + AllocateOutputBuffer(executorExecutor, *(output.get())); + + uint64 end_micros = tensorflow::Env::Default()->NowMicros(); + + { + tensorflow::mutex_lock lock(mutex_); + const double nanoseconds = (end_micros - start_micros) * 1000.0; + execution_profile_.set_compute_time_ns(std::max(nanoseconds, 1.0)); + } + + return ret; +} + +StatusOr> ExecutorExecutable::ExecuteOnStream( + const ServiceExecutableRunOptions* run_options, + tensorflow::gtl::ArraySlice arguments, + HloExecutionProfile* hlo_execution_profile) { + return tensorflow::errors::Unimplemented( + "ExecuteOnStream is not yet supported on Executor."); +} + +StatusOr ExecutorExecutable::ExecuteAsyncOnStream( + const ServiceExecutableRunOptions* run_options, + tensorflow::gtl::ArraySlice arguments) { + return tensorflow::errors::Unimplemented( + "ExecuteAsyncOnStream is not yet supported on Executor."); +} + +/*static*/ int64 ExecutorExecutable::ShapeSizeBytes(const Shape& shape) { + if (ShapeUtil::IsOpaque(shape)) { + return sizeof(void*); + } + return ShapeUtil::ByteSizeOf(shape, sizeof(void*)); +} + +std::unique_ptr ExecutorExecutable::CreateCostAnalysis() + const { + return MakeUnique(ShapeSizeBytes); +} + +} // namespace executorplugin +} // namespace xla diff --git a/tensorflow/compiler/plugin/executor/executable.h b/tensorflow/compiler/plugin/executor/executable.h new file mode 100644 index 0000000000000000000000000000000000000000..8572aba43d9a18604422eaae683f78c3cf056292 --- /dev/null +++ b/tensorflow/compiler/plugin/executor/executable.h @@ -0,0 +1,67 @@ +/* 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_EXECUTOR_DRIVER_EXECUTOR_EXECUTABLE_H_ +#define TENSORFLOW_COMPILER_EXECUTOR_DRIVER_EXECUTOR_EXECUTABLE_H_ + +#include +#include +#include +#include +#include + +#include "tensorflow/compiler/xla/service/executable.h" +#include "tensorflow/compiler/xla/service/hlo_module.h" +#include "tensorflow/compiler/xla/service/hlo_module_config.h" + +#include "tensorflow/stream_executor/lib/status.h" +#include "tensorflow/stream_executor/lib/statusor.h" + +namespace xla { +namespace executorplugin { + +class ExecutorExecutable : public Executable { + public: + ExecutorExecutable(std::unique_ptr hlo_module); + ~ExecutorExecutable() override; + + StatusOr ExecuteOnStream( + const ServiceExecutableRunOptions* run_options, + tensorflow::gtl::ArraySlice + arguments, + HloExecutionProfile* hlo_execution_profile) override; + + StatusOr> ExecuteOnStream( + const ServiceExecutableRunOptions* run_options, + tensorflow::gtl::ArraySlice arguments, + HloExecutionProfile* hlo_execution_profile) override; + + StatusOr ExecuteAsyncOnStream( + const ServiceExecutableRunOptions* run_options, + tensorflow::gtl::ArraySlice + arguments) override; + + static int64 ShapeSizeBytes(const Shape& shape); + + std::unique_ptr CreateCostAnalysis() const override; + + private: + TF_DISALLOW_COPY_AND_ASSIGN(ExecutorExecutable); +}; + +} // namespace executorplugin +} // namespace xla + +#endif // TENSORFLOW_COMPILER_EXECUTOR_DRIVER_EXECUTOR_EXECUTABLE_H_ diff --git a/tensorflow/compiler/plugin/executor/executor.cc b/tensorflow/compiler/plugin/executor/executor.cc new file mode 100644 index 0000000000000000000000000000000000000000..908b996bc95ac8d36f6c5577857b1a3a3826c3d4 --- /dev/null +++ b/tensorflow/compiler/plugin/executor/executor.cc @@ -0,0 +1,129 @@ +/* Copyright 2017 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "tensorflow/compiler/plugin/executor/executor.h" + +#include +#include + +#include "tensorflow/compiler/plugin/executor/platform_id.h" +#include "tensorflow/compiler/xla/status_macros.h" + +namespace perftools { +namespace gputools { +namespace executorplugin { + +host::HostStream *AsExecutorStream(Stream *stream) { + DCHECK(stream != nullptr); + return dynamic_cast(stream->implementation()); +} + +ExecutorExecutor::ExecutorExecutor(const PluginConfig &plugin_config) + : plugin_config_(plugin_config) {} + +ExecutorExecutor::~ExecutorExecutor() {} + +void *ExecutorExecutor::Allocate(uint64 size) { return new char[size]; } + +void *ExecutorExecutor::AllocateSubBuffer(DeviceMemoryBase *parent, + uint64 offset_bytes, + uint64 size_bytes) { + return parent + offset_bytes; +} + +void ExecutorExecutor::Deallocate(DeviceMemoryBase *mem) { + if (!mem->is_sub_buffer()) { + delete[] static_cast(mem->opaque()); + } +} + +bool ExecutorExecutor::Memcpy(Stream *stream, void *host_dst, + const DeviceMemoryBase &dev_src, uint64 size) { + AsExecutorStream(stream)->EnqueueTask([this, host_dst, dev_src, size]() { + port::Status ok = SynchronousMemcpy(host_dst, dev_src, size); + }); + return true; +} + +bool ExecutorExecutor::Memcpy(Stream *stream, DeviceMemoryBase *dev_dst, + const void *host_src, uint64 size) { + AsExecutorStream(stream)->EnqueueTask([this, dev_dst, host_src, size]() { + port::Status ok = SynchronousMemcpy(dev_dst, host_src, size); + }); + return true; +} + +port::Status ExecutorExecutor::SynchronousMemcpy(DeviceMemoryBase *dev_dst, + const void *host_src, + uint64 size) { + memcpy(dev_dst->opaque(), host_src, size); + return port::Status::OK(); +} + +port::Status ExecutorExecutor::SynchronousMemcpy(void *host_dst, + const DeviceMemoryBase &dev_src, + uint64 size) { + memcpy(host_dst, dev_src.opaque(), size); + return port::Status::OK(); +} + +bool ExecutorExecutor::HostCallback(Stream *stream, + std::function callback) { + AsExecutorStream(stream)->EnqueueTask(callback); + return true; +} + +bool ExecutorExecutor::CreateStreamDependency(Stream *dependent, Stream *other) { + AsExecutorStream(dependent)->EnqueueTask( + [other]() { other->BlockHostUntilDone(); }); + AsExecutorStream(dependent)->BlockUntilDone(); + return true; +} + +bool ExecutorExecutor::StartTimer(Stream *stream, Timer *timer) { + dynamic_cast(timer->implementation())->Start(stream); + return true; +} + +bool ExecutorExecutor::StopTimer(Stream *stream, Timer *timer) { + dynamic_cast(timer->implementation())->Stop(stream); + return true; +} + +bool ExecutorExecutor::BlockHostUntilDone(Stream *stream) { + AsExecutorStream(stream)->BlockUntilDone(); + return true; +} + +DeviceDescription *ExecutorExecutor::PopulateDeviceDescription() const { + internal::DeviceDescriptionBuilder builder; + + builder.set_device_address_bits(64); + + builder.set_name("Executor"); + builder.set_device_vendor("VectorName"); + builder.set_platform_version("1.0"); + builder.set_driver_version("1.0"); + builder.set_runtime_version("1.0"); + builder.set_pci_bus_id("1"); + builder.set_device_memory_size(static_cast(4) * 1024 * 1024 * 1024); + builder.set_clock_rate_ghz(static_cast(CLOCKS_PER_SEC) / 1e9); + + return builder.Build().release(); +} + +} // namespace executorplugin +} // namespace gputools +} // namespace perftools diff --git a/tensorflow/compiler/plugin/executor/executor.h b/tensorflow/compiler/plugin/executor/executor.h new file mode 100644 index 0000000000000000000000000000000000000000..32fdb157e48edbeef019c75d7422f7d7f507e721 --- /dev/null +++ b/tensorflow/compiler/plugin/executor/executor.h @@ -0,0 +1,213 @@ +/* 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. +==============================================================================*/ + +// Declares the ExecutorExecutor class, which is a CPU-only implementation of +// the StreamExecutor interface. For now, this is used for testing and to +// examine the performance of host-based StreamExecutor code. +#ifndef TENSORFLOW_COMPILER_EXECUTOR_STREAM_EXECUTOR_EXECUTOR_EXECUTOR_H_ +#define TENSORFLOW_COMPILER_EXECUTOR_STREAM_EXECUTOR_EXECUTOR_EXECUTOR_H_ + +#include "tensorflow/stream_executor/host/host_stream.h" +#include "tensorflow/stream_executor/host/host_timer.h" + +#include "tensorflow/compiler/xla/shape_util.h" + +#include "tensorflow/stream_executor/blas.h" +#include "tensorflow/stream_executor/lib/error.h" +#include "tensorflow/stream_executor/lib/status.h" +#include "tensorflow/stream_executor/lib/statusor.h" +#include "tensorflow/stream_executor/rng.h" +#include "tensorflow/stream_executor/stream_executor.h" +#include "tensorflow/stream_executor/stream_executor_internal.h" + +#include +#include + +namespace perftools { +namespace gputools { +namespace executorplugin { + +using Args = tensorflow::gtl::ArraySlice; + +class ExecutorExecutor : public internal::StreamExecutorInterface { + public: + explicit ExecutorExecutor(const PluginConfig &plugin_config); + ~ExecutorExecutor() override; + + port::Status Init(int device_ordinal, DeviceOptions device_options) override { + return port::Status::OK(); + } + + bool GetKernel(const MultiKernelLoaderSpec &spec, + KernelBase *kernel) override { + return false; + } + bool Launch(Stream *stream, const ThreadDim &thread_dims, + const BlockDim &block_dims, const KernelBase &kernel, + const KernelArgsArrayBase &args) override { + return false; + } + + void *Allocate(uint64 size) override; + void *AllocateSubBuffer(DeviceMemoryBase *mem, uint64 offset_bytes, + uint64 size_bytes) override; + void Deallocate(DeviceMemoryBase *mem) override; + + void *HostMemoryAllocate(uint64 size) override { return new char[size]; } + void HostMemoryDeallocate(void *mem) override { + delete[] static_cast(mem); + } + bool HostMemoryRegister(void *mem, uint64 size) override { return true; } + bool HostMemoryUnregister(void *mem) override { return true; } + + bool Memcpy(Stream *stream, void *host_dst, const DeviceMemoryBase &pop_src, + uint64 size) override; + bool Memcpy(Stream *stream, DeviceMemoryBase *pop_dst, const void *host_src, + uint64 size) override; + bool MemcpyDeviceToDevice(Stream *stream, DeviceMemoryBase *pop_dst, + const DeviceMemoryBase &host_src, + uint64 size) override { + return false; + } + + bool MemZero(Stream *stream, DeviceMemoryBase *location, + uint64 size) override { + return false; + } + bool Memset(Stream *stream, DeviceMemoryBase *location, uint8 pattern, + uint64 size) override { + return false; + } + bool Memset32(Stream *stream, DeviceMemoryBase *location, uint32 pattern, + uint64 size) override { + return false; + } + + // No "synchronize all activity" implemented for this platform at the moment. + bool SynchronizeAllActivity() override { return false; } + bool SynchronousMemZero(DeviceMemoryBase *location, uint64 size) override { + return false; + } + + bool SynchronousMemSet(DeviceMemoryBase *location, int value, + uint64 size) override { + return false; + } + + port::Status SynchronousMemcpy(DeviceMemoryBase *pop_dst, + const void *host_src, uint64 size) override; + port::Status SynchronousMemcpy(void *host_dst, + const DeviceMemoryBase &pop_src, + uint64 size) override; + port::Status SynchronousMemcpyDeviceToDevice(DeviceMemoryBase *pop_dst, + const DeviceMemoryBase &pop_src, + uint64 size) override { + return port::Status{port::error::UNIMPLEMENTED, ""}; + } + + bool HostCallback(Stream *stream, std::function callback) override; + + port::Status AllocateEvent(Event *event) override { + return port::Status{port::error::UNIMPLEMENTED, ""}; + } + + port::Status DeallocateEvent(Event *event) override { + return port::Status{port::error::UNIMPLEMENTED, ""}; + } + + port::Status RecordEvent(Stream *stream, Event *event) override { + return port::Status{port::error::UNIMPLEMENTED, ""}; + } + + port::Status WaitForEvent(Stream *stream, Event *event) override { + return port::Status{port::error::UNIMPLEMENTED, ""}; + } + + Event::Status PollForEventStatus(Event *event) override { + return Event::Status::kError; + } + + bool AllocateStream(Stream *stream) override { return true; } + void DeallocateStream(Stream *stream) override {} + bool CreateStreamDependency(Stream *dependent, Stream *other) override; + + bool AllocateTimer(Timer *timer) override { return true; } + void DeallocateTimer(Timer *timer) override {} + bool StartTimer(Stream *stream, Timer *timer) override; + bool StopTimer(Stream *stream, Timer *timer) override; + + bool BlockHostUntilDone(Stream *stream) override; + + int PlatformDeviceCount() override { return 1; } + + bool DeviceMemoryUsage(int64 *free, int64 *total) const override { + return false; + } + + DeviceDescription *PopulateDeviceDescription() const override; + + port::Status EnablePeerAccessTo(StreamExecutorInterface *other) override { + return port::Status::OK(); + } + + bool CanEnablePeerAccessTo(StreamExecutorInterface *other) override { + return true; + } + + SharedMemoryConfig GetDeviceSharedMemoryConfig() override { + return SharedMemoryConfig::kDefault; + } + + port::Status SetDeviceSharedMemoryConfig(SharedMemoryConfig config) override { + return port::Status{port::error::UNIMPLEMENTED, + "Shared memory not supported"}; + } + + std::unique_ptr CreateEventImplementation() + override { + return nullptr; + } + + std::unique_ptr CreateKernelImplementation() + override { + return nullptr; + } + + std::unique_ptr GetStreamImplementation() + override { + return std::unique_ptr(new host::HostStream()); + } + + std::unique_ptr GetTimerImplementation() override { + return std::unique_ptr(new host::HostTimer()); + } + + port::StatusOr ExecuteGraph(const xla::Shape &shape, + Args args); + + private: + DeviceMemoryBase AllocateSingleOutput(const xla::Shape &shape); + + port::StatusOr AllocateOutputBuffer( + const xla::Shape &shape); + + const PluginConfig plugin_config_; +}; + +} // namespace executorplugin +} // namespace gputools +} // namespace perftools + +#endif // TENSORFLOW_COMPILER_EXECUTOR_STREAM_EXECUTOR_EXECUTOR_EXECUTOR_H_ diff --git a/tensorflow/compiler/plugin/executor/platform.cc b/tensorflow/compiler/plugin/executor/platform.cc new file mode 100644 index 0000000000000000000000000000000000000000..404e1c3da349107a74228ed8e975a9fdf9597019 --- /dev/null +++ b/tensorflow/compiler/plugin/executor/platform.cc @@ -0,0 +1,110 @@ +/* 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/plugin/executor/platform.h" +#include "tensorflow/compiler/plugin/executor/executor.h" +#include "tensorflow/compiler/plugin/executor/platform_id.h" + +#include "tensorflow/stream_executor/lib/error.h" +#include "tensorflow/stream_executor/lib/initialize.h" +#include "tensorflow/stream_executor/lib/ptr_util.h" +#include "tensorflow/stream_executor/lib/status.h" +#include "tensorflow/stream_executor/lib/status_macros.h" +#include "tensorflow/stream_executor/lib/stringprintf.h" + +namespace se = ::perftools::gputools; +namespace sep = ::perftools::gputools::executorplugin; + +namespace perftools { +namespace gputools { +namespace executorplugin { + +PLATFORM_DEFINE_ID(kExecutorPlatformId); + +ExecutorPlatform::ExecutorPlatform() : name_("Executor") {} + +ExecutorPlatform::~ExecutorPlatform() {} + +Platform::Id ExecutorPlatform::id() const { return kExecutorPlatformId; } + +int ExecutorPlatform::VisibleDeviceCount() const { return 1; } + +const string& ExecutorPlatform::Name() const { return name_; } + +port::StatusOr ExecutorPlatform::ExecutorForDevice( + int ordinal) { + StreamExecutorConfig config; + config.ordinal = ordinal; + config.plugin_config = PluginConfig(); + config.device_options = DeviceOptions::Default(); + return GetExecutor(config); +} + +port::StatusOr +ExecutorPlatform::ExecutorForDeviceWithPluginConfig( + int device_ordinal, const PluginConfig& plugin_config) { + StreamExecutorConfig config; + config.ordinal = device_ordinal; + config.plugin_config = plugin_config; + config.device_options = DeviceOptions::Default(); + return GetExecutor(config); +} + +port::StatusOr ExecutorPlatform::GetExecutor( + const StreamExecutorConfig& config) { + return executor_cache_.GetOrCreate( + config, [&]() { return GetUncachedExecutor(config); }); +} + +port::StatusOr> +ExecutorPlatform::GetUncachedExecutor(const StreamExecutorConfig& config) { + auto executor = port::MakeUnique( + this, port::MakeUnique(config.plugin_config)); + auto init_status = executor->Init(config.ordinal, config.device_options); + if (!init_status.ok()) { + return port::Status{ + port::error::INTERNAL, + port::Printf( + "failed initializing StreamExecutor for device ordinal %d: %s", + config.ordinal, init_status.ToString().c_str())}; + } + + return std::move(executor); +} + +void ExecutorPlatform::RegisterTraceListener( + std::unique_ptr listener) { + LOG(FATAL) << "not yet implemented: register executor trace listener"; +} + +void ExecutorPlatform::UnregisterTraceListener(TraceListener* listener) { + LOG(FATAL) << "not yet implemented: unregister executor trace listener"; +} + +static void InitializeExecutorPlatform() { + std::unique_ptr platform(new sep::ExecutorPlatform); + SE_CHECK_OK(se::MultiPlatformManager::RegisterPlatform(std::move(platform))); +} + +} // namespace executorplugin +} // namespace gputools +} // namespace perftools + +REGISTER_MODULE_INITIALIZER(executor_platform, sep::InitializeExecutorPlatform()); + +DECLARE_MODULE_INITIALIZER(multi_platform_manager); +// Note that module initialization sequencing is not supported in the +// open-source project, so this will be a no-op there. +REGISTER_MODULE_INITIALIZER_SEQUENCE(executor_platform, multi_platform_manager); diff --git a/tensorflow/compiler/plugin/executor/platform.h b/tensorflow/compiler/plugin/executor/platform.h new file mode 100644 index 0000000000000000000000000000000000000000..624bcd5a4eb5878962ab90cf1a57a5909a8f5dff --- /dev/null +++ b/tensorflow/compiler/plugin/executor/platform.h @@ -0,0 +1,80 @@ +/* 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_EXECUTOR_STREAM_EXECUTOR_EXECUTOR_PLATFORM_H_ +#define TENSORFLOW_COMPILER_EXECUTOR_STREAM_EXECUTOR_EXECUTOR_PLATFORM_H_ + +#include +#include +#include + +#include "tensorflow/stream_executor/executor_cache.h" +#include "tensorflow/stream_executor/lib/statusor.h" +#include "tensorflow/stream_executor/multi_platform_manager.h" +#include "tensorflow/stream_executor/platform.h" +#include "tensorflow/stream_executor/platform/mutex.h" +#include "tensorflow/stream_executor/platform/port.h" +#include "tensorflow/stream_executor/platform/thread_annotations.h" +#include "tensorflow/stream_executor/stream_executor_pimpl.h" +#include "tensorflow/stream_executor/trace_listener.h" + +namespace perftools { +namespace gputools { +namespace executorplugin { + +class ExecutorPlatform : public Platform { + public: + ExecutorPlatform(); + ~ExecutorPlatform() override; + + Platform::Id id() const override; + + // Device count is less clear-cut for CPUs than accelerators. This call + // currently returns the number of thread units in the host, as reported by + // base::NumCPUs(). + int VisibleDeviceCount() const override; + + const string& Name() const override; + + port::StatusOr ExecutorForDevice(int ordinal) override; + + port::StatusOr ExecutorForDeviceWithPluginConfig( + int ordinal, const PluginConfig& config) override; + + port::StatusOr GetExecutor( + const StreamExecutorConfig& config) override; + + port::StatusOr> GetUncachedExecutor( + const StreamExecutorConfig& config) override; + + void RegisterTraceListener(std::unique_ptr listener) override; + + void UnregisterTraceListener(TraceListener* listener) override; + + private: + // This platform's name. + string name_; + + // Cache of created StreamExecutors. + ExecutorCache executor_cache_; + + SE_DISALLOW_COPY_AND_ASSIGN(ExecutorPlatform); +}; + +} // namespace executorplugin +} // namespace gputools +} // namespace perftools + +#endif // TENSORFLOW_COMPILER_EXECUTOR_STREAM_EXECUTOR_EXECUTOR_PLATFORM_H_ diff --git a/tensorflow/compiler/plugin/executor/platform_id.h b/tensorflow/compiler/plugin/executor/platform_id.h new file mode 100644 index 0000000000000000000000000000000000000000..8d2b29a3e4e260b86b19d286e13f4066b7037ec4 --- /dev/null +++ b/tensorflow/compiler/plugin/executor/platform_id.h @@ -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. +==============================================================================*/ + +#ifndef TENSORFLOW_STREAM_EXECUTOR_EXECUTOR_PLATFORM_ID_H_ +#define TENSORFLOW_STREAM_EXECUTOR_EXECUTOR_PLATFORM_ID_H_ + +#include "tensorflow/stream_executor/platform.h" + +namespace perftools { +namespace gputools { +namespace executorplugin { + +extern const Platform::Id kExecutorPlatformId; + +} // namespace executorplugin +} // namespace gputools +} // namespace perftools + +#endif // TENSORFLOW_STREAM_EXECUTOR_EXECUTOR_PLATFORM_ID_H_ diff --git a/tensorflow/compiler/plugin/executor/transfer_manager.cc b/tensorflow/compiler/plugin/executor/transfer_manager.cc new file mode 100644 index 0000000000000000000000000000000000000000..51c5deeea5d5fd03d0fb99d4f33413c7bf4abe0f --- /dev/null +++ b/tensorflow/compiler/plugin/executor/transfer_manager.cc @@ -0,0 +1,187 @@ +/* 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/plugin/executor/transfer_manager.h" +#include "tensorflow/compiler/plugin/executor/platform_id.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/compiler/xla/types.h" +#include "tensorflow/compiler/xla/util.h" +#include "tensorflow/compiler/xla/xla_data.pb.h" +#include "tensorflow/core/lib/core/errors.h" +#include "tensorflow/core/platform/logging.h" +#include "tensorflow/core/platform/stream_executor_no_cuda.h" + +#include +#include +#include + +namespace sep = ::perftools::gputools::executorplugin; + +namespace xla { +namespace executorplugin { + +ExecutorTransferManager::ExecutorTransferManager() {} + +se::Platform::Id ExecutorTransferManager::PlatformId() const { + return se::executorplugin::kExecutorPlatformId; +} + +Status ExecutorTransferManager::TransferLiteralFromDevice( + se::StreamExecutor* executor, const se::DeviceMemoryBase& source, + const Shape& device_shape, const Shape& literal_shape, Literal* literal) { + TF_RET_CHECK(ShapeUtil::Compatible(device_shape, literal_shape)); + + // Tuples are a special case and contain one or more shapes inside of them to + // an arbitrary nesting depth. + if (device_shape.element_type() == TUPLE) { + *literal->mutable_shape() = literal_shape; + TF_ASSIGN_OR_RETURN( + std::vector element_buffers, + ShallowCopyTupleFromDevice(executor, source, device_shape)); + TF_RET_CHECK(element_buffers.size() == + ShapeUtil::TupleElementCount(device_shape)); + for (int64 i = 0; i < element_buffers.size(); ++i) { + const Shape& element_device_shape = device_shape.tuple_shapes(i); + const Shape& element_literal_shape = literal_shape.tuple_shapes(i); + Literal* element_literal = literal->add_tuple_literals(); + // Recursively call TransferFromDevice to copy over the data in the + // element array. + TF_RETURN_IF_ERROR(TransferLiteralFromDevice( + executor, element_buffers[i], element_device_shape, + element_literal_shape, element_literal)); + } + return Status::OK(); + } + + *literal->mutable_shape() = device_shape; + literal->Reserve(ShapeUtil::ElementsIn(device_shape)); + TF_RETURN_IF_ERROR(TransferBufferFromDevice( + executor, source, ShapeUtil::ByteSizeOf(device_shape), + literal->MutableInternalData())); + if (!ShapeUtil::Equal(literal_shape, device_shape)) { + literal->Swap( + literal->Relayout(literal_shape.layout()).get()); + } + TF_RET_CHECK(ShapeUtil::Equal(literal_shape, literal->shape())); + return Status::OK(); +} + +StatusOr> +ExecutorTransferManager::ShallowCopyTupleFromDevice( + se::StreamExecutor* executor, const se::DeviceMemoryBase& source, + const Shape& shape) { + TF_RET_CHECK(ShapeUtil::IsTuple(shape)); + + std::vector element_pointers(ShapeUtil::TupleElementCount(shape), + nullptr); + int64 tuple_size = ShapeUtil::ByteSizeOf(shape, sizeof(void*)); + auto copy_status = executor->SynchronousMemcpyD2H(source, tuple_size, + element_pointers.data()); + if (!copy_status.ok()) { + return AddStatus( + Status(static_cast(copy_status.code()), + copy_status.error_message()), + "failed transfer of tuple buffer " + ShapeUtil::HumanString(shape)); + } + + // Create a DeviceMemoryBase from each void* pointer. + std::vector destination; + for (int i = 0; i < element_pointers.size(); ++i) { + if (element_pointers[i] == nullptr && + !ShapeUtil::HasZeroElements(shape.tuple_shapes(i))) { + return FailedPrecondition("tuple contains nullptr at element %d", i); + } + int64 buffer_size = + ShapeUtil::ByteSizeOf(shape.tuple_shapes(i), sizeof(void*)); + destination.emplace_back(element_pointers[i], buffer_size); + } + return std::move(destination); +} + +Status ExecutorTransferManager::TransferLiteralToDevice( + se::StreamExecutor* executor, const Literal& literal, + se::DeviceMemoryBase* destination) { + const Shape& shape = literal.shape(); + + if (ShapeUtil::IsTuple(literal.shape())) { + std::vector tuple_elements_on_device; + for (const Literal& tuple_element : literal.tuple_literals()) { + se::DeviceMemoryBase allocation = executor->AllocateArray( + GetByteSizeRequirement(tuple_element.shape())); + TF_RETURN_IF_ERROR( + TransferLiteralToDevice(executor, tuple_element, &allocation)); + tuple_elements_on_device.push_back(allocation.opaque()); + } + return TransferBufferToDevice( + executor, tuple_elements_on_device.size() * sizeof(void*), + tuple_elements_on_device.data(), destination); + } + + return TransferBufferToDevice(executor, GetByteSizeRequirement(shape), + literal.InternalData(), + destination); +} + +Status ExecutorTransferManager::TransferLiteralToInfeed( + se::StreamExecutor* executor, const Literal& literal) { + const Shape& shape = literal.shape(); + VLOG(1) << "transferring literal shape to infeed: " + << ShapeUtil::HumanString(shape); + + return Status::OK(); +} + +Status ExecutorTransferManager::TransferBufferToInfeed( + se::StreamExecutor* executor, int64 size, const void* source) { + return Unimplemented("Transfer to Infeed"); +} + +Status ExecutorTransferManager::TransferLiteralFromOutfeed( + perftools::gputools::StreamExecutor* executor, const Shape& literal_shape, + Literal* literal) { + const Shape& shape = literal->shape(); + VLOG(1) << "transferring literal shape from outfeed: " + << ShapeUtil::HumanString(shape); + + return Status::OK(); +} + +Status ExecutorTransferManager::ResetDevices( + tensorflow::gtl::ArraySlice + executors) { + return Unimplemented("Device reset not supported"); +} + +int64 ExecutorTransferManager::GetByteSizeRequirement(const Shape& shape) { + return ShapeUtil::ByteSizeOf(shape, sizeof(void*)); +} + +} // namespace executorplugin +} // namespace xla + +static std::unique_ptr CreateExecutorTransferManager() { + return xla::MakeUnique(); +} + +static bool InitModule() { + xla::TransferManager::RegisterTransferManager(sep::kExecutorPlatformId, + &CreateExecutorTransferManager); + return true; +} +static bool module_initialized = InitModule(); diff --git a/tensorflow/compiler/plugin/executor/transfer_manager.h b/tensorflow/compiler/plugin/executor/transfer_manager.h new file mode 100644 index 0000000000000000000000000000000000000000..7a42e5a2d7542eaad7f8f90f011c65a9c526cc11 --- /dev/null +++ b/tensorflow/compiler/plugin/executor/transfer_manager.h @@ -0,0 +1,77 @@ +/* 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_EXECUTOR_DRIVER_EXECUTOR_TRANSFER_MANAGER_H_ +#define TENSORFLOW_COMPILER_EXECUTOR_DRIVER_EXECUTOR_TRANSFER_MANAGER_H_ + +#include "tensorflow/compiler/xla/service/transfer_manager.h" +#include "tensorflow/compiler/xla/statusor.h" +#include "tensorflow/compiler/xla/xla_data.pb.h" +#include "tensorflow/core/platform/macros.h" +#include "tensorflow/core/platform/stream_executor_no_cuda.h" +#include "tensorflow/core/platform/types.h" + +#include + +namespace se = ::perftools::gputools; + +namespace xla { +namespace executorplugin { + +class ExecutorTransferManager : public TransferManager { + public: + ExecutorTransferManager(); + + ~ExecutorTransferManager() override {} + + se::Platform::Id PlatformId() const override; + + StatusOr> ShallowCopyTupleFromDevice( + se::StreamExecutor* executor, const se::DeviceMemoryBase& source, + const Shape& shape) override; + + Status TransferLiteralFromDevice(se::StreamExecutor* executor, + const se::DeviceMemoryBase& source, + const Shape& device_shape, + const Shape& literal_shape, + Literal* literal) override; + + Status TransferLiteralToDevice(se::StreamExecutor* executor, + const Literal& literal, + se::DeviceMemoryBase* destination) override; + + Status TransferLiteralToInfeed(se::StreamExecutor* executor, + const Literal& literal) override; + + Status TransferBufferToInfeed(se::StreamExecutor* executor, + int64 size, const void* source) override; + + Status TransferLiteralFromOutfeed(se::StreamExecutor* executor, + const Shape& literal_shape, + Literal* literal) override; + + Status ResetDevices( + tensorflow::gtl::ArraySlice executors) override; + + int64 GetByteSizeRequirement(const Shape& shape) override; + + private: + TF_DISALLOW_COPY_AND_ASSIGN(ExecutorTransferManager); +}; + +} // namespace executorplugin +} // namespace xla + +#endif // TENSORFLOW_COMPILER_EXECUTOR_DRIVER_EXECUTOR_TRANSFER_MANAGER_H_ diff --git a/tensorflow/compiler/tests/BUILD b/tensorflow/compiler/tests/BUILD index 03e255e6b842668a491d254953926500ce3a50ec..a54d1f54f9533509534505228e66315f78f1bbfa 100644 --- a/tensorflow/compiler/tests/BUILD +++ b/tensorflow/compiler/tests/BUILD @@ -40,17 +40,13 @@ py_library( "//tensorflow/python:client_testlib", "//tensorflow/python:framework_for_generated_wrappers", "//tensorflow/python:platform", + "//tensorflow/python:random_seed", + "//tensorflow/python:session", "//tensorflow/python:variables", + "//third_party/py/numpy", ], ) -cc_library( - name = "depthwise_conv2d_test_kernel", - testonly = 1, - srcs = ["depthwise_conv2d_test_kernel.cc"], - deps = ["//tensorflow/core:framework_lite"], -) - tf_xla_py_test( name = "adagrad_test", size = "small", @@ -65,6 +61,20 @@ tf_xla_py_test( ], ) +tf_xla_py_test( + name = "adam_test", + size = "small", + srcs = ["adam_test.py"], + deps = [ + ":xla_test", + "//tensorflow/python:array_ops", + "//tensorflow/python:framework_for_generated_wrappers", + "//tensorflow/python:math_ops", + "//tensorflow/python:platform_test", + "//tensorflow/python:training", + ], +) + tf_xla_py_test( name = "binary_ops_test", size = "small", @@ -143,6 +153,27 @@ tf_xla_py_test( ], ) +tf_xla_py_test( + name = "depthwise_conv_op_test", + size = "medium", + srcs = ["depthwise_conv_op_test.py"], + shard_count = 5, + tags = [ + "noasan", + "nomsan", + "notsan", + ], + deps = [ + ":xla_test", + "//tensorflow/python:array_ops", + "//tensorflow/python:framework_for_generated_wrappers", + "//tensorflow/python:nn", + "//tensorflow/python:nn_ops", + "//tensorflow/python:nn_ops_gen", + "//tensorflow/python:platform_test", + ], +) + tf_xla_py_test( name = "dynamic_stitch_test", size = "small", @@ -156,6 +187,38 @@ tf_xla_py_test( ], ) +tf_xla_py_test( + name = "slice_ops_test", + size = "small", + srcs = ["slice_ops_test.py"], + # TODO(b/62962492): Test fails with assertion error. + tags = [ + "manual", + "notap", + ], + deps = [ + ":xla_test", + "//tensorflow/python:array_ops", + "//tensorflow/python:data_flow_ops", + "//tensorflow/python:framework_for_generated_wrappers", + "//tensorflow/python:platform_test", + ], +) + +tf_xla_py_test( + name = "ftrl_test", + size = "small", + srcs = ["ftrl_test.py"], + deps = [ + ":xla_test", + "//tensorflow/python:array_ops", + "//tensorflow/python:framework_for_generated_wrappers", + "//tensorflow/python:math_ops", + "//tensorflow/python:platform_test", + "//tensorflow/python:training", + ], +) + tf_xla_py_test( name = "function_test", size = "small", @@ -282,7 +345,7 @@ tf_xla_py_test( tf_xla_py_test( name = "reverse_ops_test", - size = "small", + size = "medium", srcs = ["reverse_ops_test.py"], deps = [ ":xla_test", @@ -305,6 +368,66 @@ tf_xla_py_test( ], ) +tf_xla_py_test( + name = "segment_reduction_ops_test", + size = "medium", + srcs = ["segment_reduction_ops_test.py"], + deps = [ + ":xla_test", + "//tensorflow/python:array_ops", + "//tensorflow/python:framework_for_generated_wrappers", + "//tensorflow/python:math_ops", + "//tensorflow/python:math_ops_gen", + "//tensorflow/python:platform_test", + ], +) + +tf_xla_py_test( + name = "spacetobatch_op_test", + size = "medium", + srcs = ["spacetobatch_op_test.py"], + shard_count = 3, + deps = [ + ":xla_test", + "//tensorflow/python:array_ops", + "//tensorflow/python:framework_for_generated_wrappers", + "//tensorflow/python:math_ops", + "//tensorflow/python:platform_test", + ], +) + +tf_xla_py_test( + name = "stack_ops_test", + size = "small", + srcs = ["stack_ops_test.py"], + deps = [ + ":xla_test", + "//tensorflow/python:array_ops", + "//tensorflow/python:data_flow_ops", + "//tensorflow/python:framework_for_generated_wrappers", + "//tensorflow/python:platform_test", + ], +) + +tf_xla_py_test( + name = "tensor_array_ops_test", + size = "small", + srcs = ["tensor_array_ops_test.py"], + deps = [ + ":xla_test", + "//tensorflow/python:array_ops", + "//tensorflow/python:framework_for_generated_wrappers", + "//tensorflow/python:math_ops", + "//tensorflow/python:math_ops_gen", + "//tensorflow/python:nn_ops", + "//tensorflow/python:nn_ops_gen", + "//tensorflow/python:platform_test", + "//tensorflow/python:tensor_array_grad", + "//tensorflow/python:tensor_array_ops", + "//tensorflow/python:training", + ], +) + tf_xla_py_test( name = "ternary_ops_test", size = "small", @@ -333,6 +456,23 @@ tf_xla_py_test( ], ) +tf_xla_py_test( + name = "fused_batchnorm_test", + size = "small", + srcs = ["fused_batchnorm_test.py"], + deps = [ + ":xla_test", + "//tensorflow/python:framework_for_generated_wrappers", + "//tensorflow/python:math_ops", + "//tensorflow/python:math_ops_gen", + "//tensorflow/python:nn", + "//tensorflow/python:nn_ops", + "//tensorflow/python:nn_ops_gen", + "//tensorflow/python:platform_test", + "//tensorflow/python:training", + ], +) + tf_xla_py_test( name = "variable_ops_test", size = "small", @@ -350,6 +490,23 @@ tf_xla_py_test( ], ) +tf_xla_py_test( + name = "gather_test", + size = "small", + srcs = ["gather_test.py"], + # Gather needs CustomCall on CPU, which is not available in normal + # (not precompiled) TensorFlow. The flag below excludes the CPU + # backend. + disabled_backends = "cpu", + deps = [ + ":xla_test", + "//tensorflow/python:array_ops", + "//tensorflow/python:data_flow_ops", + "//tensorflow/python:framework_for_generated_wrappers", + "//tensorflow/python:platform_test", + ], +) + cuda_py_test( name = "xla_device_test", size = "small", @@ -381,6 +538,11 @@ cuda_py_test( "//tensorflow/python:math_ops", "//tensorflow/python:nn_ops", ], + # TODO(b/62961789): Test fails with SIGABRT + tags = [ + "manual", + "notap", + ], ) cc_library( @@ -449,8 +611,12 @@ cuda_py_test( # --dump_graph_dir, and the config file was written by hand. # # Run the following to build a minimal benchmark of the computation on Android: -# $ bazel build -c opt --config=android_arm \ -# third_party/tensorflow/compiler/tests:lstm_layer_inference_benchmark +# $ bazel build -c opt --cxxopt='-std=c++11' --linkopt='-lm' \ +# --cpu=armeabi-v7a \ +# --host_crosstool_top=@bazel_tools//tools/cpp:toolchain \ +# --crosstool_top=//external:android/crosstool \ +# //tensorflow/compiler/tests:lstm_layer_inference_benchmark + # # Currently the resulting binary size is ~190KB tf_library( diff --git a/tensorflow/compiler/tests/adagrad_test.py b/tensorflow/compiler/tests/adagrad_test.py index 0a2c9e26c6fbd827d5ab669dea5419f9fa50025b..9a93b3216404d8ed21fd6c57757bec1730c119b4 100644 --- a/tensorflow/compiler/tests/adagrad_test.py +++ b/tensorflow/compiler/tests/adagrad_test.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""Functional tests for aggregate operations.""" +"""Tests for Adagrad.""" from __future__ import absolute_import from __future__ import division @@ -49,9 +49,11 @@ class AdagradOptimizerTest(XLATestCase): ada_update.run() # Validate updated params self.assertAllCloseAccordingToType( - np.array([-1.6026098728179932, -0.6026098728179932]), var0.eval()) + np.array([-1.6026098728179932, -0.6026098728179932]), var0.eval(), + float_rtol=1e-5) self.assertAllCloseAccordingToType( - np.array([2.715679168701172, 3.715679168701172]), var1.eval()) + np.array([2.715679168701172, 3.715679168701172]), var1.eval(), + float_rtol=1e-5) def testTensorLearningRate(self): for dtype in self.float_types: @@ -73,9 +75,11 @@ class AdagradOptimizerTest(XLATestCase): ada_update.run() # Validate updated params self.assertAllCloseAccordingToType( - np.array([-1.6026098728179932, -0.6026098728179932]), var0.eval()) + np.array([-1.6026098728179932, -0.6026098728179932]), var0.eval(), + float_rtol=1e-5) self.assertAllCloseAccordingToType( - np.array([2.715679168701172, 3.715679168701172]), var1.eval()) + np.array([2.715679168701172, 3.715679168701172]), var1.eval(), + float_rtol=1e-5) def testSharing(self): for dtype in self.float_types: @@ -107,9 +111,11 @@ class AdagradOptimizerTest(XLATestCase): 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]), var0.eval(), + float_rtol=1e-5) self.assertAllCloseAccordingToType( - np.array([2.715679168701172, 3.715679168701172]), var1.eval()) + np.array([2.715679168701172, 3.715679168701172]), var1.eval(), + float_rtol=1e-5) if __name__ == "__main__": diff --git a/tensorflow/compiler/tests/adam_test.py b/tensorflow/compiler/tests/adam_test.py new file mode 100644 index 0000000000000000000000000000000000000000..3215dc36e5b2d517aa951db1b0d41188185ef93a --- /dev/null +++ b/tensorflow/compiler/tests/adam_test.py @@ -0,0 +1,176 @@ +# 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 Adam.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import numpy as np + +from tensorflow.compiler.tests.xla_test import XLATestCase +from tensorflow.python.framework import constant_op +from tensorflow.python.ops import array_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.platform import test +from tensorflow.python.training import adam + + +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 AdamOptimizerTest(XLATestCase): + + def testBasic(self): + for dtype in self.float_types: + with self.test_session(), self.test_scope(): + variable_scope.get_variable_scope().set_use_resource(True) + + # 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) + grads0_np = np.array([0.1, 0.1], dtype=dtype) + var1_np = np.array([3.0, 4.0], dtype=dtype) + grads1_np = np.array([0.01, 0.01], dtype=dtype) + + var0 = resource_variable_ops.ResourceVariable(var0_np) + var1 = resource_variable_ops.ResourceVariable(var1_np) + grads0 = array_ops.placeholder(dtype) + grads1 = array_ops.placeholder(dtype) + opt = adam.AdamOptimizer() + 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()) + + 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(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()) + + def testTensorLearningRate(self): + for dtype in self.float_types: + with self.test_session(), self.test_scope(): + variable_scope.get_variable_scope().set_use_resource(True) + + # 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) + grads0_np = np.array([0.1, 0.1], dtype=dtype) + var1_np = np.array([3.0, 4.0], dtype=dtype) + grads1_np = np.array([0.01, 0.01], dtype=dtype) + + var0 = resource_variable_ops.ResourceVariable(var0_np) + var1 = resource_variable_ops.ResourceVariable(var1_np) + grads0 = array_ops.placeholder(dtype) + grads1 = array_ops.placeholder(dtype) + opt = adam.AdamOptimizer(constant_op.constant(0.001)) + 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()) + + 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(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()) + + def testSharing(self): + for dtype in self.float_types: + with self.test_session(), self.test_scope(): + variable_scope.get_variable_scope().set_use_resource(True) + + # 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) + grads0_np = np.array([0.1, 0.1], dtype=dtype) + var1_np = np.array([3.0, 4.0], dtype=dtype) + grads1_np = np.array([0.01, 0.01], dtype=dtype) + + var0 = resource_variable_ops.ResourceVariable(var0_np) + var1 = resource_variable_ops.ResourceVariable(var1_np) + grads0 = array_ops.placeholder(dtype) + grads1 = array_ops.placeholder(dtype) + opt = adam.AdamOptimizer() + update1 = opt.apply_gradients(zip([grads0, grads1], [var0, var1])) + update2 = opt.apply_gradients(zip([grads0, grads1], [var0, var1])) + 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(feed_dict={grads0: grads0_np, grads1: grads1_np}) + else: + update2.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()) + + +if __name__ == "__main__": + test.main() diff --git a/tensorflow/compiler/tests/binary_ops_test.py b/tensorflow/compiler/tests/binary_ops_test.py index 9efdaee7ab66f7cfc84bc1c30a9ba700e268abe2..e6862f0d9dd7ec05b4e0c4ba26ab5f16a7aa9ad7 100644 --- a/tensorflow/compiler/tests/binary_ops_test.py +++ b/tensorflow/compiler/tests/binary_ops_test.py @@ -52,6 +52,12 @@ class BinaryOpsTest(XLATestCase): def testFloatOps(self): for dtype in self.float_types: + self._testBinary( + lambda x, y: math_ops.approximate_equal(x, y, tolerance=0.0001), + np.array([[[[-1, 2.00009999], [-3, 4.01]]]], dtype=dtype), + np.array([[[[-1.001, 2], [-3.00009, 4]]]], dtype=dtype), + expected=np.array([[[[False, True], [True, False]]]], dtype=dtype)) + self._testBinary( gen_math_ops._real_div, np.array([3, 3, -1.5, -8, 44], dtype=dtype), @@ -82,6 +88,12 @@ class BinaryOpsTest(XLATestCase): dtype(4), expected=np.array([[16], [81]], dtype=dtype)) + self._testBinary( + gen_math_ops._reciprocal_grad, + np.array([4, -3, -2, 1], dtype=dtype), + np.array([5, -6, 7, -8], dtype=dtype), + expected=np.array([-80, 54, -28, 8], dtype=dtype)) + self._testBinary( gen_math_ops._sigmoid_grad, np.array([4, 3, 2, 1], dtype=dtype), @@ -94,6 +106,12 @@ class BinaryOpsTest(XLATestCase): np.array([5, 6, 7, 8], dtype=dtype), expected=np.array([-160, -81, -28, -4], dtype=dtype)) + self._testBinary( + gen_math_ops._sqrt_grad, + np.array([4, 3, 2, 1], dtype=dtype), + np.array([5, 6, 7, 8], dtype=dtype), + expected=np.array([0.625, 1, 1.75, 4], dtype=dtype)) + self._testBinary( gen_nn_ops._softplus_grad, np.array([4, 3, 2, 1], dtype=dtype), @@ -101,12 +119,33 @@ class BinaryOpsTest(XLATestCase): expected=np.array( [3.97322869, 2.99258232, 1.99817801, 0.99966466], dtype=dtype)) + self._testBinary( + gen_nn_ops._softsign_grad, + np.array([4, 3, 2, 1], dtype=dtype), + np.array([5, 6, 7, 8], dtype=dtype), + expected=np.array( + [0.11111111, 0.06122449, 0.03125, 0.01234568], dtype=dtype)) + self._testBinary( gen_math_ops._tanh_grad, np.array([4, 3, 2, 1], dtype=dtype), np.array([5, 6, 7, 8], dtype=dtype), expected=np.array([-75, -48, -21, 0], dtype=dtype)) + self._testBinary( + gen_nn_ops._elu_grad, + np.array([1, 2, 3, 4, 5, 6], dtype=dtype), + np.array([-.6, -.4, -.2, 0, .2, .4], dtype=dtype), + expected=np.array([0.4, 1.2, 2.4, 4, 5, 6], dtype=dtype)) + + self._testBinary( + gen_nn_ops._selu_grad, + np.array([1, 2, 3, 4, 5, 6], dtype=dtype), + np.array([-.6, -.4, -.2, .2, .4, .6], dtype=dtype), + expected=np.array( + [1.158099340847, 2.7161986816948, 4.67429802254, + 4.202803949422, 5.2535049367774, 6.30420592413], dtype=dtype)) + self._testBinary( gen_nn_ops._relu_grad, np.array([1, 2, 3, 4, 5, 6, 7, 8, 9, 10], dtype=dtype), @@ -549,17 +588,18 @@ class BinaryOpsTest(XLATestCase): self._testBinary( math_ops.matmul, np.array( - [[[[1000, 100], [10, 1]], [[2000, 200], [20, 2]]], - [[[3000, 300], [30, 3]], [[4000, 400], [40, 4]]]], + [[[[7, 13], [10, 1]], [[2, 0.25], [20, 2]]], + [[[3, 5], [30, 3]], [[0.75, 1], [40, 4]]]], dtype=np.float32), np.array( [[[[1, 2], [3, 4]], [[5, 6], [7, 8]]], [[[11, 22], [33, 44]], [[55, 66], [77, 88]]]], dtype=np.float32), expected=np.array( - [[[[1300, 2400], [13, 24]], [[11400, 13600], [114, 136]]], - [[[42900, 79200], [429, 792]], [[250800, 299200], [2508, 2992]]]], + [[[[46, 66], [13, 24]], [[11.75, 14], [114, 136]]], + [[[198, 286], [429, 792]], [[118.25, 137.5], [2508, 2992]]]], dtype=np.float32)) + self._testBinary( math_ops.matmul, np.array([], dtype=np.float32).reshape((2, 2, 0)), @@ -575,7 +615,7 @@ class BinaryOpsTest(XLATestCase): # Regression test for b/31472796. if hasattr(np, "matmul"): - x = np.arange(0, 3 * 5 * 16 * 7, dtype=np.float32).reshape((3, 5, 16, 7)) + x = np.arange(0, 3 * 5 * 2 * 7, dtype=np.float32).reshape((3, 5, 2, 7)) self._testBinary( lambda x, y: math_ops.matmul(x, y, adjoint_b=True), x, x, @@ -635,6 +675,80 @@ class BinaryOpsTest(XLATestCase): [0, 0, 0, 0, 0, 0]], dtype=dtype)) + def testMirrorPad(self): + mirror_pad = lambda t, paddings: array_ops.pad(t, paddings, "REFLECT") + for dtype in self.numeric_types: + self._testBinary( + mirror_pad, + np.array( + [ + [1, 2, 3], # + [4, 5, 6], # + ], + dtype=dtype), + np.array([[ + 1, + 1, + ], [2, 2]], dtype=np.int32), + expected=np.array( + [ + [6, 5, 4, 5, 6, 5, 4], # + [3, 2, 1, 2, 3, 2, 1], # + [6, 5, 4, 5, 6, 5, 4], # + [3, 2, 1, 2, 3, 2, 1] + ], + dtype=dtype)) + self._testBinary( + mirror_pad, + np.array([[1, 2, 3], [4, 5, 6]], dtype=dtype), + np.array([[0, 0], [0, 0]], dtype=np.int32), + expected=np.array([[1, 2, 3], [4, 5, 6]], dtype=dtype)) + self._testBinary( + mirror_pad, + np.array( + [ + [1, 2, 3], # + [4, 5, 6], # + [7, 8, 9] + ], + dtype=dtype), + np.array([[2, 2], [0, 0]], dtype=np.int32), + expected=np.array( + [ + [7, 8, 9], # + [4, 5, 6], # + [1, 2, 3], # + [4, 5, 6], # + [7, 8, 9], # + [4, 5, 6], # + [1, 2, 3] + ], + dtype=dtype)) + self._testBinary( + mirror_pad, + np.array( + [ + [[1, 2, 3], [4, 5, 6]], + [[7, 8, 9], [10, 11, 12]], + ], dtype=dtype), + np.array([[0, 0], [1, 1], [1, 1]], dtype=np.int32), + expected=np.array( + [ + [ + [5, 4, 5, 6, 5], # + [2, 1, 2, 3, 2], # + [5, 4, 5, 6, 5], # + [2, 1, 2, 3, 2], # + ], + [ + [11, 10, 11, 12, 11], # + [8, 7, 8, 9, 8], # + [11, 10, 11, 12, 11], # + [8, 7, 8, 9, 8], # + ] + ], + dtype=dtype)) + def testReshape(self): for dtype in self.numeric_types: self._testBinary( @@ -758,6 +872,24 @@ class BinaryOpsTest(XLATestCase): np.array([1, 0], dtype=np.int32), expected=np.array([[1, 3], [2, 4]], dtype=dtype)) + def testCross(self): + for dtype in self.float_types: + self._testBinary( + gen_math_ops.cross, + np.zeros((4, 3), dtype=dtype), + np.zeros((4, 3), dtype=dtype), + expected=np.zeros((4, 3), dtype=dtype)) + self._testBinary( + gen_math_ops.cross, + np.array([1, 2, 3], dtype=dtype), + np.array([4, 5, 6], dtype=dtype), + expected=np.array([-3, 6, -3], dtype=dtype)) + self._testBinary( + gen_math_ops.cross, + np.array([[1, 2, 3], [10, 11, 12]], dtype=dtype), + np.array([[4, 5, 6], [40, 50, 60]], dtype=dtype), + expected=np.array([[-3, 6, -3], [60, -120, 60]], dtype=dtype)) + if __name__ == "__main__": googletest.main() diff --git a/tensorflow/compiler/tests/build_defs.bzl b/tensorflow/compiler/tests/build_defs.bzl index 820db13d0b123a465b1599eec486af4c86ccc376..a56c53de0fb5f76c94064e2bdc2f1a543a207b09 100644 --- a/tensorflow/compiler/tests/build_defs.bzl +++ b/tensorflow/compiler/tests/build_defs.bzl @@ -1,12 +1,14 @@ """Build rules for Tensorflow/XLA testing.""" load("@local_config_cuda//cuda:build_defs.bzl", "cuda_is_configured") +load("//tensorflow/compiler/tests:plugin.bzl", "plugins") def all_backends(): + b = ["cpu"] + plugins.keys() if cuda_is_configured(): - return ["cpu", "gpu"] + return b + ["gpu"] else: - return ["cpu"] + return b def tf_xla_py_test(name, srcs=[], deps=[], tags=[], data=[], main=None, disabled_backends=None, **kwargs): @@ -53,6 +55,13 @@ def tf_xla_py_test(name, srcs=[], deps=[], tags=[], data=[], main=None, backend_args += ["--test_device=XLA_GPU", "--types=DT_FLOAT,DT_DOUBLE,DT_INT32,DT_INT64,DT_BOOL"] backend_tags += ["requires-gpu-sm35"] + elif backend in plugins: + backend_args += ["--test_device=" + plugins[backend]["device"], + "--types=" + plugins[backend]["types"]] + backend_tags += plugins[backend]["tags"] + backend_args += plugins[backend]["args"] + backend_deps += plugins[backend]["deps"] + backend_data += plugins[backend]["data"] else: fail("Unknown backend {}".format(backend)) diff --git a/tensorflow/compiler/tests/conv2d_test.py b/tensorflow/compiler/tests/conv2d_test.py index 4bc118b5bdb370258a6113db28a0fd63b9965354..0d617eb37c5d92c87abb0f996b731112257a2b80 100644 --- a/tensorflow/compiler/tests/conv2d_test.py +++ b/tensorflow/compiler/tests/conv2d_test.py @@ -26,10 +26,8 @@ import numpy as np from tensorflow.compiler.tests.xla_test import XLATestCase 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.ops import nn_impl from tensorflow.python.ops import nn_ops from tensorflow.python.platform import googletest @@ -447,80 +445,5 @@ class Conv2DBackpropFilterTest(XLATestCase): expected=expected_output) -class DepthwiseConv2DTest(XLATestCase): - - CPU_DEVICE = "/job:localhost/replica:0/task:0/cpu:0" - - def ConfigsToTest(self): - input_sizes = [[4, 35, 35, 2], [4, 147, 147, 2], [3, 299, 299, 3], - [5, 183, 183, 1]] - filter_sizes = [[5, 5, 2, 1], [3, 3, 2, 8], [2, 2, 3, 8], [5, 5, 1, 2]] - strides = [1, 3, 2, 2] - # pylint: disable=invalid-name - VALID = "VALID" - SAME = "SAME" - # pylint: enable=invalid-name - paddings = [SAME, VALID, SAME, SAME, SAME] - for i, f, s, p in zip(input_sizes, filter_sizes, strides, paddings): - yield i, f, s, p - - def _VerifyValues(self, input_size, filter_size, stride, padding): - imag = np.random.rand(*input_size).astype(np.float32) - filt = np.random.rand(*filter_size).astype(np.float32) - strides = [1, stride, stride, 1] - - with self.test_session(): - with self.test_scope(): - imag_ph = array_ops.placeholder(dtypes.float32, shape=input_size) - filt_ph = array_ops.placeholder(dtypes.float32, shape=filter_size) - feed_dict = {imag_ph: imag, filt_ph: filt} - xla_out = nn_impl.depthwise_conv2d(imag_ph, filt_ph, strides, - padding).eval(feed_dict=feed_dict) - - with self.test_session(): - with ops.device(self.CPU_DEVICE): - imag_ph = array_ops.placeholder(dtypes.float32, shape=input_size) - filt_ph = array_ops.placeholder(dtypes.float32, shape=filter_size) - feed_dict = {imag_ph: imag, filt_ph: filt} - cpu_out = nn_impl.depthwise_conv2d(imag_ph, filt_ph, strides, - padding).eval(feed_dict=feed_dict) - - self.assertAllClose(xla_out, cpu_out) - - # This is disabled because we need a mechanism to set command-line flags, - # i.e. an implementation of SetCommandLineOption() below. - # - # def _VerifyDummy(self, input_size, filter_size, stride, padding): - # imag = np.random.rand(*input_size).astype(np.float32) - # filt = np.random.rand(*filter_size).astype(np.float32) - # strides = [1, stride, stride, 1] - # - # with self.test_session(): - # with self.test_scope(): - # imag_ph = tf.placeholder(tf.float32, shape=input_size) - # filt_ph = tf.placeholder(tf.float32, shape=filter_size) - # feed_dict = {imag_ph: imag, filt_ph: filt} - # SetCommandLineOption( - # "tf_tla_depthwise_conv2d_custom_func", - # "DummyDepthwiseConv2dKernel") - # xla_out = tf.nn.depthwise_conv2d( - # imag_ph, filt_ph, strides, padding).eval(feed_dict=feed_dict) - # SetCommandLineOption( - # "tf_tla_depthwise_conv2d_custom_func", "") - # - # expected = np.array(range(np.ravel(xla_out).shape[0]), dtype=np.float32) - # self.assertAllClose(np.ravel(xla_out), expected) - - def testBasic(self): - for i, f, s, p in self.ConfigsToTest(): - self._VerifyValues(i, f, s, p) - - # Test disabled until _VerifyDummy(), above can be implemented. - # def testCustomFunc(self): - # if self.has_custom_call: - # for i, f, s, p in self.ConfigsToTest(): - # self._VerifyDummy(i, f, s, p) - - if __name__ == "__main__": googletest.main() diff --git a/tensorflow/compiler/tests/depthwise_conv2d_test_kernel.cc b/tensorflow/compiler/tests/depthwise_conv2d_test_kernel.cc deleted file mode 100644 index 97b71c02288aa10ff9c2c8fac3eab0fcf7ed613a..0000000000000000000000000000000000000000 --- a/tensorflow/compiler/tests/depthwise_conv2d_test_kernel.cc +++ /dev/null @@ -1,30 +0,0 @@ -/* Copyright 2017 The TensorFlow Authors. All Rights Reserved. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -==============================================================================*/ - -#include "tensorflow/core/platform/types.h" - -using tensorflow::int64; - -// A dummy implementation that fills the output with 0, 1, 2,... -// to test the custom call implementation of DepthwiseConv2dNative op. -// TODO(keveman): Test this after adding a real implementation for the kernel. -extern "C" void DummyDepthwiseConv2dKernel(float* output, void** inputs) { - const int64* output_size = reinterpret_cast(inputs[4]); - const int64 total_size = - output_size[0] * output_size[1] * output_size[2] * output_size[3]; - for (int64 i = 0; i < total_size; ++i) { - *(output + i) = i; - } -} diff --git a/tensorflow/compiler/tests/depthwise_conv_op_test.py b/tensorflow/compiler/tests/depthwise_conv_op_test.py new file mode 100644 index 0000000000000000000000000000000000000000..0a0d335ca76dd7ec7ca3b12f9e8a83b596daa07e --- /dev/null +++ b/tensorflow/compiler/tests/depthwise_conv_op_test.py @@ -0,0 +1,389 @@ +# 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. +# ============================================================================== +"""Functional tests for depthwise convolutional operations.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import numpy as np +from six.moves import xrange # pylint: disable=redefined-builtin + +from tensorflow.compiler.tests.xla_test import XLATestCase +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 nn_ops +import tensorflow.python.ops.nn_grad # pylint: disable=unused-import +from tensorflow.python.platform import test + + +# Reference implementation of depthwise_conv2d +def ReferenceDepthwiseConv2D(input_tensor, filter_tensor, strides, padding, + data_format=None): + # Reference implementation of depthwise convolution that uses regular + # convolution. + convs = [] + in_channels = filter_tensor.shape[2] + # Use a custom implementation of depthwise conv2d using slicing. + for channel in xrange(in_channels): + # Slice the input along channel + if data_format == "NCHW": + input_slice = input_tensor[:, channel:channel+1, :, :] + else: + input_slice = input_tensor[:, :, :, channel:channel+1] + + # Slice the filters. Filters are H, W, InC, DepthMultiplier + filter_slice = filter_tensor[:, :, channel:channel+1, :] + # Do conv + convs.append(nn_ops.conv2d(input_slice, filter_slice, + strides, padding, + data_format=data_format, + name="depthwise_slice_%d" % channel)) + + # Concat along dimension. + if data_format == "NCHW": + return array_ops.concat(convs, 1) + else: + return array_ops.concat(convs, 3) + + +def ConfigsToTest(): + """Iterator for different convolution shapes, strides and paddings. + + Yields: + Tuple (input_size, filter_size, out_size, stride, padding), the depthwise + convolution parameters. + """ + input_sizes = [[4, 5, 5, 48], [4, 8, 8, 84], [4, 17, 17, 48], [4, 9, 27, 8], + [4, 31, 31, 7], [4, 35, 35, 2], [4, 147, 147, 2], + [3, 299, 299, 3], [5, 183, 183, 1]] + filter_sizes = [[1, 1, 48, 2], [1, 3, 84, 1], [3, 1, 48, 4], [3, 3, 8, 1], + [3, 3, 7, 1], [5, 5, 2, 1], [3, 3, 2, 8], [2, 2, 3, + 8], [5, 5, 1, 2]] + out_sizes = [[4, 5, 5, 96], [4, 8, 8, 84], [4, 17, 17, 192], [4, 9, 27, 8], + [4, 31, 31, 7], [4, 35, 35, 2], [4, 49, 49, 16], + [3, 150, 150, 24], [5, 92, 92, 2]] + strides = [1, 1, 1, 1, 1, 1, 3, 2, 2] + # pylint: disable=invalid-name + VALID = "VALID" + SAME = "SAME" + # pylint: enable=invalid-name + paddings = [SAME, SAME, SAME, SAME, SAME, SAME, VALID, SAME, SAME, SAME] + for i, f, o, s, p in zip(input_sizes, filter_sizes, out_sizes, strides, + paddings): + yield i, f, o, s, p + + +def CheckGradConfigsToTest(): + """Iterator for different convolution shapes, strides and paddings. + + compute_gradient_error() is very expensive. So the configs should be + relatively small. + + Yields: + Tuple (input_size, filter_size, out_size, stride, padding), the depthwise + convolution parameters. + """ + input_sizes = [[2, 5, 8, 1], [4, 5, 5, 1], [2, 4, 4, 2], [1, 15, 15, 2], + [2, 15, 16, 1]] + filter_sizes = [[4, 4, 1, 2], [2, 2, 1, 2], [3, 1, 2, 2], [1, 3, 2, 1], + [3, 3, 1, 2]] + out_sizes = [[2, 5, 8, 2], [4, 2, 2, 2], [2, 4, 4, 4], [1, 15, 15, 2], + [2, 5, 5, 2]] + strides = [1, 2, 1, 1, 3] + # pylint: disable=invalid-name + VALID = "VALID" + SAME = "SAME" + # pylint: enable=invalid-name + paddings = [SAME, VALID, SAME, SAME, VALID] + for i, f, o, s, p in zip(input_sizes, filter_sizes, out_sizes, strides, + paddings): + yield i, f, o, s, p + + +class DepthwiseConv2DTest(XLATestCase): + + # This is testing that depthwise_conv2d and depthwise_conv2d_native + # produce the same results. It also tests that NCHW and NWHC + # formats agree, by comparing the depthwise_conv2d_native with + # 'NCHW' format (with transposition) matches the 'NHWC' format using + # the higher level interface. + def _VerifyValues(self, + tensor_in_sizes, + filter_in_sizes, + stride, + padding, + data_type, + data_format="NHWC"): + """Verifies the output values of the convolution function. + + Args: + tensor_in_sizes: Input tensor dimensions in + [batch, input_rows, input_cols, input_depth]. + filter_in_sizes: Filter tensor dimensions in + [filter_rows, filter_cols, input_depth, depth_multiplier]. + stride: Stride. + padding: Padding type. + data_type: The data type to use. + data_format: The data_format of the input. "NHWC" or "NCHW". + """ + total_size_1 = 1 + total_size_2 = 1 + for s in tensor_in_sizes: + total_size_1 *= s + for s in filter_in_sizes: + total_size_2 *= s + # Initializes the input and filter tensor with numbers incrementing from 1. + x1 = np.array([f * 1.0 for f in range(1, total_size_1 + 1)], + dtype=data_type).reshape(tensor_in_sizes) + x2 = np.array([f * 1.0 for f in range(1, total_size_2 + 1)], + dtype=data_type).reshape(filter_in_sizes) + with self.test_session() as sess: + if data_type == np.float32: + tolerance = 1e-5 + else: + self.assertEqual(data_type, np.float64) + tolerance = 1e-8 + + t1 = array_ops.placeholder(shape=tensor_in_sizes, dtype=data_type) + t2 = array_ops.placeholder(shape=filter_in_sizes, dtype=data_type) + + native_t1 = t1 + strides = [1, stride, stride, 1] + if data_format == "NCHW": + # Transpose from NWHC input to NCHW + # Ex. [4, 5, 5, 48] to [4, 48, 5, 5] + native_t1 = array_ops.transpose(t1, [0, 3, 1, 2]) + strides = [1, 1, stride, stride] + + with self.test_scope(): + conv_native = nn_ops.depthwise_conv2d_native( + native_t1, + t2, + strides=strides, + data_format=data_format, + padding=padding) + + if data_format == "NCHW": + # Transpose back from NCHW to NHWC + conv_native = array_ops.transpose(conv_native, [0, 2, 3, 1]) + + with ops.device("CPU"): + conv_interface = ReferenceDepthwiseConv2D( + t1, t2, strides=[1, stride, stride, 1], padding=padding) + + native_result = sess.run(conv_native, {t1: x1, t2: x2}) + interface_result = sess.run(conv_interface, {t1: x1, t2: x2}) + + print("data_type:", data_type, "max diff = ", + np.amax(np.absolute(native_result - interface_result))) + self.assertAllClose( + np.ravel(native_result), np.ravel(interface_result), rtol=tolerance) + + def testDepthwiseConv2D(self): + for index, (input_size, filter_size, _, stride, + padding) in enumerate(ConfigsToTest()): + print("Testing DepthwiseConv2D,", index, "th config:", input_size, "*", + filter_size, "stride:", stride, "padding:", padding) + for data_type in self.float_types: + # TODO(phawkins): the reference implementation only supports float32. + if data_type == np.float32: + self._VerifyValues( + input_size, filter_size, stride, padding, data_type) + + def testDepthwiseConv2DFormat(self): + for index, (input_size, filter_size, _, stride, + padding) in enumerate(ConfigsToTest()): + print("Testing DepthwiseConv2DFormat,", index, "th config:", input_size, + "*", filter_size, "stride:", stride, "padding:", padding) + for data_type in self.float_types: + # TODO(phawkins): the reference implementation only supports float32. + if data_type == np.float32: + self._VerifyValues( + input_size, + filter_size, + stride, + padding, + data_type, + data_format="NCHW") + +# This is testing against hand calculated results. + + def _VerifyHandValues(self, tensor_in_sizes, filter_in_sizes, stride, padding, + expected): + """Verifies the output values of the depthwise convolution function. + + Args: + tensor_in_sizes: Input tensor dimensions in + [batch, input_rows, input_cols, input_depth]. + filter_in_sizes: Filter tensor dimensions in + [filter_rows, filter_cols, input_depth, depth_multiplier]. + stride: Stride. + padding: Padding type. + expected: An array containing the expected operation outputs. + """ + total_size_1 = 1 + total_size_2 = 1 + for s in tensor_in_sizes: + total_size_1 *= s + for s in filter_in_sizes: + total_size_2 *= s + # Initializes the input tensor with array containing incrementing + # numbers from 1. + x1 = np.array([f * 1.0 for f in range(1, total_size_1 + 1)], + dtype=np.float32).reshape(tensor_in_sizes) + x2 = np.array([f * 1.0 for f in range(1, total_size_2 + 1)], + dtype=np.float32).reshape(filter_in_sizes) + with self.test_session() as sess: + t1 = array_ops.placeholder(shape=tensor_in_sizes, dtype=np.float32) + t2 = array_ops.placeholder(shape=filter_in_sizes, dtype=np.float32) + with self.test_scope(): + conv = nn_ops.depthwise_conv2d_native( + 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.assertShapeEqual(value, conv) + + def testConv2D2x2Filter(self): + # The inputs look like this (it's a 3 x 2 matrix, each of depth 2): + # + # [ (1.0, 2.0), (3.0, 4.0), ( 5.0, 6.0) ] + # [ (7.0, 8.0), (9.0, 10.0), (11.0, 12.0) ] + # We can view this as two inputs + # + # input depth 0: + # + # [ 1.0, 3.0, 5.0 ] + # [ 7.0, 9.0, 11.0 ] + # + # input depth 1: + # + # [ 2.0, 4.0, 6.0 ] + # [ 8.0, 10.0, 12.0 ] + # + # The filter looks like this (it has two 2 x 2 patches, each generating 2 + # depths): + # + # filter #0: + # + # [ (1.0, 3.0), ( 5.0, 7.0)] + # [ (9.0, 11.0), (13.0, 15.0)] + # + # filter #1: + # + # [ ( 2.0, 4.0), ( 6.0, 8.0)] + # [ (10.0, 12.0), (14.0, 16.0)] + # + # So the outputs are: + # + # (position 0, 0: in_depth 0, output_depth 0 -- using filter #0) + # 1.0 * 1.0 + 7.0 * 9.0 + 3.0 * 5.0 + 9.0 * 13.0 = 196 + # (position 0, 0: in_depth 0, output_depth 1 -- using filter #1) + # 1.0 * 2.0 + 7.0 * 10.0 + 3.0 * 6.0 + 9.0 * 14.0 = 216 + # (position 0, 0: in_depth 1, output_depth 2 -- using filter #0) + # 2.0 * 3.0 + 8.0 * 11.0 + 4.0 * 7.0 + 10.0 * 15.0 = 272 + # (position 0, 0: in_depth 1, output_depth 3 -- using filter #1) + # 2.0 * 4.0 + 8.0 * 12.0 + 4.0 * 8.0 + 10.0 * 16.0 = 296 + # + # (position 1, 0: in_depth 0, output_depth 0 -- using filter #0) + # 3.0 * 1.0 + 9.0 * 9.0 + 5.0 * 5.0 + 11.0 * 13.0 = 252 + # (position 1, 0: in_depth 0, output_depth 1 -- using filter #1) + # 3.0 * 2.0 + 9.0 * 10.0 + 5.0 * 6.0 + 11.0 * 14.0 = 280 + # (position 1, 0: in_depth 1, output_depth 2 -- using filter #0) + # 4.0 * 3.0 + 10.0 * 11.0 + 6.0 * 7.0 + 12.0 * 15.0 = 344 + # (position 1, 0: in_depth 1, output_depth 3 -- using filter #1) + # 4.0 * 4.0 + 10.0 * 12.0 + 6.0 * 8.0 + 12.0 * 16.0 = 376 + expected_output = [196, 216, 272, 296, 252, 280, 344, 376] + self._VerifyHandValues( + tensor_in_sizes=[1, 2, 3, 2], + filter_in_sizes=[2, 2, 2, 2], + stride=1, + padding="VALID", + expected=expected_output) + + def _CompareBackpropInput(self, input_sizes, filter_sizes, output_sizes, + stride, padding): + x1 = np.random.rand(*filter_sizes).astype(np.float32) + x2 = np.random.rand(*output_sizes).astype(np.float32) + + def _GetVal(use_xla): + with self.test_session(): + t0 = constant_op.constant(input_sizes, shape=[len(input_sizes)]) + t1 = array_ops.placeholder(np.float32, shape=filter_sizes) + t2 = array_ops.placeholder(np.float32, shape=output_sizes) + if use_xla: + with self.test_scope(): + backprop = nn_ops.depthwise_conv2d_native_backprop_input( + t0, t1, t2, strides=[1, stride, stride, 1], padding=padding) + else: + backprop = nn_ops.depthwise_conv2d_native_backprop_input( + t0, t1, t2, strides=[1, stride, stride, 1], padding=padding) + + ret = backprop.eval({t1: x1, t2: x2}) + self.assertShapeEqual(ret, backprop) + return ret + + gpu_value = _GetVal(use_xla=True) + cpu_value = _GetVal(use_xla=False) + self.assertAllClose(cpu_value, gpu_value, rtol=1e-4, atol=1e-4) + + def testDepthwiseConv2DInputGradCompare(self): + for index, (input_size, filter_size, output_size, stride, + padding) in enumerate(ConfigsToTest()): + print("Testing DepthwiseConv2DInputGradCompare,", index, "th config:", + input_size, "*", filter_size, "stride:", stride, "padding:", + padding) + self._CompareBackpropInput(input_size, filter_size, output_size, stride, + padding) + + def _CompareBackpropFilter(self, input_sizes, filter_sizes, output_sizes, + stride, padding): + x0 = np.random.rand(*input_sizes).astype(np.float32) + x2 = np.random.rand(*output_sizes).astype(np.float32) + + def _GetVal(use_xla): + with self.test_session(): + 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) + if use_xla: + with self.test_scope(): + backprop = nn_ops.depthwise_conv2d_native_backprop_filter( + t0, t1, t2, strides=[1, stride, stride, 1], padding=padding) + else: + backprop = nn_ops.depthwise_conv2d_native_backprop_filter( + t0, t1, t2, strides=[1, stride, stride, 1], padding=padding) + ret = backprop.eval({t0: x0, t2: x2}) + self.assertShapeEqual(ret, backprop) + return ret + + gpu_value = _GetVal(use_xla=True) + cpu_value = _GetVal(use_xla=False) + self.assertAllClose(cpu_value, gpu_value, rtol=1e-4, atol=1e-4) + + def testDepthwiseConv2DFilterGradCompare(self): + 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) + self._CompareBackpropFilter(input_size, filter_size, output_size, + stride, padding) + + +if __name__ == "__main__": + test.main() diff --git a/tensorflow/compiler/tests/ftrl_test.py b/tensorflow/compiler/tests/ftrl_test.py new file mode 100644 index 0000000000000000000000000000000000000000..7e3871312c86530b6d3cb0bbacc16c25d3469832 --- /dev/null +++ b/tensorflow/compiler/tests/ftrl_test.py @@ -0,0 +1,292 @@ +# 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 Ftrl optimizer.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import numpy as np + +from tensorflow.compiler.tests.xla_test import XLATestCase +from tensorflow.python.framework import constant_op +from tensorflow.python.ops import resource_variable_ops +from tensorflow.python.ops import variables +from tensorflow.python.platform import test +from tensorflow.python.training import adagrad +from tensorflow.python.training import ftrl +from tensorflow.python.training import gradient_descent + + +class FtrlOptimizerTest(XLATestCase): + + def initVariableAndGradient(self, dtype): + var0 = resource_variable_ops.ResourceVariable([0.0, 0.0], dtype=dtype) + var1 = resource_variable_ops.ResourceVariable([0.0, 0.0], dtype=dtype) + grads0 = constant_op.constant([0.1, 0.2], dtype=dtype) + grads1 = constant_op.constant([0.02, 0.04], dtype=dtype) + + return var0, var1, grads0, grads1 + + def equivAdagradTest_FtrlPart(self, steps, dtype): + var0, var1, grads0, grads1 = self.initVariableAndGradient(dtype) + opt = ftrl.FtrlOptimizer( + 3.0, + learning_rate_power=-0.5, # using Adagrad learning rate + initial_accumulator_value=0.1, + l1_regularization_strength=0.0, + l2_regularization_strength=0.0) + 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()) + + # Run Ftrl for a few steps + for _ in range(steps): + ftrl_update.run() + + return var0.eval(), var1.eval() + + def equivAdagradTest_AdagradPart(self, steps, dtype): + var0, var1, grads0, grads1 = self.initVariableAndGradient(dtype) + opt = adagrad.AdagradOptimizer(3.0, initial_accumulator_value=0.1) + 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()) + + # Run Adagrad for a few steps + for _ in range(steps): + adagrad_update.run() + + return var0.eval(), var1.eval() + + def equivGradientDescentTest_FtrlPart(self, steps, dtype): + var0, var1, grads0, grads1 = self.initVariableAndGradient(dtype) + opt = ftrl.FtrlOptimizer( + 3.0, + learning_rate_power=-0.0, # using Fixed learning rate + initial_accumulator_value=0.1, + l1_regularization_strength=0.0, + l2_regularization_strength=0.0) + 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()) + + # Run Ftrl for a few steps + for _ in range(steps): + ftrl_update.run() + + return var0.eval(), var1.eval() + + def equivGradientDescentTest_GradientDescentPart(self, steps, dtype): + var0, var1, grads0, grads1 = self.initVariableAndGradient(dtype) + opt = gradient_descent.GradientDescentOptimizer(3.0, name="sgd") + 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()) + + # Run GradientDescent for a few steps + for _ in range(steps): + sgd_update.run() + + return var0.eval(), var1.eval() + + def testFtrlwithoutRegularization(self): + for dtype in self.float_types: + with self.test_session(), self.test_scope(): + var0 = resource_variable_ops.ResourceVariable([0.0, 0.0], dtype=dtype) + var1 = resource_variable_ops.ResourceVariable([0.0, 0.0], dtype=dtype) + grads0 = constant_op.constant([0.1, 0.2], dtype=dtype) + grads1 = constant_op.constant([0.01, 0.02], dtype=dtype) + opt = ftrl.FtrlOptimizer( + 3.0, + initial_accumulator_value=0.1, + l1_regularization_strength=0.0, + l2_regularization_strength=0.0) + 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()) + + # Run 3 steps FTRL + for _ in range(3): + ftrl_update.run() + + # Validate updated params + self.assertAllCloseAccordingToType( + np.array([-2.60260963, -4.29698515]), var0.eval(), float_rtol=1e-5) + self.assertAllCloseAccordingToType( + np.array([-0.28432083, -0.56694895]), var1.eval(), float_rtol=1e-5) + + def testFtrlwithoutRegularization2(self): + for dtype in self.float_types: + with self.test_session(), self.test_scope(): + var0 = resource_variable_ops.ResourceVariable([1.0, 2.0], dtype=dtype) + var1 = resource_variable_ops.ResourceVariable([4.0, 3.0], dtype=dtype) + grads0 = constant_op.constant([0.1, 0.2], dtype=dtype) + grads1 = constant_op.constant([0.01, 0.02], dtype=dtype) + opt = ftrl.FtrlOptimizer( + 3.0, + initial_accumulator_value=0.1, + l1_regularization_strength=0.0, + l2_regularization_strength=0.0) + 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()) + + # Run 3 steps FTRL + for _ in range(3): + ftrl_update.run() + + # Validate updated params + self.assertAllClose( + np.array([-2.55607247, -3.98729396]), var0.eval(), 1e-5, 1e-5) + self.assertAllClose( + np.array([-0.28232238, -0.56096673]), var1.eval(), 1e-5, 1e-5) + + def testFtrlWithL1(self): + for dtype in self.float_types: + with self.test_session(), self.test_scope(): + var0 = resource_variable_ops.ResourceVariable([1.0, 2.0], dtype=dtype) + var1 = resource_variable_ops.ResourceVariable([4.0, 3.0], dtype=dtype) + grads0 = constant_op.constant([0.1, 0.2], dtype=dtype) + grads1 = constant_op.constant([0.01, 0.02], dtype=dtype) + opt = ftrl.FtrlOptimizer( + 3.0, + initial_accumulator_value=0.1, + l1_regularization_strength=0.001, + l2_regularization_strength=0.0) + 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()) + + # Run 10 steps FTRL + for _ in range(10): + ftrl_update.run() + + # Validate updated params + self.assertAllClose(np.array([-7.66718769, -10.91273689]), var0.eval(), + rtol=1e-4) + self.assertAllClose(np.array([-0.93460727, -1.86147261]), var1.eval(), + rtol=1e-4) + + def testFtrlWithL1_L2(self): + for dtype in self.float_types: + with self.test_session(), self.test_scope(): + var0 = resource_variable_ops.ResourceVariable([1.0, 2.0], dtype=dtype) + var1 = resource_variable_ops.ResourceVariable([4.0, 3.0], dtype=dtype) + grads0 = constant_op.constant([0.1, 0.2], dtype=dtype) + grads1 = constant_op.constant([0.01, 0.02], dtype=dtype) + opt = ftrl.FtrlOptimizer( + 3.0, + initial_accumulator_value=0.1, + l1_regularization_strength=0.001, + l2_regularization_strength=2.0) + 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()) + + # Run 10 steps FTRL + for _ in range(10): + ftrl_update.run() + + # Validate updated params + self.assertAllClose(np.array([-0.24059935, -0.46829352]), var0.eval(), + rtol=1e-5) + self.assertAllClose(np.array([-0.02406147, -0.04830509]), var1.eval(), + rtol=1e-5) + + def testFtrlWithL1_L2_L2Shrinkage(self): + """Test the new FTRL op with support for l2 shrinkage. + + The addition of this parameter which places a constant pressure on weights + towards the origin causes the gradient descent trajectory to differ. The + weights will tend to have smaller magnitudes with this parameter set. + """ + for dtype in self.float_types: + with self.test_session(), self.test_scope(): + var0 = resource_variable_ops.ResourceVariable([1.0, 2.0], dtype=dtype) + var1 = resource_variable_ops.ResourceVariable([4.0, 3.0], dtype=dtype) + grads0 = constant_op.constant([0.1, 0.2], dtype=dtype) + grads1 = constant_op.constant([0.01, 0.02], dtype=dtype) + opt = ftrl.FtrlOptimizer( + 3.0, + initial_accumulator_value=0.1, + l1_regularization_strength=0.001, + l2_regularization_strength=2.0, + l2_shrinkage_regularization_strength=0.1) + 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()) + + # Run 10 steps FTRL + for _ in range(10): + ftrl_update.run() + + # Validate updated params + self.assertAllClose(np.array([-0.21931979, -0.40642974]), var0.eval(), + rtol=1e-4) + self.assertAllClose(np.array([-0.0282721, -0.07188385]), var1.eval(), + rtol=1e-4) + + # When variables are initialized with Zero, FTRL-Proximal has two properties: + # 1. Without L1&L2 but with fixed learning rate, FTRL-Proximal is identical + # with GradientDescent. + # 2. Without L1&L2 but with adaptive learning rate, FTRL-Proximal is idential + # with Adagrad. + # So, basing on these two properties, we test if our implementation of + # FTRL-Proximal performs same updates as Adagrad or GradientDescent. + def testEquivAdagradwithoutRegularization(self): + steps = 5 + for dtype in self.float_types: + with self.test_session(), self.test_scope(): + val0, val1 = self.equivAdagradTest_FtrlPart(steps, dtype) + with self.test_session(), self.test_scope(): + val2, val3 = self.equivAdagradTest_AdagradPart(steps, dtype) + + self.assertAllClose(val0, val2, rtol=1e-4) + self.assertAllClose(val1, val3, rtol=1e-4) + + def testEquivGradientDescentwithoutRegularization(self): + steps = 5 + for dtype in self.float_types: + with self.test_session(), self.test_scope(): + val0, val1 = self.equivGradientDescentTest_FtrlPart(steps, dtype) + with self.test_session(), self.test_scope(): + val2, val3 = self.equivGradientDescentTest_GradientDescentPart( + steps, dtype) + + self.assertAllClose(val0, val2, rtol=1e-5) + self.assertAllClose(val1, val3, rtol=1e-5) + + +if __name__ == "__main__": + test.main() diff --git a/tensorflow/compiler/tests/fused_batchnorm_test.py b/tensorflow/compiler/tests/fused_batchnorm_test.py new file mode 100644 index 0000000000000000000000000000000000000000..936fcf8b6be0f8cd67ba07a8bef9d35a732d30ba --- /dev/null +++ b/tensorflow/compiler/tests/fused_batchnorm_test.py @@ -0,0 +1,193 @@ +# 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 tests for fused batch norm operations.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import numpy as np + +from tensorflow.compiler.tests.xla_test import XLATestCase +from tensorflow.python.ops import array_ops +from tensorflow.python.ops import gen_nn_ops +from tensorflow.python.ops import gradient_checker +from tensorflow.python.ops import nn +from tensorflow.python.platform import test + + +class FusedBatchNormTest(XLATestCase): + + def _reference_training(self, x, scale, offset, epsilon, data_format): + if data_format != "NHWC": + raise ValueError("data_format must be NHWC, got %s." % data_format) + x_square = x * x + x_square_sum = np.sum(x_square, (0, 1, 2)) + x_sum = np.sum(x, axis=(0, 1, 2)) + element_count = np.size(x) / int(np.shape(x)[0]) + mean = x_sum / element_count + var = x_square_sum / element_count - mean * mean + normalized = (x - mean) / np.sqrt(var + epsilon) + return (normalized * scale + offset), mean, var + + def _reference_grad(self, x, grad_y, scale, mean, var, epsilon, data_format): + # Use the following formulas to calculate gradients: + # grad_scale = + # sum(grad_y * (x - mean)) * rsqrt(var + epsilon) + # + # grad_offset = sum(output_y) + # + # grad_x = + # 1/N * scale * rsqrt(var + epsilon) * (N * grad_y - sum(grad_y) - + # (x - mean) * sum(grad_y * (x - mean)) / (var + epsilon)) + if data_format != "NHWC": + raise ValueError("data_format must be NHWC, got %s." % data_format) + grad_x = scale * (grad_y - np.mean(grad_y, axis=(0, 1, 2)) - + (x - mean) * np.mean(grad_y * + (x - mean), axis=(0, 1, 2)) / + (var + epsilon)) / np.sqrt(var + epsilon) + grad_scale = np.sum( + grad_y * (x - mean) / np.sqrt(var + epsilon), axis=(0, 1, 2)) + grad_offset = np.sum(grad_y, axis=(0, 1, 2)) + return grad_x, grad_scale, grad_offset + + def testInference(self): + x_shape = [2, 2, 6, 2] + scale_shape = [2] + x_val = np.random.random_sample(x_shape).astype(np.float32) + scale_val = np.random.random_sample(scale_shape).astype(np.float32) + + offset_val = np.random.random_sample(scale_shape).astype(np.float32) + data_format = "NHWC" + with self.test_session() as sess, self.test_scope(): + # To avoid constant folding + t_val = array_ops.placeholder(np.float32, shape=x_shape, name="x") + scale = array_ops.placeholder(np.float32, shape=[2], name="scale") + offset = array_ops.placeholder(np.float32, shape=[2], name="offset") + epsilon = 0.001 + y_ref, mean_ref, var_ref = self._reference_training( + x_val, scale_val, offset_val, epsilon, data_format) + y, mean, variance = nn.fused_batch_norm( + t_val, + scale, + offset, + mean=mean_ref, + variance=var_ref, + epsilon=epsilon, + data_format=data_format, + is_training=False) + + y_val, _, _ = sess.run( + [y, mean, + variance], {t_val: x_val, + scale: scale_val, + offset: offset_val}) + self.assertAllClose(y_val, y_ref, atol=1e-3) + + def _testLearning(self, use_gradient_checker): + x_shape = [2, 2, 6, 2] + scale_shape = [2] + x_val = np.random.random_sample(x_shape).astype(np.float32) + scale_val = np.random.random_sample(scale_shape).astype(np.float32) + + offset_val = np.random.random_sample(scale_shape).astype(np.float32) + mean_val = np.random.random_sample(scale_shape).astype(np.float32) + var_val = np.random.random_sample(scale_shape).astype(np.float32) + data_format = "NHWC" + with self.test_session() as sess, self.test_scope(): + # To avoid constant folding + t_val = array_ops.placeholder(np.float32, shape=x_shape, name="x") + scale = array_ops.placeholder(np.float32, shape=[2], name="scale") + offset = array_ops.placeholder(np.float32, shape=[2], name="offset") + epsilon = 0.001 + y, mean, var = nn.fused_batch_norm( + t_val, + scale, + offset, + mean=None, + variance=None, + epsilon=epsilon, + data_format=data_format, + is_training=True) + # Check gradient. + if use_gradient_checker: + err = gradient_checker.compute_gradient_error( + t_val, + x_shape, + y, + x_shape, + extra_feed_dict={ + t_val: x_val, + scale: scale_val, + offset: offset_val + }) + self.assertLess(err, 1e-3) + + y_val, mean_val, var_val = sess.run( + [y, mean, var], {t_val: x_val, + scale: scale_val, + offset: offset_val}) + y_ref, mean_ref, var_ref = self._reference_training( + x_val, scale_val, offset_val, epsilon, data_format) + self.assertAllClose(mean_val, mean_ref, atol=1e-3) + self.assertAllClose(y_val, y_ref, atol=1e-3) + self.assertAllClose(var_val, var_ref, atol=1e-3) + + def testLearning(self): + self._testLearning(False) + + def testLearningWithGradientChecker(self): + self._testLearning(True) + + def testGradient(self): + # TODO(b/64270657): Use gradient_checker here in addition to comparing with + # this reference implementation. + x_shape = [2, 2, 6, 2] + scale_shape = [2] + grad_val = np.random.random_sample(x_shape).astype(np.float32) + x_val = np.random.random_sample(x_shape).astype(np.float32) + scale_val = np.random.random_sample(scale_shape).astype(np.float32) + mean_val = np.random.random_sample(scale_shape).astype(np.float32) + var_val = np.random.random_sample(scale_shape).astype(np.float32) + epsilon = 0.001 + + with self.test_session() as sess, self.test_scope(): + grad = array_ops.placeholder(np.float32, shape=x_shape, name="grad") + x = array_ops.placeholder(np.float32, shape=x_shape, name="x") + mean = array_ops.placeholder(np.float32, shape=scale_shape, name="mean") + var = array_ops.placeholder(np.float32, shape=scale_shape, name="var") + scale = array_ops.placeholder(np.float32, shape=scale_shape, name="scale") + grad_x, grad_scale, grad_offset, _, _ = gen_nn_ops.fused_batch_norm_grad( + grad, x, scale, mean, var, data_format="NHWC") + + grad_x_val, grad_scale_val, grad_offset_val = sess.run( + [grad_x, grad_scale, grad_offset], { + grad: grad_val, + x: x_val, + mean: mean_val, + var: var_val, + scale: scale_val + }) + + grad_x_ref, grad_scale_ref, grad_offset_ref = self._reference_grad( + x_val, grad_val, scale_val, mean_val, var_val, epsilon, "NHWC") + + self.assertAllClose(grad_x_val, grad_x_ref, atol=1e-2) + self.assertAllClose(grad_scale_val, grad_scale_ref, atol=1e-2) + self.assertAllClose(grad_offset_val, grad_offset_ref, atol=1e-3) + + +if __name__ == "__main__": + test.main() diff --git a/tensorflow/compiler/tests/gather_test.py b/tensorflow/compiler/tests/gather_test.py new file mode 100644 index 0000000000000000000000000000000000000000..9f752dd072bd90b02c2ab801a09a6d17f8ea0e58 --- /dev/null +++ b/tensorflow/compiler/tests/gather_test.py @@ -0,0 +1,105 @@ +# 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 tests for XLA Gather Op.""" + +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.python.framework import constant_op +from tensorflow.python.framework import dtypes +from tensorflow.python.ops import array_ops +from tensorflow.python.platform import test + +_TEST_TYPES = [dtypes.float32] + + +class GatherTest(xla_test.XLATestCase): + + def _buildParams(self, data, dtype): + data = data.astype(dtype.as_numpy_dtype) + # For complex types, adds an index-dependent imaginary component so we can + # tell we got the right value. + if dtype.is_complex: + return data + 10j * data + return data + + def testScalar1D(self): + with self.test_session() as session, self.test_scope(): + data = np.array([0, 1, 2, 3, 7, 5]) + for dtype in _TEST_TYPES: + for indices in 4, [1, 2, 2, 4, 5]: + params_np = self._buildParams(data, dtype) + params = array_ops.placeholder(dtype=dtype) + indices_tf = constant_op.constant(indices) + gather_t = array_ops.gather(params, indices_tf) + gather_val = session.run(gather_t, feed_dict={params: params_np}) + np_val = params_np[indices] + self.assertAllEqual(np_val, gather_val) + self.assertEqual(np_val.shape, gather_val.shape) + + def testScalar2D(self): + with self.test_session() as session, self.test_scope(): + data = np.array([[0, 1, 2], [3, 4, 5], [6, 7, 8], [9, 10, 11], + [12, 13, 14]]) + for dtype in _TEST_TYPES: + params_np = self._buildParams(data, dtype) + params = array_ops.placeholder(dtype=dtype) + indices = constant_op.constant(2) + gather_t = array_ops.gather(params, indices) + gather_val = session.run(gather_t, feed_dict={params: params_np}) + self.assertAllEqual(np.take(params_np, 2, axis=0), gather_val) + expected_shape = data.shape[:0] + data.shape[1:] + self.assertEqual(expected_shape, gather_val.shape) + + def testSimpleTwoD32(self): + with self.test_session() as session, self.test_scope(): + data = np.array([[0, 1, 2], [3, 4, 5], [6, 7, 8], [9, 10, 11], + [12, 13, 14]]) + for dtype in _TEST_TYPES: + params_np = self._buildParams(data, dtype) + params = array_ops.placeholder(dtype=dtype) + # The indices must be in bounds for any axis. + indices = constant_op.constant([0, 1, 0, 2]) + gather_t = array_ops.gather(params, indices) + gather_val = session.run(gather_t, feed_dict={params: params_np}) + self.assertAllEqual( + np.take(params_np, [0, 1, 0, 2], axis=0), gather_val) + expected_shape = data.shape[:0] + (4,) + data.shape[1:] + self.assertEqual(expected_shape, gather_val.shape) + + def testHigherRank(self): + # Check that scalar and empty indices shapes work as well. + shape = (2, 1, 3, 2) + for indices_shape in (), (0,), (2, 0), (2, 3): + for dtype in _TEST_TYPES: + params = self._buildParams(np.random.randn(*shape), dtype) + indices = np.random.randint(shape[0], size=indices_shape) + with self.test_session() as sess, self.test_scope(): + tf_params = array_ops.placeholder(dtype=dtype) + tf_indices = constant_op.constant(indices, dtype=dtypes.int32) + gather = array_ops.gather(tf_params, tf_indices) + gather_value = sess.run(gather, feed_dict={tf_params: params}) + gather_np = np.take(params, indices, 0) + self.assertAllEqual(gather_np, gather_value) + expected_shape = (params.shape[:0] + indices.shape + params.shape[1:]) + self.assertEqual(expected_shape, gather_value.shape) + + +if __name__ == "__main__": + test.main() diff --git a/tensorflow/compiler/tests/plugin.bzl b/tensorflow/compiler/tests/plugin.bzl new file mode 100644 index 0000000000000000000000000000000000000000..fbc8781a3e59faecf985cde5114bf56a041c4be0 --- /dev/null +++ b/tensorflow/compiler/tests/plugin.bzl @@ -0,0 +1,30 @@ +# 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. +# ============================================================================== +"""Additional XLA devices to be included in the unit test suite.""" + +# If you wish to edit this file without checking it into the repo, consider: +# git update-index --assume-unchanged tensorflow/compiler/tests/plugin.bzl + +plugins = { + #"example": { + # "device":"XLA_MY_DEVICE", + # "types":"DT_FLOAT,DT_HALF,DT_INT32", + # "tags":[], + # "args":["--disabled_manifest=tensorflow/compiler/plugin/example/disabled_manifest.txt"], + # "data":["//tensorflow/compiler/plugin/example:disabled_manifest.txt"], + # "deps":[], + #}, +} + diff --git a/tensorflow/compiler/tests/pooling_ops_test.py b/tensorflow/compiler/tests/pooling_ops_test.py index 52290e63548910309a9b6b75b7a4642ebeed1efa..7c19a99c4eb4be3ca34b3ce949216e557b0a681d 100644 --- a/tensorflow/compiler/tests/pooling_ops_test.py +++ b/tensorflow/compiler/tests/pooling_ops_test.py @@ -376,7 +376,7 @@ class PoolGradTest(XLATestCase): self.assertAllClose( expected_input_gradient_vals.flatten(), actual.flatten(), - rtol=1e-5, + rtol=1e-4, atol=1e-6) self.assertShapeEqual(actual, inputs) diff --git a/tensorflow/compiler/tests/randomized_tests.cc b/tensorflow/compiler/tests/randomized_tests.cc index a0cd905f1731ab5aef884a69119d7bb8d3279edd..49c1699b6edc9d16bbba4578fdadd86a14d8c56c 100644 --- a/tensorflow/compiler/tests/randomized_tests.cc +++ b/tensorflow/compiler/tests/randomized_tests.cc @@ -76,6 +76,7 @@ namespace { // Command line flags: see main() below. int64 tf_xla_random_seed = 0; int32 tf_xla_test_repetitions = 20; +int64 tf_xla_max_tensor_size = 100000LL; string* tf_xla_test_device_ptr; // initial value set in main() bool tf_xla_test_use_jit = true; @@ -94,7 +95,12 @@ class OpTestBuilder { explicit OpTestBuilder(const string& op_name); // Adds an input 'tensor'. - OpTestBuilder& Input(Tensor tensor); + OpTestBuilder& Input(const Tensor& tensor); + + // Adds a random input tensor with 'type'. If 'dims' is not provided, + // RandomDims() is used. + OpTestBuilder& RandomInput(DataType type); + OpTestBuilder& RandomInput(DataType type, std::vector dims); // Sets an attribute. template @@ -111,25 +117,54 @@ class OpTestBuilder { // sets it to the NodeDef of the operator under test. Fills 'inputs' and // 'outputs' with the names of the input placeholder nodes and the output // identity nodes, respectively. - Status BuildGraph(string name_prefix, string device, bool use_jit, - GraphDef* graphdef, NodeDef** test_node_def, + Status BuildGraph(const string& name_prefix, const string& device, + bool use_jit, GraphDef* graphdef, NodeDef** test_node_def, std::vector* inputs, std::vector* outputs) const; - const std::vector& inputs() const { return inputs_; } + struct InputDescription { + Tensor tensor; + + DataType type = DT_INVALID; + bool has_dims = false; + std::vector dims; + }; + + const std::vector& inputs() const { return inputs_; } private: NodeDef node_def_; - std::vector inputs_; + std::vector inputs_; }; OpTestBuilder::OpTestBuilder(const string& op_name) { node_def_.set_op(op_name); } -OpTestBuilder& OpTestBuilder::Input(Tensor tensor) { +OpTestBuilder& OpTestBuilder::Input(const Tensor& tensor) { VLOG(1) << "Adding input: " << tensor.DebugString(); - inputs_.push_back(tensor); + InputDescription input; + input.tensor = tensor; + inputs_.push_back(input); + return *this; +} + +OpTestBuilder& OpTestBuilder::RandomInput(DataType type) { + VLOG(1) << "Adding random input: " << type; + InputDescription input; + input.type = type; + inputs_.push_back(input); + return *this; +} + +OpTestBuilder& OpTestBuilder::RandomInput(DataType type, + std::vector dims) { + VLOG(1) << "Adding input: " << type << " " << TensorShape(dims).DebugString(); + InputDescription input; + input.type = type; + input.has_dims = true; + input.dims = std::move(dims); + inputs_.push_back(input); return *this; } @@ -146,9 +181,9 @@ OpTestBuilder& OpTestBuilder::Attr(StringPiece attr_name, return *this; } -Status OpTestBuilder::BuildGraph(string name_prefix, string device, - bool use_jit, GraphDef* graphdef, - NodeDef** test_node_def, +Status OpTestBuilder::BuildGraph(const string& name_prefix, + const string& device, bool use_jit, + GraphDef* graphdef, NodeDef** test_node_def, std::vector* inputs, std::vector* outputs) const { OpRegistryInterface* op_registry = OpRegistry::Global(); @@ -207,23 +242,36 @@ class OpTest : public ::testing::Test { public: OpTest(); - // Runs 'fn' up to --tf_xla_test_repetitions times, or until a failure occurs; - // whichever happens first. - void Repeatedly(std::function fn); + enum TestResult { + // The test saw an unrecoverable error. Don't try any more runs. + kFatalError, + // The parameters of the test were invalid (e.g., the "golden" + // implementation failed, or the parameters are oversize). Reruns are ok. + kInvalid, + // The test ran successfully, and we have a verdict. Does *not* mean the + // test passed. + kOk, + }; + + // Runs 'fn' up to --tf_xla_test_repetitions times, or until a test failure + // occurs; whichever happens first. Reruns if the TestResult is kInvalid. + void Repeatedly(const std::function& fn); // Select a random element from 'candidates'. template T Choose(gtl::ArraySlice candidates); static constexpr int kDefaultMaxRank = 5; - static constexpr int64 kDefaultMaxDimensionSize = 20LL; + static constexpr int64 kDefaultMaxDimensionSize = 256LL; + + // Returns true if 'dims' have a size less than tf_xla_max_tensor_size. + bool TensorSizeIsOk(gtl::ArraySlice dims); - // Returns a random dimension size. + // Returns a random dimension size, in the range [min, max). int64 RandomDim(int64 min = 0, int64 max = kDefaultMaxDimensionSize); // Returns a random shape. The tensor has rank in the range [min_rank, - // max_rank). - // Each dimension has size [0, kDefaultMaxDimensionSize]. + // max_rank). Each dimension has size [min_size, max_size). std::vector RandomDims(int min_rank = 0, int max_rank = kDefaultMaxRank, int64 min_size = 0, @@ -279,8 +327,9 @@ class OpTest : public ::testing::Test { // element-wise difference between x and y must no more than // atol + rtol * abs(x); or both elements may be NaN or infinity. For // non-floating-point tensors the element values must match exactly. - void ExpectTfAndXlaOutputsAreClose(const OpTestBuilder& builder, - double atol = 1e-2, double rtol = 1e-2); + TestResult ExpectTfAndXlaOutputsAreClose(const OpTestBuilder& builder, + double atol = 1e-2, + double rtol = 1e-2); protected: // Per-test state: @@ -303,10 +352,10 @@ OpTest::OpTest() { } else { seed = static_cast(s); } - LOG(INFO) << "Random seed for test case: " << seed - << ". To reproduce the " - "results of this test, pass flag --tf_xla_random_seed=" - << seed; + LOG(ERROR) << "Random seed for test case: " << seed + << ". To reproduce the " + "results of this test, pass flag --tf_xla_random_seed=" + << seed; generator_.reset(new std::mt19937(seed)); // Create a session with an empty graph. @@ -316,10 +365,35 @@ OpTest::OpTest() { TF_CHECK_OK(session_->Create(def)); } -void OpTest::Repeatedly(std::function fn) { +void OpTest::Repeatedly(const std::function& fn) { int const max_repetitions = tf_xla_test_repetitions; - for (int i = 0; !HasFailure() && i < max_repetitions; ++i) { - fn(); + int valid_test_runs = 0; + // We run up to 20 * max_repetitions times; the idea is that if we roll the + // dice enough times we will find some valid parameters. We want to put an + // upper limit on the number iterations just in case the probability of + // finding feasible parameters is very low. + for (int i = 0; !HasFailure() && i < max_repetitions * 20 && + valid_test_runs < max_repetitions; + ++i) { + TestResult result = fn(); + switch (result) { + case kOk: + ++valid_test_runs; + break; + + case kFatalError: + ASSERT_TRUE(false) << "Test had fatal failure"; + return; + + case kInvalid: + break; + } + } + if (!HasFailure()) { + EXPECT_GE(valid_test_runs, max_repetitions) + << "Not enough test instances passed; this means that either the " + "golden implementation is buggy or the operator harness is not " + "producing well-formed test cases with a high probability."; } } @@ -334,6 +408,14 @@ int64 OpTest::RandomDim(int64 min, int64 max) { return size_distribution(generator()); } +bool OpTest::TensorSizeIsOk(gtl::ArraySlice dims) { + int64 size = 1LL; + for (int64 dim : dims) { + size *= dim; + } + return size < tf_xla_max_tensor_size; +} + std::vector OpTest::RandomDims(int min_rank, int max_rank, int64 min_size, int64 max_size) { CHECK_LE(0, min_rank); @@ -341,9 +423,13 @@ std::vector OpTest::RandomDims(int min_rank, int max_rank, std::uniform_int_distribution rank_distribution(min_rank, max_rank); int rank = rank_distribution(generator()); std::vector dims(rank); - std::generate(dims.begin(), dims.end(), [this, min_size, max_size]() { - return RandomDim(min_size, max_size); - }); + // TODO(phawkins): too small a maximum tensor size could lead to an infinite + // loop here. + do { + std::generate(dims.begin(), dims.end(), [this, min_size, max_size]() { + return RandomDim(min_size, max_size); + }); + } while (!TensorSizeIsOk(dims)); return dims; } @@ -527,6 +613,8 @@ std::vector OpTest::ImageDims(TensorFormat format, int batch, dims.push_back(dim); } break; + case FORMAT_NCHW_VECT_C: + LOG(FATAL) << "FORMAT_NCHW_VECT_C not supported."; } return dims; } @@ -607,53 +695,84 @@ Status TensorsAreClose(const Tensor& a, const Tensor& b, double atol, } } -void OpTest::ExpectTfAndXlaOutputsAreClose(const OpTestBuilder& builder, - double atol, double rtol) { +OpTest::TestResult OpTest::ExpectTfAndXlaOutputsAreClose( + const OpTestBuilder& builder, double atol, double rtol) { + const std::vector& inputs = builder.inputs(); + std::vector input_tensors; + input_tensors.reserve(inputs.size()); + for (const OpTestBuilder::InputDescription& input : inputs) { + if (input.type == DT_INVALID) { + input_tensors.push_back(input.tensor); + } else { + std::vector dims; + if (input.has_dims) { + dims = input.dims; + } else { + dims = RandomDims(); + } + if (!TensorSizeIsOk(dims)) { + VLOG(1) << "Input: " << input.type << " " + << TensorShape(input.dims).DebugString(); + VLOG(1) << "Ignoring oversize dims."; + return kInvalid; + } + input_tensors.push_back(RandomTensor(input.type, dims)); + } + VLOG(1) << "Input: " << input_tensors.back().DebugString(); + } + string cpu_device = LocalDeviceToFullDeviceName(strings::StrCat(DEVICE_CPU, ":0")); string test_device = LocalDeviceToFullDeviceName(*tf_xla_test_device_ptr); DeviceNameUtils::ParsedName parsed_name; - ASSERT_TRUE( - DeviceNameUtils::ParseLocalName(*tf_xla_test_device_ptr, &parsed_name)); + if (!DeviceNameUtils::ParseLocalName(*tf_xla_test_device_ptr, &parsed_name)) { + LOG(ERROR) << "Could not parse device name: " << *tf_xla_test_device_ptr; + return kFatalError; + } DeviceType test_device_type(parsed_name.type); ++num_tests_; GraphDef graph; std::vector expected_inputs, test_inputs; std::vector expected_fetches, test_fetches; - TF_ASSERT_OK(builder.BuildGraph( + Status status = builder.BuildGraph( strings::StrCat("test", num_tests_, "_expected"), cpu_device, /* use_jit= */ false, &graph, /* test_node_def= */ nullptr, - &expected_inputs, &expected_fetches)); + &expected_inputs, &expected_fetches); + if (!status.ok()) { + LOG(ERROR) << "Expected graph construction failed: " << status; + return kFatalError; + } NodeDef* node_def; - TF_ASSERT_OK(builder.BuildGraph(strings::StrCat("test", num_tests_, "_test"), - test_device, tf_xla_test_use_jit, &graph, - &node_def, &test_inputs, &test_fetches)); + status = builder.BuildGraph(strings::StrCat("test", num_tests_, "_test"), + test_device, tf_xla_test_use_jit, &graph, + &node_def, &test_inputs, &test_fetches); + if (!status.ok()) { + LOG(ERROR) << "Test graph construction failed: " << status; + return kFatalError; + } // Check that there's a kernel corresponding to 'node_def' on the device under // test. - Status status = FindKernelDef(test_device_type, *node_def, nullptr, nullptr); + status = FindKernelDef(test_device_type, *node_def, nullptr, nullptr); if (!status.ok()) { VLOG(1) << "Skipping test because there is no corresponding registered " << "kernel on the test device: " << status; - return; + return kInvalid; } - TF_ASSERT_OK(session_->Extend(graph)); - - const std::vector& input_tensors = builder.inputs(); - if (VLOG_IS_ON(1)) { - for (const Tensor& input : input_tensors) { - VLOG(1) << "Input: " << input.DebugString(); - } + status = session_->Extend(graph); + if (!status.ok()) { + LOG(ERROR) << "Session::Extend() failed: " << status; + return kFatalError; } std::vector> expected_feeds(expected_inputs.size()); std::vector> test_feeds(test_inputs.size()); - ASSERT_EQ(input_tensors.size(), expected_inputs.size()); - ASSERT_EQ(input_tensors.size(), test_inputs.size()); + CHECK_EQ(input_tensors.size(), expected_inputs.size()); + CHECK_EQ(input_tensors.size(), test_inputs.size()); for (int i = 0; i < input_tensors.size(); ++i) { expected_feeds[i] = {expected_inputs[i], input_tensors[i]}; @@ -665,18 +784,27 @@ void OpTest::ExpectTfAndXlaOutputsAreClose(const OpTestBuilder& builder, Status s = session_->Run(expected_feeds, expected_fetches, {}, &expected_outputs); if (!s.ok()) { - VLOG(1) << "Expected graph failed with status: " << s << ". Skipping test"; - return; + VLOG(1) << "Expected graph failed with status: " << s << ". Ignoring test"; + return kInvalid; + } + for (const Tensor& expected : expected_outputs) { + VLOG(1) << "Expected: " << expected.DebugString(); } VLOG(1) << "Running test graph"; - TF_ASSERT_OK(session_->Run(test_feeds, test_fetches, {}, &test_outputs)); + status = session_->Run(test_feeds, test_fetches, {}, &test_outputs); + if (!status.ok()) { + LOG(ERROR) << "Test graph failed: " << status; + return kFatalError; + } - ASSERT_EQ(expected_outputs.size(), test_outputs.size()); + CHECK_EQ(expected_outputs.size(), test_outputs.size()); for (int j = 0; s.ok() && j < test_outputs.size(); ++j) { s = TensorsAreClose(expected_outputs[j], test_outputs[j], atol, rtol); } TF_EXPECT_OK(s); + + return kOk; } // Helper that converts 'values' to an int32 or int64 Tensor. @@ -696,8 +824,15 @@ Tensor AsIntTensor(DataType dtype, const std::vector& values) { TEST_F(OpTest, Abs) { Repeatedly([this]() { DataType type = Choose({DT_INT32, DT_FLOAT}); - ExpectTfAndXlaOutputsAreClose( - OpTestBuilder("Abs").Input(RandomTensor(type)).Attr("T", type)); + return ExpectTfAndXlaOutputsAreClose( + OpTestBuilder("Abs").RandomInput(type).Attr("T", type)); + }); +} + +TEST_F(OpTest, Acosh) { + Repeatedly([this]() { + return ExpectTfAndXlaOutputsAreClose( + OpTestBuilder("Acosh").RandomInput(DT_FLOAT).Attr("T", DT_FLOAT)); }); } @@ -705,10 +840,10 @@ TEST_F(OpTest, Add) { Repeatedly([this]() { DataType type = Choose({DT_INT32, DT_FLOAT}); auto dims = BroadcastableDims(); - ExpectTfAndXlaOutputsAreClose(OpTestBuilder("Add") - .Input(RandomTensor(type, dims.first)) - .Input(RandomTensor(type, dims.second)) - .Attr("T", type)); + return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("Add") + .RandomInput(type, dims.first) + .RandomInput(type, dims.second) + .Attr("T", type)); }); } @@ -723,49 +858,74 @@ TEST_F(OpTest, AddN) { builder.Attr("T", type); builder.Attr("N", n); for (int i = 0; i < n; ++i) { - builder.Input(RandomTensor(type, shape)); + builder.RandomInput(type, shape); } - ExpectTfAndXlaOutputsAreClose(builder); + return ExpectTfAndXlaOutputsAreClose(builder); }); } TEST_F(OpTest, All) { Repeatedly([this]() { - Tensor data = RandomTensor(DT_BOOL); - Tensor indices = RandomReductionIndices(data.dims()); + std::vector data_dims = RandomDims(); + Tensor indices = RandomReductionIndices(data_dims.size()); bool keep_dims = Choose({false, true}); - ExpectTfAndXlaOutputsAreClose( - OpTestBuilder("All").Input(data).Input(indices).Attr("keep_dims", - keep_dims)); + return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("All") + .RandomInput(DT_BOOL, data_dims) + .Input(indices) + .Attr("keep_dims", keep_dims)); }); } TEST_F(OpTest, Any) { Repeatedly([this]() { - Tensor data = RandomTensor(DT_BOOL); - Tensor indices = RandomReductionIndices(data.dims()); + std::vector data_dims = RandomDims(); + Tensor indices = RandomReductionIndices(data_dims.size()); bool keep_dims = Choose({false, true}); - ExpectTfAndXlaOutputsAreClose( - OpTestBuilder("Any").Input(data).Input(indices).Attr("keep_dims", - keep_dims)); + return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("Any") + .RandomInput(DT_BOOL, data_dims) + .Input(indices) + .Attr("keep_dims", keep_dims)); + }); +} + +TEST_F(OpTest, ApproximateEqual) { + Repeatedly([this]() { + auto dims = RandomDims(); + return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("ApproximateEqual") + .RandomInput(DT_FLOAT, dims) + .RandomInput(DT_FLOAT, dims) + .Attr("T", DT_FLOAT)); + }); +} + +TEST_F(OpTest, Asinh) { + Repeatedly([this]() { + return ExpectTfAndXlaOutputsAreClose( + OpTestBuilder("Asinh").RandomInput(DT_FLOAT).Attr("T", DT_FLOAT)); + }); +} + +TEST_F(OpTest, Atanh) { + Repeatedly([this]() { + return ExpectTfAndXlaOutputsAreClose( + OpTestBuilder("Atanh").RandomInput(DT_FLOAT).Attr("T", DT_FLOAT)); }); } TEST_F(OpTest, AvgPool) { Repeatedly([this]() { - WindowedSpatialDims d = ChooseWindowedSpatialDims(2); std::uniform_int_distribution random_int(1, 5); - - int kernel_rows = random_int(generator()), - kernel_cols = random_int(generator()); + std::vector dims = RandomDims(4, 4, 1); + int kernel_rows = + std::uniform_int_distribution(1, dims[1])(generator()); + int kernel_cols = + std::uniform_int_distribution(1, dims[2])(generator()); int stride_rows = random_int(generator()), stride_cols = random_int(generator()); string padding = Choose({"SAME", "VALID"}); - ExpectTfAndXlaOutputsAreClose( + return ExpectTfAndXlaOutputsAreClose( OpTestBuilder("AvgPool") - .Input( - RandomTensor(DT_FLOAT, {RandomDim(1), RandomDim(kernel_rows), - RandomDim(kernel_cols), RandomDim(1)})) + .RandomInput(DT_FLOAT, dims) .Attr("T", DT_FLOAT) .Attr("ksize", {1, kernel_rows, kernel_cols, 1}) .Attr("strides", {1, stride_rows, stride_cols, 1}) @@ -779,23 +939,28 @@ TEST_F(OpTest, AvgPool) { TEST_F(OpTest, AvgPool3D) { Repeatedly([this]() { std::uniform_int_distribution random_int(1, 5); + std::vector dims = RandomDims(5, 5, 1); + std::vector input_dims, kernel_dims, stride_dims; for (int i = 0; i < 3; ++i) { - kernel_dims.push_back(random_int(generator())); - input_dims.push_back(RandomDim(kernel_dims.back())); + kernel_dims.push_back( + std::uniform_int_distribution(1, dims[i])(generator())); + input_dims.push_back(dims[i]); stride_dims.push_back(random_int(generator())); } + int64 batch = dims[3]; + int64 feature = dims[4]; string padding = Choose({"SAME", "VALID"}); - ExpectTfAndXlaOutputsAreClose( + return ExpectTfAndXlaOutputsAreClose( OpTestBuilder("AvgPool3D") - .Input(RandomTensor(DT_FLOAT, ImageDims(FORMAT_NHWC, RandomDim(1), - RandomDim(1), input_dims))) + .RandomInput(DT_FLOAT, + ImageDims(FORMAT_NHWC, batch, feature, input_dims)) .Attr("T", DT_FLOAT) .Attr("ksize", ImageDims(FORMAT_NHWC, 1, 1, kernel_dims)) .Attr("strides", ImageDims(FORMAT_NHWC, 1, 1, stride_dims)) .Attr("padding", padding) - .Attr("data_format", "NHWC")); + .Attr("data_format", "NDHWC")); }); // TODO(phawkins): test NCHW format (not supported by CPU) } @@ -808,10 +973,10 @@ TEST_F(OpTest, AvgPoolGrad) { AsInt32s(ImageDims(FORMAT_NHWC, batch, features, d.input_dims)); std::vector output_dims = ImageDims(FORMAT_NHWC, batch, features, d.output_dims); - ExpectTfAndXlaOutputsAreClose( + return ExpectTfAndXlaOutputsAreClose( OpTestBuilder("AvgPoolGrad") .Input(test::AsTensor(input_dims)) - .Input(RandomTensor(DT_FLOAT, output_dims)) + .RandomInput(DT_FLOAT, output_dims) .Attr("T", DT_FLOAT) .Attr("ksize", ImageDims(FORMAT_NHWC, 1, 1, d.kernel_dims)) .Attr("strides", ImageDims(FORMAT_NHWC, 1, 1, d.stride_dims)) @@ -828,15 +993,15 @@ TEST_F(OpTest, AvgPool3DGrad) { AsInt32s(ImageDims(FORMAT_NHWC, batch, features, d.input_dims)); std::vector output_dims = ImageDims(FORMAT_NHWC, batch, features, d.output_dims); - ExpectTfAndXlaOutputsAreClose( + return ExpectTfAndXlaOutputsAreClose( OpTestBuilder("AvgPool3DGrad") .Input(test::AsTensor(input_dims)) - .Input(RandomTensor(DT_FLOAT, output_dims)) + .RandomInput(DT_FLOAT, output_dims) .Attr("T", DT_FLOAT) .Attr("ksize", ImageDims(FORMAT_NHWC, 1, 1, d.kernel_dims)) .Attr("strides", ImageDims(FORMAT_NHWC, 1, 1, d.stride_dims)) .Attr("padding", d.padding == SAME ? "SAME" : "VALID") - .Attr("data_format", "NHWC")); + .Attr("data_format", "NDHWC")); }); } @@ -848,60 +1013,127 @@ TEST_F(OpTest, BatchMatMul) { std::vector x_dims(output_dims), y_dims(output_dims); x_dims[ndims - 1] = inner_dim; y_dims[ndims - 2] = inner_dim; - ExpectTfAndXlaOutputsAreClose(OpTestBuilder("BatchMatMul") - .Input(RandomTensor(DT_FLOAT, x_dims)) - .Input(RandomTensor(DT_FLOAT, y_dims)) - .Attr("T", DT_FLOAT)); - - std::swap(x_dims[ndims - 1], x_dims[ndims - 2]); - ExpectTfAndXlaOutputsAreClose(OpTestBuilder("BatchMatMul") - .Input(RandomTensor(DT_FLOAT, x_dims)) - .Input(RandomTensor(DT_FLOAT, y_dims)) - .Attr("T", DT_FLOAT) - .Attr("adj_x", true)); - - std::swap(y_dims[ndims - 1], y_dims[ndims - 2]); - ExpectTfAndXlaOutputsAreClose(OpTestBuilder("BatchMatMul") - .Input(RandomTensor(DT_FLOAT, x_dims)) - .Input(RandomTensor(DT_FLOAT, y_dims)) - .Attr("T", DT_FLOAT) - .Attr("adj_x", true) - .Attr("adj_y", true)); - - std::swap(x_dims[ndims - 1], x_dims[ndims - 2]); - ExpectTfAndXlaOutputsAreClose(OpTestBuilder("BatchMatMul") - .Input(RandomTensor(DT_FLOAT, x_dims)) - .Input(RandomTensor(DT_FLOAT, y_dims)) - .Attr("T", DT_FLOAT) - .Attr("adj_y", true)); + + std::bernoulli_distribution random_bool; + bool adj_x = random_bool(generator()); + bool adj_y = random_bool(generator()); + if (adj_x) { + std::swap(x_dims[ndims - 1], x_dims[ndims - 2]); + } + if (adj_y) { + std::swap(y_dims[ndims - 1], y_dims[ndims - 2]); + } + + return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("BatchMatMul") + .RandomInput(DT_FLOAT, x_dims) + .RandomInput(DT_FLOAT, y_dims) + .Attr("T", DT_FLOAT) + .Attr("adj_x", adj_x) + .Attr("adj_y", adj_y)); + }); +} + +TEST_F(OpTest, BatchToSpace) { + Repeatedly([this]() { + const int num_block_dims = 2; + std::vector block_dims = + RandomDims(num_block_dims, num_block_dims, 0, 5); + int64 block_size = RandomDim(2, 5); + + std::vector input_dims(1 + num_block_dims + 1); + input_dims[0] = RandomDim(); + for (int i = 0; i < num_block_dims; ++i) { + input_dims[0] *= block_size; + input_dims[1 + i] = block_dims[i]; + } + input_dims[1 + num_block_dims] = RandomDim(); + + std::vector crop_vals; + std::uniform_int_distribution distribution(0, 4); + for (int i = 0; i < num_block_dims; ++i) { + // Chooses crop values; does not always choose legal values. + crop_vals.push_back(distribution(generator())); + crop_vals.push_back(distribution(generator())); + } + Tensor crops; + CHECK(crops.CopyFrom(AsIntTensor(DT_INT32, crop_vals), + TensorShape({num_block_dims, 2}))); + + return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("BatchToSpace") + .RandomInput(DT_FLOAT, input_dims) + .Input(crops) + .Attr("T", DT_FLOAT) + .Attr("block_size", block_size)); + }); +} + +TEST_F(OpTest, BatchToSpaceND) { + Repeatedly([this]() { + std::vector block_dims = RandomDims(1, 3, 0, 5); + int num_block_dims = block_dims.size(); + std::vector remaining_dims = RandomDims(0, 3); + std::vector block_multipliers = + RandomDims(block_dims.size(), block_dims.size(), 0, 4); + + std::vector input_dims(1 + num_block_dims + remaining_dims.size()); + input_dims[0] = RandomDim(); + for (int i = 0; i < num_block_dims; ++i) { + input_dims[0] *= block_dims[i]; + } + std::copy(block_multipliers.begin(), block_multipliers.end(), + input_dims.begin() + 1); + std::copy(remaining_dims.begin(), remaining_dims.end(), + input_dims.begin() + 1 + num_block_dims); + + std::vector crop_vals; + std::uniform_int_distribution distribution(0, 3); + for (int i = 0; i < num_block_dims; ++i) { + // Chooses crop values; does not always choose legal values. + crop_vals.push_back(distribution(generator())); + crop_vals.push_back(distribution(generator())); + } + Tensor crops; + CHECK(crops.CopyFrom(AsIntTensor(DT_INT32, crop_vals), + TensorShape({num_block_dims, 2}))); + + return ExpectTfAndXlaOutputsAreClose( + OpTestBuilder("BatchToSpaceND") + .RandomInput(DT_FLOAT, input_dims) + .Input(test::AsTensor( + std::vector(block_dims.begin(), block_dims.end()))) + .Input(crops) + .Attr("T", DT_FLOAT)); }); } TEST_F(OpTest, BiasAdd) { Repeatedly([this]() { - auto x = RandomTensor(DT_FLOAT, RandomDims(2, kDefaultMaxRank)); - auto y = RandomTensor(DT_FLOAT, {x.dim_size(x.dims() - 1)}); + auto x_dims = RandomDims(2, kDefaultMaxRank); + auto y_dims = {x_dims[x_dims.size() - 1]}; // TODO(phawkins): test both data formats. - ExpectTfAndXlaOutputsAreClose( - OpTestBuilder("BiasAdd").Input(x).Input(y).Attr("T", DT_FLOAT)); + return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("BiasAdd") + .RandomInput(DT_FLOAT, x_dims) + .RandomInput(DT_FLOAT, y_dims) + .Attr("T", DT_FLOAT)); }); } TEST_F(OpTest, BiasAddGrad) { Repeatedly([this]() { - auto x = RandomTensor(DT_FLOAT); // TODO(phawkins): test both data formats. - ExpectTfAndXlaOutputsAreClose( - OpTestBuilder("BiasAddGrad").Input(x).Attr("T", DT_FLOAT)); + return ExpectTfAndXlaOutputsAreClose( + OpTestBuilder("BiasAddGrad").RandomInput(DT_FLOAT).Attr("T", DT_FLOAT)); }); } TEST_F(OpTest, BiasAddV1) { Repeatedly([this]() { - auto x = RandomTensor(DT_FLOAT, RandomDims(2, kDefaultMaxRank)); - auto y = RandomTensor(DT_FLOAT, {x.dim_size(x.dims() - 1)}); - ExpectTfAndXlaOutputsAreClose( - OpTestBuilder("BiasAddV1").Input(x).Input(y).Attr("T", DT_FLOAT)); + auto x_dims = RandomDims(2, kDefaultMaxRank); + auto y_dims = {x_dims[x_dims.size() - 1]}; + return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("BiasAddV1") + .RandomInput(DT_FLOAT, x_dims) + .RandomInput(DT_FLOAT, y_dims) + .Attr("T", DT_FLOAT)); }); } @@ -911,10 +1143,11 @@ TEST_F(OpTest, BroadcastGradientArgs) { // DataType type = Choose({DT_INT32, DT_INT64}); DataType type = DT_INT32; auto dims = BroadcastableDims(); - ExpectTfAndXlaOutputsAreClose(OpTestBuilder("BroadcastGradientArgs") - .Input(AsIntTensor(type, dims.first)) - .Input(AsIntTensor(type, dims.second)) - .Attr("T", type)); + return ExpectTfAndXlaOutputsAreClose( + OpTestBuilder("BroadcastGradientArgs") + .Input(AsIntTensor(type, dims.first)) + .Input(AsIntTensor(type, dims.second)) + .Attr("T", type)); }); } @@ -923,18 +1156,17 @@ TEST_F(OpTest, Cast) { DataType src_type, dst_type; src_type = Choose({DT_INT32, DT_FLOAT, DT_BOOL}); dst_type = Choose({DT_INT32, DT_FLOAT, DT_BOOL}); - ExpectTfAndXlaOutputsAreClose(OpTestBuilder("Cast") - .Input(RandomTensor(src_type)) - .Attr("SrcT", src_type) - .Attr("DstT", dst_type)); + return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("Cast") + .RandomInput(src_type) + .Attr("SrcT", src_type) + .Attr("DstT", dst_type)); }); } TEST_F(OpTest, Ceil) { Repeatedly([this]() { - ExpectTfAndXlaOutputsAreClose(OpTestBuilder("Ceil") - .Input(RandomTensor(DT_FLOAT)) - .Attr("T", DT_FLOAT)); + return ExpectTfAndXlaOutputsAreClose( + OpTestBuilder("Ceil").RandomInput(DT_FLOAT).Attr("T", DT_FLOAT)); }); } @@ -954,9 +1186,9 @@ TEST_F(OpTest, Concat) { for (int i = 0; i < n; ++i) { std::vector shape = dims; shape[concat_dim] = RandomDim(); - builder.Input(RandomTensor(type, shape)); + builder.RandomInput(type, shape); } - ExpectTfAndXlaOutputsAreClose(builder); + return ExpectTfAndXlaOutputsAreClose(builder); }); } @@ -976,7 +1208,7 @@ TEST_F(OpTest, ConcatOffset) { shape[concat_dim] = RandomDim(); builder.Input(test::AsTensor(shape)); } - ExpectTfAndXlaOutputsAreClose(builder); + return ExpectTfAndXlaOutputsAreClose(builder); }); } @@ -989,15 +1221,15 @@ TEST_F(OpTest, Conv2D) { int64 batch = RandomDim(); - Tensor data = RandomTensor( - DT_FLOAT, ImageDims(FORMAT_NHWC, batch, features_in, d.input_dims)); + std::vector data_dims = + ImageDims(FORMAT_NHWC, batch, features_in, d.input_dims); - Tensor kernel = RandomTensor(DT_FLOAT, {d.kernel_dims[0], d.kernel_dims[1], - features_in, features_out}); - ExpectTfAndXlaOutputsAreClose( + std::vector kernel_dims = {d.kernel_dims[0], d.kernel_dims[1], + features_in, features_out}; + return ExpectTfAndXlaOutputsAreClose( OpTestBuilder("Conv2D") - .Input(data) - .Input(kernel) + .RandomInput(DT_FLOAT, data_dims) + .RandomInput(DT_FLOAT, kernel_dims) .Attr("T", DT_FLOAT) .Attr("strides", ImageDims(FORMAT_NHWC, 1, 1, d.stride_dims)) .Attr("padding", d.padding == SAME ? "SAME" : "VALID") @@ -1012,17 +1244,17 @@ TEST_F(OpTest, Conv2DBackpropFilter) { int features_in = random_int(generator()); int features_out = random_int(generator()); int32 batch = RandomDim(); - Tensor activations = RandomTensor( - DT_FLOAT, ImageDims(FORMAT_NHWC, batch, features_in, d.input_dims)); - Tensor backprop = RandomTensor( - DT_FLOAT, ImageDims(FORMAT_NHWC, batch, features_out, d.output_dims)); + std::vector activations = + ImageDims(FORMAT_NHWC, batch, features_in, d.input_dims); + std::vector backprop = + ImageDims(FORMAT_NHWC, batch, features_out, d.output_dims); Tensor kernel_shape = test::AsTensor(AsInt32s( {d.kernel_dims[0], d.kernel_dims[1], features_in, features_out})); - ExpectTfAndXlaOutputsAreClose( + return ExpectTfAndXlaOutputsAreClose( OpTestBuilder("Conv2DBackpropFilter") - .Input(activations) + .RandomInput(DT_FLOAT, activations) .Input(kernel_shape) - .Input(backprop) + .RandomInput(DT_FLOAT, backprop) .Attr("T", DT_FLOAT) .Attr("strides", ImageDims(FORMAT_NHWC, 1, 1, d.stride_dims)) .Attr("padding", d.padding == SAME ? "SAME" : "VALID") @@ -1039,15 +1271,15 @@ TEST_F(OpTest, Conv2DBackpropInput) { int32 batch = RandomDim(); Tensor in_shape = test::AsTensor( AsInt32s(ImageDims(FORMAT_NHWC, batch, features_in, d.input_dims))); - Tensor backprop = RandomTensor( - DT_FLOAT, ImageDims(FORMAT_NHWC, batch, features_out, d.output_dims)); - Tensor kernel = RandomTensor(DT_FLOAT, {d.kernel_dims[0], d.kernel_dims[1], - features_in, features_out}); - ExpectTfAndXlaOutputsAreClose( + std::vector backprop = + ImageDims(FORMAT_NHWC, batch, features_out, d.output_dims); + std::vector kernel = {d.kernel_dims[0], d.kernel_dims[1], + features_in, features_out}; + return ExpectTfAndXlaOutputsAreClose( OpTestBuilder("Conv2DBackpropInput") .Input(in_shape) - .Input(kernel) - .Input(backprop) + .RandomInput(DT_FLOAT, kernel) + .RandomInput(DT_FLOAT, backprop) .Attr("T", DT_FLOAT) .Attr("strides", ImageDims(FORMAT_NHWC, 1, 1, d.stride_dims)) .Attr("padding", d.padding == SAME ? "SAME" : "VALID") @@ -1061,17 +1293,15 @@ TEST_F(OpTest, Conv3D) { std::uniform_int_distribution random_int(1, 5); int features_in = random_int(generator()); int features_out = random_int(generator()); - Tensor data = - RandomTensor(DT_FLOAT, {RandomDim(), d.input_dims[0], d.input_dims[1], - d.input_dims[2], features_in}); - - Tensor kernel = - RandomTensor(DT_FLOAT, {d.kernel_dims[0], d.kernel_dims[1], - d.kernel_dims[2], features_in, features_out}); - ExpectTfAndXlaOutputsAreClose( + std::vector data = {RandomDim(), d.input_dims[0], d.input_dims[1], + d.input_dims[2], features_in}; + + std::vector kernel = {d.kernel_dims[0], d.kernel_dims[1], + d.kernel_dims[2], features_in, features_out}; + return ExpectTfAndXlaOutputsAreClose( OpTestBuilder("Conv3D") - .Input(data) - .Input(kernel) + .RandomInput(DT_FLOAT, data) + .RandomInput(DT_FLOAT, kernel) .Attr("T", DT_FLOAT) .Attr("strides", ImageDims(FORMAT_NHWC, 1, 1, d.stride_dims)) .Attr("padding", d.padding == SAME ? "SAME" : "VALID")); @@ -1085,18 +1315,18 @@ TEST_F(OpTest, Conv3DBackpropFilter) { int features_in = random_int(generator()); int features_out = random_int(generator()); int32 batch = RandomDim(1); - Tensor activations = RandomTensor( - DT_FLOAT, ImageDims(FORMAT_NHWC, batch, features_in, d.input_dims)); - Tensor backprop = RandomTensor( - DT_FLOAT, ImageDims(FORMAT_NHWC, batch, features_out, d.output_dims)); + std::vector activations = + ImageDims(FORMAT_NHWC, batch, features_in, d.input_dims); + std::vector backprop = + ImageDims(FORMAT_NHWC, batch, features_out, d.output_dims); Tensor kernel_shape = test::AsTensor( AsInt32s({d.kernel_dims[0], d.kernel_dims[1], d.kernel_dims[2], features_in, features_out})); - ExpectTfAndXlaOutputsAreClose( + return ExpectTfAndXlaOutputsAreClose( OpTestBuilder("Conv3DBackpropFilterV2") - .Input(activations) + .RandomInput(DT_FLOAT, activations) .Input(kernel_shape) - .Input(backprop) + .RandomInput(DT_FLOAT, backprop) .Attr("T", DT_FLOAT) .Attr("strides", ImageDims(FORMAT_NHWC, 1, 1, d.stride_dims)) .Attr("padding", d.padding == SAME ? "SAME" : "VALID")); @@ -1112,28 +1342,118 @@ TEST_F(OpTest, Conv3DBackpropInput) { int32 batch = RandomDim(1); Tensor in_shape = test::AsTensor( AsInt32s(ImageDims(FORMAT_NHWC, batch, features_in, d.input_dims))); - Tensor backprop = RandomTensor( - DT_FLOAT, ImageDims(FORMAT_NHWC, batch, features_out, d.output_dims)); - Tensor kernel = - RandomTensor(DT_FLOAT, {d.kernel_dims[0], d.kernel_dims[1], - d.kernel_dims[2], features_in, features_out}); - ExpectTfAndXlaOutputsAreClose( + std::vector backprop = + ImageDims(FORMAT_NHWC, batch, features_out, d.output_dims); + std::vector kernel = {d.kernel_dims[0], d.kernel_dims[1], + d.kernel_dims[2], features_in, features_out}; + return ExpectTfAndXlaOutputsAreClose( OpTestBuilder("Conv3DBackpropInputV2") .Input(in_shape) - .Input(kernel) - .Input(backprop) + .RandomInput(DT_FLOAT, kernel) + .RandomInput(DT_FLOAT, backprop) + .Attr("T", DT_FLOAT) + .Attr("strides", ImageDims(FORMAT_NHWC, 1, 1, d.stride_dims)) + .Attr("padding", d.padding == SAME ? "SAME" : "VALID")); + }); +} + +TEST_F(OpTest, DepthwiseConv2DNative) { + Repeatedly([this]() { + WindowedSpatialDims d = ChooseWindowedSpatialDims(2); + std::uniform_int_distribution random_int(1, 5); + int features_in = random_int(generator()); + int depth_multiplier = random_int(generator()); + std::vector input_dims = {RandomDim(), d.input_dims[0], + d.input_dims[1], features_in}; + + std::vector kernel_dims = {d.kernel_dims[0], d.kernel_dims[1], + features_in, depth_multiplier}; + return ExpectTfAndXlaOutputsAreClose( + OpTestBuilder("DepthwiseConv2dNative") + .RandomInput(DT_FLOAT, input_dims) + .RandomInput(DT_FLOAT, kernel_dims) .Attr("T", DT_FLOAT) .Attr("strides", ImageDims(FORMAT_NHWC, 1, 1, d.stride_dims)) .Attr("padding", d.padding == SAME ? "SAME" : "VALID")); }); } +TEST_F(OpTest, DepthwiseConv2DBackpropFilter) { + Repeatedly([this]() { + WindowedSpatialDims d = ChooseWindowedSpatialDims(2); + std::uniform_int_distribution random_int(1, 5); + int features_in = random_int(generator()); + int depth_multiplier = random_int(generator()); + int32 batch = RandomDim(); + std::vector activations = + ImageDims(FORMAT_NHWC, batch, features_in, d.input_dims); + std::vector backprop = ImageDims( + FORMAT_NHWC, batch, features_in * depth_multiplier, d.output_dims); + Tensor kernel_shape = test::AsTensor(AsInt32s( + {d.kernel_dims[0], d.kernel_dims[1], features_in, depth_multiplier})); + return ExpectTfAndXlaOutputsAreClose( + OpTestBuilder("DepthwiseConv2dNativeBackpropFilter") + .RandomInput(DT_FLOAT, activations) + .Input(kernel_shape) + .RandomInput(DT_FLOAT, backprop) + .Attr("T", DT_FLOAT) + .Attr("strides", ImageDims(FORMAT_NHWC, 1, 1, d.stride_dims)) + .Attr("padding", d.padding == SAME ? "SAME" : "VALID") + .Attr("data_format", "NHWC")); + }); +} + +TEST_F(OpTest, Cos) { + Repeatedly([this]() { + return ExpectTfAndXlaOutputsAreClose( + OpTestBuilder("Cos").RandomInput(DT_FLOAT).Attr("T", DT_FLOAT)); + }); +} + +TEST_F(OpTest, Cosh) { + Repeatedly([this]() { + return ExpectTfAndXlaOutputsAreClose( + OpTestBuilder("Cosh").RandomInput(DT_FLOAT).Attr("T", DT_FLOAT)); + }); +} + +TEST_F(OpTest, DepthwiseConv2DBackpropInput) { + Repeatedly([this]() { + WindowedSpatialDims d = ChooseWindowedSpatialDims(2); + std::uniform_int_distribution random_int(1, 5); + int features_in = random_int(generator()); + int depth_multiplier = random_int(generator()); + int32 batch = RandomDim(); + Tensor in_shape = test::AsTensor( + AsInt32s(ImageDims(FORMAT_NHWC, batch, features_in, d.input_dims))); + std::vector backprop = ImageDims( + FORMAT_NHWC, batch, features_in * depth_multiplier, d.output_dims); + std::vector kernel = {d.kernel_dims[0], d.kernel_dims[1], + features_in, depth_multiplier}; + return ExpectTfAndXlaOutputsAreClose( + OpTestBuilder("DepthwiseConv2dNativeBackpropInput") + .Input(in_shape) + .RandomInput(DT_FLOAT, kernel) + .RandomInput(DT_FLOAT, backprop) + .Attr("T", DT_FLOAT) + .Attr("strides", ImageDims(FORMAT_NHWC, 1, 1, d.stride_dims)) + .Attr("padding", d.padding == SAME ? "SAME" : "VALID") + .Attr("data_format", "NHWC")); + }); +} + TEST_F(OpTest, Diag) { Repeatedly([this]() { DataType type = Choose({DT_INT32, DT_FLOAT}); - ExpectTfAndXlaOutputsAreClose(OpTestBuilder("Diag") - .Input(RandomTensor(type, RandomDims(1))) - .Attr("T", type)); + std::vector dims; + // Diag causes a quadratic blowup in output size. + int64 size; + do { + dims = RandomDims(1); + size = TensorShape(dims).num_elements(); + } while (size * size < tf_xla_max_tensor_size); + return ExpectTfAndXlaOutputsAreClose( + OpTestBuilder("Diag").RandomInput(type, dims).Attr("T", type)); }); } @@ -1145,9 +1465,9 @@ TEST_F(OpTest, DiagPart) { std::vector doubled_dims(dims.size() * 2); std::copy(dims.begin(), dims.end(), doubled_dims.begin()); std::copy(dims.begin(), dims.end(), doubled_dims.begin() + dims.size()); - ExpectTfAndXlaOutputsAreClose(OpTestBuilder("DiagPart") - .Input(RandomTensor(type, doubled_dims)) - .Attr("T", type)); + return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("DiagPart") + .RandomInput(type, doubled_dims) + .Attr("T", type)); }); } @@ -1155,10 +1475,10 @@ TEST_F(OpTest, Div) { Repeatedly([this]() { DataType type = Choose({DT_INT32, DT_FLOAT}); auto dims = BroadcastableDims(); - ExpectTfAndXlaOutputsAreClose(OpTestBuilder("Div") - .Input(RandomTensor(type, dims.first)) - .Input(RandomTensor(type, dims.second)) - .Attr("T", type)); + return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("Div") + .RandomInput(type, dims.first) + .RandomInput(type, dims.second) + .Attr("T", type)); }); } @@ -1184,11 +1504,11 @@ TEST_F(OpTest, DynamicStitch) { } while (size == 0); // Shuffle the range of indices that cover the output. - // TODO(phawkins): The documentation for DynamicStitch doesn't require that - // the indices cover all positions of the output. The XLA implementation - // does so require. However, the native TF implementation leaves undefined - // values if we don't cover everything, so we can't really test that case - // anyway. + // TODO(phawkins): The documentation for DynamicStitch doesn't require + // that the indices cover all positions of the output. The XLA + // implementation does so require. However, the native TF implementation + // leaves undefined values if we don't cover everything, so we can't + // really test that case anyway. std::vector indices(size); std::iota(indices.begin(), indices.end(), 0); std::shuffle(indices.begin(), indices.end(), generator()); @@ -1207,10 +1527,43 @@ TEST_F(OpTest, DynamicStitch) { std::vector dims(index_dims[i].begin(), index_dims[i].end()); std::copy(constant_dims.begin(), constant_dims.end(), std::back_inserter(dims)); - Tensor t = RandomTensor(type, dims); - builder.Input(t); + builder.RandomInput(type, dims); } - ExpectTfAndXlaOutputsAreClose(builder); + return ExpectTfAndXlaOutputsAreClose(builder); + }); +} + +TEST_F(OpTest, Elu) { + Repeatedly([this]() { + return ExpectTfAndXlaOutputsAreClose( + OpTestBuilder("Elu").RandomInput(DT_FLOAT).Attr("T", DT_FLOAT)); + }); +} + +TEST_F(OpTest, EluGrad) { + Repeatedly([this]() { + auto dims = RandomDims(); + return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("EluGrad") + .RandomInput(DT_FLOAT, dims) + .RandomInput(DT_FLOAT, dims) + .Attr("T", DT_FLOAT)); + }); +} + +TEST_F(OpTest, Selu) { + Repeatedly([this]() { + return ExpectTfAndXlaOutputsAreClose( + OpTestBuilder("Selu").RandomInput(DT_FLOAT).Attr("T", DT_FLOAT)); + }); +} + +TEST_F(OpTest, SeluGrad) { + Repeatedly([this]() { + auto dims = RandomDims(); + return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("SeluGrad") + .RandomInput(DT_FLOAT, dims) + .RandomInput(DT_FLOAT, dims) + .Attr("T", DT_FLOAT)); }); } @@ -1218,50 +1571,58 @@ TEST_F(OpTest, Equal) { Repeatedly([this]() { DataType type = Choose({DT_INT32, DT_FLOAT}); auto dims = BroadcastableDims(); - ExpectTfAndXlaOutputsAreClose(OpTestBuilder("Equal") - .Input(RandomTensor(type, dims.first)) - .Input(RandomTensor(type, dims.second)) - .Attr("T", type)); + return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("Equal") + .RandomInput(type, dims.first) + .RandomInput(type, dims.second) + .Attr("T", type)); }); } TEST_F(OpTest, Exp) { Repeatedly([this]() { - ExpectTfAndXlaOutputsAreClose( - OpTestBuilder("Exp").Input(RandomTensor(DT_FLOAT)).Attr("T", DT_FLOAT)); + return ExpectTfAndXlaOutputsAreClose( + OpTestBuilder("Exp").RandomInput(DT_FLOAT).Attr("T", DT_FLOAT)); + }); +} + +TEST_F(OpTest, Expm1) { + Repeatedly([this]() { + return ExpectTfAndXlaOutputsAreClose( + OpTestBuilder("Expm1").RandomInput(DT_FLOAT).Attr("T", DT_FLOAT)); }); } TEST_F(OpTest, ExpandDims) { Repeatedly([this]() { DataType type = Choose(kAllXlaTypes); - Tensor in = RandomTensor(type); + std::vector in_dims = RandomDims(); Tensor dim(DT_INT32, TensorShape()); - std::uniform_int_distribution d(-1 - in.dims(), in.dims()); + std::uniform_int_distribution d(-1 - in_dims.size(), in_dims.size()); dim.scalar()() = d(generator()); - ExpectTfAndXlaOutputsAreClose( - OpTestBuilder("ExpandDims").Input(in).Input(dim).Attr("T", type)); + return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("ExpandDims") + .RandomInput(type, in_dims) + .Input(dim) + .Attr("T", type)); }); } TEST_F(OpTest, Fill) { Repeatedly([this]() { DataType type = Choose(kAllXlaTypes); - Tensor scalar = RandomTensor(type, {}); std::vector dims = RandomDims(); std::vector shape(dims.begin(), dims.end()); - ExpectTfAndXlaOutputsAreClose(OpTestBuilder("Fill") - .Input(test::AsTensor(shape)) - .Input(scalar) - .Attr("T", type)); + return ExpectTfAndXlaOutputsAreClose( + OpTestBuilder("Fill") + .Input(test::AsTensor(shape)) + .RandomInput(type, {}) + .Attr("T", type)); }); } TEST_F(OpTest, Floor) { Repeatedly([this]() { - ExpectTfAndXlaOutputsAreClose(OpTestBuilder("Floor") - .Input(RandomTensor(DT_FLOAT)) - .Attr("T", DT_FLOAT)); + return ExpectTfAndXlaOutputsAreClose( + OpTestBuilder("Floor").RandomInput(DT_FLOAT).Attr("T", DT_FLOAT)); }); } @@ -1269,10 +1630,10 @@ TEST_F(OpTest, FloorDiv) { Repeatedly([this]() { DataType type = DT_INT32; auto dims = BroadcastableDims(); - ExpectTfAndXlaOutputsAreClose(OpTestBuilder("FloorDiv") - .Input(RandomTensor(type, dims.first)) - .Input(RandomTensor(type, dims.second)) - .Attr("T", type)); + return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("FloorDiv") + .RandomInput(type, dims.first) + .RandomInput(type, dims.second) + .Attr("T", type)); }); } @@ -1280,10 +1641,10 @@ TEST_F(OpTest, FloorMod) { Repeatedly([this]() { DataType type = Choose({DT_INT32, DT_FLOAT}); auto dims = BroadcastableDims(); - ExpectTfAndXlaOutputsAreClose(OpTestBuilder("FloorMod") - .Input(RandomTensor(type, dims.first)) - .Input(RandomTensor(type, dims.second)) - .Attr("T", type)); + return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("FloorMod") + .RandomInput(type, dims.first) + .RandomInput(type, dims.second) + .Attr("T", type)); }); } @@ -1291,10 +1652,10 @@ TEST_F(OpTest, Greater) { Repeatedly([this]() { DataType type = Choose({DT_INT32, DT_FLOAT}); auto dims = BroadcastableDims(); - ExpectTfAndXlaOutputsAreClose(OpTestBuilder("Greater") - .Input(RandomTensor(type, dims.first)) - .Input(RandomTensor(type, dims.second)) - .Attr("T", type)); + return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("Greater") + .RandomInput(type, dims.first) + .RandomInput(type, dims.second) + .Attr("T", type)); }); } @@ -1302,28 +1663,18 @@ TEST_F(OpTest, GreaterEqual) { Repeatedly([this]() { DataType type = Choose({DT_INT32, DT_FLOAT}); auto dims = BroadcastableDims(); - ExpectTfAndXlaOutputsAreClose(OpTestBuilder("GreaterEqual") - .Input(RandomTensor(type, dims.first)) - .Input(RandomTensor(type, dims.second)) - .Attr("T", type)); - }); -} - -TEST_F(OpTest, Reciprocal) { - Repeatedly([this]() { - ExpectTfAndXlaOutputsAreClose(OpTestBuilder("Reciprocal") - .Input(RandomTensor(DT_FLOAT)) - .Attr("T", DT_FLOAT)); + return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("GreaterEqual") + .RandomInput(type, dims.first) + .RandomInput(type, dims.second) + .Attr("T", type)); }); } TEST_F(OpTest, L2Loss) { Repeatedly([this]() { - DataType type = Choose({DT_INT32, DT_FLOAT}); - // TODO(b/31644876): scalars currently crash. - ExpectTfAndXlaOutputsAreClose(OpTestBuilder("L2Loss") - .Input(RandomTensor(type, RandomDims(1))) - .Attr("T", type)); + DataType type = DT_FLOAT; + return ExpectTfAndXlaOutputsAreClose( + OpTestBuilder("L2Loss").RandomInput(type).Attr("T", type)); }); } @@ -1331,10 +1682,10 @@ TEST_F(OpTest, Less) { Repeatedly([this]() { DataType type = Choose({DT_INT32, DT_FLOAT}); auto dims = BroadcastableDims(); - ExpectTfAndXlaOutputsAreClose(OpTestBuilder("Less") - .Input(RandomTensor(type, dims.first)) - .Input(RandomTensor(type, dims.second)) - .Attr("T", type)); + return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("Less") + .RandomInput(type, dims.first) + .RandomInput(type, dims.second) + .Attr("T", type)); }); } @@ -1342,10 +1693,10 @@ TEST_F(OpTest, LessEqual) { Repeatedly([this]() { DataType type = Choose({DT_INT32, DT_FLOAT}); auto dims = BroadcastableDims(); - ExpectTfAndXlaOutputsAreClose(OpTestBuilder("LessEqual") - .Input(RandomTensor(type, dims.first)) - .Input(RandomTensor(type, dims.second)) - .Attr("T", type)); + return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("LessEqual") + .RandomInput(type, dims.first) + .RandomInput(type, dims.second) + .Attr("T", type)); }); } @@ -1357,10 +1708,10 @@ TEST_F(OpTest, LinSpace) { }; std::uniform_int_distribution distribution(-50, 50); DataType type = Choose({DT_INT32, DT_INT64}); - ExpectTfAndXlaOutputsAreClose( + return ExpectTfAndXlaOutputsAreClose( OpTestBuilder("LinSpace") - .Input(RandomTensor(DT_FLOAT, {})) - .Input(RandomTensor(DT_FLOAT, {})) + .RandomInput(DT_FLOAT, {}) + .RandomInput(DT_FLOAT, {}) .Input(ToScalar(type, distribution(generator()))) .Attr("T", DT_FLOAT) .Attr("Tidx", type)); @@ -1369,62 +1720,69 @@ TEST_F(OpTest, LinSpace) { TEST_F(OpTest, Log) { Repeatedly([this]() { - ExpectTfAndXlaOutputsAreClose( - OpTestBuilder("Log").Input(RandomTensor(DT_FLOAT)).Attr("T", DT_FLOAT)); + return ExpectTfAndXlaOutputsAreClose( + OpTestBuilder("Log").RandomInput(DT_FLOAT).Attr("T", DT_FLOAT)); + }); +} + +TEST_F(OpTest, Log1p) { + Repeatedly([this]() { + return ExpectTfAndXlaOutputsAreClose( + OpTestBuilder("Log1p").RandomInput(DT_FLOAT).Attr("T", DT_FLOAT)); }); } TEST_F(OpTest, LogicalAnd) { Repeatedly([this]() { auto dims = BroadcastableDims(); - ExpectTfAndXlaOutputsAreClose( + return ExpectTfAndXlaOutputsAreClose( OpTestBuilder("LogicalAnd") - .Input(RandomTensor(DT_BOOL, dims.first)) - .Input(RandomTensor(DT_BOOL, dims.second))); + .RandomInput(DT_BOOL, dims.first) + .RandomInput(DT_BOOL, dims.second)); }); } TEST_F(OpTest, LogicalNot) { Repeatedly([this]() { - ExpectTfAndXlaOutputsAreClose( - OpTestBuilder("LogicalNot").Input(RandomTensor(DT_BOOL))); + return ExpectTfAndXlaOutputsAreClose( + OpTestBuilder("LogicalNot").RandomInput(DT_BOOL)); }); } TEST_F(OpTest, LogicalOr) { Repeatedly([this]() { auto dims = BroadcastableDims(); - ExpectTfAndXlaOutputsAreClose( + return ExpectTfAndXlaOutputsAreClose( OpTestBuilder("LogicalOr") - .Input(RandomTensor(DT_BOOL, dims.first)) - .Input(RandomTensor(DT_BOOL, dims.second))); + .RandomInput(DT_BOOL, dims.first) + .RandomInput(DT_BOOL, dims.second)); }); } TEST_F(OpTest, LogSoftmax) { Repeatedly([this]() { - ExpectTfAndXlaOutputsAreClose( + return ExpectTfAndXlaOutputsAreClose( OpTestBuilder("LogSoftmax") - .Input(RandomTensor(DT_FLOAT, RandomDims(2, 2))) + .RandomInput(DT_FLOAT, RandomDims(2, 2)) .Attr("T", DT_FLOAT)); }); } TEST_F(OpTest, LRN) { Repeatedly([this]() { - Tensor data; // TODO(b/31362467): Crashes with 0 dims on GPU. Re-enable when fixed. - data = RandomTensor(DT_FLOAT, RandomDims(4, 4, 1, 8)); + std::vector data_dims = RandomDims(4, 4, 1, 8); // CuDNN requires depth_radius > 0. - std::uniform_int_distribution radius(1, data.dim_size(3)); + std::uniform_int_distribution radius(1, data_dims[3]); std::uniform_real_distribution coeff(0.01, 2.0); - ExpectTfAndXlaOutputsAreClose(OpTestBuilder("LRN") - .Input(data) - .Attr("T", DT_FLOAT) - .Attr("depth_radius", radius(generator())) - .Attr("bias", coeff(generator())) - .Attr("alpha", coeff(generator())) - .Attr("beta", coeff(generator()))); + return ExpectTfAndXlaOutputsAreClose( + OpTestBuilder("LRN") + .RandomInput(DT_FLOAT, data_dims) + .Attr("T", DT_FLOAT) + .Attr("depth_radius", radius(generator())) + .Attr("bias", coeff(generator())) + .Attr("alpha", coeff(generator())) + .Attr("beta", coeff(generator()))); }); } @@ -1432,21 +1790,19 @@ TEST_F(OpTest, LRNGrad) { Repeatedly([this]() { // TODO(b/31362467): Crashes with 0 dims on GPU. Re-enable when fixed. std::vector dims = RandomDims(4, 4, 1, 8); - Tensor input_grads = RandomTensor(DT_FLOAT, dims); - Tensor input_image = RandomTensor(DT_FLOAT, dims); - Tensor output_image = RandomTensor(DT_FLOAT, dims); // CuDNN requires depth_radius > 0. - std::uniform_int_distribution radius(1, input_grads.dim_size(3)); + std::uniform_int_distribution radius(1, dims[3]); std::uniform_real_distribution coeff(0.0, 2.0); - ExpectTfAndXlaOutputsAreClose(OpTestBuilder("LRNGrad") - .Input(input_grads) - .Input(input_image) - .Input(output_image) - .Attr("T", DT_FLOAT) - .Attr("depth_radius", radius(generator())) - .Attr("bias", coeff(generator())) - .Attr("alpha", coeff(generator())) - .Attr("beta", coeff(generator()))); + return ExpectTfAndXlaOutputsAreClose( + OpTestBuilder("LRNGrad") + .RandomInput(DT_FLOAT, dims) + .RandomInput(DT_FLOAT, dims) + .RandomInput(DT_FLOAT, dims) + .Attr("T", DT_FLOAT) + .Attr("depth_radius", radius(generator())) + .Attr("bias", coeff(generator())) + .Attr("alpha", coeff(generator())) + .Attr("beta", coeff(generator()))); }); } @@ -1456,59 +1812,57 @@ TEST_F(OpTest, MatMul) { int64 y = RandomDim(); int64 z = RandomDim(); - ExpectTfAndXlaOutputsAreClose(OpTestBuilder("MatMul") - .Input(RandomTensor(DT_FLOAT, {x, y})) - .Input(RandomTensor(DT_FLOAT, {y, z})) - .Attr("T", DT_FLOAT)); + std::vector a_dims = {x, y}; + std::vector b_dims = {y, z}; - ExpectTfAndXlaOutputsAreClose(OpTestBuilder("MatMul") - .Input(RandomTensor(DT_FLOAT, {y, x})) - .Input(RandomTensor(DT_FLOAT, {y, z})) - .Attr("T", DT_FLOAT) - .Attr("transpose_a", true)); - - ExpectTfAndXlaOutputsAreClose(OpTestBuilder("MatMul") - .Input(RandomTensor(DT_FLOAT, {x, y})) - .Input(RandomTensor(DT_FLOAT, {z, y})) - .Attr("T", DT_FLOAT) - .Attr("transpose_b", true)); + std::bernoulli_distribution random_bool; + bool transpose_a = random_bool(generator()); + bool transpose_b = random_bool(generator()); + if (transpose_a) { + std::swap(a_dims[0], a_dims[1]); + } + if (transpose_b) { + std::swap(b_dims[0], b_dims[1]); + } - ExpectTfAndXlaOutputsAreClose(OpTestBuilder("MatMul") - .Input(RandomTensor(DT_FLOAT, {y, x})) - .Input(RandomTensor(DT_FLOAT, {z, y})) - .Attr("T", DT_FLOAT) - .Attr("transpose_a", true) - .Attr("transpose_b", true)); + return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("MatMul") + .RandomInput(DT_FLOAT, a_dims) + .RandomInput(DT_FLOAT, b_dims) + .Attr("T", DT_FLOAT) + .Attr("transpose_a", transpose_a) + .Attr("transpose_b", transpose_b)); }); } TEST_F(OpTest, MatrixDiag) { Repeatedly([this]() { - DataType type = Choose({DT_BOOL, DT_INT32, DT_FLOAT}); - ExpectTfAndXlaOutputsAreClose(OpTestBuilder("MatrixDiag") - .Input(RandomTensor(type, RandomDims(1))) - .Attr("T", type)); + DataType type = Choose({DT_INT32, DT_FLOAT}); + return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("MatrixDiag") + .RandomInput(type, RandomDims(1)) + .Attr("T", type)); }); } TEST_F(OpTest, MatrixDiagPart) { Repeatedly([this]() { - DataType type = Choose({DT_BOOL, DT_INT32, DT_FLOAT}); - ExpectTfAndXlaOutputsAreClose(OpTestBuilder("MatrixDiagPart") - .Input(RandomTensor(type, RandomDims(2))) - .Attr("T", type)); + DataType type = Choose({DT_INT32, DT_FLOAT}); + return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("MatrixDiagPart") + .RandomInput(type, RandomDims(2)) + .Attr("T", type)); }); } TEST_F(OpTest, Max) { Repeatedly([this]() { DataType type = Choose({DT_INT32, DT_FLOAT}); - Tensor data = RandomTensor(type); - Tensor indices = RandomReductionIndices(data.dims()); + std::vector data_dims = RandomDims(); + Tensor indices = RandomReductionIndices(data_dims.size()); bool keep_dims = Choose({false, true}); - ExpectTfAndXlaOutputsAreClose( - OpTestBuilder("Max").Input(data).Input(indices).Attr("T", type).Attr( - "keep_dims", keep_dims)); + return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("Max") + .RandomInput(type, data_dims) + .Input(indices) + .Attr("T", type) + .Attr("keep_dims", keep_dims)); }); } @@ -1516,26 +1870,28 @@ TEST_F(OpTest, Maximum) { Repeatedly([this]() { DataType type = Choose({DT_INT32, DT_FLOAT}); auto dims = BroadcastableDims(); - ExpectTfAndXlaOutputsAreClose(OpTestBuilder("Maximum") - .Input(RandomTensor(type, dims.first)) - .Input(RandomTensor(type, dims.second)) - .Attr("T", type)); + return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("Maximum") + .RandomInput(type, dims.first) + .RandomInput(type, dims.second) + .Attr("T", type)); }); } TEST_F(OpTest, MaxPool) { Repeatedly([this]() { std::uniform_int_distribution random_int(1, 5); - int kernel_rows = random_int(generator()), - kernel_cols = random_int(generator()); + std::vector dims = RandomDims(4, 4, 1); + int kernel_rows = + std::uniform_int_distribution(1, dims[1])(generator()); + int kernel_cols = + std::uniform_int_distribution(1, dims[2])(generator()); int stride_rows = random_int(generator()), stride_cols = random_int(generator()); + string padding = Choose({"SAME", "VALID"}); - ExpectTfAndXlaOutputsAreClose( + return ExpectTfAndXlaOutputsAreClose( OpTestBuilder("MaxPool") - .Input( - RandomTensor(DT_FLOAT, {RandomDim(1), RandomDim(kernel_rows), - RandomDim(kernel_cols), RandomDim(1)})) + .RandomInput(DT_FLOAT, dims) .Attr("T", DT_FLOAT) .Attr("ksize", {1, kernel_rows, kernel_cols, 1}) .Attr("strides", {1, stride_rows, stride_cols, 1}) @@ -1548,28 +1904,32 @@ TEST_F(OpTest, MaxPool) { TEST_F(OpTest, MaxPool3D) { Repeatedly([this]() { std::uniform_int_distribution random_int(1, 5); - std::vector input_dims; - std::vector kernel_dims, stride_dims; - input_dims.push_back(RandomDim(1)); + std::vector dims = RandomDims(5, 5, 1); + + std::vector input_dims, kernel_dims, stride_dims; kernel_dims.push_back(1); stride_dims.push_back(1); for (int i = 0; i < 3; ++i) { - kernel_dims.push_back(random_int(generator())); - input_dims.push_back(RandomDim(kernel_dims.back())); + kernel_dims.push_back( + std::uniform_int_distribution(1, dims[i])(generator())); + input_dims.push_back(dims[i]); stride_dims.push_back(random_int(generator())); } - input_dims.push_back(RandomDim(1)); kernel_dims.push_back(1); stride_dims.push_back(1); + int64 batch = dims[3]; + int64 feature = dims[4]; string padding = Choose({"SAME", "VALID"}); - ExpectTfAndXlaOutputsAreClose(OpTestBuilder("MaxPool3D") - .Input(RandomTensor(DT_FLOAT, input_dims)) - .Attr("T", DT_FLOAT) - .Attr("ksize", kernel_dims) - .Attr("strides", stride_dims) - .Attr("padding", padding) - .Attr("data_format", "NHWC")); + return ExpectTfAndXlaOutputsAreClose( + OpTestBuilder("MaxPool3D") + .RandomInput(DT_FLOAT, + ImageDims(FORMAT_NHWC, batch, feature, input_dims)) + .Attr("T", DT_FLOAT) + .Attr("ksize", kernel_dims) + .Attr("strides", stride_dims) + .Attr("padding", padding) + .Attr("data_format", "NDHWC")); }); // TODO(phawkins): test NCHW format (not supported by CPU) } @@ -1579,24 +1939,28 @@ TEST_F(OpTest, Mean) { DataType type = Choose({DT_INT32, DT_FLOAT}); // TODO(phawkins): CPU and XLA differ output for reducing across a // size-0 dimension (nan vs 0). For now, require size >= 1. - Tensor data = RandomTensor(type, RandomDims(0, kDefaultMaxRank, 1)); - Tensor indices = RandomReductionIndices(data.dims()); + std::vector data_dims = RandomDims(0, kDefaultMaxRank, 1); + Tensor indices = RandomReductionIndices(data_dims.size()); bool keep_dims = Choose({false, true}); - ExpectTfAndXlaOutputsAreClose( - OpTestBuilder("Mean").Input(data).Input(indices).Attr("T", type).Attr( - "keep_dims", keep_dims)); + return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("Mean") + .RandomInput(type, data_dims) + .Input(indices) + .Attr("T", type) + .Attr("keep_dims", keep_dims)); }); } TEST_F(OpTest, Min) { Repeatedly([this]() { DataType type = Choose({DT_INT32, DT_FLOAT}); - Tensor data = RandomTensor(type); - Tensor indices = RandomReductionIndices(data.dims()); + std::vector data_dims = RandomDims(); + Tensor indices = RandomReductionIndices(data_dims.size()); bool keep_dims = Choose({false, true}); - ExpectTfAndXlaOutputsAreClose( - OpTestBuilder("Min").Input(data).Input(indices).Attr("T", type).Attr( - "keep_dims", keep_dims)); + return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("Min") + .RandomInput(type, data_dims) + .Input(indices) + .Attr("T", type) + .Attr("keep_dims", keep_dims)); }); } @@ -1604,21 +1968,20 @@ TEST_F(OpTest, Minimum) { Repeatedly([this]() { DataType type = Choose({DT_INT32, DT_FLOAT}); auto dims = BroadcastableDims(); - ExpectTfAndXlaOutputsAreClose(OpTestBuilder("Minimum") - .Input(RandomTensor(type, dims.first)) - .Input(RandomTensor(type, dims.second)) - .Attr("T", type)); + return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("Minimum") + .RandomInput(type, dims.first) + .RandomInput(type, dims.second) + .Attr("T", type)); }); } TEST_F(OpTest, Mod) { Repeatedly([this]() { auto dims = BroadcastableDims(); - ExpectTfAndXlaOutputsAreClose( - OpTestBuilder("Mod") - .Input(RandomTensor(DT_INT32, dims.first)) - .Input(RandomTensor(DT_INT32, dims.second)) - .Attr("T", DT_INT32)); + return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("Mod") + .RandomInput(DT_INT32, dims.first) + .RandomInput(DT_INT32, dims.second) + .Attr("T", DT_INT32)); }); } @@ -1626,18 +1989,18 @@ TEST_F(OpTest, Mul) { Repeatedly([this]() { DataType type = Choose({DT_INT32, DT_FLOAT}); auto dims = BroadcastableDims(); - ExpectTfAndXlaOutputsAreClose(OpTestBuilder("Mul") - .Input(RandomTensor(type, dims.first)) - .Input(RandomTensor(type, dims.second)) - .Attr("T", type)); + return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("Mul") + .RandomInput(type, dims.first) + .RandomInput(type, dims.second) + .Attr("T", type)); }); } TEST_F(OpTest, Neg) { Repeatedly([this]() { DataType type = Choose({DT_INT32, DT_FLOAT}); - ExpectTfAndXlaOutputsAreClose( - OpTestBuilder("Neg").Input(RandomTensor(type)).Attr("T", type)); + return ExpectTfAndXlaOutputsAreClose( + OpTestBuilder("Neg").RandomInput(type).Attr("T", type)); }); } @@ -1645,10 +2008,10 @@ TEST_F(OpTest, NotEqual) { Repeatedly([this]() { DataType type = Choose({DT_INT32, DT_FLOAT}); auto dims = BroadcastableDims(); - ExpectTfAndXlaOutputsAreClose(OpTestBuilder("NotEqual") - .Input(RandomTensor(type, dims.first)) - .Input(RandomTensor(type, dims.second)) - .Attr("T", type)); + return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("NotEqual") + .RandomInput(type, dims.first) + .RandomInput(type, dims.second) + .Attr("T", type)); }); } @@ -1676,9 +2039,17 @@ TEST_F(OpTest, OneHot) { builder.Attr("axis", axis); builder.Input(indices); builder.Input(test::AsScalar(depth)); - builder.Input(RandomTensor(type, {})); - builder.Input(RandomTensor(type, {})); - ExpectTfAndXlaOutputsAreClose(builder); + builder.RandomInput(type, {}); + builder.RandomInput(type, {}); + return ExpectTfAndXlaOutputsAreClose(builder); + }); +} + +TEST_F(OpTest, OnesLike) { + Repeatedly([this]() { + DataType type = Choose({DT_INT32, DT_FLOAT}); + return ExpectTfAndXlaOutputsAreClose( + OpTestBuilder("OnesLike").RandomInput(type).Attr("T", type)); }); } @@ -1697,9 +2068,9 @@ TEST_F(OpTest, Pack) { builder.Attr("N", n); builder.Attr("axis", axis); for (int i = 0; i < n; ++i) { - builder.Input(RandomTensor(type, dims)); + builder.RandomInput(type, dims); } - ExpectTfAndXlaOutputsAreClose(builder); + return ExpectTfAndXlaOutputsAreClose(builder); }); } @@ -1707,23 +2078,26 @@ TEST_F(OpTest, Pack) { TEST_F(OpTest, Pad) { Repeatedly([this]() { DataType type = Choose(kAllXlaTypes); - Tensor t = RandomTensor(type); + 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; std::vector paddings_vec; std::uniform_int_distribution distribution(0, 7); - for (int i = 0; i < t.dims(); ++i) { + for (int i = 0; i < t_dims.size(); ++i) { paddings_vec.push_back(distribution(generator())); paddings_vec.push_back(distribution(generator())); } Tensor paddings; - CHECK(paddings.CopyFrom(AsIntTensor(tpaddings, paddings_vec), - TensorShape({t.dims(), 2}))); - ExpectTfAndXlaOutputsAreClose( - OpTestBuilder("Pad").Input(t).Input(paddings).Attr("T", type).Attr( - "Tpaddings", tpaddings)); + CHECK( + paddings.CopyFrom(AsIntTensor(tpaddings, paddings_vec), + TensorShape({static_cast(t_dims.size()), 2}))); + return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("Pad") + .RandomInput(type, t_dims) + .Input(paddings) + .Attr("T", type) + .Attr("Tpaddings", tpaddings)); }); } @@ -1732,23 +2106,24 @@ TEST_F(OpTest, Pow) { // nontermination. Repeatedly([this]() { auto dims = BroadcastableDims(); - ExpectTfAndXlaOutputsAreClose( - OpTestBuilder("Pow") - .Input(RandomTensor(DT_FLOAT, dims.first)) - .Input(RandomTensor(DT_FLOAT, dims.second)) - .Attr("T", DT_FLOAT)); + return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("Pow") + .RandomInput(DT_FLOAT, dims.first) + .RandomInput(DT_FLOAT, dims.second) + .Attr("T", DT_FLOAT)); }); } TEST_F(OpTest, Prod) { Repeatedly([this]() { DataType type = Choose({DT_INT32, DT_FLOAT}); - Tensor data = RandomTensor(type); - Tensor indices = RandomReductionIndices(data.dims()); + std::vector data_dims = RandomDims(); + Tensor indices = RandomReductionIndices(data_dims.size()); bool keep_dims = Choose({false, true}); - ExpectTfAndXlaOutputsAreClose( - OpTestBuilder("Prod").Input(data).Input(indices).Attr("T", type).Attr( - "keep_dims", keep_dims)); + return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("Prod") + .RandomInput(type, data_dims) + .Input(indices) + .Attr("T", type) + .Attr("keep_dims", keep_dims)); }); } @@ -1763,7 +2138,7 @@ TEST_F(OpTest, Range) { }; std::uniform_int_distribution distribution(-50, 50); DataType tidx = Choose({DT_INT32, DT_INT64, DT_FLOAT, DT_DOUBLE}); - ExpectTfAndXlaOutputsAreClose( + return ExpectTfAndXlaOutputsAreClose( OpTestBuilder("Range") .Input(ToScalar(tidx, distribution(generator()))) .Input(ToScalar(tidx, distribution(generator()))) @@ -1775,8 +2150,8 @@ TEST_F(OpTest, Range) { TEST_F(OpTest, Rank) { Repeatedly([this]() { DataType type = Choose({DT_INT32, DT_FLOAT}); - ExpectTfAndXlaOutputsAreClose( - OpTestBuilder("Rank").Input(RandomTensor(type)).Attr("T", type)); + return ExpectTfAndXlaOutputsAreClose( + OpTestBuilder("Rank").RandomInput(type).Attr("T", type)); }); } @@ -1784,46 +2159,60 @@ TEST_F(OpTest, RealDiv) { Repeatedly([this]() { DataType type = DT_FLOAT; auto dims = BroadcastableDims(); - ExpectTfAndXlaOutputsAreClose(OpTestBuilder("RealDiv") - .Input(RandomTensor(type, dims.first)) - .Input(RandomTensor(type, dims.second)) - .Attr("T", type)); + return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("RealDiv") + .RandomInput(type, dims.first) + .RandomInput(type, dims.second) + .Attr("T", type)); + }); +} + +TEST_F(OpTest, Reciprocal) { + Repeatedly([this]() { + return ExpectTfAndXlaOutputsAreClose( + OpTestBuilder("Reciprocal").RandomInput(DT_FLOAT).Attr("T", DT_FLOAT)); }); } +TEST_F(OpTest, ReciprocalGrad) { + Repeatedly([this]() { + std::vector dims = RandomDims(); + return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("ReciprocalGrad") + .RandomInput(DT_FLOAT, dims) + .RandomInput(DT_FLOAT, dims) + .Attr("T", DT_FLOAT)); + }); +} TEST_F(OpTest, Relu) { Repeatedly([this]() { - ExpectTfAndXlaOutputsAreClose(OpTestBuilder("Relu") - .Input(RandomTensor(DT_FLOAT)) - .Attr("T", DT_FLOAT)); + return ExpectTfAndXlaOutputsAreClose( + OpTestBuilder("Relu").RandomInput(DT_FLOAT).Attr("T", DT_FLOAT)); }); } TEST_F(OpTest, Relu6) { Repeatedly([this]() { - ExpectTfAndXlaOutputsAreClose(OpTestBuilder("Relu6") - .Input(RandomTensor(DT_FLOAT)) - .Attr("T", DT_FLOAT)); + return ExpectTfAndXlaOutputsAreClose( + OpTestBuilder("Relu6").RandomInput(DT_FLOAT).Attr("T", DT_FLOAT)); }); } TEST_F(OpTest, Relu6Grad) { Repeatedly([this]() { auto dims = RandomDims(1); - ExpectTfAndXlaOutputsAreClose(OpTestBuilder("Relu6Grad") - .Input(RandomTensor(DT_FLOAT, dims)) - .Input(RandomTensor(DT_FLOAT, dims)) - .Attr("T", DT_FLOAT)); + return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("Relu6Grad") + .RandomInput(DT_FLOAT, dims) + .RandomInput(DT_FLOAT, dims) + .Attr("T", DT_FLOAT)); }); } TEST_F(OpTest, ReluGrad) { Repeatedly([this]() { auto dims = RandomDims(1); - ExpectTfAndXlaOutputsAreClose(OpTestBuilder("ReluGrad") - .Input(RandomTensor(DT_FLOAT, dims)) - .Input(RandomTensor(DT_FLOAT, dims)) - .Attr("T", DT_FLOAT)); + return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("ReluGrad") + .RandomInput(DT_FLOAT, dims) + .RandomInput(DT_FLOAT, dims) + .Attr("T", DT_FLOAT)); }); } @@ -1845,10 +2234,9 @@ TEST_F(OpTest, Reshape) { } } } - Tensor data = RandomTensor(type, dims_before); - ExpectTfAndXlaOutputsAreClose( + return ExpectTfAndXlaOutputsAreClose( OpTestBuilder("Reshape") - .Input(data) + .RandomInput(type, dims_before) .Input(test::AsTensor( std::vector(dims_after.begin(), dims_after.end()))) .Attr("T", type)); @@ -1860,56 +2248,61 @@ TEST_F(OpTest, Reverse) { std::vector dims = RandomDims(1); DataType type = Choose({DT_INT32, DT_FLOAT}); int64 rank = dims.size(); - ExpectTfAndXlaOutputsAreClose(OpTestBuilder("Reverse") - .Input(RandomTensor(type, dims)) - .Input(RandomTensor(DT_BOOL, {rank})) - .Attr("T", DT_FLOAT)); + return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("Reverse") + .RandomInput(type, dims) + .RandomInput(DT_BOOL, {rank}) + .Attr("T", DT_FLOAT)); }); } TEST_F(OpTest, ReverseV2) { Repeatedly([this]() { DataType type = Choose({DT_INT32, DT_FLOAT}); - Tensor data = RandomTensor(type); - Tensor indices = RandomReductionIndices(data.dims()); - ExpectTfAndXlaOutputsAreClose(OpTestBuilder("ReverseV2") - .Input(data) - .Input(indices) - .Attr("T", DT_FLOAT)); + std::vector data_dims = RandomDims(); + Tensor indices = RandomReductionIndices(data_dims.size()); + return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("ReverseV2") + .RandomInput(type, data_dims) + .Input(indices) + .Attr("T", DT_FLOAT)); + }); +} + +TEST_F(OpTest, Rint) { + Repeatedly([this]() { + return ExpectTfAndXlaOutputsAreClose( + OpTestBuilder("Rint").RandomInput(DT_FLOAT).Attr("T", DT_FLOAT)); }); } TEST_F(OpTest, Round) { Repeatedly([this]() { - ExpectTfAndXlaOutputsAreClose(OpTestBuilder("Round") - .Input(RandomTensor(DT_FLOAT)) - .Attr("T", DT_FLOAT)); + return ExpectTfAndXlaOutputsAreClose( + OpTestBuilder("Round").RandomInput(DT_FLOAT).Attr("T", DT_FLOAT)); }); } TEST_F(OpTest, Rsqrt) { Repeatedly([this]() { - ExpectTfAndXlaOutputsAreClose(OpTestBuilder("Rsqrt") - .Input(RandomTensor(DT_FLOAT)) - .Attr("T", DT_FLOAT)); + return ExpectTfAndXlaOutputsAreClose( + OpTestBuilder("Rsqrt").RandomInput(DT_FLOAT).Attr("T", DT_FLOAT)); }); } TEST_F(OpTest, RsqrtGrad) { Repeatedly([this]() { auto dims = RandomDims(); - ExpectTfAndXlaOutputsAreClose(OpTestBuilder("RsqrtGrad") - .Input(RandomTensor(DT_FLOAT, dims)) - .Input(RandomTensor(DT_FLOAT, dims)) - .Attr("T", DT_FLOAT)); + return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("RsqrtGrad") + .RandomInput(DT_FLOAT, dims) + .RandomInput(DT_FLOAT, dims) + .Attr("T", DT_FLOAT)); }); } TEST_F(OpTest, Shape) { Repeatedly([this]() { DataType type = Choose(kAllXlaTypes); - ExpectTfAndXlaOutputsAreClose( - OpTestBuilder("Shape").Input(RandomTensor(type)).Attr("T", type)); + return ExpectTfAndXlaOutputsAreClose( + OpTestBuilder("Shape").RandomInput(type).Attr("T", type)); }); } @@ -1921,72 +2314,86 @@ TEST_F(OpTest, ShapeN) { builder.Attr("T", type); builder.Attr("N", n); for (int i = 0; i < n; ++i) { - builder.Input(RandomTensor(type)); + builder.RandomInput(type); } - ExpectTfAndXlaOutputsAreClose(builder); + return ExpectTfAndXlaOutputsAreClose(builder); }); } TEST_F(OpTest, Sigmoid) { Repeatedly([this]() { - ExpectTfAndXlaOutputsAreClose(OpTestBuilder("Sigmoid") - .Input(RandomTensor(DT_FLOAT)) - .Attr("T", DT_FLOAT)); + return ExpectTfAndXlaOutputsAreClose( + OpTestBuilder("Sigmoid").RandomInput(DT_FLOAT).Attr("T", DT_FLOAT)); }); } TEST_F(OpTest, SigmoidGrad) { Repeatedly([this]() { auto dims = RandomDims(); - ExpectTfAndXlaOutputsAreClose(OpTestBuilder("SigmoidGrad") - .Input(RandomTensor(DT_FLOAT, dims)) - .Input(RandomTensor(DT_FLOAT, dims)) - .Attr("T", DT_FLOAT)); + return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("SigmoidGrad") + .RandomInput(DT_FLOAT, dims) + .RandomInput(DT_FLOAT, dims) + .Attr("T", DT_FLOAT)); }); } TEST_F(OpTest, Sign) { Repeatedly([this]() { DataType type = Choose({DT_INT32, DT_FLOAT}); - ExpectTfAndXlaOutputsAreClose( - OpTestBuilder("Sign").Input(RandomTensor(type)).Attr("T", type)); + return ExpectTfAndXlaOutputsAreClose( + OpTestBuilder("Sign").RandomInput(type).Attr("T", type)); + }); +} + +TEST_F(OpTest, Sin) { + Repeatedly([this]() { + return ExpectTfAndXlaOutputsAreClose( + OpTestBuilder("Sin").RandomInput(DT_FLOAT).Attr("T", DT_FLOAT)); + }); +} + +TEST_F(OpTest, Sinh) { + Repeatedly([this]() { + return ExpectTfAndXlaOutputsAreClose( + OpTestBuilder("Sinh").RandomInput(DT_FLOAT).Attr("T", DT_FLOAT)); }); } TEST_F(OpTest, Size) { Repeatedly([this]() { DataType type = Choose({DT_INT32, DT_FLOAT}); - ExpectTfAndXlaOutputsAreClose( - OpTestBuilder("Size").Input(RandomTensor(type)).Attr("T", type)); + return ExpectTfAndXlaOutputsAreClose( + OpTestBuilder("Size").RandomInput(type).Attr("T", type)); }); } TEST_F(OpTest, Slice) { Repeatedly([this]() { DataType type = Choose(kAllXlaTypes); - Tensor data = RandomTensor(type); + std::vector data_dims = RandomDims(); - std::vector begin(data.dims()), size(data.dims()); - for (int i = 0; i < data.dims(); ++i) { - begin[i] = std::uniform_int_distribution( - 0, data.dim_size(i))(generator()); + std::vector begin(data_dims.size()), size(data_dims.size()); + for (int i = 0; i < data_dims.size(); ++i) { + begin[i] = + std::uniform_int_distribution(0, data_dims[i])(generator()); size[i] = std::uniform_int_distribution( - -1, data.dim_size(i) - begin[i])(generator()); + -1, data_dims[i] - begin[i])(generator()); } - ExpectTfAndXlaOutputsAreClose(OpTestBuilder("Slice") - .Input(data) - .Input(test::AsTensor(begin)) - .Input(test::AsTensor(size)) - .Attr("T", type) - .Attr("Index", DT_INT32)); + return ExpectTfAndXlaOutputsAreClose( + OpTestBuilder("Slice") + .RandomInput(type, data_dims) + .Input(test::AsTensor(begin)) + .Input(test::AsTensor(size)) + .Attr("T", type) + .Attr("Index", DT_INT32)); }); } TEST_F(OpTest, Softmax) { Repeatedly([this]() { - ExpectTfAndXlaOutputsAreClose( + return ExpectTfAndXlaOutputsAreClose( OpTestBuilder("Softmax") - .Input(RandomTensor(DT_FLOAT, RandomDims(2, 2))) + .RandomInput(DT_FLOAT, RandomDims(2, 2)) .Attr("T", DT_FLOAT)); }); } @@ -1994,28 +2401,126 @@ TEST_F(OpTest, Softmax) { TEST_F(OpTest, SoftmaxCrossEntropyWithLogits) { Repeatedly([this]() { std::vector dims = RandomDims(2, 2, 1); - ExpectTfAndXlaOutputsAreClose(OpTestBuilder("SoftmaxCrossEntropyWithLogits") - .Input(RandomTensor(DT_FLOAT, dims)) - .Input(RandomTensor(DT_FLOAT, dims)) - .Attr("T", DT_FLOAT)); + return ExpectTfAndXlaOutputsAreClose( + OpTestBuilder("SoftmaxCrossEntropyWithLogits") + .RandomInput(DT_FLOAT, dims) + .RandomInput(DT_FLOAT, dims) + .Attr("T", DT_FLOAT)); }); } TEST_F(OpTest, Softplus) { Repeatedly([this]() { - ExpectTfAndXlaOutputsAreClose(OpTestBuilder("Softplus") - .Input(RandomTensor(DT_FLOAT)) - .Attr("T", DT_FLOAT)); + return ExpectTfAndXlaOutputsAreClose( + OpTestBuilder("Softplus").RandomInput(DT_FLOAT).Attr("T", DT_FLOAT)); }); } TEST_F(OpTest, SoftplusGrad) { Repeatedly([this]() { std::vector dims = RandomDims(); - ExpectTfAndXlaOutputsAreClose(OpTestBuilder("SoftplusGrad") - .Input(RandomTensor(DT_FLOAT, dims)) - .Input(RandomTensor(DT_FLOAT, dims)) - .Attr("T", DT_FLOAT)); + return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("SoftplusGrad") + .RandomInput(DT_FLOAT, dims) + .RandomInput(DT_FLOAT, dims) + .Attr("T", DT_FLOAT)); + }); +} + +TEST_F(OpTest, Softsign) { + Repeatedly([this]() { + return ExpectTfAndXlaOutputsAreClose( + OpTestBuilder("Softsign").RandomInput(DT_FLOAT).Attr("T", DT_FLOAT)); + }); +} + +TEST_F(OpTest, SoftsignGrad) { + Repeatedly([this]() { + std::vector dims = RandomDims(); + return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("SoftsignGrad") + .RandomInput(DT_FLOAT, dims) + .RandomInput(DT_FLOAT, dims) + .Attr("T", DT_FLOAT)); + }); +} + +TEST_F(OpTest, SpaceToBatch) { + Repeatedly([this]() { + std::vector block_dims = RandomDims(4, 4, 0, 5); + const int num_block_dims = 2; + int64 block_size = RandomDim(2, 5); + + std::vector input_dims(1 + num_block_dims + 1); + input_dims[0] = RandomDim(); + for (int i = 0; i < num_block_dims; ++i) { + input_dims[1 + i] = block_dims[i] * block_size; + } + input_dims[1 + num_block_dims] = RandomDim(); + + std::vector padding_vals; + std::uniform_int_distribution distribution(0, 7); + for (int i = 0; i < num_block_dims; ++i) { + int64 pad_before; + int64 pad_after; + do { + pad_before = distribution(generator()); + pad_after = distribution(generator()); + } while (pad_before + pad_after > input_dims[1 + i]); + input_dims[1 + i] -= pad_before + pad_after; + padding_vals.push_back(pad_before); + padding_vals.push_back(pad_after); + } + Tensor paddings; + CHECK(paddings.CopyFrom(AsIntTensor(DT_INT32, padding_vals), + TensorShape({num_block_dims, 2}))); + + return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("SpaceToBatch") + .RandomInput(DT_FLOAT, input_dims) + .Input(paddings) + .Attr("T", DT_FLOAT) + .Attr("block_size", block_size)); + }); +} + +TEST_F(OpTest, SpaceToBatchND) { + Repeatedly([this]() { + std::vector block_dims = RandomDims(1, 3, 0, 5); + int num_block_dims = block_dims.size(); + std::vector remaining_dims = RandomDims(0, 3); + std::vector block_multipliers = + RandomDims(block_dims.size(), block_dims.size(), 0, 4); + + std::vector input_dims(1 + num_block_dims + remaining_dims.size()); + input_dims[0] = RandomDim(); + for (int i = 0; i < num_block_dims; ++i) { + input_dims[1 + i] = block_dims[i] * block_multipliers[i]; + } + std::copy(remaining_dims.begin(), remaining_dims.end(), + input_dims.begin() + 1 + num_block_dims); + + std::vector padding_vals; + std::uniform_int_distribution distribution(0, 7); + for (int i = 0; i < num_block_dims; ++i) { + int64 pad_before; + int64 pad_after; + do { + pad_before = distribution(generator()); + pad_after = distribution(generator()); + } while (pad_before + pad_after > input_dims[1 + i]); + input_dims[1 + i] -= pad_before + pad_after; + padding_vals.push_back(pad_before); + padding_vals.push_back(pad_after); + } + Tensor paddings; + CHECK(paddings.CopyFrom(AsIntTensor(DT_INT32, padding_vals), + TensorShape({num_block_dims, 2}))); + + return ExpectTfAndXlaOutputsAreClose( + OpTestBuilder("SpaceToBatchND") + .RandomInput(DT_FLOAT, input_dims) + .Input(test::AsTensor( + std::vector(block_dims.begin(), block_dims.end()))) + .Input(paddings) + .Attr("T", DT_FLOAT)); }); } @@ -2025,33 +2530,26 @@ TEST_F(OpTest, SparseMatMul) { int64 y = RandomDim(); int64 z = RandomDim(); - ExpectTfAndXlaOutputsAreClose(OpTestBuilder("SparseMatMul") - .Input(RandomTensor(DT_FLOAT, {x, y})) - .Input(RandomTensor(DT_FLOAT, {y, z})) - .Attr("Ta", DT_FLOAT) - .Attr("Tb", DT_FLOAT)); - - ExpectTfAndXlaOutputsAreClose(OpTestBuilder("SparseMatMul") - .Input(RandomTensor(DT_FLOAT, {y, x})) - .Input(RandomTensor(DT_FLOAT, {y, z})) - .Attr("Ta", DT_FLOAT) - .Attr("Tb", DT_FLOAT) - .Attr("transpose_a", true)); + std::vector a_dims = {x, y}; + std::vector b_dims = {y, z}; - ExpectTfAndXlaOutputsAreClose(OpTestBuilder("SparseMatMul") - .Input(RandomTensor(DT_FLOAT, {x, y})) - .Input(RandomTensor(DT_FLOAT, {z, y})) - .Attr("Ta", DT_FLOAT) - .Attr("Tb", DT_FLOAT) - .Attr("transpose_b", true)); + std::bernoulli_distribution random_bool; + bool transpose_a = random_bool(generator()); + bool transpose_b = random_bool(generator()); + if (transpose_a) { + std::swap(a_dims[0], a_dims[1]); + } + if (transpose_b) { + std::swap(b_dims[0], b_dims[1]); + } - ExpectTfAndXlaOutputsAreClose(OpTestBuilder("SparseMatMul") - .Input(RandomTensor(DT_FLOAT, {y, x})) - .Input(RandomTensor(DT_FLOAT, {z, y})) - .Attr("Ta", DT_FLOAT) - .Attr("Tb", DT_FLOAT) - .Attr("transpose_a", true) - .Attr("transpose_b", true)); + return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("SparseMatMul") + .RandomInput(DT_FLOAT, a_dims) + .RandomInput(DT_FLOAT, b_dims) + .Attr("Ta", DT_FLOAT) + .Attr("Tb", DT_FLOAT) + .Attr("transpose_a", transpose_a) + .Attr("transpose_b", transpose_b)); }); } @@ -2067,9 +2565,9 @@ TEST_F(OpTest, SparseSoftmaxCrossEntropyWithLogits) { std::uniform_int_distribution(0, num_classes - 1)(generator()); } - ExpectTfAndXlaOutputsAreClose( + return ExpectTfAndXlaOutputsAreClose( OpTestBuilder("SparseSoftmaxCrossEntropyWithLogits") - .Input(RandomTensor(DT_FLOAT, dims)) + .RandomInput(DT_FLOAT, dims) .Input(test::AsTensor(indices)) .Attr("T", DT_FLOAT) .Attr("Tlabels", DT_INT32)); @@ -2087,56 +2585,64 @@ TEST_F(OpTest, Split) { // Ensure 'dim' is evenly divisible by 'n'. dims[dim] /= n; dims[dim] *= n; - ExpectTfAndXlaOutputsAreClose(OpTestBuilder("Split") - .Input(test::AsScalar(dim)) - .Input(RandomTensor(type, dims)) - .Attr("T", type) - .Attr("num_split", n)); + return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("Split") + .Input(test::AsScalar(dim)) + .RandomInput(type, dims) + .Attr("T", type) + .Attr("num_split", n)); }); } TEST_F(OpTest, Sqrt) { Repeatedly([this]() { - ExpectTfAndXlaOutputsAreClose(OpTestBuilder("Sqrt") - .Input(RandomTensor(DT_FLOAT)) - .Attr("T", DT_FLOAT)); + return ExpectTfAndXlaOutputsAreClose( + OpTestBuilder("Sqrt").RandomInput(DT_FLOAT).Attr("T", DT_FLOAT)); + }); +} + +TEST_F(OpTest, SqrtGrad) { + Repeatedly([this]() { + auto dims = RandomDims(); + return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("SqrtGrad") + .RandomInput(DT_FLOAT, dims) + .RandomInput(DT_FLOAT, dims) + .Attr("T", DT_FLOAT)); }); } TEST_F(OpTest, SquaredDifference) { Repeatedly([this]() { auto dims = BroadcastableDims(); - ExpectTfAndXlaOutputsAreClose( - OpTestBuilder("SquaredDifference") - .Input(RandomTensor(DT_FLOAT, dims.first)) - .Input(RandomTensor(DT_FLOAT, dims.second)) - .Attr("T", DT_FLOAT)); + return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("SquaredDifference") + .RandomInput(DT_FLOAT, dims.first) + .RandomInput(DT_FLOAT, dims.second) + .Attr("T", DT_FLOAT)); }); } TEST_F(OpTest, Square) { Repeatedly([this]() { DataType type = Choose({DT_INT32, DT_FLOAT}); - ExpectTfAndXlaOutputsAreClose( - OpTestBuilder("Square").Input(RandomTensor(type)).Attr("T", type)); + return ExpectTfAndXlaOutputsAreClose( + OpTestBuilder("Square").RandomInput(type).Attr("T", type)); }); } TEST_F(OpTest, Squeeze) { Repeatedly([this]() { DataType type = Choose(kAllXlaTypes); - Tensor t = RandomTensor(type, RandomDims(0, kDefaultMaxRank, 0, 5)); + std::vector t_dims = RandomDims(0, kDefaultMaxRank, 0, 5); std::bernoulli_distribution random_bool; std::vector squeeze_dims; - for (int i = 0; i < t.dims(); ++i) { - if (t.dim_size(i) == 1 && random_bool(generator())) { + for (int i = 0; i < t_dims.size(); ++i) { + if (t_dims[i] == 1 && random_bool(generator())) { squeeze_dims.push_back(i); } } - ExpectTfAndXlaOutputsAreClose(OpTestBuilder("Squeeze") - .Input(t) - .Attr("squeeze_dims", squeeze_dims) - .Attr("T", type)); + return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("Squeeze") + .RandomInput(type, t_dims) + .Attr("squeeze_dims", squeeze_dims) + .Attr("T", type)); }); } @@ -2144,58 +2650,59 @@ TEST_F(OpTest, Sub) { Repeatedly([this]() { DataType type = Choose({DT_INT32, DT_FLOAT}); auto dims = BroadcastableDims(); - ExpectTfAndXlaOutputsAreClose(OpTestBuilder("Sub") - .Input(RandomTensor(type, dims.first)) - .Input(RandomTensor(type, dims.second)) - .Attr("T", type)); + return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("Sub") + .RandomInput(type, dims.first) + .RandomInput(type, dims.second) + .Attr("T", type)); }); } TEST_F(OpTest, Sum) { Repeatedly([this]() { DataType type = Choose({DT_INT32, DT_FLOAT}); - Tensor data = RandomTensor(type); - Tensor indices = RandomReductionIndices(data.dims()); + std::vector data_dims = RandomDims(); + Tensor indices = RandomReductionIndices(data_dims.size()); bool keep_dims = Choose({false, true}); - ExpectTfAndXlaOutputsAreClose( - OpTestBuilder("Sum").Input(data).Input(indices).Attr("T", type).Attr( - "keep_dims", keep_dims)); + return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("Sum") + .RandomInput(type, data_dims) + .Input(indices) + .Attr("T", type) + .Attr("keep_dims", keep_dims)); }); } TEST_F(OpTest, StridedSlice) { Repeatedly([this]() { DataType type = Choose(kAllXlaTypes); - Tensor data = RandomTensor(type); - - std::vector begin(data.dims()), end(data.dims()); - std::vector strides(data.dims()); - for (int i = 0; i < data.dims(); ++i) { + std::vector data_dims = RandomDims(); + std::vector begin(data_dims.size()), end(data_dims.size()); + std::vector strides(data_dims.size()); + for (int i = 0; i < data_dims.size(); ++i) { begin[i] = std::uniform_int_distribution( - -2 * data.dim_size(i), 2 * data.dim_size(i))(generator()); + -2 * data_dims[i], 2 * data_dims[i])(generator()); end[i] = std::uniform_int_distribution( - -2 * data.dim_size(i), 2 * data.dim_size(i))(generator()); + -2 * data_dims[i], 2 * data_dims[i])(generator()); // TODO(b/31360685): support strides other than 1 or -1 strides[i] = std::bernoulli_distribution()(generator()) ? 1 : -1; } - int64 max_bitmask = (1LL << data.dims()) - 1; + int64 max_bitmask = (1LL << data_dims.size()) - 1; std::uniform_int_distribution bitmask_distribution(0, max_bitmask); int64 begin_mask = bitmask_distribution(generator()); int64 end_mask = bitmask_distribution(generator()); // Create a ellipsis bitmask with at most one 1 bit set. int64 ellipsis_mask = 0; - if (data.dims() > 0 && std::bernoulli_distribution()(generator())) { - int ellipsis_pos = - std::uniform_int_distribution(0, data.dims() - 1)(generator()); + if (!data_dims.empty() && std::bernoulli_distribution()(generator())) { + int ellipsis_pos = std::uniform_int_distribution( + 0, data_dims.size() - 1)(generator()); ellipsis_mask = 1LL << ellipsis_pos; } int64 new_axis_mask = bitmask_distribution(generator()); int64 shrink_axis_mask = bitmask_distribution(generator()); - ExpectTfAndXlaOutputsAreClose( + return ExpectTfAndXlaOutputsAreClose( OpTestBuilder("StridedSlice") - .Input(data) + .RandomInput(type, data_dims) .Input(test::AsTensor(begin)) .Input(test::AsTensor(end)) .Input(test::AsTensor(strides)) @@ -2245,13 +2752,13 @@ TEST_F(OpTest, StridedSliceGrad) { // TODO(phawkins): use shape inference for the forward op to compute the // gradient shape for the backward op. At present, there is a low // probability of the golden op succeeding. - ExpectTfAndXlaOutputsAreClose( + return ExpectTfAndXlaOutputsAreClose( OpTestBuilder("StridedSliceGrad") .Input(test::AsTensor(dims)) .Input(test::AsTensor(begin)) .Input(test::AsTensor(end)) .Input(test::AsTensor(strides)) - .Input(RandomTensor(type, RandomDims(1))) + .RandomInput(type, RandomDims(1)) .Attr("T", type) .Attr("Index", DT_INT64) .Attr("begin_mask", begin_mask) @@ -2262,50 +2769,57 @@ TEST_F(OpTest, StridedSliceGrad) { }); } +TEST_F(OpTest, Tan) { + Repeatedly([this]() { + return ExpectTfAndXlaOutputsAreClose( + OpTestBuilder("Tan").RandomInput(DT_FLOAT).Attr("T", DT_FLOAT)); + }); +} + TEST_F(OpTest, Tanh) { Repeatedly([this]() { - ExpectTfAndXlaOutputsAreClose(OpTestBuilder("Tanh") - .Input(RandomTensor(DT_FLOAT)) - .Attr("T", DT_FLOAT)); + return ExpectTfAndXlaOutputsAreClose( + OpTestBuilder("Tanh").RandomInput(DT_FLOAT).Attr("T", DT_FLOAT)); }); } TEST_F(OpTest, TanhGrad) { Repeatedly([this]() { auto dims = RandomDims(); - ExpectTfAndXlaOutputsAreClose(OpTestBuilder("TanhGrad") - .Input(RandomTensor(DT_FLOAT, dims)) - .Input(RandomTensor(DT_FLOAT, dims)) - .Attr("T", DT_FLOAT)); + return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("TanhGrad") + .RandomInput(DT_FLOAT, dims) + .RandomInput(DT_FLOAT, dims) + .Attr("T", DT_FLOAT)); }); } TEST_F(OpTest, Tile) { Repeatedly([this]() { DataType type = Choose(kAllXlaTypes); - Tensor t = RandomTensor(type, RandomDims(1)); - std::vector multiples(t.dims()); - for (int i = 0; i < t.dims(); ++i) { + std::vector t_dims = RandomDims(1); + std::vector multiples(t_dims.size()); + for (int i = 0; i < t_dims.size(); ++i) { multiples[i] = std::uniform_int_distribution(1, 3)(generator()); } - ExpectTfAndXlaOutputsAreClose(OpTestBuilder("Tile") - .Input(t) - .Input(test::AsTensor(multiples)) - .Attr("T", type)); + return ExpectTfAndXlaOutputsAreClose( + OpTestBuilder("Tile") + .RandomInput(type, t_dims) + .Input(test::AsTensor(multiples)) + .Attr("T", type)); }); } TEST_F(OpTest, Transpose) { Repeatedly([this]() { DataType type = Choose(kAllXlaTypes); - Tensor data = RandomTensor(type); - std::vector perm(data.dims()); + std::vector data_dims = RandomDims(); + std::vector perm(data_dims.size()); std::iota(perm.begin(), perm.end(), 0); std::shuffle(perm.begin(), perm.end(), generator()); - ExpectTfAndXlaOutputsAreClose(OpTestBuilder("Transpose") - .Input(data) - .Input(test::AsTensor(perm)) - .Attr("T", type)); + return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("Transpose") + .RandomInput(type, data_dims) + .Input(test::AsTensor(perm)) + .Attr("T", type)); }); } @@ -2313,10 +2827,10 @@ TEST_F(OpTest, TruncateDiv) { Repeatedly([this]() { DataType type = DT_INT32; auto dims = BroadcastableDims(); - ExpectTfAndXlaOutputsAreClose(OpTestBuilder("TruncateDiv") - .Input(RandomTensor(type, dims.first)) - .Input(RandomTensor(type, dims.second)) - .Attr("T", type)); + return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("TruncateDiv") + .RandomInput(type, dims.first) + .RandomInput(type, dims.second) + .Attr("T", type)); }); } @@ -2324,26 +2838,18 @@ TEST_F(OpTest, TruncateMod) { Repeatedly([this]() { DataType type = Choose({DT_INT32, DT_FLOAT}); auto dims = BroadcastableDims(); - ExpectTfAndXlaOutputsAreClose(OpTestBuilder("TruncateMod") - .Input(RandomTensor(type, dims.first)) - .Input(RandomTensor(type, dims.second)) - .Attr("T", type)); + return ExpectTfAndXlaOutputsAreClose(OpTestBuilder("TruncateMod") + .RandomInput(type, dims.first) + .RandomInput(type, dims.second) + .Attr("T", type)); }); } TEST_F(OpTest, ZerosLike) { Repeatedly([this]() { DataType type = Choose({DT_INT32, DT_FLOAT}); - ExpectTfAndXlaOutputsAreClose( - OpTestBuilder("ZerosLike").Input(RandomTensor(type)).Attr("T", type)); - }); -} - -TEST_F(OpTest, OnesLike) { - Repeatedly([this]() { - DataType type = Choose({DT_INT32, DT_FLOAT}); - ExpectTfAndXlaOutputsAreClose( - OpTestBuilder("OnesLike").Input(RandomTensor(type)).Attr("T", type)); + return ExpectTfAndXlaOutputsAreClose( + OpTestBuilder("ZerosLike").RandomInput(type).Attr("T", type)); }); } @@ -2362,6 +2868,9 @@ int main(int argc, char** argv) { tensorflow::Flag("tf_xla_test_repetitions", &tensorflow::tf_xla_test_repetitions, "Number of repetitions for each test."), + tensorflow::Flag("tf_xla_max_tensor_size", + &tensorflow::tf_xla_max_tensor_size, + "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_test_use_jit", &tensorflow::tf_xla_test_use_jit, diff --git a/tensorflow/compiler/tests/segment_reduction_ops_test.py b/tensorflow/compiler/tests/segment_reduction_ops_test.py new file mode 100644 index 0000000000000000000000000000000000000000..260a04421b62310c109d8f0ea72875a50c234bb0 --- /dev/null +++ b/tensorflow/compiler/tests/segment_reduction_ops_test.py @@ -0,0 +1,139 @@ +# 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 cases for segment reduction ops.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import functools +import numpy as np + +from tensorflow.compiler.tests.xla_test import XLATestCase +from tensorflow.python.ops import array_ops +from tensorflow.python.ops import math_ops +from tensorflow.python.platform import googletest + + +class SegmentReductionOpsTest(XLATestCase): + """Test cases for segment reduction ops.""" + + def UnsortedSegmentSum(self, data, indices, num_segments): + with self.test_session() as sess, self.test_scope(): + d = array_ops.placeholder(data.dtype, shape=data.shape) + if isinstance(indices, int): + i = array_ops.placeholder(np.int32, shape=[]) + else: + i = array_ops.placeholder(indices.dtype, shape=indices.shape) + return sess.run( + math_ops.unsorted_segment_sum(d, i, num_segments), + {d: data, + i: indices}) + + def testUnsortedSegmentSum0DIndices1DData(self): + for dtype in self.numeric_types: + self.assertAllClose( + np.array( + [[0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0], [0, 1, 2, 3, 4, 5], + [0, 0, 0, 0, 0, 0]], + dtype=dtype), + self.UnsortedSegmentSum( + np.array([0, 1, 2, 3, 4, 5], dtype=dtype), 2, 4)) + + def testUnsortedSegmentSum1DIndices1DData(self): + for dtype in self.numeric_types: + self.assertAllClose( + np.array([1, 3, 2, 9], dtype=dtype), + self.UnsortedSegmentSum( + np.array([0, 1, 2, 3, 4, 5], dtype=dtype), + np.array([3, 0, 2, 1, 3, 3], dtype=np.int32), 4)) + + def testUnsortedSegmentSum1DIndices2DDataDisjoint(self): + for dtype in self.numeric_types: + data = np.array( + [[0, 1, 2, 3], [20, 21, 22, 23], [30, 31, 32, 33], [40, 41, 42, 43], + [50, 51, 52, 53]], + dtype=dtype) + indices = np.array([8, 1, 0, 3, 7], dtype=np.int32) + num_segments = 10 + y = self.UnsortedSegmentSum(data, indices, num_segments) + self.assertAllClose( + np.array( + [[30, 31, 32, 33], [20, 21, 22, 23], [0, 0, 0, 0], + [40, 41, 42, 43], [0, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0], + [50, 51, 52, 53], [0, 1, 2, 3], [0, 0, 0, 0]], + dtype=dtype), y) + + def testUnsortedSegmentSum1DIndices2DDataNonDisjoint(self): + for dtype in self.numeric_types: + data = np.array( + [[0, 1, 2, 3], [20, 21, 22, 23], [30, 31, 32, 33], [40, 41, 42, 43], + [50, 51, 52, 53]], + dtype=dtype) + indices = np.array([0, 1, 2, 0, 1], dtype=np.int32) + num_segments = 4 + y = self.UnsortedSegmentSum(data, indices, num_segments) + self.assertAllClose( + np.array( + [[40, 42, 44, 46], [70, 72, 74, 76], [30, 31, 32, 33], + [0, 0, 0, 0]], + dtype=dtype), y) + + def testUnsortedSegmentSum2DIndices3DData(self): + for dtype in self.numeric_types: + data = np.array( + [[[0, 1, 2], [10, 11, 12]], [[100, 101, 102], [110, 111, 112]], + [[200, 201, 202], [210, 211, 212]], [[300, 301, 302], + [310, 311, 312]]], + dtype=dtype) + indices = np.array([[3, 5], [3, 1], [5, 0], [6, 2]], dtype=np.int32) + num_segments = 8 + y = self.UnsortedSegmentSum(data, indices, num_segments) + self.assertAllClose( + np.array( + [[210, 211, 212], [110, 111, 112], [310, 311, 312], + [100, 102, 104], [0, 0, 0.], [210, 212, 214], [300, 301, + 302], [0, 0, 0]], + dtype=dtype), y) + + def testUnsortedSegmentSum1DIndices3DData(self): + for dtype in self.numeric_types: + data = np.array( + [[[0, 1, 2], [10, 11, 12]], [[100, 101, 102], [110, 111, 112]], + [[200, 201, 202], [210, 211, 212]], [[300, 301, 302], + [310, 311, 312]]], + dtype=dtype) + indices = np.array([3, 0, 2, 5], dtype=np.int32) + num_segments = 6 + y = self.UnsortedSegmentSum(data, indices, num_segments) + self.assertAllClose( + np.array( + [[[100, 101, 102.], [110, 111, 112]], [[0, 0, 0], [0, 0, 0]], + [[200, 201, 202], [210, 211, 212]], [[0, 1, 2.], [10, 11, 12]], + [[0, 0, 0], [0, 0, 0]], [[300, 301, 302], [310, 311, 312]]], + dtype=dtype), y) + + def testUnsortedSegmentSumShapeError(self): + for dtype in self.numeric_types: + data = np.ones((4, 8, 7), dtype=dtype) + indices = np.ones((3, 2), dtype=np.int32) + num_segments = 4 + self.assertRaises(ValueError, + functools.partial(self.UnsortedSegmentSum, data, + indices, num_segments)) + + +if __name__ == '__main__': + googletest.main() diff --git a/tensorflow/compiler/tests/slice_ops_test.py b/tensorflow/compiler/tests/slice_ops_test.py new file mode 100644 index 0000000000000000000000000000000000000000..4ddf2ee0dcb2b5f514ff9820c07f7cc10609ff66 --- /dev/null +++ b/tensorflow/compiler/tests/slice_ops_test.py @@ -0,0 +1,145 @@ +# 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 slicing.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import numpy as np + +from tensorflow.compiler.tests.xla_test import XLATestCase +from tensorflow.python.framework import dtypes +from tensorflow.python.ops import array_ops +from tensorflow.python.platform import googletest + + + +class SliceTest(XLATestCase): + + def test1D(self): + for dtype in self.numeric_types: + with self.test_session(): + i = array_ops.placeholder(dtype, shape=[10]) + with self.test_scope(): + o = array_ops.slice(i, [2], [4]) + params = { + i: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9], + } + result = o.eval(feed_dict=params) + + self.assertAllEqual([2, 3, 4, 5], result) + + def test3D(self): + for dtype in self.numeric_types: + with self.test_session(): + i = array_ops.placeholder(dtype, shape=[3, 3, 10]) + with self.test_scope(): + o = array_ops.slice(i, [1, 2, 2], [1, 1, 4]) + params = { + i: [[[0, 1, 2, 3, 4, 5, 6, 7, 8, 9], + [9, 8, 7, 6, 5, 4, 3, 2, 1, 0], + [5, 3, 1, 7, 9, 2, 4, 6, 8, 0]], + [[5, 5, 5, 5, 5, 5, 5, 5, 5, 5], + [1, 1, 1, 1, 1, 1, 1, 1, 1, 1], + [8, 7, 6, 5, 4, 3, 2, 1, 8, 7]], + [[7, 5, 7, 5, 7, 5, 7, 5, 7, 5], + [1, 2, 1, 2, 1, 2, 1, 2, 1, 2], + [9, 8, 7, 9, 8, 7, 9, 8, 7, 9]]] + } + result = o.eval(feed_dict=params) + + self.assertAllEqual([[[6, 5, 4, 3]]], result) + + + +class StridedSliceTest(XLATestCase): + + def test1D(self): + for dtype in self.numeric_types: + with self.test_session(): + i = array_ops.placeholder(dtype, shape=[10]) + with self.test_scope(): + o = array_ops.strided_slice(i, [2], [6], [2]) + params = { + i: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9], + } + result = o.eval(feed_dict=params) + + self.assertAllEqual([2, 4], result) + + def test1DNegtiveStride(self): + for dtype in self.numeric_types: + with self.test_session(): + i = array_ops.placeholder(dtype, shape=[10]) + with self.test_scope(): + o = array_ops.strided_slice(i, [6], [2], [-2]) + params = { + i: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9], + } + result = o.eval(feed_dict=params) + + self.assertAllEqual([6, 4], result) + + def test3D(self): + for dtype in self.numeric_types: + with self.test_session(): + i = array_ops.placeholder(dtype, shape=[3, 3, 10]) + with self.test_scope(): + o = array_ops.strided_slice(i, [0, 2, 2], [2, 3, 6], [1, 1, 2]) + params = { + i: [[[0, 1, 2, 3, 4, 5, 6, 7, 8, 9], + [9, 8, 7, 6, 5, 4, 3, 2, 1, 0], + [5, 3, 1, 7, 9, 2, 4, 6, 8, 0]], + [[5, 5, 5, 5, 5, 5, 5, 5, 5, 5], + [1, 1, 1, 1, 1, 1, 1, 1, 1, 1], + [8, 7, 6, 5, 4, 3, 2, 1, 8, 7]], + [[7, 5, 7, 5, 7, 5, 7, 5, 7, 5], + [1, 2, 1, 2, 1, 2, 1, 2, 1, 2], + [9, 8, 7, 9, 8, 7, 9, 8, 7, 9]]] + } + result = o.eval(feed_dict=params) + + self.assertAllEqual([[[1, 9]], [[6, 4]]], result) + + def test3DNegativeStride(self): + for dtype in self.numeric_types: + with self.test_session(): + i = array_ops.placeholder(dtype, shape=[3, 4, 10]) + with self.test_scope(): + o = array_ops.strided_slice(i, [2, 2, 6], [0, 0, 2], [-1, -1, -2]) + params = { + i: [[[0, 1, 2, 3, 4, 5, 6, 7, 8, 9], + [9, 8, 7, 6, 5, 4, 3, 2, 1, 0], + [5, 3, 1, 7, 9, 2, 4, 6, 8, 0], + [4, 5, 2, 4, 3, 7, 6, 8, 9, 4]], + [[5, 5, 5, 5, 5, 5, 5, 5, 5, 5], + [4, 3, 4, 5, 7, 6, 5, 3, 4, 5], + [8, 7, 6, 5, 4, 3, 2, 1, 8, 7], + [7, 1, 7, 1, 8, 1, 8, 1, 3, 1]], + [[7, 5, 7, 5, 7, 5, 7, 5, 7, 5], + [1, 2, 1, 2, 1, 2, 1, 2, 1, 2], + [9, 8, 7, 9, 8, 7, 9, 8, 7, 9], + [9, 9, 5, 5, 6, 6, 3, 3, 6, 6]]] + } + result = o.eval(feed_dict=params) + + self.assertAllEqual([[[9, 8], + [1, 1]], + [[2, 4], + [5, 7]]], result) + +if __name__ == "__main__": + googletest.main() diff --git a/tensorflow/compiler/tests/spacetobatch_op_test.py b/tensorflow/compiler/tests/spacetobatch_op_test.py new file mode 100644 index 0000000000000000000000000000000000000000..c013f4b50a4cf95be8028248c52b10b1c3be2bd3 --- /dev/null +++ b/tensorflow/compiler/tests/spacetobatch_op_test.py @@ -0,0 +1,272 @@ +# 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 tests for SpaceToBatch and BatchToSpace ops.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import numpy as np + +from tensorflow.compiler.tests.xla_test import XLATestCase +from tensorflow.python.ops import array_ops +from tensorflow.python.ops import gen_array_ops +from tensorflow.python.platform import test + + +def space_to_batch_direct(input_array, block_shape, paddings): + """Direct Python implementation of space-to-batch conversion. + + This is used for tests only. + + Args: + input_array: N-D array + block_shape: 1-D array of shape [num_block_dims]. + paddings: 2-D array of shape [num_block_dims, 2]. + + Returns: + Converted tensor. + """ + input_array = np.array(input_array) + block_shape = np.array(block_shape) + num_block_dims = len(block_shape) + paddings = np.array(paddings).reshape((len(block_shape), 2)) + + padded = np.pad(input_array, + pad_width=([[0, 0]] + list(paddings) + [[0, 0]] * + (input_array.ndim - 1 - num_block_dims)), + mode="constant") + reshaped_padded_shape = [input_array.shape[0]] + output_shape = [input_array.shape[0] * np.prod(block_shape)] + for block_dim, block_shape_value in enumerate(block_shape): + reduced_size = padded.shape[block_dim + 1] // block_shape_value + reshaped_padded_shape.append(reduced_size) + output_shape.append(reduced_size) + reshaped_padded_shape.append(block_shape_value) + reshaped_padded_shape.extend(input_array.shape[num_block_dims + 1:]) + output_shape.extend(input_array.shape[num_block_dims + 1:]) + + reshaped_padded = padded.reshape(reshaped_padded_shape) + permuted_reshaped_padded = np.transpose(reshaped_padded, ( + list(np.arange(num_block_dims) * 2 + 2) + [0] + + list(np.arange(num_block_dims) * 2 + 1) + list( + np.arange(input_array.ndim - num_block_dims - 1) + 1 + num_block_dims + * 2))) + return permuted_reshaped_padded.reshape(output_shape) + + +class SpaceToBatchTest(XLATestCase): + """Tests input-output pairs for the SpaceToBatch and BatchToSpace ops.""" + + def _testPad(self, inputs, paddings, block_size, outputs): + with self.test_session() as sess, self.test_scope(): + for dtype in self.float_types: + # outputs = space_to_batch(inputs) + placeholder = array_ops.placeholder(dtype) + x_tf = gen_array_ops._space_to_batch( + placeholder, paddings, block_size=block_size) + self.assertAllEqual(sess.run(x_tf, {placeholder: inputs}), outputs) + # inputs = batch_to_space(outputs) + x_tf = gen_array_ops._batch_to_space( + placeholder, paddings, block_size=block_size) + self.assertAllEqual(sess.run(x_tf, {placeholder: outputs}), inputs) + + def _testOne(self, inputs, block_size, outputs): + paddings = np.zeros((2, 2), dtype=np.int32) + self._testPad(inputs, paddings, block_size, outputs) + + # [1, 2, 2, 1] <-> [4, 1, 1, 1] + def testSmallInput2x2(self): + x_np = [[[[1], [2]], [[3], [4]]]] + block_size = 2 + x_out = [[[[1]]], [[[2]]], [[[3]]], [[[4]]]] + self._testOne(x_np, block_size, x_out) + + # [1, 2, 2, 1] <-> [1, 3, 3, 1] (padding) <-> [9, 1, 1, 1] + def testSmallInput2x2Pad1x0(self): + x_np = [[[[1], [2]], [[3], [4]]]] + paddings = np.array([[1, 0], [1, 0]], dtype=np.int32) + block_size = 3 + x_out = [[[[0]]], [[[0]]], [[[0]]], [[[0]]], [[[1]]], [[[2]]], [[[0]]], + [[[3]]], [[[4]]]] + self._testPad(x_np, paddings, block_size, x_out) + + # Test with depth larger than 1. + # [1, 2, 2, 3] <-> [4, 1, 1, 3] + def testDepthInput2x2(self): + x_np = [[[[1, 2, 3], [4, 5, 6]], [[7, 8, 9], [10, 11, 12]]]] + block_size = 2 + x_out = [[[[1, 2, 3]]], [[[4, 5, 6]]], [[[7, 8, 9]]], [[[10, 11, 12]]]] + self._testOne(x_np, block_size, x_out) + + # Test for larger input dimensions. + # [1, 4, 4, 1] <-> [4, 2, 2, 1] + def testLargerInput2x2(self): + x_np = [[[[1], [2], [3], [4]], [[5], [6], [7], [8]], + [[9], [10], [11], [12]], [[13], [14], [15], [16]]]] + block_size = 2 + x_out = [[[[1], [3]], [[9], [11]]], [[[2], [4]], [[10], [12]]], + [[[5], [7]], [[13], [15]]], [[[6], [8]], [[14], [16]]]] + self._testOne(x_np, block_size, x_out) + + # Test with batch larger than 1. + # [2, 2, 4, 1] <-> [8, 1, 2, 1] + def testBatchInput2x2(self): + x_np = [[[[1], [2], [3], [4]], [[5], [6], [7], [8]]], + [[[9], [10], [11], [12]], [[13], [14], [15], [16]]]] + block_size = 2 + x_out = [[[[1], [3]]], [[[9], [11]]], [[[2], [4]]], [[[10], [12]]], + [[[5], [7]]], [[[13], [15]]], [[[6], [8]]], [[[14], [16]]]] + self._testOne(x_np, block_size, x_out) + + # Tests for larger input spatial dimensions AND batch larger than 1, to ensure + # that elements are correctly laid out spatially and properly interleaved + # along the batch dimension. + # [2, 4, 4, 1] <-> [8, 2, 2, 1] + def testLargerInputBatch2x2(self): + x_np = [[[[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]]]] + x_out = [[[[1], [3]], [[9], [11]]], [[[17], [19]], [[25], [27]]], + [[[2], [4]], [[10], [12]]], [[[18], [20]], [[26], [28]]], + [[[5], [7]], [[13], [15]]], [[[21], [23]], [[29], [31]]], + [[[6], [8]], [[14], [16]]], [[[22], [24]], [[30], [32]]]] + block_size = 2 + self._testOne(x_np, block_size, x_out) + + +class SpaceToBatchNDTest(XLATestCase): + """Tests input-output pairs for the SpaceToBatchND and BatchToSpaceND ops.""" + + def _testPad(self, inputs, block_shape, paddings, outputs): + block_shape = np.array(block_shape) + paddings = np.array(paddings).reshape((len(block_shape), 2)) + with self.test_session() as sess, self.test_scope(): + for dtype in self.float_types: + placeholder = array_ops.placeholder(dtype) + # outputs = space_to_batch(inputs) + x_tf = array_ops.space_to_batch_nd(placeholder, block_shape, paddings) + self.assertAllEqual(sess.run(x_tf, {placeholder: inputs}), outputs) + # inputs = batch_to_space(outputs) + placeholder = array_ops.placeholder(dtype) + x_tf = array_ops.batch_to_space_nd(placeholder, block_shape, paddings) + self.assertAllEqual(sess.run(x_tf, {placeholder: outputs}), inputs) + + def _testDirect(self, input_shape, block_shape, paddings): + inputs = np.arange(np.prod(input_shape), dtype=np.float32) + inputs = inputs.reshape(input_shape) + self._testPad(inputs, block_shape, paddings, + space_to_batch_direct(inputs, block_shape, paddings)) + + def testZeroBlockDimsZeroRemainingDims(self): + self._testPad( + inputs=[1, 2], + block_shape=[], + paddings=[], + outputs=[1, 2],) + + def testZeroBlockDimsOneRemainingDim(self): + self._testPad( + inputs=[[1, 2], [3, 4]], + block_shape=[], + paddings=[], + outputs=[[1, 2], [3, 4]]) + + # Same thing, but with a no-op block dim. + self._testPad( + inputs=[[1, 2], [3, 4]], + block_shape=[1], + paddings=[[0, 0]], + outputs=[[1, 2], [3, 4]]) + + def testZeroBlockDimsTwoRemainingDims(self): + self._testPad( + inputs=[[[1, 2], [3, 4]], [[5, 6], [7, 8]]], + block_shape=[], + paddings=[], + outputs=[[[1, 2], [3, 4]], [[5, 6], [7, 8]]]) + + # Same thing, but with a no-op block dim. + self._testPad( + inputs=[[[1, 2], [3, 4]], [[5, 6], [7, 8]]], + block_shape=[1], + paddings=[[0, 0]], + outputs=[[[1, 2], [3, 4]], [[5, 6], [7, 8]]]) + + # Same thing, but with two no-op block dims. + self._testPad( + inputs=[[[1, 2], [3, 4]], [[5, 6], [7, 8]]], + block_shape=[1, 1], + paddings=[[0, 0], [0, 0]], + outputs=[[[1, 2], [3, 4]], [[5, 6], [7, 8]]]) + + def testOneBlockDimZeroRemainingDims(self): + self._testPad( + inputs=[[1, 2, 3], [4, 5, 6]], + block_shape=[2], + paddings=[1, 0], + outputs=[[0, 2], [0, 5], [1, 3], [4, 6]]) + + def testOneBlockDimOneRemainingDim(self): + self._testPad( + inputs=[[[1, 11], [2, 21], [3, 31]], [[4, 41], [5, 51], [6, 61]]], + block_shape=[2], + paddings=[1, 0], + outputs=[[[0, 0], [2, 21]], [[0, 0], [5, 51]], [[1, 11], [3, 31]], + [[4, 41], [6, 61]]]) + + def testDirect0(self): + # Test with zero-size remaining dimension. + self._testDirect( + input_shape=[3, 1, 2, 0], block_shape=[3], paddings=[[0, 2]]) + + def testDirect1(self): + # Test with zero-size blocked dimension. + self._testDirect( + input_shape=[3, 0, 2, 5], block_shape=[3], paddings=[[0, 0]]) + + def testDirect2(self): + # Test with padding up from zero size. + self._testDirect( + input_shape=[3, 0, 2, 5], block_shape=[3], paddings=[[1, 2]]) + + def testDirect3(self): + self._testDirect( + input_shape=[3, 3, 4, 5, 2], + block_shape=[3, 4, 2], + paddings=[[1, 2], [0, 0], [3, 0]]) + + def testDirect4(self): + self._testDirect( + input_shape=[3, 3, 4, 5, 2], + block_shape=[3, 4, 2, 2], + paddings=[[1, 2], [0, 0], [3, 0], [0, 0]]) + + def testDirect5(self): + self._testDirect( + input_shape=[3, 2, 2, 3, 4, 5, 2, 5], + block_shape=[1, 1, 3, 4, 2, 2], + paddings=[[0, 0], [0, 0], [1, 2], [0, 0], [3, 0], [0, 0]]) + + def testDirect6(self): + self._testDirect( + input_shape=[3, 2, 2, 3, 4, 5, 2, 5], + block_shape=[1, 1, 3, 4, 2, 2, 1], + paddings=[[0, 0], [0, 0], [1, 2], [0, 0], [3, 0], [0, 0], [0, 0]]) + + +if __name__ == "__main__": + test.main() diff --git a/tensorflow/compiler/tests/stack_ops_test.py b/tensorflow/compiler/tests/stack_ops_test.py new file mode 100644 index 0000000000000000000000000000000000000000..2b9c2279737ccee531d488d27ccdb0cafa1dc8fc --- /dev/null +++ b/tensorflow/compiler/tests/stack_ops_test.py @@ -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. +# ============================================================================== +"""Tests for tensorflow.ops.stack_ops.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import numpy as np + +from tensorflow.compiler.tests.xla_test import XLATestCase +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_data_flow_ops +from tensorflow.python.platform import test + + +class StackOpTest(XLATestCase): + + def testStackPushPop(self): + with self.test_session(), self.test_scope(): + size = array_ops.placeholder(dtypes.int32) + v = array_ops.placeholder(dtypes.float32) + h = gen_data_flow_ops._stack_v2(size, dtypes.float32, stack_name="foo") + c = gen_data_flow_ops._stack_push_v2(h, v) + with ops.control_dependencies([c]): + c1 = gen_data_flow_ops._stack_pop_v2(h, dtypes.float32) + self.assertAllClose([[4.0, 5.0]], c1.eval({size: 5, v: [[4.0, 5.0]]})) + + def testStackPushPopSwap(self): + with self.test_session(), self.test_scope(): + a = np.arange(2000) + x = array_ops.placeholder(dtypes.float32) + h = gen_data_flow_ops._stack_v2(5, dtypes.float32, stack_name="foo") + c = gen_data_flow_ops._stack_push_v2(h, x, swap_memory=True) + with ops.control_dependencies([c]): + c1 = gen_data_flow_ops._stack_pop_v2(h, dtypes.float32) + self.assertAllClose(a, c1.eval({x: a})) + + def testMultiStack(self): + with self.test_session(), self.test_scope(): + v = array_ops.placeholder(dtypes.float32) + h1 = gen_data_flow_ops._stack_v2(5, dtypes.float32, stack_name="foo") + c1 = gen_data_flow_ops._stack_push_v2(h1, v) + with ops.control_dependencies([c1]): + c1 = gen_data_flow_ops._stack_pop_v2(h1, dtypes.float32) + h2 = gen_data_flow_ops._stack_v2(5, dtypes.float32, stack_name="bar") + c2 = gen_data_flow_ops._stack_push_v2(h2, 5.0) + with ops.control_dependencies([c2]): + c2 = gen_data_flow_ops._stack_pop_v2(h2, dtypes.float32) + r = c1 + c2 + self.assertAllClose(9.0, r.eval({v: 4.0})) + + def testSameNameStacks(self): + """Different stacks with the same name do not interfere.""" + with self.test_session() as sess, self.test_scope(): + v1 = array_ops.placeholder(dtypes.float32) + v2 = array_ops.placeholder(dtypes.float32) + h1 = gen_data_flow_ops._stack_v2(5, dtypes.float32, stack_name="foo") + h2 = gen_data_flow_ops._stack_v2(5, dtypes.float32, stack_name="foo") + + c1 = gen_data_flow_ops._stack_push_v2(h1, v1) + with ops.control_dependencies([c1]): + c2 = gen_data_flow_ops._stack_push_v2(h2, v2) + with ops.control_dependencies([c2]): + pop1 = gen_data_flow_ops._stack_pop_v2(h1, dtypes.float32) + pop2 = gen_data_flow_ops._stack_pop_v2(h2, dtypes.float32) + + out1, out2 = sess.run([pop1, pop2], {v1: 4.0, v2: 5.0}) + self.assertAllClose(out1, 4.0) + self.assertAllClose(out2, 5.0) + + def testCloseStack(self): + with self.test_session() as sess, self.test_scope(): + size = array_ops.placeholder(dtypes.int32) + h = gen_data_flow_ops._stack_v2(size, dtypes.float32, stack_name="foo") + c1 = gen_data_flow_ops._stack_close_v2(h) + sess.run(c1, {size: 5}) + + def testPushCloseStack(self): + with self.test_session() as sess, self.test_scope(): + v = array_ops.placeholder(dtypes.float32) + h = gen_data_flow_ops._stack_v2(5, dtypes.float32, stack_name="foo") + c = gen_data_flow_ops._stack_push_v2(h, v) + with ops.control_dependencies([c]): + c1 = gen_data_flow_ops._stack_close_v2(h) + sess.run(c1, {v: [[4.0, 5.0]]}) + + +if __name__ == "__main__": + test.main() diff --git a/tensorflow/compiler/tests/tensor_array_ops_test.py b/tensorflow/compiler/tests/tensor_array_ops_test.py new file mode 100644 index 0000000000000000000000000000000000000000..ac039e01623b954e291760fb9b50ef8eae3da7c1 --- /dev/null +++ b/tensorflow/compiler/tests/tensor_array_ops_test.py @@ -0,0 +1,1022 @@ +# 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 tests for XLA TensorArray 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.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 gen_data_flow_ops +from tensorflow.python.ops import gradients_impl +from tensorflow.python.ops import resource_variable_ops +from tensorflow.python.ops import tensor_array_grad # pylint: disable=unused-import +from tensorflow.python.ops import tensor_array_ops +from tensorflow.python.ops import variables +from tensorflow.python.platform import test + + +def _make_converter(dtype): + def _converter(x): + return np.asarray(x).astype(dtype.as_numpy_dtype) + return _converter + + +class TensorArrayTest(xla_test.XLATestCase): + + def testTensorArrayWriteRead(self): + with self.test_session() as session, self.test_scope(): + ta = tensor_array_ops.TensorArray( + dtype=dtypes.float32, + tensor_array_name="foo", + size=3) + + w0 = ta.write(0, [[4.0, 5.0]]) + w1 = w0.write(1, [[1.0, 3.0]]) + w2 = w1.write(2, [[7.0, -8.5]]) + + r0 = w2.read(0) + r1 = w2.read(1) + r2 = w2.read(2) + flow = w2.flow + + d0, d1, d2, flow_val = session.run([r0, r1, r2, flow]) + self.assertAllEqual([[4.0, 5.0]], d0) + self.assertAllEqual([[1.0, 3.0]], d1) + self.assertAllEqual([[7.0, -8.5]], d2) + self.assertAllEqual([], flow_val.shape) + + def _testTensorArrayWritePack(self, tf_dtype): + with self.test_session(), self.test_scope(): + ta = tensor_array_ops.TensorArray( + dtype=tf_dtype, tensor_array_name="foo", size=3) + + convert = _make_converter(tf_dtype) + + w0 = ta.write(0, convert([[4.0, 5.0]])) + w1 = w0.write(1, convert([[6.0, 7.0]])) + w2 = w1.write(2, convert([[8.0, 9.0]])) + + c0 = w2.stack() + + self.assertAllEqual( + convert([[[4.0, 5.0]], [[6.0, 7.0]], [[8.0, 9.0]]]), c0.eval()) + + def testTensorArrayWritePack(self): + for dtype in self.numeric_tf_types: + self._testTensorArrayWritePack(dtype) + + def testEmptyTensorArrayPack(self): + with self.test_session(), self.test_scope(): + ta = tensor_array_ops.TensorArray( + dtype=dtypes.float32, tensor_array_name="foo", size=3) + + empty_element = np.zeros((0, 1), dtype=np.float32) + w0 = ta.write(0, empty_element) + w1 = w0.write(1, empty_element) + w2 = w1.write(2, empty_element) + + c0 = w2.stack() + + self.assertAllEqual([3, 0, 1], c0.eval().shape) + + def _testTensorArrayWriteConcat(self, tf_dtype): + with self.test_session(), self.test_scope(): + ta = tensor_array_ops.TensorArray( + dtype=tf_dtype, tensor_array_name="foo", size=3) + + convert = _make_converter(tf_dtype) + + w0 = ta.write(0, convert([[4.0, 5.0], [104.0, 105.0]])) + w1 = w0.write(1, convert([[6.0, 7.0], [106.0, 107.0]])) + w2 = w1.write(2, convert([[8.0, 9.0], [204.0, 205.0]])) + + 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()) + + def testTensorArrayWriteConcat(self): + for dtype in self.numeric_tf_types: + self._testTensorArrayWriteConcat(dtype) + + def _testTensorArrayUnpackRead(self, tf_dtype): + with self.test_session() as session, self.test_scope(): + ta = tensor_array_ops.TensorArray( + dtype=tf_dtype, tensor_array_name="foo", size=3) + + convert = _make_converter(tf_dtype) + + # Unpack a vector into scalars + w0 = ta.unstack(convert([1.0, 2.0, 3.0])) + r0 = w0.read(0) + r1 = w0.read(1) + r2 = w0.read(2) + + d0, d1, d2 = session.run([r0, r1, r2]) + self.assertAllEqual(convert(1.0), d0) + self.assertAllEqual(convert(2.0), d1) + self.assertAllEqual(convert(3.0), d2) + + ta = tensor_array_ops.TensorArray( + dtype=tf_dtype, tensor_array_name="foo", size=3) + + # Unpack a matrix into vectors. + w1 = ta.unstack(convert([[1.0, 1.1], [2.0, 2.1], [3.0, 3.1]])) + r0 = w1.read(0) + r1 = w1.read(1) + r2 = w1.read(2) + + d0, d1, d2 = session.run([r0, r1, r2]) + self.assertAllEqual(convert([1.0, 1.1]), d0) + self.assertAllEqual(convert([2.0, 2.1]), d1) + self.assertAllEqual(convert([3.0, 3.1]), d2) + + # Reset ta because we're going to change the shape, else shape + # inference will throw an error. + ta = tensor_array_ops.TensorArray( + dtype=tf_dtype, tensor_array_name="foo", size=3) + + # Try unpacking an empty matrix, which should not cause an error. + w2 = ta.unstack(convert([[], [], []])) + r0 = w2.read(0) + r1 = w2.read(1) + r2 = w2.read(2) + + d0, d1, d2 = session.run([r0, r1, r2]) + self.assertAllEqual(convert([]), d0) + self.assertAllEqual(convert([]), d1) + self.assertAllEqual(convert([]), d2) + + def _testTensorArrayUnpackReadMaybeLegacy(self): + for dtype in self.numeric_tf_types: + self._testTensorArrayUnpackRead(dtype) + + def testTensorArrayUnpackRead(self): + self._testTensorArrayUnpackReadMaybeLegacy() + + def _testTensorArraySplitRead(self, tf_dtype): + with self.test_session() as session, self.test_scope(): + ta = tensor_array_ops.TensorArray( + dtype=tf_dtype, tensor_array_name="foo", size=3) + + convert = _make_converter(tf_dtype) + + # Split an empty vector. + lengths = constant_op.constant([0, 0, 0]) + w0 = ta.split(convert([]), lengths=lengths) + r0 = w0.read(0) + r1 = w0.read(1) + r2 = w0.read(2) + + d0, d1, d2 = session.run([r0, r1, r2]) + self.assertAllEqual(convert([]), d0) + self.assertAllEqual(convert([]), d1) + self.assertAllEqual(convert([]), d2) + + # Split a vector. + ta = tensor_array_ops.TensorArray( + dtype=tf_dtype, tensor_array_name="foo", size=3) + lengths = constant_op.constant([1, 1, 1]) + w0 = ta.split(convert([1.0, 2.0, 3.0]), lengths=lengths) + r0 = w0.read(0) + r1 = w0.read(1) + r2 = w0.read(2) + + d0, d1, d2 = session.run([r0, r1, r2]) + self.assertAllEqual(convert([1.0]), d0) + self.assertAllEqual(convert([2.0]), d1) + self.assertAllEqual(convert([3.0]), d2) + + # Split a matrix. + ta = tensor_array_ops.TensorArray( + dtype=tf_dtype, tensor_array_name="foo", size=3) + lengths = constant_op.constant([1, 1, 1]) + w0 = ta.split( + convert([[1.0, 101.0], [2.0, 201.0], [3.0, 301.0]]), lengths=lengths) + r0 = w0.read(0) + r1 = w0.read(1) + r2 = w0.read(2) + + d0, d1, d2 = session.run([r0, r1, r2]) + self.assertAllEqual(convert([[1.0, 101.0]]), d0) + self.assertAllEqual(convert([[2.0, 201.0]]), d1) + self.assertAllEqual(convert([[3.0, 301.0]]), d2) + + def testTensorArraySplitRead(self): + for dtype in self.numeric_tf_types: + self._testTensorArraySplitRead(dtype) + + def testTensorGradArrayWriteRead(self): + with self.test_session() as session, self.test_scope(): + ta = tensor_array_ops.TensorArray( + dtype=dtypes.float32, + tensor_array_name="foo", + size=3) + + w0 = ta.write(0, [[4.0]]) + w1 = w0.write(1, [[1.0]]) + w2 = w1.write(2, [[-3.0]]) + + g_ta = w2.grad("grad") + + g_w0 = g_ta.write(0, [[5.0]]) + g_w1 = g_w0.write(1, [[2.0]]) + g_w2 = g_w1.write(2, [[-2.0]]) + + r0 = w2.read(0) + r1 = w2.read(1) + r2 = w2.read(2) + + g_r0 = g_w2.read(0) + g_r1 = g_w2.read(1) + g_r2 = g_w2.read(2) + + d0, d1, d2, g_d0, g_d1, g_d2 = session.run([r0, r1, r2, g_r0, g_r1, g_r2]) + self.assertAllEqual([[4.0]], d0) + self.assertAllEqual([[1.0]], d1) + self.assertAllEqual([[-3.0]], d2) + self.assertAllEqual([[5.0]], g_d0) + self.assertAllEqual([[2.0]], g_d1) + self.assertAllEqual([[-2.0]], g_d2) + + def testTensorGradArrayDynamicWriteRead(self): + with self.test_session() as session, self.test_scope(): + ta = tensor_array_ops.TensorArray( + dtype=dtypes.float32, + tensor_array_name="foo", + size=3) + + w0 = ta.write(0, [[4.0]]) + w1 = w0.write(1, [[1.0]]) + w2 = w1.write(2, [[-3.0]]) + + g_ta = w2.grad("grad") # Get gradient array here so we know the shape + + s = w2.size() + g_s = g_ta.size() + + g_w0 = g_ta.write(0, [[5.0]]) + g_w1 = g_w0.write(1, [[2.0]]) + g_w2 = g_w1.write(2, [[-2.0]]) + + r0 = w2.read(0) + r1 = w2.read(1) + r2 = w2.read(2) + + g_r0 = g_w2.read(0) + g_r1 = g_w2.read(1) + g_r2 = g_w2.read(2) + + d0, d1, d2, g_d0, g_d1, g_d2, vs, g_vs = session.run( + [r0, r1, r2, g_r0, g_r1, g_r2, s, g_s]) + self.assertAllEqual([[4.0]], d0) + self.assertAllEqual([[1.0]], d1) + self.assertAllEqual([[-3.0]], d2) + self.assertAllEqual([[5.0]], g_d0) + self.assertAllEqual([[2.0]], g_d1) + self.assertAllEqual([[-2.0]], g_d2) + self.assertAllEqual(3, vs) + self.assertAllEqual(3, g_vs) + + def testTensorGradAccessTwiceReceiveSameObject(self): + with self.test_session() as session, self.test_scope(): + ta = tensor_array_ops.TensorArray( + dtype=dtypes.float32, tensor_array_name="foo", size=3, + element_shape=[1, 2]) + g_ta_0 = ta.grad("grad") + g_ta_1 = ta.grad("grad") + + with ops.control_dependencies([g_ta_0.write(0, [[4.0, 5.0]]).flow]): + # Write with one gradient handle, read with another copy of it + r1_0 = g_ta_1.read(0) + + t_g_ta_0, t_g_ta_1, d_r1_0 = session.run( + [g_ta_0.handle.op, g_ta_1.handle.op, r1_0]) + self.assertAllEqual(t_g_ta_0, t_g_ta_1) + self.assertAllEqual([[4.0, 5.0]], d_r1_0) + + def testTensorArrayWriteWrongIndexOrDataTypeFails(self): + with self.test_session(), self.test_scope(): + ta = tensor_array_ops.TensorArray( + dtype=dtypes.float32, tensor_array_name="foo", size=3) + + # Test writing the wrong datatype. + with self.assertRaisesOpError( + "TensorArray dtype is float but op has dtype int32"): + ta.write(-1, np.int32(7)).flow.eval() + + def testTensorArrayReadWrongIndexOrDataTypeFails(self): + # Find two different floating point types, create an array of + # the first type, but try to read the other type. + if len(self.float_types) > 1: + dtype1 = self.float_types[0] + dtype2 = self.float_types[1] + with self.test_session(), self.test_scope(): + ta = tensor_array_ops.TensorArray( + dtype=dtype1, tensor_array_name="foo", size=3) + + w0 = ta.write(0, [[4.0, 5.0]]) + + # Test reading wrong datatype. + 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() + + # Test reading from a different index than the one we wrote to + w0.read(1) + + def testTensorArraySplitIncompatibleShapesFails(self): + with self.test_session(), self.test_scope(): + ta = tensor_array_ops.TensorArray( + dtype=dtypes.float32, + tensor_array_name="foo", + size=3, + infer_shape=False) + + with self.assertRaisesOpError( + r"value is not 1D"): + lengths = array_ops.placeholder(dtypes.int64) + ta.split([1.0, 2.0, 3.0], lengths).flow.eval(feed_dict={lengths: 1}) + + with self.assertRaisesOpError( + r"lengths must be equal: 1 vs. 2"): + ta.split([1.0, 2.0, 3.0], [1, 2, 3]).flow.eval() + + with self.assertRaisesOpError( + r"value must have rank >= 1"): + ta.split(1.0, [1]).flow.eval() + + ta = tensor_array_ops.TensorArray( + dtype=dtypes.float32, + tensor_array_name="foo", + size=2, + infer_shape=False) + + with self.assertRaisesOpError( + r"TensorArray's size is not equal to the size of lengths " + r"\(1 vs. 2\)"): + ta.split([1.0], [1]).flow.eval() + + def _testTensorArrayWriteGradientAddMultipleAdds(self, dtype): + with self.test_session(), self.test_scope(): + ta = tensor_array_ops.TensorArray( + dtype=dtype, tensor_array_name="foo", size=3, infer_shape=False) + + c = lambda x: np.asarray(x, dtype=dtype.as_numpy_dtype) + + w0 = ta.write(2, c(3.0)) + w1 = w0.write(2, c(4.0)) + + ta_grad = w1.grad("grad") + + w0_grad = ta_grad.write(2, c(3.0)) + w1_grad = w0_grad.write(2, c(4.0)) + w2_grad = w1_grad.write(2, c(5.0)) + + # Assert that aggregation works correctly + self.assertAllEqual(c(12.00), w2_grad.read(2).eval()) + + # Using differing shapes causes an exception + wb0_grad = ta_grad.write(1, c(1.0)) + wb1_grad = wb0_grad.write(1, c([1.0])) + + with self.assertRaisesOpError( + r"Mismatched TensorArray sizes"): + wb1_grad.flow.eval() + + def testTensorArrayWriteGradientAddMultipleAdds(self): + for dtype in self.numeric_tf_types: + self._testTensorArrayWriteGradientAddMultipleAdds(dtype) + + def testMultiTensorArray(self): + with self.test_session(), self.test_scope(): + h1 = tensor_array_ops.TensorArray( + size=1, dtype=dtypes.float32, tensor_array_name="foo") + w1 = h1.write(0, 4.0) + r1 = w1.read(0) + + h2 = tensor_array_ops.TensorArray( + size=1, dtype=dtypes.float32, tensor_array_name="bar") + + w2 = h2.write(0, 5.0) + r2 = w2.read(0) + r = r1 + r2 + self.assertAllClose(9.0, r.eval()) + + def _testTensorArrayGradientWriteReadType(self, dtype): + with self.test_session() as session, self.test_scope(): + ta = tensor_array_ops.TensorArray( + dtype=dtypes.as_dtype(dtype), + tensor_array_name="foo", + size=3, + infer_shape=False) + + c = lambda x: np.array(x, dtype=dtype) + + value_0 = constant_op.constant(c([[4.0, 5.0]])) + value_1 = constant_op.constant(c([[3.0, 3.5]])) + + w0 = ta.write(0, value_0) + w1 = w0.write(1, value_1) + r0 = w1.read(0) + r1 = w1.read(1) + r0_2 = w1.read(0) + + # Test individual components' gradients + grad_just_r0 = gradients_impl.gradients( + ys=[r0], xs=[value_0], grad_ys=[c([[2.0, 3.0]])]) + grad_just_r0_vals = session.run(grad_just_r0) + self.assertAllEqual(c([[2.0, 3.0]]), grad_just_r0_vals[0]) + + grad_r0_r0_2 = gradients_impl.gradients( + ys=[r0, r0_2], + xs=[value_0], + grad_ys=[c([[2.0, 3.0]]), c([[1.0, -1.0]])]) + grad_r0_r0_2_vals = session.run(grad_r0_r0_2) + self.assertAllEqual(c([[3.0, 2.0]]), grad_r0_r0_2_vals[0]) + + grad_just_r1 = gradients_impl.gradients( + ys=[r1], xs=[value_1], grad_ys=[c([[-2.0, -4.0]])]) + grad_just_r1_vals = session.run(grad_just_r1) + self.assertAllEqual(c([[-2.0, -4.0]]), grad_just_r1_vals[0]) + + # Test combined gradients + grad = gradients_impl.gradients( + ys=[r0, r0_2, r1], + xs=[value_0, value_1], + grad_ys=[c([[2.0, 3.0]]), c([[1.0, -1.0]]), c([[-2.0, -10.0]])]) + grad_vals = session.run(grad) + self.assertEqual(len(grad_vals), 2) + self.assertAllEqual(c([[3.0, 2.0]]), grad_vals[0]) + self.assertAllEqual(c([[-2.0, -10.0]]), grad_vals[1]) + + def testTensorArrayGradientWriteRead(self): + for dtype in self.numeric_types: + self._testTensorArrayGradientWriteReadType(dtype) + + def _testTensorArrayGradientWritePackConcatAndRead(self): + with self.test_session() as sess, self.test_scope(): + ta = tensor_array_ops.TensorArray( + dtype=dtypes.float32, + tensor_array_name="foo", + size=2, + clear_after_read=False) + + value_0 = constant_op.constant([-1.0, 1.0]) + value_1 = constant_op.constant([-10.0, 10.0]) + + w0 = ta.write(0, value_0) + w1 = w0.write(1, value_1) + p0 = w1.stack() + r0 = w1.read(0) + s0 = w1.concat() + + # Test gradient accumulation between read(0), pack(), and concat(). + with ops.control_dependencies([p0, r0, s0]): + grad_r = gradients_impl.gradients( + ys=[p0, r0, s0], + xs=[value_0, value_1], + grad_ys=[ + [[2.0, 3.0], [4.0, 5.0]], # stack gradient + [-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 + + 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]) + + def testTensorArrayGradientWritePackConcatAndRead(self): + self._testTensorArrayGradientWritePackConcatAndRead() + + def testTensorArrayReadTwice(self): + with self.test_session(), self.test_scope(): + value = constant_op.constant([[1.0, -1.0], [10.0, -10.0]]) + + ta_readtwice = tensor_array_ops.TensorArray( + dtype=dtypes.float32, + tensor_array_name="foo", + size=2, + clear_after_read=False) + w_readtwice = ta_readtwice.unstack(value) + r0_readtwice = w_readtwice.read(0) + with ops.control_dependencies([r0_readtwice]): + r1_readtwice = w_readtwice.read(0) + + self.assertAllEqual([1.0, -1.0], r1_readtwice.eval()) + + def _testTensorArrayGradientUnpackRead(self): + with self.test_session() as session, self.test_scope(): + ta = tensor_array_ops.TensorArray( + dtype=dtypes.float32, + tensor_array_name="foo", + size=2, + clear_after_read=False) + + value = constant_op.constant([[1.0, -1.0], [10.0, -10.0]]) + + w = ta.unstack(value) + r0 = w.read(0) + r0_1 = w.read(0) + r1 = w.read(1) + + # Test combined gradients + aggregation of read(0). + grad = gradients_impl.gradients( + ys=[r0, r0_1, r1], + xs=[value], + grad_ys=[[2.0, 3.0], [-1.5, 1.5], [4.0, 5.0]]) + grad_vals = session.run(grad) + + self.assertEqual(len(grad_vals), 1) + self.assertAllEqual([[2.0 - 1.5, 3.0 + 1.5], [4.0, 5.0]], grad_vals[0]) + + def testTensorArrayGradientUnpackRead(self): + self._testTensorArrayGradientUnpackRead() + + def testTensorArrayGradientSplitConcat(self): + with self.test_session() as session, self.test_scope(): + ta = tensor_array_ops.TensorArray( + dtype=dtypes.float32, tensor_array_name="foo", size=2) + + value = constant_op.constant( + [[1.0, -1.0], [10.0, -10.0], [100.0, -100.0], [1000.0, -1000.0]]) + + w = ta.split(value, [2, 2]) + r = w.concat() + + # Test combined gradients + grad = gradients_impl.gradients( + ys=[r], + xs=[value], + grad_ys=[[[2.0, -2.0], [20.0, -20.0], [200.0, -200.0], + [2000.0, -2000.0]]]) + grad_vals = session.run(grad) + + self.assertEqual(len(grad_vals), 1) + self.assertAllEqual([[2.0, -2.0], [20.0, -20.0], [200.0, -200.0], + [2000.0, -2000.0]], + grad_vals[0]) + + def testCloseTensorArray(self): + with self.test_session() as session, self.test_scope(): + ta = tensor_array_ops.TensorArray( + dtype=dtypes.float32, tensor_array_name="foo", size=3) + c1 = ta.close() + session.run(c1) + + def testSizeTensorArray(self): + with self.test_session(), self.test_scope(): + ta = tensor_array_ops.TensorArray( + dtype=dtypes.float32, tensor_array_name="foo", size=3) + s = ta.size() + self.assertAllEqual(3, s.eval()) + + def testWriteCloseTensorArray(self): + with self.test_session(), self.test_scope(): + ta = tensor_array_ops.TensorArray( + dtype=dtypes.float32, + tensor_array_name="foo", + size=3, + infer_shape=False) + w0 = ta.write(0, [[4.0, 5.0]]) + w1 = w0.write(1, [3.0]) + w1.close().run() # Expected to run without problems + + # TODO(phawkins): implement while loops. + # def _testWhileLoopWritePackGradients(self, dynamic_size, dtype): + # np_dtype = dtype.as_numpy_dtype + # with self.test_session() as session, self.test_scope(): + # v0 = array_ops.identity(np.arange(3 * 5, dtype=np_dtype).reshape(3, 5)) + # var = variables.Variable(np.arange(100, 105, dtype=np_dtype)) + # state0 = array_ops.identity(np.array([1] * 5, dtype=np_dtype)) + # ta = tensor_array_ops.TensorArray( + # dtype=dtype, + # tensor_array_name="foo", + # size=0 if dynamic_size else 3, + # dynamic_size=dynamic_size) + # time_0 = array_ops.identity(0) + + # def body(time, ta_t, state): + # sliced = array_ops.slice( + # v0, begin=array_ops.stack([time, 0]), size=[1, -1]) + # sliced = array_ops.squeeze(sliced) + # out = sliced + var + state + # state += sliced + # ta_t = ta_t.write(time, out) + # return (time + 1, ta_t, state) + + # (unused_0, h_final, unused_2) = control_flow_ops.while_loop( + # cond=lambda time, unused_1, unused_2: time < 3, + # body=body, + # loop_vars=(time_0, ta, state0), + # shape_invariants=(time_0.get_shape(), tensor_shape.unknown_shape(), + # tensor_shape.unknown_shape()), + # parallel_iterations=3) + # vout = h_final.stack() + + # grad_val = -np.arange(3 * 5, dtype=np_dtype).reshape(3, 5) + # v0_grad = gradients_impl.gradients([vout], [v0], [grad_val])[0] + # state0_grad = gradients_impl.gradients([vout], [state0], [grad_val])[0] + # var_grad = gradients_impl.gradients([vout], [var], [grad_val])[0] + + # variables.global_variables_initializer().run() + # state0_t, var_t, v0_t, vout_t, v0_grad_t, var_grad_t, state0_grad_t = ( + # session.run([state0, var, v0, vout, v0_grad, var_grad, state0_grad]) + # ) + # just_v0_grad_t, = session.run([v0_grad]) + + # # state = [ state0 | state0 + v0[0] | state0 + v0[0] + v0[1] ] + # # vout = [ v0[0] + var + state[0] | + # # v0[1] + var + state[1] | + # # v0[2] + var + state[2] ] + # # = [ v0[0] + var + state0 | + # # v0[1] + var + state0 + v0[0] | + # # v0[2] + var + state0 + v0[0] + v0[1] ] + # # + # # d(vout[0])/d(v0) = [1 | 0 | 0 ] + # # d(vout[1])/d(v0) = [1 | 1 | 0 ] + # # d(vout[2])/d(v0) = [1 | 1 | 1 ] + # # d(vout)/d(var) = [1 | 1 | 1] + # # d(vout)/d(state0) = [ 1 | 1 | 1 ] + + # state_per_time = np.array( + # [state0_t, state0_t + v0_t[0, :], + # state0_t + v0_t[0, :] + v0_t[1, :]]) + + # # Compare forward prop + # self.assertAllClose(v0_t + var_t + state_per_time, vout_t) + + # # Compare backward prop + # expected_v0_grad_t = np.array([ + # grad_val[0, :] + grad_val[1, :] + grad_val[2, :], + # grad_val[1, :] + grad_val[2, :], grad_val[2, :] + # ]) + + # self.assertAllEqual(expected_v0_grad_t, v0_grad_t) + # self.assertAllEqual(expected_v0_grad_t, just_v0_grad_t) + # self.assertAllClose(grad_val.sum(axis=0), var_grad_t) + # self.assertAllClose(grad_val.sum(axis=0), state0_grad_t) + + # def testWhileLoopWritePackGradients(self): + # self._testWhileLoopWritePackGradients( + # dynamic_size=False, dtype=dtypes.float32) + # # TODO(ebrevdo): re-enable when While supports non-float32 gradients. + # # self._testWhileLoopWritePackGradients( + # # dynamic_size=False, dtype=tf.int64) + + # def testWhileLoopDynamicWritePackGradients(self): + # self._testWhileLoopWritePackGradients( + # dynamic_size=True, dtype=dtypes.float32) + + # def testGradSerialTwoLoops(self): + # with self.test_session(), self.test_scope(): + # num_steps = 100 + # acc = tensor_array_ops.TensorArray( + # dtype=dtypes.float32, + # size=num_steps, + # clear_after_read=False, + # element_shape=tensor_shape.scalar()) + # i = constant_op.constant(0, name="i") + # x = constant_op.constant(2.0, name="x") + + # c = lambda i, acc: i < 5 + + # def b(i, acc): + # x1 = control_flow_ops.cond( + # math_ops.equal(i, 0), lambda: x, + # lambda: math_ops.multiply(acc.read(i - 1), 2.0)) + # return i + 1, acc.write(i, x1) + + # i1, acc1 = control_flow_ops.while_loop(c, b, [i, acc]) + + # z = constant_op.constant(0.0) + + # def fn(i, acc): + # return i + 1, acc.write(i, z) + + # _, acc2 = control_flow_ops.while_loop(lambda i, acc: i < num_steps, fn, + # [i1, acc1]) + + # r = acc2.stack() + # grad = gradients_impl.gradients(r, [x])[0] + # self.assertAllClose(31.0, grad.eval()) + + def testSumOfTwoReadVariablesWithoutRepeatGrad(self): + with self.test_session() as session, self.test_scope(): + a = array_ops.identity( + np.arange( + 3 * 5, dtype=np.float32).reshape(3, 5) + 1) + b = array_ops.identity( + np.arange( + 3 * 5, dtype=np.float32).reshape(3, 5) + 1 + 3 * 5) + ta = tensor_array_ops.TensorArray(dtype=dtypes.float32, size=2) + ta = ta.write(0, a, name="write_a") + ta = ta.write(1, b, name="write_b") + c = ( + ta.read( + 0, name="read_a_0") + # a + b + ta.read( + 1, name="read_b_0")) + g0 = -(np.arange(3 * 5, dtype=np.float32).reshape(3, 5) + 1) + grad_a = gradients_impl.gradients([c], [a], [g0])[0] # d(a+b)/da = 1 + grad_b = gradients_impl.gradients([c], [b], [g0])[0] # d(a+b)/db = 1 + + # Test gradients calculated individually + grad_a_t, = session.run([grad_a]) + self.assertAllEqual(grad_a_t, g0) + + grad_b_t, = session.run([grad_b]) + self.assertAllEqual(grad_b_t, g0) + + # Test gradients calculated jointly. + joint_grad_a_t, joint_grad_b_t = session.run([grad_a, grad_b]) + self.assertAllEqual(joint_grad_a_t, g0) + self.assertAllEqual(joint_grad_b_t, g0) + + def testWriteShape(self): + with self.test_session(), self.test_scope(): + ta = tensor_array_ops.TensorArray( + dtype=dtypes.float32, tensor_array_name="foo", size=3) + c0 = constant_op.constant([4.0, 5.0]) + w0 = ta.write(0, c0) + r0 = w0.read(0) + self.assertAllEqual(c0.get_shape(), r0.get_shape()) + + ta = tensor_array_ops.TensorArray( + dtype=dtypes.float32, tensor_array_name="foo", size=3) + c1 = constant_op.constant([6.0, 7.0]) + w1 = w0.write(1, c1) + r0 = w1.read(0) + r1 = w1.read(1) + self.assertAllEqual(c0.get_shape(), r0.get_shape()) + self.assertAllEqual(c1.get_shape(), r1.get_shape()) + + ta = tensor_array_ops.TensorArray( + dtype=dtypes.float32, tensor_array_name="foo", size=3) + c2 = constant_op.constant([4.0, 5.0, 6.0]) + with self.assertRaises(ValueError): + w0.write(0, c2) + + def testPartlyUnknownShape(self): + with self.test_session(), self.test_scope(): + ta = tensor_array_ops.TensorArray( + dtype=dtypes.float32, tensor_array_name="foo", size=6) + + c0 = array_ops.placeholder(dtypes.float32, [None, None, None, 3]) + w0 = ta.write(0, c0) + r0 = w0.read(0) + self.assertAllEqual([None, None, None, 3], r0.get_shape().as_list()) + + c1 = array_ops.placeholder(dtypes.float32, [None, None, None, 3]) + w1 = w0.write(1, c1) + r1 = w1.read(0) + self.assertAllEqual([None, None, None, 3], r1.get_shape().as_list()) + + # Writing less specific shape (doesn't change type.) + c2 = array_ops.placeholder(dtypes.float32, [None, None, None, None]) + w2 = w1.write(2, c2) + r2 = w2.read(0) + self.assertAllEqual([None, None, None, 3], r2.get_shape().as_list()) + + # Writing more specific shape in one dimension and less specific in + # another. + c3 = array_ops.placeholder(dtypes.float32, [None, None, 2, None]) + w3 = w2.write(3, c3) + r3 = w3.read(0) + self.assertAllEqual([None, None, 2, 3], r3.get_shape().as_list()) + + # Writing partly defined shape using TensorArray.scatter. + c4 = array_ops.placeholder(dtypes.float32, [2, None, 4, 2, 3]) + w4 = w3.scatter([4, 5], c4) + r4 = w4.read(0) + self.assertAllEqual([None, 4, 2, 3], r4.get_shape().as_list()) + + # Writing fully defined shape using TensorArray.split. + c5 = array_ops.placeholder(dtypes.float32, [10, 4, 2, 3]) + w5 = w4.split(c5, constant_op.constant([5, 5])) + r5 = w5.read(0) + self.assertAllEqual([5, 4, 2, 3], r5.get_shape().as_list()) + + def _testUnpackShape(self): + with self.test_session(), self.test_scope(): + ta = tensor_array_ops.TensorArray( + dtype=dtypes.float32, + tensor_array_name="foo", + size=0, + infer_shape=True) + value = constant_op.constant( + [[1.0, -1.0], [10.0, -10.0], [100.0, -100.0]]) + w0 = ta.unstack(value) + r0 = w0.read(0) + self.assertAllEqual((2,), r0.get_shape()) + + c1 = constant_op.constant([4.0, 5.0]) + w1 = w0.write(3, c1) + r1 = w1.read(0) + self.assertAllEqual(c1.get_shape(), r1.get_shape()) + + c2 = constant_op.constant([4.0, 5.0, 6.0]) + with self.assertRaises(ValueError): + w1.write(4, c2) + + def testUnpackShape(self): + self._testUnpackShape() + + def testSplitShape(self): + with self.test_session(), self.test_scope(): + ta = tensor_array_ops.TensorArray( + dtype=dtypes.float32, + tensor_array_name="foo", + size=0, + infer_shape=True) + value = constant_op.constant([[1.0, -1.0], [2.0, -2.0], [3.0, -3.0]]) + w0 = ta.split(value, [1, 1, 1]) + r0 = w0.read(0) + self.assertAllEqual((1, 2), r0.get_shape()) + + ta1 = tensor_array_ops.TensorArray( + dtype=dtypes.float32, + tensor_array_name="foo1", + size=0, + infer_shape=True) + w0 = ta1.split(value, [1, 2]) + r0 = w0.read(0) + self.assertAllEqual(r0.get_shape(), tensor_shape.unknown_shape()) + + def testWriteUnknownShape(self): + with self.test_session(), self.test_scope(): + ta = tensor_array_ops.TensorArray( + dtype=dtypes.float32, + tensor_array_name="foo", + size=3, + infer_shape=True) + c0 = array_ops.placeholder(dtypes.float32) + w0 = ta.write(0, c0) + r0 = w0.read(0) + self.assertAllEqual(r0.get_shape(), tensor_shape.unknown_shape()) + + def _testGradientWhenNotAllComponentsRead(self): + with self.test_session() as session, self.test_scope(): + ta = tensor_array_ops.TensorArray(dtype=dtypes.float32, size=2) + x = constant_op.constant([2.0, 3.0]) + w = ta.unstack(x) + r0 = w.read(0) + # Calculate (dr0/dx0, dr0/dx1). since r0 = x0, gradients are (1, 0). + grad_r0 = gradients_impl.gradients(ys=[r0], xs=[x], grad_ys=[1.0]) + grad_r0_vals = session.run(grad_r0)[0] + self.assertAllEqual(grad_r0_vals, [1.0, 0.0]) + + def testGradientWhenNotAllComponentsRead(self): + self._testGradientWhenNotAllComponentsRead() + + def _testTensorArrayEvalEmpty(self): + with self.test_session(), self.test_scope(): + ta = tensor_array_ops.TensorArray( + dtype=dtypes.float32, size=0, infer_shape=False) + with self.assertRaisesOpError( + "TensorArray has size zero, but element shape is not fully " + "defined. Currently only static shapes are supported when packing " + "zero-size TensorArrays."): + ta.stack().eval() + + def testTensorArrayEvalEmpty(self): + self._testTensorArrayEvalEmpty() + + def _testTensorArrayEvalEmptyWithDefault(self): + with self.test_session(), self.test_scope(): + ta = tensor_array_ops.TensorArray( + dtype=dtypes.float32, size=0, infer_shape=True) + 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) + # Concatenating zero tensors along their first dimension gives a + # first dimension of zero + self.assertAllEqual([0, 5], ta.concat().eval().shape) + + def testTensorArrayEvalEmptyWithDefault(self): + self._testTensorArrayEvalEmptyWithDefault() + + def testTensorArrayScatterReadAndGradients(self): + with self.test_session() as session, self.test_scope(): + ta = tensor_array_ops.TensorArray( + dtype=dtypes.float32, + tensor_array_name="foo", + size=10) + + indices = constant_op.constant([1, 8]) + value = constant_op.constant([[1.0, -1.0], [10.0, -10.0]]) + + w = ta.scatter(indices, value) + r0 = w.read(1) + r1 = w.read(8) + + # Test combined gradients + aggregation of read(0). + grad = gradients_impl.gradients( + ys=[r0, r1], xs=[value], grad_ys=[[2.0, 3.0], [4.0, 5.0]]) + read_vals, grad_vals = session.run([[r0, r1], grad]) + + self.assertEqual(len(read_vals), 2) + self.assertEqual(len(grad_vals), 1) + self.assertAllEqual([1.0, -1.0], read_vals[0]) + self.assertAllEqual([10.0, -10.0], read_vals[1]) + self.assertAllEqual([[2.0, 3.0], [4.0, 5.0]], grad_vals[0]) + + def testTensorArrayWriteGatherAndGradients(self): + with self.test_session() as session, self.test_scope(): + ta = tensor_array_ops.TensorArray( + dtype=dtypes.float32, + tensor_array_name="foo", + size=10) + + values = constant_op.constant([[1.0 * x, -1.0 * x] for x in range(10)]) + indices = constant_op.constant([1, 8]) + + w = ta.unstack(values) + g = w.gather(indices) + + # Test combined gradients + aggregation of read(0). + grad = gradients_impl.gradients( + ys=[g], xs=[values], grad_ys=[[[2.0, 3.0], [4.0, 5.0]]]) + g_vals, grad_vals = session.run([[g], grad]) + + # Gradients for 8 of the 10 unread components are zero. + expected_grad = np.zeros((10, 2)) + expected_grad[1] = [2.0, 3.0] + expected_grad[8] = [4.0, 5.0] + + self.assertEqual(len(g_vals), 1) + self.assertEqual(len(grad_vals), 1) + self.assertAllEqual([[1.0, -1.0], [8.0, -8.0]], g_vals[0]) + self.assertAllEqual(expected_grad, grad_vals[0]) + + def testTensorArrayIdentity(self): + with self.test_session() as session, self.test_scope(): + ta0 = tensor_array_ops.TensorArray(dtype=dtypes.float32, size=2, + infer_shape=False) + ta1 = tensor_array_ops.TensorArray(dtype=dtypes.int32, size=4, + infer_shape=True) + + ta0 = ta0.write(0, 0.) + ta1 = ta1.write(0, 1) + + v0 = resource_variable_ops.ResourceVariable(0) + v1 = resource_variable_ops.ResourceVariable(0) + + with ops.control_dependencies([v0.assign_add(1)]): + ta0 = ta0.identity() + + with ops.control_dependencies([v1.assign_add(1)]): + ta1 = ta1.identity() + + read0 = ta0.read(0) + read1 = ta1.read(0) + + size0 = ta0.size() + size1 = ta1.size() + + # Tests correct properties on new TensorArrays. + self.assertEqual(dtypes.float32, ta0.dtype) + self.assertEqual(dtypes.int32, ta1.dtype) + self.assertEqual(tensor_shape.unknown_shape(), read0.get_shape()) + self.assertEqual(tensor_shape.scalar(), read1.get_shape()) + + variables.global_variables_initializer().run() + + read0_v, read1_v, size0_v, size1_v = session.run( + (read0, read1, size0, size1)) + + # Tests that the control dependencies was added and executed. + self.assertEqual(1, v0.eval()) + self.assertEqual(1, v1.eval()) + + # Tests correct TensorArray. + self.assertEqual(read0_v, 0) + self.assertEqual(read1_v, 1) + self.assertEqual(size0_v, 2) + self.assertEqual(size1_v, 4) + +if __name__ == "__main__": + test.main() diff --git a/tensorflow/compiler/tests/unary_ops_test.py b/tensorflow/compiler/tests/unary_ops_test.py index 1e85d3a2c8b06ee956ccd844c1724ebbf1fd7495..b21f1998a5d351d4a86438236441be541eef42b0 100644 --- a/tensorflow/compiler/tests/unary_ops_test.py +++ b/tensorflow/compiler/tests/unary_ops_test.py @@ -18,6 +18,8 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function +import unittest + import numpy as np from six.moves import xrange # pylint: disable=redefined-builtin @@ -117,27 +119,61 @@ class UnaryOpsTest(XLATestCase): def testFloatOps(self): for dtype in self.float_types: + self._assertOpOutputMatchesExpected( + math_ops.acosh, + np.array([1, 2, 3, 4], dtype=dtype), + expected=np.array([0, 1.3169579, 1.76274717, 2.06343707], + dtype=dtype)) + + self._assertOpOutputMatchesExpected( + math_ops.asinh, + np.array([1, 2, 3, 4], dtype=dtype), + expected=np.array([0.88137359, 1.44363548, 1.81844646, 2.09471255], + dtype=dtype)) + + self._assertOpOutputMatchesExpected( + math_ops.atanh, + np.array([0.1, 0.2, 0.3, 0.4], dtype=dtype), + expected=np.array([0.10033535, 0.20273255, 0.3095196, 0.42364893], + dtype=dtype)) + self._assertOpOutputMatchesExpected( math_ops.ceil, np.array([[-1.7, 1.2]], dtype=dtype), expected=np.array([[-1, 2]], dtype=dtype)) + self._assertOpOutputMatchesExpected( + math_ops.cosh, + np.array([1, 2, 3, 4], dtype=dtype), + expected=np.array([1.54308063, 3.76219569, 10.067662, 27.30823284], + dtype=dtype)) + self._assertOpOutputMatchesExpected( math_ops.exp, np.array([[-1, 1]], dtype=dtype), expected=np.array([[0.36787945, 2.7182817]], dtype=dtype)) + self._assertOpOutputMatchesExpected( + math_ops.expm1, + np.array([[-1, 1]], dtype=dtype), + expected=np.array([[-0.63212056, 1.71828183]], dtype=dtype)) + self._assertOpOutputMatchesExpected( math_ops.floor, np.array([[-1.7, 1.2]], dtype=dtype), expected=np.array([[-2, 1]], dtype=dtype)) + self._assertOpOutputMatchesExpected( + math_ops.is_finite, + np.array([[np.NINF, -2, -1, 0, 0.5, 1, 2, np.inf, np.nan]], + dtype=dtype), + expected=np.array([[0, 1, 1, 1, 1, 1, 1, 0, 0]], dtype=np.bool)) + # Tests for tf.nn ops. self._assertOpOutputMatchesExpected( nn_ops.l2_loss, np.array([[[]]], dtype=dtype), expected=dtype(0)) - # TODO(b/31644876): enable this test case when fixed. - # self._assertOpOutputMatchesExpected(tf.nn.l2_loss, dtype(4), dtype(10)) + self._assertOpOutputMatchesExpected(nn_ops.l2_loss, dtype(4), dtype(8)) self._assertOpOutputMatchesExpected( nn_ops.l2_loss, np.array([[-2, 4]], dtype=dtype), expected=dtype(10)) @@ -152,6 +188,16 @@ class UnaryOpsTest(XLATestCase): np.array([[1, 2]], dtype=dtype), expected=np.array([[0, 0.69314718]], dtype=dtype)) + self._assertOpOutputMatchesExpected( + math_ops.sin, + np.array([[1, 2]], dtype=dtype), + expected=np.array([[0.841478, 0.909302]], dtype=dtype)) + + self._assertOpOutputMatchesExpected( + math_ops.cos, + np.array([[1, 2]], dtype=dtype), + expected=np.array([[0.540297, -0.41614]], dtype=dtype)) + # TODO(b/34703906): improve log1p implementation and make tolerance # tighter. self._assertOpOutputMatchesExpected( @@ -159,6 +205,12 @@ class UnaryOpsTest(XLATestCase): np.array([[1e-14, 1e-15, 0.6]], dtype=dtype), expected=np.log1p(np.array([[1e-14, 1e-15, 0.6]], dtype=dtype))) + self._assertOpOutputMatchesExpected( + math_ops.rint, + np.array([[-1.7, 1.2, 4.0, 0.0], [-3.5, -2.5, -1.5, -0.5], + [0.5, 1.5, 2.5, 3.5]], dtype=dtype), + expected=np.array([[-2, 1, 4, 0], [-4, -2, -2, 0], [0, 2, 2, 4]], + dtype=dtype)) self._assertOpOutputMatchesExpected( math_ops.round, np.array([[-1.7, 1.2, 4.0, 0.0], [-3.5, -2.5, -1.5, -0.5], @@ -182,11 +234,28 @@ class UnaryOpsTest(XLATestCase): [0.7310586, 0.880797, 0.95257413, 0.98201376]], dtype=dtype)) + self._assertOpOutputMatchesExpected( + math_ops.sigmoid, + np.array([-300, -150, 0, 150, 300], dtype=dtype), + expected=np.array([0, 0, 0.5, 1, 1], dtype=dtype)) + + self._assertOpOutputMatchesExpected( + math_ops.sinh, + np.array([1, 2, 3, 4], dtype=dtype), + expected=np.array([1.17520119, 3.62686041, 10.01787493, 27.2899172], + dtype=dtype)) + self._assertOpOutputMatchesExpected( math_ops.sqrt, np.array([[4, 9]], dtype=dtype), expected=np.array([[2, 3]], dtype=dtype)) + self._assertOpOutputMatchesExpected( + math_ops.tan, + np.array([1, 2, 3, 4], dtype=dtype), + expected=np.array([1.55740772, -2.18503986, -0.14254654, 1.15782128], + dtype=dtype)) + self._assertOpOutputMatchesExpected( math_ops.tanh, np.array( @@ -209,6 +278,16 @@ class UnaryOpsTest(XLATestCase): [-3.4401896, -2.4401896, -1.4401897, -0.44018969]], dtype=dtype)) + self._assertOpOutputMatchesExpected( + nn_ops.elu, + np.array([[-1, 0, 1]], dtype=dtype), + expected=np.array([[-0.63212056, 0, 1]], dtype=dtype)) + + self._assertOpOutputMatchesExpected( + nn_ops.selu, + np.array([[-1, 0, 1]], dtype=dtype), + expected=np.array([[-1.11133074, 0., 1.05070099]], dtype=dtype)) + self._assertOpOutputMatchesExpected( nn_ops.relu, np.array([[-1, 1]], dtype=dtype), @@ -235,6 +314,19 @@ class UnaryOpsTest(XLATestCase): np.array([[-2, 0, 8]], dtype=dtype), expected=np.array([[0.126928, 0.6931472, 8.0003354]], dtype=dtype)) + self._assertOpOutputMatchesExpected( + nn_ops.softsign, + np.array([[-2, -1, 0, 1, 2]], dtype=dtype), + expected=np.array([[-0.66666669, -0.5, 0, 0.5, 0.66666669]], + dtype=dtype)) + + self._assertOpOutputMatchesExpected( + math_ops.is_finite, + np.array( + [[42, float("inf"), -123], [float("nan"), 0, -0.0]], dtype=dtype), + expected=np.array( + [[True, False, True], [False, True, True]], dtype=np.bool)) + def testNumericOps(self): for dtype in self.numeric_types: self._assertOpOutputMatchesExpected( @@ -262,6 +354,23 @@ class UnaryOpsTest(XLATestCase): np.array([[4, 3], [2, 1]], dtype=dtype), expected=np.array([[1, 1], [1, 1]], dtype=dtype)) + # TODO(phawkins): these tests fail unless fastmath optimizations + # are disabled. Use more robust IsInf/IsNaN detection and enable these + # tests. + @unittest.skip("test case fails in fast-math mode") + def testIsInfAndIsNan(self): + for dtype in self.float_types: + self._assertOpOutputMatchesExpected( + math_ops.is_inf, + np.array([[np.NINF, -2, -1, 0, 0.5, 1, 2, np.inf, np.nan]], + dtype=dtype), + expected=np.array([[1, 0, 0, 0, 0, 0, 0, 1, 0]], dtype=np.bool)) + self._assertOpOutputMatchesExpected( + math_ops.is_nan, + np.array([[np.NINF, -2, -1, 0, 0.5, 1, 2, np.inf, np.nan]], + dtype=dtype), + expected=np.array([[0, 0, 0, 0, 0, 0, 0, 0, 1]], dtype=np.bool)) + def testLogicalOps(self): self._assertOpOutputMatchesExpected( math_ops.logical_not, diff --git a/tensorflow/compiler/tests/variable_ops_test.py b/tensorflow/compiler/tests/variable_ops_test.py index dcb9e2db2f8ca7ef6e89cb9c6493d15dcaacd46e..a6b59fc731e7556cbfa6e0c2c4f889b58568e622 100644 --- a/tensorflow/compiler/tests/variable_ops_test.py +++ b/tensorflow/compiler/tests/variable_ops_test.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""Tests for XLA JIT compiler.""" +"""Tests for reading and writing variables.""" from __future__ import absolute_import from __future__ import division @@ -21,11 +21,14 @@ from __future__ import print_function import numpy as np from tensorflow.compiler.tests.xla_test import XLATestCase +from tensorflow.python.framework import constant_op from tensorflow.python.framework import dtypes +from tensorflow.python.framework import errors from tensorflow.python.framework import ops from tensorflow.python.ops import array_ops from tensorflow.python.ops import init_ops from tensorflow.python.ops import math_ops +from tensorflow.python.ops import resource_variable_ops from tensorflow.python.ops import state_ops from tensorflow.python.ops import variable_scope from tensorflow.python.ops import variables @@ -36,6 +39,69 @@ from tensorflow.python.training.gradient_descent import GradientDescentOptimizer class VariableOpsTest(XLATestCase): """Test cases for resource variable operators.""" + def testOneWriteOneOutput(self): + # Regression test for a bug where computations with one non-constant + # output and one variable update were mishandled. + for dtype in self.numeric_types: + init = np.array([[1, 2], [3, 4]], dtype=dtype) + with self.test_session() as sess, self.test_scope(): + v = resource_variable_ops.ResourceVariable(init) + sess.run(variables.variables_initializer([v])) + p = array_ops.placeholder(dtype) + x = v.assign_add(p) + with ops.control_dependencies([x]): + y = v.read_value() + self.assertAllClose(np.array([[2, 3], [4, 5]], dtype=dtype), + sess.run(y, {p: 1})) + + def testSparseRead0DIndices(self): + for dtype in self.numeric_types: + init = np.array([[0, 1, 2, 3], [4, 5, 6, 7], [8, 9, 10, 11]], dtype=dtype) + with self.test_session() as sess, self.test_scope(): + v = resource_variable_ops.ResourceVariable(init) + sess.run(variables.variables_initializer([v])) + x = v.sparse_read(2) + self.assertAllClose(np.array([8, 9, 10, 11], dtype=dtype), sess.run(x)) + + def testSparseRead1DIndices(self): + for dtype in self.numeric_types: + init = np.array([[0, 1, 2, 3], [4, 5, 6, 7], [8, 9, 10, 11]], dtype=dtype) + with self.test_session() as sess, self.test_scope(): + v = resource_variable_ops.ResourceVariable(init) + sess.run(variables.variables_initializer([v])) + x = v.sparse_read([2, 1]) + self.assertAllClose( + np.array([[8, 9, 10, 11], [4, 5, 6, 7]], dtype=dtype), sess.run(x)) + + def testSparseRead2DIndices(self): + for dtype in self.numeric_types: + init = np.array([[0, 1, 2, 3], [4, 5, 6, 7], [8, 9, 10, 11]], dtype=dtype) + with self.test_session() as sess, self.test_scope(): + v = resource_variable_ops.ResourceVariable(init) + sess.run(variables.variables_initializer([v])) + x = v.sparse_read([[2, 1], [0, 2]]) + self.assertAllClose( + np.array( + [[[8, 9, 10, 11], [4, 5, 6, 7]], [[0, 1, 2, 3], [8, 9, 10, + 11]]], + dtype=dtype), sess.run(x)) + + def testSparseRead2DIndices3DTensor(self): + for dtype in self.numeric_types: + init = np.array( + [[[0, 1, 2], [3, 4, 5]], [[10, 11, 12], [13, 14, 15]], + [[20, 21, 22], [23, 24, 25]], [[30, 31, 32], [33, 34, 35]]], + dtype=dtype) + with self.test_session() as sess, self.test_scope(): + v = resource_variable_ops.ResourceVariable(init) + sess.run(variables.variables_initializer([v])) + x = v.sparse_read([[2, 1], [3, 0]]) + self.assertAllClose( + np.array( + [[[[20, 21, 22], [23, 24, 25]], [[10, 11, 12], [13, 14, 15]]], + [[[30, 31, 32], [33, 34, 35]], [[0, 1, 2], [3, 4, 5]]]], + dtype=dtype), sess.run(x)) + def testReadWrite(self): """Tests initialization, reading, and writing a resource variable.""" with self.test_session() as session: @@ -98,5 +164,68 @@ class VariableOpsTest(XLATestCase): self.assertAllClose(np.array([1.9, 2.9], dtype=np.float32), vb, rtol=1e-4) +class StridedSliceAssignChecker(object): + """Compares the results of a slice assignment using Tensorflow and numpy.""" + + def __init__(self, test, x, dtype): + self.dtype = dtype + self.test = test + self.x_np = np.array(x).astype(dtype) + + def __setitem__(self, index, value): + value = np.array(value).astype(self.dtype) + + with self.test.test_session() as sess, self.test.test_scope(): + x = constant_op.constant(self.x_np, dtype=self.dtype) + var = resource_variable_ops.ResourceVariable(x) + sess.run(variables.variables_initializer([var])) + val = sess.run(var[index].assign(value)) + # val_copy is used to check that tf.assign works equivalently to the + # assign method above. + val_copy = sess.run(state_ops.assign(var[index], value)) + valnp = np.copy(self.x_np) + valnp[index] = np.array(value) + self.test.assertAllEqual(val, valnp) + self.test.assertAllEqual(val_copy, valnp) + + +class SliceAssignTest(XLATestCase): + + def testSliceAssign(self): + for dtype in self.numeric_types: + checker = StridedSliceAssignChecker(self, [[1, 2, 3], [4, 5, 6]], + dtype=dtype) + # No-op assignment + checker[:] = [[10, 20, 30], [40, 50, 60]] + # Checks trivial (1,1) shape tensor + checker[1:2, 1:2] = [[66]] + # shrink shape changes + checker[1:2, 1] = [66] + checker[1, 1:2] = [66] + checker[1, 1] = 66 + # newaxis shape changes + checker[:, None, :] = [[[10, 20, 30]], [[40, 50, 50]]] + # shrink and newaxis + checker[None, None, 0, 0:1] = [[[99]]] + # Non unit strides + checker[::1, 1::-1] = [[3, 33], [4, 44]] + # degenerate interval + checker[8:10, 0] = [] + checker[8:10, 8:10] = [[]] + + # Assign vector to scalar (rank-0) using newaxis + checker2 = StridedSliceAssignChecker(self, 222, dtype=dtype) + checker2[()] = 6 # no indices + checker2[...] = 6 # ellipsis + checker2[None] = [6] # new axis + + def testUninitialized(self): + with self.assertRaisesRegexp(errors.InvalidArgumentError, + "uninitialized variable"): + with self.test_session() as sess, self.test_scope(): + v = resource_variable_ops.ResourceVariable([1, 2]) + sess.run(v[:].assign([1, 2])) + + if __name__ == "__main__": googletest.main() diff --git a/tensorflow/compiler/tests/xla_device_test.py b/tensorflow/compiler/tests/xla_device_test.py index 1388a892ba5a1d07c05eedf277085099923ae901..f5c228f8305d740b994dadc34c93b4e0ae32d785 100644 --- a/tensorflow/compiler/tests/xla_device_test.py +++ b/tensorflow/compiler/tests/xla_device_test.py @@ -18,15 +18,10 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function -import numpy as np - from tensorflow.python.client import session as session_lib -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.platform import test @@ -48,34 +43,6 @@ class XlaDeviceTest(test.TestCase): result = sess.run(w, {x: [1.5, 0.5]}) self.assertAllClose(result, [12., 2.], rtol=1e-3) - def testLoops(self): - """Tests that loops work on XLA devices.""" - - with session_lib.Session() as session: - x = array_ops.placeholder(dtypes.float32) - with ops.device("device:XLA_CPU:0"): - c = lambda i, _: math_ops.less(i, 5) - b = lambda i, x: (i + 1, x * 2.0 + 1.0) - _, y = control_flow_ops.while_loop(c, b, (constant_op.constant(0), x)) - - result = session.run(y, {x: np.float32(2)}) - self.assertAllClose(result, np.float32(95), rtol=1e-3) - - def testCond(self): - """Tests that tf.cond works on XLA devices.""" - - with session_lib.Session() as session: - x = array_ops.placeholder(dtypes.float32) - y = array_ops.placeholder(dtypes.float32) - c = array_ops.placeholder(dtypes.bool) - with ops.device("device:XLA_CPU:0"): - z = x + 1.0 - w = control_flow_ops.cond(c, lambda: z, lambda: y) - t = math_ops.add(z, w) - - result = session.run(t, {x: np.float32(2), y: np.float32(4), c: True}) - self.assertAllClose(result, np.float32(6), rtol=1e-3) - if __name__ == "__main__": test.main() diff --git a/tensorflow/compiler/tests/xla_test.py b/tensorflow/compiler/tests/xla_test.py index f7fe186cf8b35974e72fed6ab6a45a191b898763..79549644ea03575f84d08ba2f114ae60df2a14da 100644 --- a/tensorflow/compiler/tests/xla_test.py +++ b/tensorflow/compiler/tests/xla_test.py @@ -54,16 +54,20 @@ class XLATestCase(test.TestCase): self.device = FLAGS.test_device self.has_custom_call = (self.device == 'XLA_CPU') self.all_tf_types = [ - dtypes.DType(types_pb2.DataType.Value(name)) + dtypes.as_dtype(types_pb2.DataType.Value(name)) for name in FLAGS.types.split(',') ] - self.all_types = [dtype.as_numpy_dtype for dtype in self.all_tf_types] - self.int_types = [ - dtype.as_numpy_dtype for dtype in self.all_tf_types if dtype.is_integer + self.int_tf_types = [ + dtype for dtype in self.all_tf_types if dtype.is_integer ] - self.float_types = [ - dtype.as_numpy_dtype for dtype in self.all_tf_types if dtype.is_floating + self.float_tf_types = [ + dtype for dtype in self.all_tf_types if dtype.is_floating ] + self.numeric_tf_types = self.int_tf_types + self.float_tf_types + + self.all_types = [dtype.as_numpy_dtype for dtype in self.all_tf_types] + self.int_types = [dtype.as_numpy_dtype for dtype in self.int_tf_types] + self.float_types = [dtype.as_numpy_dtype for dtype in self.float_tf_types] self.numeric_types = self.int_types + self.float_types # Parse the manifest file, if any, into a regex identifying tests to diff --git a/tensorflow/compiler/tf2xla/BUILD b/tensorflow/compiler/tf2xla/BUILD index 7a18c1e3750afa276d6721ffea9a4d481cb37136..a8d743c0716ea50ac7882b97c12ac886cf304bee 100644 --- a/tensorflow/compiler/tf2xla/BUILD +++ b/tensorflow/compiler/tf2xla/BUILD @@ -20,6 +20,8 @@ package( default_visibility = [":internal"], ) +load("@local_config_cuda//cuda:build_defs.bzl", "if_cuda_is_configured") + cc_library( name = "xla_compiler", srcs = [ @@ -29,7 +31,10 @@ cc_library( "xla_helpers.cc", "xla_op_kernel.cc", "xla_op_registry.cc", - ], + "xla_cpu_backend.cc", + ] + if_cuda_is_configured([ + "xla_gpu_backend.cc", + ]), hdrs = [ "xla_compilation_device.h", "xla_compiler.h", @@ -42,15 +47,16 @@ cc_library( deps = [ ":common", ":dump_graph", + ":functionalize_control_flow", "//tensorflow/compiler/xla:literal_util", "//tensorflow/compiler/xla:shape_util", "//tensorflow/compiler/xla:statusor", + "//tensorflow/compiler/xla:types", "//tensorflow/compiler/xla:xla_data_proto", "//tensorflow/compiler/xla/client:client_library", "//tensorflow/compiler/xla/client:computation", "//tensorflow/compiler/xla/client:computation_builder", "//tensorflow/compiler/xla/client:local_client", - "//tensorflow/compiler/xla/service:cpu_plugin", "//tensorflow/core:core_cpu", "//tensorflow/core:core_cpu_internal", "//tensorflow/core:framework", @@ -92,6 +98,7 @@ cc_library( cc_test( name = "xla_compiler_test", + size = "small", srcs = ["xla_compiler_test.cc"], deps = [ ":xla_compiler", @@ -100,6 +107,7 @@ cc_test( "//tensorflow/cc:ops", "//tensorflow/compiler/tf2xla/kernels:xla_ops", "//tensorflow/compiler/xla:literal_util", + "//tensorflow/compiler/xla:shape_util", "//tensorflow/compiler/xla/client:client_library", "//tensorflow/compiler/xla/client:local_client", "//tensorflow/compiler/xla/tests:literal_test_util", @@ -113,6 +121,7 @@ cc_test( cc_test( name = "str_util_test", + size = "small", srcs = [ "str_util_test.cc", ], @@ -126,6 +135,7 @@ cc_test( cc_test( name = "literal_util_test", + size = "small", srcs = [ "literal_util_test.cc", ], @@ -149,7 +159,6 @@ cc_library( "//tensorflow/core:framework", "//tensorflow/core:lib", "//tensorflow/core:lib_internal", - "//tensorflow/core:protos_all_cc", ], ) @@ -162,13 +171,10 @@ cc_test( "//tensorflow/cc:cc_ops", "//tensorflow/cc:function_ops", "//tensorflow/cc:ops", - "//tensorflow/core:core_cpu", "//tensorflow/core:core_cpu_internal", - "//tensorflow/core:framework", "//tensorflow/core:ops", "//tensorflow/core:test", "//tensorflow/core:test_main", - "//tensorflow/core:testlib", ], ) @@ -200,6 +206,59 @@ cc_library( ], ) +cc_library( + name = "functionalize_control_flow", + srcs = ["functionalize_control_flow.cc"], + hdrs = ["functionalize_control_flow.h"], + deps = [ + "//tensorflow/compiler/jit:graph_to_functiondef", + "//tensorflow/compiler/tf2xla:dump_graph", + "//tensorflow/compiler/tf2xla/ops:functional_ops", + "//tensorflow/compiler/xla:status_macros", + "//tensorflow/compiler/xla:util", + "//tensorflow/core:core_cpu", + "//tensorflow/core:core_cpu_internal", + "//tensorflow/core:framework", + ], +) + +cc_test( + name = "functionalize_control_flow_test", + srcs = ["functionalize_control_flow_test.cc"], + deps = [ + ":functionalize_control_flow", + ":test_util", + "//tensorflow/cc:cc_ops", + "//tensorflow/cc:cc_ops_internal", + "//tensorflow/cc:function_ops", + "//tensorflow/cc:ops", + "//tensorflow/cc:resource_variable_ops", + "//tensorflow/compiler/tf2xla/cc:functional_ops", + "//tensorflow/compiler/xla:status_macros", + "//tensorflow/core:core_cpu", + "//tensorflow/core:framework", + "//tensorflow/core:framework_internal", + "//tensorflow/core:ops", + "//tensorflow/core:resource_variable_ops_op_lib", + "//tensorflow/core:test", + "//tensorflow/core:test_main", + ], +) + +cc_library( + name = "test_util", + testonly = 1, + srcs = ["test_util.cc"], + hdrs = ["test_util.h"], + deps = [ + "//tensorflow/compiler/xla:status_macros", + "//tensorflow/core:core_cpu", + "//tensorflow/core:framework", + "//tensorflow/core:lib", + "//tensorflow/core:protos_all_cc", + ], +) + # ----------------------------------------------------------------------------- filegroup( diff --git a/tensorflow/compiler/tf2xla/cc/BUILD b/tensorflow/compiler/tf2xla/cc/BUILD new file mode 100644 index 0000000000000000000000000000000000000000..311dddca94c458a60fd00afe5532840e0dbf0437 --- /dev/null +++ b/tensorflow/compiler/tf2xla/cc/BUILD @@ -0,0 +1,67 @@ +package( + default_visibility = ["//tensorflow/compiler/tf2xla:internal"], +) + +licenses(["notice"]) # Apache 2.0 + +load("//tensorflow:tensorflow.bzl", "tf_gen_op_wrapper_cc") + +tf_gen_op_wrapper_cc( + name = "functional_ops_gen", + include_internal_ops = 1, + out_ops_file = "ops/functional_ops", + deps = ["//tensorflow/compiler/tf2xla/ops:functional_ops"], +) + +cc_library( + name = "functional_ops", + srcs = ["ops/functional_ops.cc"], + hdrs = ["ops/functional_ops.h"], + deps = [ + "//tensorflow/cc:const_op", + "//tensorflow/cc:ops", + "//tensorflow/cc:scope", + "//tensorflow/compiler/tf2xla/ops:functional_ops", + "//tensorflow/core:core_cpu", + "//tensorflow/core:framework", + "//tensorflow/core:lib", + "//tensorflow/core:protos_all_cc", + ], +) + +tf_gen_op_wrapper_cc( + name = "sendrecv_ops_gen", + include_internal_ops = 1, + out_ops_file = "ops/sendrecv_ops", + deps = ["//tensorflow/compiler/tf2xla/ops:sendrecv_ops"], +) + +cc_library( + name = "sendrecv_ops", + srcs = ["ops/sendrecv_ops.cc"], + hdrs = ["ops/sendrecv_ops.h"], + deps = [ + "//tensorflow/cc:const_op", + "//tensorflow/cc:ops", + "//tensorflow/cc:scope", + "//tensorflow/compiler/tf2xla/ops:sendrecv_ops", + "//tensorflow/core:core_cpu", + "//tensorflow/core:framework", + "//tensorflow/core:lib", + "//tensorflow/core:protos_all_cc", + ], +) + +# ----------------------------------------------------------------------------- + +filegroup( + name = "all_files", + srcs = glob( + ["**/*"], + exclude = [ + "**/METADATA", + "**/OWNERS", + ], + ), + visibility = ["//tensorflow:__subpackages__"], +) diff --git a/tensorflow/compiler/tf2xla/const_analysis.cc b/tensorflow/compiler/tf2xla/const_analysis.cc index 53aa749a0a90bf3fad06ed4bc57c4327c5d24dcc..e4e1689a2de5780525a1e20c6a22911633845fdf 100644 --- a/tensorflow/compiler/tf2xla/const_analysis.cc +++ b/tensorflow/compiler/tf2xla/const_analysis.cc @@ -35,6 +35,10 @@ Status BackwardsConstAnalysis(const Graph& g, {"Any", "reduction_indices"}, {"ArgMax", "dimension"}, {"AvgPoolGrad", "orig_input_shape"}, + {"AvgPool3DGrad", "orig_input_shape"}, + {"BatchToSpace", "crops"}, + {"BatchToSpaceND", "block_shape"}, + {"BatchToSpaceND", "crops"}, {"BroadcastGradientArgs", "s0"}, {"BroadcastGradientArgs", "s1"}, {"Concat", "concat_dim"}, @@ -45,6 +49,8 @@ Status BackwardsConstAnalysis(const Graph& g, {"Conv2DBackpropInput", "input_sizes"}, {"Conv3DBackpropFilterV2", "filter_sizes"}, {"Conv3DBackpropInputV2", "input_sizes"}, + {"DepthwiseConv2dNativeBackpropFilter", "filter_sizes"}, + {"DepthwiseConv2dNativeBackpropInput", "input_sizes"}, {"DynamicStitch", "indices"}, {"ExpandDims", "dim"}, {"Fill", "dims"}, @@ -57,6 +63,7 @@ Status BackwardsConstAnalysis(const Graph& g, {"Min", "reduction_indices"}, {"OneHot", "depth"}, {"Pad", "paddings"}, + {"MirrorPad", "paddings"}, {"Prod", "reduction_indices"}, {"RandomStandardNormal", "shape"}, {"RandomUniform", "shape"}, @@ -65,13 +72,20 @@ Status BackwardsConstAnalysis(const Graph& g, {"Range", "limit"}, {"Range", "delta"}, {"Reshape", "shape"}, + {"ResourceStridedSliceAssign", "begin"}, + {"ResourceStridedSliceAssign", "end"}, + {"ResourceStridedSliceAssign", "strides"}, {"Reverse", "dims"}, {"ReverseV2", "axis"}, {"Slice", "begin"}, {"Slice", "size"}, + {"SpaceToBatch", "paddings"}, + {"SpaceToBatchND", "block_shape"}, + {"SpaceToBatchND", "paddings"}, {"Split", "split_dim"}, {"SplitV", "split_dim"}, {"SplitV", "size_splits"}, + {"StackV2", "max_size"}, {"StridedSlice", "begin"}, {"StridedSlice", "end"}, {"StridedSlice", "strides"}, @@ -80,6 +94,8 @@ Status BackwardsConstAnalysis(const Graph& g, {"StridedSliceGrad", "end"}, {"StridedSliceGrad", "strides"}, {"Sum", "reduction_indices"}, + {"TensorArrayV3", "size"}, + {"TensorArraySplitV3", "lengths"}, {"Tile", "multiples"}, {"Transpose", "perm"}}; @@ -102,7 +118,7 @@ Status BackwardsConstAnalysis(const Graph& g, if (must_be_const.find(node) != must_be_const.end()) { if (node->type_string() == "_Arg") { int index; - status = GetNodeAttr(node->def(), "index", &index); + status = GetNodeAttr(node->attrs(), "index", &index); if (!status.ok()) return; compile_time_const_args->at(index) = true; return; @@ -118,8 +134,8 @@ Status BackwardsConstAnalysis(const Graph& g, if (range.first == range.second) return; NameRangeMap input_name_ranges; - status = NameRangesForNode(node->def(), node->op_def(), &input_name_ranges, - nullptr); + status = + NameRangesForNode(*node, node->op_def(), &input_name_ranges, nullptr); if (!status.ok()) return; for (auto it = range.first; it != range.second; ++it) { diff --git a/tensorflow/compiler/tf2xla/dump_graph.cc b/tensorflow/compiler/tf2xla/dump_graph.cc index af5753c2600f0b064b6aea4eba556054c38d8d9c..ddd912b87315f7943915153b5bf73531107af54d 100644 --- a/tensorflow/compiler/tf2xla/dump_graph.cc +++ b/tensorflow/compiler/tf2xla/dump_graph.cc @@ -38,7 +38,8 @@ string MakeUniquePath(string name) { // Remove illegal characters from `name`. for (int i = 0; i < name.size(); ++i) { - if (name[i] == '/') { + char ch = name[i]; + if (ch == '/' || ch == '[' || ch == ']' || ch == '*' || ch == '?') { name[i] = '_'; } } diff --git a/tensorflow/compiler/tf2xla/functionalize_control_flow.cc b/tensorflow/compiler/tf2xla/functionalize_control_flow.cc new file mode 100644 index 0000000000000000000000000000000000000000..1c7a2046aa549beb2de58d21f517363d4fe8aea7 --- /dev/null +++ b/tensorflow/compiler/tf2xla/functionalize_control_flow.cc @@ -0,0 +1,583 @@ +/* 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/functionalize_control_flow.h" + +#include +#include +#include +#include + +#include "tensorflow/compiler/jit/graph_to_functiondef.h" +#include "tensorflow/compiler/tf2xla/dump_graph.h" +#include "tensorflow/compiler/xla/ptr_util.h" +#include "tensorflow/compiler/xla/status_macros.h" +#include "tensorflow/core/framework/node_def_builder.h" +#include "tensorflow/core/graph/control_flow.h" + +namespace tensorflow { + +namespace { + +const char* const kArgOp = "_Arg"; +const char* const kRetValOp = "_Retval"; + +// Information about a loop argument. +struct Arg { + // Every loop argument has an Enter node. + Node* enter; + + // Is the loop argument a loop-invariant value? Taken from the `is_constant` + // attribute on the Enter node. + bool is_loop_invariant; + + // If 'is_loop_invariant' is true, the following are all nullptr. Non-constant + // arguments must have all of the following nodes: + Node* merge = nullptr; + Node* switch_node = nullptr; + Node* next_iteration = nullptr; + Node* exit = nullptr; +}; + +// Information about a loop frame. +struct Frame { + string name; + + // Pointer to the parent frame. The root frame has a pointer to itself. + Frame* parent = nullptr; + int num_children = 0; + + // Arguments to this loop. + std::vector args; + + // The loop condition of the loop. There should be exactly one loop condition + // in every loop. + Node* loop_cond = nullptr; + + // Set of nodes that belong to the loop frame. + std::unordered_set nodes; +}; + +// Copies a subgraph from `graph` to `output` by performing a reverse DFS +// starting at nodes in vector `stack`. +// `node_map` is a vector indexed by source node ID to dest nodes. +// Does not traverse into nodes in `node_map`, so by adding nodes to `node_map` +// before the traversal clients can cut the graph. Returns an error if the +// traversal leaves 'frame'; the client must add enough nodes to `node_map` to +// cut the graph and prevent the traversal from escaping. +// +// `squash_src_outputs` contains a bool for each source node ID. If true, then +// the source output on that node will be replaced by zero when copied. This is +// used when replacing a Switch node with an _Arg node. The output we are +// taking from the Switch node was not necessarily the first output, but _Arg +// nodes only have one output. By adding the Switch node to `squash_src_outputs` +// we rewrite the src_output of the corresponding edge to be 0. +Status CopySubgraph(const Graph& graph, const Frame& frame, + std::vector stack, + const std::vector& squash_src_outputs, + std::vector* node_map, Graph* output) { + std::vector visited(graph.num_node_ids(), false); + while (!stack.empty()) { + Node* n = stack.back(); + stack.pop_back(); + + VLOG(3) << "Copying node " << n->name(); + + if (visited[n->id()]) continue; + visited[n->id()] = true; + + for (const Edge* e : n->in_edges()) { + Node* src = e->src(); + if (frame.nodes.find(src) == frame.nodes.end()) { + // We traversed out of the loop frame, without encountering a cut node. + return errors::Internal("Graph traversal of loop frame ", frame.name, + " escaped frame at ", src->name(), + " without encountering an argument node."); + } + if ((*node_map)[src->id()] == nullptr) { + (*node_map)[src->id()] = output->CopyNode(src); + stack.push_back(src); + } + Node* src_copy = (*node_map)[e->src()->id()]; + int src_output = squash_src_outputs[e->src()->id()] ? 0 : e->src_output(); + Node* dst_copy = (*node_map)[e->dst()->id()]; + output->AddEdge(src_copy, src_output, dst_copy, e->dst_input()); + } + } + return Status::OK(); +} + +Status BuildArgNode(Graph* graph, DataType type, int index, Node** arg_node) { + NodeDef arg_def; + NodeDefBuilder builder(strings::StrCat("_Arg", index), kArgOp); + builder.Attr("T", type); + builder.Attr("index", index); + TF_RETURN_IF_ERROR(builder.Finalize(&arg_def)); + Status status; + *arg_node = graph->AddNode(arg_def, &status); + return status; +} + +Status BuildRetvalNode(Graph* graph, DataType type, int index, + Node** retval_node) { + NodeDef ret_def; + ret_def.set_op(kRetValOp); + ret_def.set_name(strings::StrCat("_Retval", index)); + AddNodeAttr("T", type, &ret_def); + AddNodeAttr("index", index, &ret_def); + Status status; + *retval_node = graph->AddNode(ret_def, &status); + return status; +} + +// Builds a graph for the loop condition. +Status BuildLoopCondition(const Graph& graph, Frame* frame, + std::unique_ptr* cond_output) { + VLOG(2) << "Building loop condition for " << frame->name; + *cond_output = xla::MakeUnique(graph.op_registry()); + Graph* output = cond_output->get(); + + // Map from nodes in the original graph to the condition graph. + std::vector node_map(graph.num_node_ids(), nullptr); + std::vector squash_src_outputs(graph.num_node_ids(), false); + + // Build one _Arg node for each Enter node. + for (int i = 0; i < frame->args.size(); ++i) { + const Arg& arg = frame->args[i]; + + Node* arg_node; + TF_RETURN_IF_ERROR( + BuildArgNode(output, arg.enter->input_type(0), i, &arg_node)); + if (arg.is_loop_invariant) { + node_map[arg.enter->id()] = arg_node; + } else { + node_map[arg.merge->id()] = arg_node; + } + } + + // Build a Retval node for the loop condition. The LoopCond nodes are always + // boolean because of the type constraints on the LoopCond op. + TF_RETURN_IF_ERROR( + BuildRetvalNode(output, DT_BOOL, 0, &node_map[frame->loop_cond->id()])); + + // Performs a reverse DFS, copying nodes and edges to the output graph. + // The _Arg and _Retval nodes were added unconditionally above, so we are + // guaranteed to get the correct function signature. + TF_RETURN_IF_ERROR(CopySubgraph(graph, *frame, {frame->loop_cond}, + squash_src_outputs, &node_map, output)); + + return Status::OK(); +} + +// Builds a graph for the loop body. +Status BuildLoopBody(const Graph& graph, Frame* frame, + DataTypeVector* arg_types, + std::unique_ptr* body_output) { + VLOG(2) << "Building loop body for " << frame->name; + *body_output = xla::MakeUnique(graph.op_registry()); + Graph* output = body_output->get(); + + // Map from nodes in the original graph to the condition graph. + std::vector node_map(graph.num_node_ids(), nullptr); + std::vector squash_src_outputs(graph.num_node_ids(), false); + + // Build one _Arg node for each Enter node. + std::vector next_iterations; + next_iterations.reserve(frame->args.size()); + arg_types->reserve(frame->args.size()); + for (int i = 0; i < frame->args.size(); ++i) { + const Arg& arg = frame->args[i]; + + DataType dtype = arg.enter->input_type(0); + arg_types->push_back(dtype); + Node* arg_node; + TF_RETURN_IF_ERROR(BuildArgNode(output, dtype, i, &arg_node)); + + if (dtype == DT_RESOURCE) { + // The convention of the XLA bridge is that resource variable arguments + // are only inputs to the loop body and have no corresponding output. + // TODO(b/37741920): change the convention so that DT_RESOURCE variables + // are both inputs and outputs, and then remove this case. + TF_RET_CHECK(arg.is_loop_invariant); + node_map[arg.enter->id()] = arg_node; + } else { + Node* retval_node; + TF_RETURN_IF_ERROR(BuildRetvalNode(output, dtype, i, &retval_node)); + + if (arg.is_loop_invariant) { + // Argument is loop-invariant. Forward it from the Arg to the Retval. + node_map[arg.enter->id()] = arg_node; + output->AddEdge(arg_node, 0, retval_node, 0); + } else { + // Argument is loop-varying. + node_map[arg.switch_node->id()] = arg_node; + // The Switch node has two outputs, but _Arg only has one. This tells + // the CopySubgraph function to rewrite the output number of edges from + // the _Arg node to be 0 rather than copying the output number from the + // Switch node. + squash_src_outputs[arg.switch_node->id()] = true; + node_map[arg.next_iteration->id()] = retval_node; + next_iterations.push_back(arg.next_iteration); + } + } + } + + // Performs a reverse DFS, copying nodes and edges to the output graph. + // The _Arg and _Retval nodes were added unconditionally above, so we are + // guaranteed to get the correct function signature. + TF_RETURN_IF_ERROR(CopySubgraph(graph, *frame, std::move(next_iterations), + squash_src_outputs, &node_map, output)); + + return Status::OK(); +} + +Status FunctionalizeLoop(Graph* graph, Frame* frame, + FunctionLibraryDefinition* library) { + VLOG(2) << "Frame " << frame->name << " before: " + << dump_graph::DumpGraphToFile("functionalize_before", *graph); + + // Split loop-varying Enter nodes with multiple successors. If the same + // Tensor is fed as input to multiple loop arguments, we may end up with a + // shared Enter node. We clone Enter nodes with multiple successors to + // maintain the invariant of a unique Enter node per argument of the final + // loop. + std::vector args; + for (const Arg& arg : frame->args) { + if (arg.is_loop_invariant) { + args.push_back(arg); + } else { + std::vector edges(arg.enter->out_edges().begin(), + arg.enter->out_edges().end()); + for (int i = 0; i < edges.size(); ++i) { + if (edges[i]->IsControlEdge() && edges[i]->dst()->IsSink()) { + continue; + } + TF_RET_CHECK(!edges[i]->IsControlEdge()) << edges[i]->src()->name(); + Arg new_arg; + new_arg.is_loop_invariant = false; + if (i == 0) { + new_arg.enter = arg.enter; + } else { + new_arg.enter = graph->CopyNode(arg.enter); + frame->nodes.insert(new_arg.enter); + for (Edge const* e : arg.enter->in_edges()) { + graph->AddEdge(e->src(), e->src_output(), new_arg.enter, + e->IsControlEdge() ? Graph::kControlSlot : 0); + } + Node* dst = edges[i]->dst(); + int dst_input = edges[i]->dst_input(); + graph->RemoveEdge(edges[i]); + graph->AddEdge(new_arg.enter, 0, dst, dst_input); + } + args.push_back(new_arg); + } + } + } + frame->args = std::move(args); + + // Order the arguments so that: + // a) resource variables are last, and + // b) sort lexicographically by name (for deterministic output). + std::sort(frame->args.begin(), frame->args.end(), + [](const Arg& a, const Arg& b) { + bool a_is_resource = (a.enter->input_type(0) == DT_RESOURCE); + bool b_is_resource = (b.enter->input_type(0) == DT_RESOURCE); + return std::tie(a_is_resource, a.enter->name()) < + std::tie(b_is_resource, b.enter->name()); + }); + + if (frame->loop_cond == nullptr) { + return errors::InvalidArgument("Loop ", frame->name, + " has no LoopCond node"); + } + + // Find the set of Switch nodes that are successors of the LoopCond. + std::unordered_set switches; + for (const Edge* edge : frame->loop_cond->out_edges()) { + if (!edge->IsControlEdge() && IsSwitch(edge->dst()) && + edge->dst_input() == 1) { + switches.insert(edge->dst()); + } + } + + // For each non-constant argument, looks for the following pattern of nodes: + // Enter ----> Merge --------> Switch --> Exit + // ^ ^ + // | | + // NextIteration LoopCond + // ^ ^ + // | | + // ... ... + for (Arg& arg : frame->args) { + if (!arg.is_loop_invariant) { + // Follow the edge from the Enter to Merge. + const Edge* enter_merge = nullptr; + for (const Edge* e : arg.enter->out_edges()) { + // Ignore control-edges to the sink node. These are allowed by the + // graph invariants, although probably they should have been stripped + // off earlier. + if (e->IsControlEdge() && e->dst()->IsSink()) { + continue; + } + if (enter_merge != nullptr) { + return errors::Internal( + "Enter node for loop-varying argument ", arg.enter->name(), + " has multiple successors: ", enter_merge->dst()->name(), " and ", + e->dst()->name()); + } + enter_merge = e; + } + if (enter_merge == nullptr) { + return errors::Internal("Enter node for loop-varying argument ", + arg.enter->name(), " has zero successors"); + } + arg.merge = enter_merge->dst(); + if (!IsMerge(arg.merge)) { + return errors::InvalidArgument( + "Successor of Enter node for loop-varying argument ", + arg.merge->name(), + " is not a Merge node; got: ", arg.merge->type_string()); + } + + // Find the NextIteration from the merge. There should be two inputs to + // the Merge and the NextIteration should be the other input. + if (arg.merge->input_types().size() != 2) { + return errors::InvalidArgument( + "Unexpected number of inputs to Merge node for loop-varying " + "argument ", + arg.merge->name(), "; expected 2, got ", + arg.merge->input_types().size()); + } + TF_RETURN_IF_ERROR(arg.merge->input_node(1 - enter_merge->dst_input(), + &arg.next_iteration)); + if (!IsNextIteration(arg.next_iteration)) { + return errors::InvalidArgument( + "Expected NextIteration node as input to Merge node; got node ", + arg.next_iteration->name(), " with kind ", + arg.next_iteration->type_string()); + } + + // Find the Switch successor of the Merge. There should be exactly one + // Switch node that is a successor of both the Merge and the LoopCond. + for (const Edge* edge : arg.merge->out_edges()) { + if (edge->dst_input() == 0 && IsSwitch(edge->dst()) && + switches.find(edge->dst()) != switches.end()) { + if (arg.switch_node != nullptr) { + return errors::InvalidArgument("Duplicate Switch successors to ", + arg.merge->name()); + } + arg.switch_node = edge->dst(); + } + } + if (arg.switch_node == nullptr) { + return errors::InvalidArgument("Missing Switch successor to ", + arg.merge->name()); + } + + // Find the Exit successor of the Switch. + for (const Edge* edge : arg.switch_node->out_edges()) { + if (edge->src_output() == 0 && IsExit(edge->dst())) { + if (arg.exit != nullptr) { + return errors::InvalidArgument("Duplicate Exit successors to ", + arg.switch_node->name()); + } + arg.exit = edge->dst(); + } + } + if (arg.exit == nullptr) { + return errors::InvalidArgument("Missing Exit successor to ", + arg.switch_node->name()); + } + } + } + + // Builds the condition and body functions. + std::unique_ptr cond_graph; + TF_RETURN_IF_ERROR(BuildLoopCondition(*graph, frame, &cond_graph)); + DataTypeVector arg_types; + std::unique_ptr body_graph; + TF_RETURN_IF_ERROR(BuildLoopBody(*graph, frame, &arg_types, &body_graph)); + + VLOG(2) << "Frame " << frame->name << " condition: " + << dump_graph::DumpGraphToFile("loop_condition", *cond_graph) + << " body: " << dump_graph::DumpGraphToFile("loop_body", *body_graph); + + static std::atomic sequence_num(0LL); + int64 id = ++sequence_num; + NameAttrList cond_name; + cond_name.set_name(strings::StrCat("_functionalize_cond_", id)); + NameAttrList body_name; + body_name.set_name(strings::StrCat("_functionalize_body_", id)); + FunctionDef cond_fdef; + TF_RETURN_IF_ERROR( + GraphToFunctionDef(*cond_graph, cond_name.name(), &cond_fdef)); + FunctionDef body_fdef; + TF_RETURN_IF_ERROR( + GraphToFunctionDef(*body_graph, body_name.name(), &body_fdef)); + + TF_RETURN_IF_ERROR(library->AddFunctionDef(cond_fdef)); + TF_RETURN_IF_ERROR(library->AddFunctionDef(body_fdef)); + + // Builds a While operator. + NodeDef while_def; + NodeDefBuilder builder(frame->loop_cond->name(), "XlaWhile"); + builder.Attr("T", arg_types); + builder.Attr("cond", cond_name); + builder.Attr("body", body_name); + std::vector inputs; + for (int i = 0; i < frame->args.size(); ++i) { + const Arg& arg = frame->args[i]; + const Edge* in_edge; + TF_RETURN_IF_ERROR(arg.enter->input_edge(0, &in_edge)); + if (in_edge->IsControlEdge()) { + builder.ControlInput(in_edge->src()->name()); + } else { + inputs.push_back(NodeDefBuilder::NodeOut( + in_edge->src()->name(), in_edge->src_output(), arg_types[i])); + } + } + builder.Input(inputs); + TF_RETURN_IF_ERROR(builder.Finalize(&while_def)); + + Status status; + Node* while_node = graph->AddNode(while_def, &status); + if (!status.ok()) { + return status; + } + + // Copies edges to the Enter nodes and from the Exit nodes onto the While. + for (int i = 0; i < frame->args.size(); ++i) { + const Arg& arg = frame->args[i]; + const Edge* in_edge; + TF_RETURN_IF_ERROR(arg.enter->input_edge(0, &in_edge)); + if (in_edge->IsControlEdge()) { + graph->AddControlEdge(in_edge->src(), while_node); + } else { + graph->AddEdge(in_edge->src(), in_edge->src_output(), while_node, i); + } + + if (!arg.is_loop_invariant) { + std::vector edges(arg.exit->out_edges().begin(), + arg.exit->out_edges().end()); + for (const Edge* edge : edges) { + Node* dst = edge->dst(); + int dst_input = edge->dst_input(); + graph->RemoveEdge(edge); + + int src_output = + dst_input == Graph::kControlSlot ? Graph::kControlSlot : i; + graph->AddEdge(while_node, src_output, dst, dst_input); + } + } + } + + // Remove the old nodes from the graph, and add the while node to the parent + // frame. + for (Node* node : frame->nodes) { + graph->RemoveNode(node); + } + frame->parent->nodes.insert(while_node); + + VLOG(2) << "Frame " << frame->name << " after: " + << dump_graph::DumpGraphToFile("functionalize_after", *graph); + + return Status::OK(); +} + +} // namespace + +// Transformation that converts Tensorflow's graph control flow constructs into +// functional equivalents. +Status FunctionalizeControlFlow(Graph* graph, + FunctionLibraryDefinition* library) { + VLOG(2) << "FunctionalizeControlFlow: " + << dump_graph::DumpGraphToFile("functionalize_initial", *graph); + // Note: BuildControlFlowInfo() requires that the graph's source node is + // connected to all source nodes in the graph. Many graphs violate this + // invariant. + std::vector cf_info; + TF_RETURN_IF_ERROR(BuildControlFlowInfo(graph, &cf_info)); + + // Builds Frames, indexed by name. + std::unordered_map frames; + for (Node* node : graph->op_nodes()) { + const ControlFlowInfo& cf = cf_info[node->id()]; + + VLOG(2) << "node: " << node->name() << " frame_name: " << cf.frame_name + << " frame: " << (cf.frame ? cf.frame->name() : "---") + << " parent_frame: " + << (cf.parent_frame ? cf.parent_frame->name() : "---"); + TF_RET_CHECK(cf.frame != nullptr && cf.parent_frame != nullptr); + + Frame& frame = frames[cf.frame_name]; + Frame* parent = &frames[cf_info[cf.parent_frame->id()].frame_name]; + if (frame.parent == nullptr) { + frame.parent = parent; + frame.name = cf.frame_name; + ++parent->num_children; + } else if (frame.parent != parent) { + return errors::InvalidArgument("Mismatched parent frames for ", + cf.frame->id(), ": ", parent->name, " vs ", + frame.parent->name); + } + + if (IsEnter(node)) { + Arg arg; + arg.enter = node; + TF_RETURN_IF_ERROR(GetNodeAttr(arg.enter->attrs(), "is_constant", + &arg.is_loop_invariant)); + frame.args.push_back(arg); + } else if (IsLoopCond(node)) { + if (frame.loop_cond) { + return errors::InvalidArgument( + "Loop ", cf.frame_name, + " has more than one LoopCond node: ", node->name(), " and ", + frame.loop_cond->name()); + } + frame.loop_cond = node; + } + frame.nodes.insert(node); + } + + // Adds frames with no children (i.e., the innermost frames) to a worklist. + std::deque worklist; + for (auto& frame : frames) { + if (frame.second.num_children == 0) { + worklist.push_back(&frame.second); + } + } + + // Eliminate loops from innermost to outermost. + while (!worklist.empty()) { + Frame* frame = worklist.front(); + worklist.pop_front(); + if (frame->parent == frame) { + // Skip the root frame. + continue; + } + + TF_RETURN_IF_ERROR(FunctionalizeLoop(graph, frame, library)); + + // If the parent has no remaining children, add it to the worklist. + --frame->parent->num_children; + if (frame->parent->num_children == 0) { + worklist.push_back(frame->parent); + } + } + + return Status::OK(); +} + +} // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/functionalize_control_flow.h b/tensorflow/compiler/tf2xla/functionalize_control_flow.h new file mode 100644 index 0000000000000000000000000000000000000000..1535dc80b0ccdba38c57b534ed7473fc8632e33f --- /dev/null +++ b/tensorflow/compiler/tf2xla/functionalize_control_flow.h @@ -0,0 +1,32 @@ +/* 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_FUNCTIONALIZE_CONTROL_FLOW_H_ +#define TENSORFLOW_COMPILER_TF2XLA_FUNCTIONALIZE_CONTROL_FLOW_H_ + +#include "tensorflow/core/framework/function.h" +#include "tensorflow/core/graph/graph.h" + +namespace tensorflow { + +// Transformation that converts tf.while_loop() loops into functional While +// operators, suitable for XLA compilation. +// TODO(b/36470387): add support for conditionals. +Status FunctionalizeControlFlow(Graph* graph, + FunctionLibraryDefinition* library); + +} // namespace tensorflow + +#endif // TENSORFLOW_COMPILER_TF2XLA_FUNCTIONALIZE_CONTROL_FLOW_H_ diff --git a/tensorflow/compiler/tf2xla/functionalize_control_flow_test.cc b/tensorflow/compiler/tf2xla/functionalize_control_flow_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..914c8999a6f13f5f2dc4e3cecc38c91afd432131 --- /dev/null +++ b/tensorflow/compiler/tf2xla/functionalize_control_flow_test.cc @@ -0,0 +1,658 @@ +/* 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/functionalize_control_flow.h" + +#include "tensorflow/cc/framework/ops.h" +#include "tensorflow/cc/ops/control_flow_ops_internal.h" +#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/cc/ops/functional_ops.h" +#include "tensorflow/compiler/tf2xla/test_util.h" +#include "tensorflow/compiler/xla/status_macros.h" +#include "tensorflow/core/framework/function.h" +#include "tensorflow/core/framework/node_def_util.h" +#include "tensorflow/core/framework/op.h" +#include "tensorflow/core/graph/graph_constructor.h" +#include "tensorflow/core/graph/graph_def_builder.h" +#include "tensorflow/core/lib/core/status_test_util.h" +#include "tensorflow/core/platform/test.h" +#include "tensorflow/core/util/equal_graph_def.h" + +namespace tensorflow { +namespace { + +// Returns the names of the "cond" and "body" functions for the While node +// in a graph. +Status FindWhileCondAndBody(const GraphDef& graph, NameAttrList* cond, + NameAttrList* body) { + for (const NodeDef& node : graph.node()) { + if (node.op() == "XlaWhile") { + const NameAttrList* result; + TF_RETURN_IF_ERROR(GetNodeAttr(node, "cond", &result)); + *cond = *result; + TF_RETURN_IF_ERROR(GetNodeAttr(node, "body", &result)); + *body = *result; + return Status::OK(); + } + } + return errors::NotFound("No XlaWhile node found in graph"); +} + +// Graph: +// x = array_ops.placeholder(dtypes.int32) +// y = control_flow_ops.while_loop(lambda i: i < 10, lambda i: i + 1, [x]) +TEST(FunctionalizeControlFlow, OneLoopVar) { + Graph graph(OpRegistry::Global()); + { + Scope scope = Scope::NewRootScope().ExitOnError(); + + auto dummy = ops::Placeholder(scope.WithOpName("Dummy"), DT_INT32); + + auto source = ops::Placeholder(scope.WithOpName("source"), DT_INT32); + auto enter = + ops::internal::Enter(scope.WithOpName("while/Enter"), source, "aloop"); + // Add an unused Enter node. These should be ignored. + auto enter2 = + ops::internal::Enter(scope.WithOpName("while/Enter2"), source, "aloop"); + auto merge = ops::Merge(scope.WithOpName("while/Merge"), + std::initializer_list{enter, dummy}); + auto ten = ops::Const( + scope.WithOpName("while/Less/y").WithControlDependencies(merge.output), + 10); + auto less = ops::Less(scope.WithOpName("while/Less"), merge.output, ten); + auto loop_cond = ops::LoopCond(scope.WithOpName("while/LoopCond"), less); + auto switch_ = + ops::Switch(scope.WithOpName("while/Switch"), merge.output, loop_cond); + auto exit = ops::internal::Exit(scope.WithOpName("while/Exit"), + switch_.output_false); + auto identity = + ops::Identity(scope.WithOpName("while/Identity"), switch_.output_true); + auto one = ops::Const( + scope.WithOpName("while/add/y").WithControlDependencies(identity), 1); + auto add = ops::Add(scope.WithOpName("while/add"), identity, one); + auto next_iteration = + ops::NextIteration(scope.WithOpName("while/NextIteration"), add); + + auto sink = ops::Identity(scope.WithOpName("sink"), exit); + + // Remove the dummy node and add the loop backedge. + scope.graph()->RemoveNode(dummy.node()); + scope.graph()->AddEdge(next_iteration.node(), 0, merge.output.node(), 1); + + TF_EXPECT_OK(scope.ToGraph(&graph)); + } + + // Regression test: control edges from an Enter node to the graph sink should + // be ignored. + for (Node* n : graph.nodes()) { + if (n->name() == "while/Enter") { + graph.AddControlEdge(n, graph.sink_node()); + } + } + + FunctionLibraryDefinition library(OpRegistry::Global(), {}); + TF_ASSERT_OK(FunctionalizeControlFlow(&graph, &library)); + + 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::XlaWhile(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)); + + EXPECT_EQ(DataTypeVector{DT_INT32}, result.arg_types); + EXPECT_EQ(DataTypeVector{DT_INT32}, result.ret_types); + TF_EXPECT_GRAPH_EQ(expected, result.gdef); + } +} + +// Graph: +// x = array_ops.placeholder(dtypes.int32) +// y = array_ops.placeholder(dtypes.int32) +// cond = lambda (i, j): i + 3 < 10 +// body = lambda (i, j): (i < 10, j * 2) +// z = control_flow_ops.while_loop(cond, body, [x, y]) +TEST(FunctionalizeControlFlow, TwoLoopVars) { + Graph graph(OpRegistry::Global()); + { + Scope scope = Scope::NewRootScope().ExitOnError(); + + auto dummy = ops::Placeholder(scope.WithOpName("Dummy"), DT_INT32); + + auto x = ops::Placeholder(scope.WithOpName("Placeholder/x"), DT_INT32); + auto y = ops::Placeholder(scope.WithOpName("Placeholder/y"), DT_INT32); + auto enter_x = + ops::internal::Enter(scope.WithOpName("while/Enter/x"), x, "aloop"); + auto enter_y = + ops::internal::Enter(scope.WithOpName("while/Enter/y"), y, "aloop"); + auto merge_x = ops::Merge(scope.WithOpName("while/Merge/x"), + std::initializer_list{enter_x, dummy}); + auto merge_y = ops::Merge(scope.WithOpName("while/Merge/y"), + std::initializer_list{enter_y, dummy}); + + // Loop condition + auto three = ops::Const(scope.WithOpName("while/cond/three") + .WithControlDependencies(merge_x.output), + 3); + auto cond_add = + ops::Add(scope.WithOpName("while/cond/Add"), merge_x.output, three); + auto ten = ops::Const(scope.WithOpName("while/cond/ten") + .WithControlDependencies(merge_x.output), + 10); + auto less = ops::Less(scope.WithOpName("while/cond/Less"), cond_add, ten); + auto loop_cond = ops::LoopCond(scope.WithOpName("while/LoopCond"), less); + + auto switch_x = ops::Switch(scope.WithOpName("while/Switch/x"), + merge_x.output, loop_cond); + auto switch_y = ops::Switch(scope.WithOpName("while/Switch/y"), + merge_y.output, loop_cond); + + auto exit_x = ops::internal::Exit(scope.WithOpName("while/Exit/x"), + switch_x.output_false); + auto exit_y = ops::internal::Exit(scope.WithOpName("while/Exit/y"), + switch_y.output_false); + + auto identity_x = ops::Identity(scope.WithOpName("while/Identity/x"), + switch_x.output_true); + auto identity_y = ops::Identity(scope.WithOpName("while/Identity/y"), + switch_y.output_true); + + 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 next_iteration_x = + ops::NextIteration(scope.WithOpName("while/NextIteration/x"), add); + auto next_iteration_y = + ops::NextIteration(scope.WithOpName("while/NextIteration/y"), mul); + + auto sink_x = ops::Identity(scope.WithOpName("sink_x"), exit_x); + auto sink_y = ops::Identity(scope.WithOpName("sink_y"), exit_y); + + // Remove the dummy node and add the loop backedges. + scope.graph()->RemoveNode(dummy.node()); + scope.graph()->AddEdge(next_iteration_x.node(), 0, merge_x.output.node(), + 1); + scope.graph()->AddEdge(next_iteration_y.node(), 0, merge_y.output.node(), + 1); + + TF_EXPECT_OK(scope.ToGraph(&graph)); + } + + FunctionLibraryDefinition library(OpRegistry::Global(), {}); + TF_ASSERT_OK(FunctionalizeControlFlow(&graph, &library)); + + 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::XlaWhile(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") + .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); + + 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); + } +} + +// Example with nesting, loop-invariant arguments, and resource variables. +// +// accum = resource_variable_ops.ResourceVariable(1) +// x = array_ops.placeholder(2, dtype=dtypes.int32) +// y = 3 + x +// +// def inner_body(j, k): +// add = state_ops.assign_add(accum, k * j + x) +// with ops.control_dependencies([add]): +// return [j + 1, k] +// +// def body(i): +// m = control_flow_ops.while_loop(lambda j, k: j < 5, inner_body, +// [1, y], name="inner") +// with ops.control_dependencies(m): +// return [i + 1] +// +// z = control_flow_ops.while_loop(lambda i: i < 10, body, [0], name="outer") +TEST(FunctionalizeControlFlow, Complex) { + Graph graph(OpRegistry::Global()); + { + Scope scope = Scope::NewRootScope().ExitOnError(); + + auto dummy = ops::Placeholder(scope.WithOpName("Dummy"), DT_INT32); + + 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({})); + + // Outer loop + auto zero = ops::Const(scope.WithOpName("outer/Const"), 0); + auto enter_i = + ops::internal::Enter(scope.WithOpName("outer/Enter_i"), zero, "outer"); + auto merge_i = ops::Merge(scope.WithOpName("outer/Merge_i"), + std::initializer_list{enter_i, dummy}); + auto ten = ops::Const(scope.WithOpName("outer/Less/y") + .WithControlDependencies(merge_i.output), + 10); + auto less_i = + ops::Less(scope.WithOpName("outer/Less_i"), merge_i.output, ten); + auto outer_loop_cond = + ops::LoopCond(scope.WithOpName("outer/LoopCond"), less_i); + auto switch_i = ops::Switch(scope.WithOpName("outer/Switch"), + merge_i.output, outer_loop_cond); + auto exit_i = ops::internal::Exit(scope.WithOpName("outer/Exit"), + switch_i.output_false); + auto identity_i = + ops::Identity(scope.WithOpName("outer/Identity"), switch_i.output_true); + + auto enter_x_outer = + ops::internal::Enter(scope.WithOpName("outer/Enter_x"), x, "outer", + ops::internal::Enter::Attrs().IsConstant(true)); + auto enter_k_outer = + ops::internal::Enter(scope.WithOpName("outer/Enter_k"), y, "outer", + ops::internal::Enter::Attrs().IsConstant(true)); + auto enter_var_outer = + ops::internal::Enter(scope.WithOpName("outer/Enter_var"), var, "outer", + ops::internal::Enter::Attrs().IsConstant(true)); + + // Inner loop + auto one_j = ops::Const( + scope.WithOpName("outer/j").WithControlDependencies(identity_i), 1); + auto enter_j = ops::internal::Enter(scope.WithOpName("outer/inner/Enter_j"), + one_j, "inner"); + auto enter_k = + ops::internal::Enter(scope.WithOpName("outer/inner/Enter_k") + .WithControlDependencies(identity_i), + enter_k_outer, "inner"); + auto enter_x = ops::internal::Enter( + scope.WithOpName("outer/inner/Enter_x"), enter_x_outer, "inner", + ops::internal::Enter::Attrs().IsConstant(true)); + auto enter_var = ops::internal::Enter( + scope.WithOpName("outer/inner/Enter_var"), enter_var_outer, "inner", + ops::internal::Enter::Attrs().IsConstant(true)); + + auto merge_j = ops::Merge(scope.WithOpName("outer/inner/Merge_j"), + std::initializer_list{enter_j, dummy}); + auto merge_k = ops::Merge(scope.WithOpName("outer/inner/Merge_k"), + std::initializer_list{enter_k, dummy}); + + auto five = ops::Const(scope.WithOpName("outer/inner/Five") + .WithControlDependencies(merge_j.output), + 5); + auto less_j = + ops::Less(scope.WithOpName("outer/inner/Less_j"), merge_j.output, five); + auto loop_cond = ops::LoopCond(scope.WithOpName("outer/LoopCond"), less_j); + + auto switch_j = ops::Switch(scope.WithOpName("outer/inner/Switch_j"), + merge_j.output, loop_cond); + auto switch_k = ops::Switch(scope.WithOpName("outer/inner/Switch_k"), + merge_k.output, loop_cond); + auto exit_j = ops::internal::Exit(scope.WithOpName("outer/inner/Exit_j"), + switch_j.output_false); + auto exit_k = ops::internal::Exit(scope.WithOpName("outer/inner/Exit_k"), + switch_k.output_false); + auto identity_j = ops::Identity(scope.WithOpName("outer/inner/Identity_j"), + switch_j.output_true); + auto identity_k = ops::Identity(scope.WithOpName("outer/inner/Identity_k"), + switch_k.output_true); + + // Variable update + 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, enter_x); + auto assign = ops::AssignAddVariableOp( + scope.WithOpName("outer/inner/assign_add"), enter_var, add_jkx); + + auto one = + ops::Const(scope.WithOpName("outer/inner/One") + .WithControlDependencies( + gtl::ArraySlice{assign.operation}), + 1); + auto add_j = + ops::Add(scope.WithOpName("outer/inner/add_j"), identity_j, one); + + auto next_iteration_j = ops::NextIteration( + scope.WithOpName("outer/inner/NextIteration_j"), add_j); + auto next_iteration_k = ops::NextIteration( + scope.WithOpName("outer/inner/NextIteration_k"), identity_k); + + // Body and backedge for outer loop. + auto one_outer = ops::Const( + scope.WithOpName("outer/add/y").WithControlDependencies(identity_i), 1); + auto add_i = + ops::Add(scope.WithOpName("outer/add") + .WithControlDependencies(gtl::ArraySlice{ + exit_j.output.op(), exit_k.output.op()}), + identity_i, one_outer); + auto next_iteration_i = + ops::NextIteration(scope.WithOpName("outer/NextIteration"), add_i); + + auto sink = ops::Identity(scope.WithOpName("sink"), exit_i); + + // Remove the dummy node and add the loop backedge. + scope.graph()->RemoveNode(dummy.node()); + scope.graph()->AddEdge(next_iteration_i.node(), 0, merge_i.output.node(), + 1); + scope.graph()->AddEdge(next_iteration_j.node(), 0, merge_j.output.node(), + 1); + scope.graph()->AddEdge(next_iteration_k.node(), 0, merge_k.output.node(), + 1); + + TF_EXPECT_OK(scope.ToGraph(&graph)); + } + + FunctionLibraryDefinition library(OpRegistry::Global(), {}); + TF_ASSERT_OK(FunctionalizeControlFlow(&graph, &library)); + + 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::XlaWhile(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::XlaWhile(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(gtl::ArraySlice{ + 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; + 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( + gtl::ArraySlice{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); + } +} + +} // namespace +} // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/kernels/BUILD b/tensorflow/compiler/tf2xla/kernels/BUILD index 2ee80a41e820b5ecc92816c84b6de9625f319b19..6e6c5dc17f5364bf6623dd07b57cf4797442bc3b 100644 --- a/tensorflow/compiler/tf2xla/kernels/BUILD +++ b/tensorflow/compiler/tf2xla/kernels/BUILD @@ -1,10 +1,7 @@ licenses(["notice"]) # Apache 2.0 package( - default_visibility = [ - "//tensorflow/compiler/tf2xla:internal", - ], - features = ["no_layering_check"], + default_visibility = ["//tensorflow/compiler/tf2xla:internal"], ) load("//tensorflow:tensorflow.bzl", "tf_kernel_library") @@ -14,24 +11,30 @@ tf_kernel_library( name = "xla_ops", srcs = [ "aggregate_ops.cc", + "arg_op.cc", "batch_matmul_op.cc", + "batch_norm_op.cc", + "batchtospace_op.cc", "bcast_ops.cc", "bias_ops.cc", "binary_ops.cc", "cast_op.cc", "concat_op.cc", + "const_op.cc", "conv_ops.cc", + "cross_op.cc", "cwise_ops.cc", - "declaration_op.cc", - "depthwise_conv_ops.cc", "diag_op.cc", "dynamic_stitch_op.cc", + "elu_op.cc", "fill_op.cc", "function_ops.cc", + "gather_op.cc", "identity_op.cc", "l2loss_op.cc", "lrn_ops.cc", "matmul_op.cc", + "mirror_pad_op.cc", "no_op.cc", "one_hot_op.cc", "pack_op.cc", @@ -44,13 +47,18 @@ tf_kernel_library( "reshape_op.cc", "retval_op.cc", "reverse_op.cc", + "segment_reduction_ops.cc", "select_op.cc", + "sendrecv_ops.cc", "sequence_ops.cc", "shape_op.cc", "slice_op.cc", "softmax_op.cc", + "spacetobatch_op.cc", "split_op.cc", + "stack_ops.cc", "strided_slice_op.cc", + "tensor_array_ops.cc", "tile_ops.cc", "training_ops.cc", "transpose_op.cc", @@ -60,24 +68,29 @@ tf_kernel_library( ], hdrs = [ "cwise_ops.h", + "gather_op.h", + "gather_op_helpers.h", "reduction_ops.h", ], deps = [ + ":while_op", "//tensorflow/compiler/tf2xla:common", "//tensorflow/compiler/tf2xla:xla_compiler", + "//tensorflow/compiler/tf2xla/ops:sendrecv_ops", "//tensorflow/compiler/xla:literal_util", + "//tensorflow/compiler/xla:util", + "//tensorflow/compiler/xla:xla_data_proto", "//tensorflow/compiler/xla/client:client_library", "//tensorflow/compiler/xla/client:computation_builder", "//tensorflow/compiler/xla/client/lib:arithmetic", "//tensorflow/core:framework", "//tensorflow/core:lib", + "//tensorflow/core:protos_all_cc", "//tensorflow/core:tensorflow_opensource", + "//tensorflow/core/kernels:bounds_check", "//tensorflow/core/kernels:concat_lib", - "//tensorflow/core/kernels:conv_2d", "//tensorflow/core/kernels:conv_ops", "//tensorflow/core/kernels:cwise_op", - "//tensorflow/core/kernels:depthwise_conv_op", - "//tensorflow/core/kernels:matmul_op", "//tensorflow/core/kernels:no_op", "//tensorflow/core/kernels:ops_util", "//tensorflow/core/kernels:pooling_ops", @@ -86,6 +99,23 @@ tf_kernel_library( ], ) +tf_kernel_library( + name = "while_op", + srcs = ["while_op.cc"], + hdrs = ["while_op.h"], + deps = [ + "//tensorflow/compiler/tf2xla:common", + "//tensorflow/compiler/tf2xla:xla_compiler", + "//tensorflow/compiler/tf2xla/ops:functional_ops", + "//tensorflow/compiler/xla:literal_util", + "//tensorflow/compiler/xla/client:computation_builder", + "//tensorflow/core:framework", + "//tensorflow/core:lib", + "//tensorflow/core:protos_all_cc", + "//tensorflow/core:tensorflow_opensource", + ], +) + # Kernels that only work on CPU, because they use XLA custom calls. # Only link this when using the CPU backend for XLA. # @@ -94,7 +124,6 @@ tf_kernel_library( tf_kernel_library( name = "xla_cpu_only_ops", srcs = [ - "gather_op.cc", "index_ops.cc", ], deps = [ @@ -111,59 +140,58 @@ tf_kernel_library( "//tensorflow/core:framework", "//tensorflow/core:lib", "//tensorflow/core:tensorflow_opensource", + "//tensorflow/core/kernels:bounds_check", ], ) -tf_kernel_library( +cc_library( name = "gather_op_kernel_float_int32", srcs = ["gather_op_kernel_float_int32.cc"], - # Makes the custom-call function visible to LLVM during JIT. - linkopts = export_dynamic_linkopts, visibility = ["//visibility:public"], deps = [ "//tensorflow/compiler/tf2xla:xla_local_runtime_context", "//tensorflow/core:framework_lite", - "//tensorflow/core/kernels:gather_functor", + "//tensorflow/core/kernels:bounds_check", + "//tensorflow/core/kernels:gather_functor_hdr", "//third_party/eigen3", ], + alwayslink = 1, ) -tf_kernel_library( +cc_library( name = "gather_op_kernel_float_int64", srcs = ["gather_op_kernel_float_int64.cc"], - # Makes the custom-call function visible to LLVM during JIT. - linkopts = export_dynamic_linkopts, visibility = ["//visibility:public"], deps = [ "//tensorflow/compiler/tf2xla:xla_local_runtime_context", "//tensorflow/core:framework_lite", - "//tensorflow/core/kernels:gather_functor", + "//tensorflow/core/kernels:bounds_check", + "//tensorflow/core/kernels:gather_functor_hdr", "//third_party/eigen3", ], + alwayslink = 1, ) -tf_kernel_library( +cc_library( name = "index_ops_kernel_argmax_float_1d", srcs = ["index_ops_kernel_argmax_float_1d.cc"], - # Makes the custom-call function visible to LLVM during JIT. - linkopts = export_dynamic_linkopts, visibility = ["//visibility:public"], deps = [ "//tensorflow/core:framework_lite", "//third_party/eigen3", ], + alwayslink = 1, ) -tf_kernel_library( +cc_library( name = "index_ops_kernel_argmax_float_2d", srcs = ["index_ops_kernel_argmax_float_2d.cc"], - # Makes the custom-call function visible to LLVM during JIT. - linkopts = export_dynamic_linkopts, visibility = ["//visibility:public"], deps = [ "//tensorflow/core:framework_lite", "//third_party/eigen3", ], + alwayslink = 1, ) # ----------------------------------------------------------------------------- diff --git a/tensorflow/compiler/tf2xla/kernels/arg_op.cc b/tensorflow/compiler/tf2xla/kernels/arg_op.cc new file mode 100644 index 0000000000000000000000000000000000000000..1156546512952871fafe93e4b5a42308322671df --- /dev/null +++ b/tensorflow/compiler/tf2xla/kernels/arg_op.cc @@ -0,0 +1,94 @@ +/* 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/type_util.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" + +namespace tensorflow { +namespace { + +// This OpKernel implements the _Arg Op for XLA JIT devices. It +// associates its output with one of the arguments to a +// subcomputation. +class ArgOp : public XlaOpKernel { + public: + explicit ArgOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) { + OP_REQUIRES_OK(ctx, ctx->GetAttr("T", &dtype_)); + OP_REQUIRES_OK(ctx, ctx->GetAttr("index", &index_)); + } + + void Compile(XlaOpKernelContext* ctx) override { + // If 'frame' is non-null, this is a function call inside an outer JIT + // compilation. Use the usual implementation of _Arg. + auto frame = ctx->call_frame(); + if (frame != nullptr) { + Tensor val; + OP_REQUIRES_OK(ctx, frame->GetArg(index_, &val)); + OP_REQUIRES(ctx, val.dtype() == dtype_, + errors::InvalidArgument( + "Type mismatch: actual ", DataTypeString(val.dtype()), + " vs. expect ", DataTypeString(dtype_))); + // Forwards the argument from the frame. + ctx->op_kernel_context()->set_output(0, val); + return; + } + + XlaContext& xc = XlaContext::Get(ctx); + const XlaContext::Argument& arg = xc.args()[index_]; + if (arg.is_resource) { + XlaResource::Kind kind; + switch (arg.kind) { + case XlaCompiler::Argument::kVariable: + kind = XlaResource::kVariable; + break; + case XlaCompiler::Argument::kTensorArray: + kind = XlaResource::kTensorArray; + break; + case XlaCompiler::Argument::kStack: + kind = XlaResource::kStack; + break; + default: + CHECK(false); + } + + // TODO(phawkins): this code assumes that variables do not alias. + XlaResource* resource; + OP_REQUIRES_OK(ctx, + xc.CreateResource(kind, index_, arg.name, arg.value.type, + arg.value.handle, &resource)); + resource->tensor_array_size = arg.tensor_array_size; + ctx->SetResourceOutput(0, resource); + } else if (arg.value.is_constant) { + ctx->SetConstantOutput(0, arg.value.constant_value); + } else { + ctx->SetOutput(0, arg.value.handle); + } + } + + private: + int index_; + DataType dtype_; + + TF_DISALLOW_COPY_AND_ASSIGN(ArgOp); +}; + +REGISTER_XLA_OP(Name("_Arg").AllowResourceTypes(), ArgOp); + +} // namespace +} // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/kernels/batch_matmul_op.cc b/tensorflow/compiler/tf2xla/kernels/batch_matmul_op.cc index f752fb3ae2b9264f1b86ed2028640bb57dcf22bf..16b778bca439b9236498945f132e8095baeb71c1 100644 --- a/tensorflow/compiler/tf2xla/kernels/batch_matmul_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/batch_matmul_op.cc @@ -94,12 +94,14 @@ class BatchMatMulOp : public XlaOpKernel { // Slice off individual matrices and reshape to 2D tensors. auto x_slice = builder->Slice( x_flat, {i, 0, 0}, - {i + 1, x_shape.dim_size(ndims - 2), x_shape.dim_size(ndims - 1)}); + {i + 1, x_shape.dim_size(ndims - 2), x_shape.dim_size(ndims - 1)}, + {1, 1, 1}); x_slice = builder->Reshape( x_slice, {x_shape.dim_size(ndims - 2), x_shape.dim_size(ndims - 1)}); auto y_slice = builder->Slice( y_flat, {i, 0, 0}, - {i + 1, y_shape.dim_size(ndims - 2), y_shape.dim_size(ndims - 1)}); + {i + 1, y_shape.dim_size(ndims - 2), y_shape.dim_size(ndims - 1)}, + {1, 1, 1}); y_slice = builder->Reshape( y_slice, {y_shape.dim_size(ndims - 2), y_shape.dim_size(ndims - 1)}); diff --git a/tensorflow/compiler/tf2xla/kernels/batch_norm_op.cc b/tensorflow/compiler/tf2xla/kernels/batch_norm_op.cc new file mode 100644 index 0000000000000000000000000000000000000000..9d2703bf9523ce603a2d1e5042c53521f35d9c2c --- /dev/null +++ b/tensorflow/compiler/tf2xla/kernels/batch_norm_op.cc @@ -0,0 +1,123 @@ +/* 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. +==============================================================================*/ + +// XLA implementation of BatchNorm operations. +#include "tensorflow/compiler/tf2xla/literal_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/core/util/tensor_format.h" + +namespace tensorflow { +namespace { + +class FusedBatchNormOp : public XlaOpKernel { + public: + explicit FusedBatchNormOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) { + string data_format; + OP_REQUIRES_OK(ctx, ctx->GetAttr("epsilon", &epsilon_)); + OP_REQUIRES_OK(ctx, ctx->GetAttr("is_training", &is_training_)); + OP_REQUIRES_OK(ctx, ctx->GetAttr("data_format", &data_format)); + TensorFormat tensor_format; + if (ctx->GetAttr("data_format", &data_format).ok()) { + OP_REQUIRES(ctx, FormatFromString(data_format, &tensor_format), + errors::InvalidArgument("Invalid data format")); + OP_REQUIRES( + ctx, (tensor_format == FORMAT_NHWC || tensor_format == FORMAT_NCHW), + errors::InvalidArgument("Not supported format")); + feature_index_ = GetTensorFeatureDimIndex(/*num_dims=*/4, tensor_format); + } + } + + void Compile(XlaOpKernelContext* ctx) override { + if (is_training_) { + xla::ComputationDataHandle output = ctx->builder()->BatchNormTraining( + ctx->Input(0), ctx->Input(1), ctx->Input(2), epsilon_, + feature_index_); + + // In training mode, outputs the normalized value as well as the + // calculated mean and variance. + for (int i = 0; i < 3; i++) { + ctx->SetOutput(i, ctx->builder()->GetTupleElement(output, i)); + } + // Output 3 and 4 for "FusedBatchNorm" are currently marked as "reserved + // space 1 & 2". They are used to pass the per-batch mean and + // variance to the gradient. Here we maintain the same behavior by setting + // them to the mean and variance calculated by BatchNormTraining. + ctx->SetOutput(3, ctx->builder()->GetTupleElement(output, 1)); + ctx->SetOutput(4, ctx->builder()->GetTupleElement(output, 2)); + } else { + xla::ComputationDataHandle output = ctx->builder()->BatchNormInference( + ctx->Input(0), ctx->Input(1), ctx->Input(2), ctx->Input(3), + ctx->Input(4), epsilon_, feature_index_); + ctx->SetOutput(0, output); + // Directly send input to output as mean and variance in inference mode. + ctx->SetOutput(1, ctx->Input(3)); + ctx->SetOutput(2, ctx->Input(4)); + ctx->SetOutput(3, ctx->Input(3)); + ctx->SetOutput(4, ctx->Input(4)); + } + } + + private: + float epsilon_; + int64 feature_index_; + bool is_training_; +}; + +REGISTER_XLA_OP(Name("FusedBatchNorm"), FusedBatchNormOp); + +class FusedBatchNormGradOp : public XlaOpKernel { + public: + explicit FusedBatchNormGradOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) { + string data_format; + OP_REQUIRES_OK(ctx, ctx->GetAttr("epsilon", &epsilon_)); + OP_REQUIRES_OK(ctx, ctx->GetAttr("data_format", &data_format)); + TensorFormat tensor_format; + if (ctx->GetAttr("data_format", &data_format).ok()) { + OP_REQUIRES(ctx, FormatFromString(data_format, &tensor_format), + errors::InvalidArgument("Invalid data format")); + OP_REQUIRES( + ctx, (tensor_format == FORMAT_NHWC || tensor_format == FORMAT_NCHW), + errors::InvalidArgument("Not supported format")); + feature_index_ = GetTensorFeatureDimIndex(4, tensor_format); + } + } + + void Compile(XlaOpKernelContext* ctx) override { + auto grad_output = ctx->Input(0); + auto activation = ctx->Input(1); + auto scale = ctx->Input(2); + auto mean = ctx->Input(3); + auto var = ctx->Input(4); + xla::ComputationDataHandle output = ctx->builder()->BatchNormGrad( + activation, scale, mean, var, grad_output, epsilon_, feature_index_); + + for (int i = 0; i < 3; i++) { + ctx->SetOutput(i, ctx->builder()->GetTupleElement(output, i)); + } + ctx->SetOutput(3, ctx->builder()->GetTupleElement(output, 1)); + ctx->SetOutput(4, ctx->builder()->GetTupleElement(output, 2)); + } + + private: + float epsilon_; + int64 feature_index_; +}; + +REGISTER_XLA_OP(Name("FusedBatchNormGrad"), FusedBatchNormGradOp); + +} // namespace +} // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/kernels/batchtospace_op.cc b/tensorflow/compiler/tf2xla/kernels/batchtospace_op.cc new file mode 100644 index 0000000000000000000000000000000000000000..21d3e64872e19109852297838043975cea6d7921 --- /dev/null +++ b/tensorflow/compiler/tf2xla/kernels/batchtospace_op.cc @@ -0,0 +1,187 @@ +/* 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_helpers.h" +#include "tensorflow/compiler/tf2xla/xla_op_kernel.h" +#include "tensorflow/compiler/tf2xla/xla_op_registry.h" + +namespace tensorflow { +namespace { + +void BatchToSpace(XlaOpKernelContext* ctx, + const xla::ComputationDataHandle& input, DataType input_dtype, + const TensorShape& input_tensor_shape, + gtl::ArraySlice block_shape, + const xla::Literal& crops) { + const int input_rank = input_tensor_shape.dims(); + const gtl::InlinedVector input_shape = + input_tensor_shape.dim_sizes(); + const int block_rank = block_shape.size(); + + OP_REQUIRES( + ctx, input_rank >= 1 + block_rank, + errors::InvalidArgument("input rank should be >= ", 1 + block_rank, + " instead of ", input_rank)); + gtl::ArraySlice remainder_shape(input_shape); + remainder_shape.remove_prefix(1 + block_rank); + + OP_REQUIRES( + ctx, + xla::ShapeUtil::Rank(crops.shape()) == 2 && + block_rank == xla::ShapeUtil::GetDimension(crops.shape(), 0) && + 2 == xla::ShapeUtil::GetDimension(crops.shape(), 1), + errors::InvalidArgument("crops should have shape [", block_rank, + ", 2] instead of ", + xla::ShapeUtil::HumanString(crops.shape()))); + + xla::ComputationBuilder* b = ctx->builder(); + const int64 batch_size = input_shape[0]; + + // Compute the product of the block_shape values. + int64 block_num_elems = 1; + for (int i = 0; i < block_rank; ++i) { + block_num_elems *= block_shape[i]; + } + OP_REQUIRES(ctx, block_num_elems > 0, + errors::InvalidArgument( + "The product of the block dimensions must be positive")); + + // 1. Reshape `input` to `reshaped` of shape: + // [block_shape[0], ..., block_shape[M-1], + // batch / prod(block_shape), + // input_shape[1], ..., input_shape[N-1]] + + OP_REQUIRES( + ctx, batch_size % block_num_elems == 0, + errors::InvalidArgument("Input batch dimension (", batch_size, + ") is not divisible by product of block sizes (", + block_num_elems, ")")); + std::vector reshaped_shape(input_rank + block_rank); + std::copy(block_shape.begin(), block_shape.end(), reshaped_shape.begin()); + reshaped_shape[block_rank] = batch_size / block_num_elems; + std::copy(input_shape.begin() + 1, input_shape.end(), + reshaped_shape.begin() + block_rank + 1); + xla::ComputationDataHandle reshaped = b->Reshape(input, reshaped_shape); + + // 2. Permute dimensions of `reshaped` to produce `permuted` of shape + // [batch / prod(block_shape), + // + // input_shape[1], block_shape[0], + // ..., + // input_shape[M], block_shape[M-1], + // + // input_shape[M+1], ..., input_shape[N-1]] + std::vector permutation(reshaped_shape.size()); + permutation[0] = block_rank; + for (int i = 0; i < block_rank; ++i) { + permutation[1 + 2 * i] = block_rank + 1 + i; + permutation[1 + 2 * i + 1] = i; + } + std::iota(permutation.begin() + 1 + block_rank * 2, permutation.end(), + 1 + block_rank * 2); + xla::ComputationDataHandle permuted = b->Transpose(reshaped, permutation); + + // 3. Reshape `permuted` to produce `reshaped_permuted` of shape + // [batch / prod(block_shape), + // + // input_shape[1] * block_shape[0], + // ..., + // input_shape[M] * block_shape[M-1], + // + // input_shape[M+1], + // ..., + // input_shape[N-1]] + std::vector reshaped_permuted_shape(input_rank); + reshaped_permuted_shape[0] = batch_size / block_num_elems; + for (int i = 0; i < block_rank; ++i) { + reshaped_permuted_shape[1 + i] = block_shape[i] * input_shape[1 + i]; + } + std::copy(remainder_shape.begin(), remainder_shape.end(), + reshaped_permuted_shape.begin() + 1 + block_rank); + + xla::ComputationDataHandle reshaped_permuted = + b->Reshape(permuted, reshaped_permuted_shape); + + // 4. Crop the start and end of dimensions `[1, ..., M]` of + // `reshaped_permuted` according to `crops` to produce the output of shape: + // [batch / prod(block_shape), + // + // input_shape[1] * block_shape[0] - crops[0,0] - crops[0,1], + // ..., + // input_shape[M] * block_shape[M-1] - crops[M-1,0] - crops[M-1,1], + // + // input_shape[M+1], ..., input_shape[N-1]] + std::vector start_indices(input_rank, 0); + std::vector end_indices = reshaped_permuted_shape; + std::vector strides(input_rank, 1); + for (int i = 0; i < block_rank; ++i) { + int64 crop_start = crops.Get({i, 0}); + int64 crop_end = crops.Get({i, 1}); + OP_REQUIRES(ctx, crop_start >= 0 && crop_end >= 0, + errors::InvalidArgument("Crops must be non-negative")); + start_indices[1 + i] = crop_start; + end_indices[1 + i] -= crop_end; + OP_REQUIRES( + ctx, start_indices[1 + i] <= end_indices[1 + i], + errors::InvalidArgument( + "Cropped size must be non-negative: start: ", crop_start, + " end: ", crop_end, " size ", reshaped_permuted_shape[1 + i])); + } + xla::ComputationDataHandle output = + b->Slice(reshaped_permuted, start_indices, end_indices, strides); + ctx->SetOutput(0, output); +} + +class BatchToSpaceNDOp : public XlaOpKernel { + public: + explicit BatchToSpaceNDOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {} + + void Compile(XlaOpKernelContext* ctx) override { + std::vector block_shape; + OP_REQUIRES_OK(ctx, ctx->ConstantInputAsIntVector(1, &block_shape)); + + xla::Literal crops; + OP_REQUIRES_OK(ctx, ctx->ConstantInputAsInt64Literal(2, &crops)); + + BatchToSpace(ctx, ctx->Input(0), input_type(0), ctx->InputShape(0), + block_shape, crops); + } +}; +REGISTER_XLA_OP(Name("BatchToSpaceND"), BatchToSpaceNDOp); + +class BatchToSpaceOp : public XlaOpKernel { + public: + explicit BatchToSpaceOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) { + OP_REQUIRES_OK(ctx, ctx->GetAttr("block_size", &block_size_)); + OP_REQUIRES( + ctx, block_size_ > 1, + errors::InvalidArgument("Block size should be > 1: ", block_size_)); + } + + void Compile(XlaOpKernelContext* ctx) override { + xla::Literal crops; + OP_REQUIRES_OK(ctx, ctx->ConstantInputAsInt64Literal(1, &crops)); + + BatchToSpace(ctx, ctx->Input(0), input_type(0), ctx->InputShape(0), + {block_size_, block_size_}, crops); + } + + private: + int block_size_; +}; +REGISTER_XLA_OP(Name("BatchToSpace"), BatchToSpaceOp); + +} // namespace +} // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/kernels/bcast_ops.cc b/tensorflow/compiler/tf2xla/kernels/bcast_ops.cc index b0fee5e4bca502a7abb4613b58ecdd2ffca2206d..bc2cd31230dfe9ca35540341d225dcb768fa34f6 100644 --- a/tensorflow/compiler/tf2xla/kernels/bcast_ops.cc +++ b/tensorflow/compiler/tf2xla/kernels/bcast_ops.cc @@ -55,7 +55,7 @@ class BCastGradArgsOp : public XlaOpKernel { BCast::Vec vec; for (int64 i = 0; i < in_shape.num_elements(); ++i) { - vec.push_back(xla::LiteralUtil::Get(literal, {i})); + vec.push_back(literal.Get({i})); } shapes.push_back(vec); } diff --git a/tensorflow/compiler/tf2xla/kernels/binary_ops.cc b/tensorflow/compiler/tf2xla/kernels/binary_ops.cc index ded20a9a3cefd12c1f1fee0593c82165a8129f40..58538b45137b26ed5aa296eb6c1077e88aea72b9 100644 --- a/tensorflow/compiler/tf2xla/kernels/binary_ops.cc +++ b/tensorflow/compiler/tf2xla/kernels/binary_ops.cc @@ -102,11 +102,16 @@ XLA_MAKE_BINARY(Mod, b->Rem(lhs, rhs, extend_dimensions)); XLA_MAKE_BINARY(Maximum, b->Max(lhs, rhs, extend_dimensions)); XLA_MAKE_BINARY(Minimum, b->Min(lhs, rhs, extend_dimensions)); XLA_MAKE_BINARY(RealDiv, b->Div(lhs, rhs, extend_dimensions)); +XLA_MAKE_BINARY(ReciprocalGrad, b->Neg(b->Mul(rhs, b->Mul(lhs, lhs)))); XLA_MAKE_BINARY( RsqrtGrad, b->Mul(b->Pow(lhs, XlaHelpers::IntegerLiteral(b, input_type(0), 3)), b->Div(rhs, XlaHelpers::IntegerLiteral(b, input_type(0), -2)), extend_dimensions)); +XLA_MAKE_BINARY(SqrtGrad, + b->Div(b->Mul(rhs, + XlaHelpers::FloatLiteral(b, input_type(0), 0.5)), + lhs, extend_dimensions)); static xla::ComputationDataHandle Square(xla::ComputationBuilder* builder, const xla::ComputationDataHandle& x) { @@ -136,6 +141,11 @@ XLA_MAKE_BINARY(SoftplusGrad, b->Div(lhs, b->Add(b->Exp(b->Neg(rhs)), XlaHelpers::One(b, input_type(1))))); +// softsigngrad(gradients, features) = gradients / (1 + abs(features)) ** 2 +XLA_MAKE_BINARY(SoftsignGrad, + b->Div(lhs, Square(b, b->Add(XlaHelpers::One(b, input_type(0)), + b->Abs(rhs))))); + XLA_MAKE_BINARY(TanhGrad, b->Mul(rhs, b->Sub(XlaHelpers::One(b, input_type(0)), b->Mul(lhs, lhs)))); @@ -143,5 +153,24 @@ XLA_MAKE_BINARY(Pow, b->Pow(lhs, rhs, extend_dimensions)); #undef XLA_MAKE_BINARY +class ApproximateEqualOp : public XlaOpKernel { + public: + explicit ApproximateEqualOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) { + OP_REQUIRES_OK(ctx, ctx->GetAttr("tolerance", &tolerance_)); + } + + // Computes the max of the scalar input x and 0. + void Compile(XlaOpKernelContext* ctx) override { + xla::ComputationBuilder* b = ctx->builder(); + auto result = b->Lt(b->Abs(b->Sub(ctx->Input(0), ctx->Input(1))), + XlaHelpers::FloatLiteral(b, input_type(0), tolerance_)); + ctx->SetOutput(0, result); + } + + private: + float tolerance_; +}; +REGISTER_XLA_OP(Name("ApproximateEqual"), ApproximateEqualOp); + } // namespace } // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/kernels/cast_op.cc b/tensorflow/compiler/tf2xla/kernels/cast_op.cc index 124e33d7935ce19ced72d1c84521ffda1090bc86..2331520230176fce7646d89140851fe37aee5fda 100644 --- a/tensorflow/compiler/tf2xla/kernels/cast_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/cast_op.cc @@ -38,17 +38,6 @@ class CastOp : public XlaOpKernel { if (src_dtype_ == dst_dtype_) { output = input; - } else if (src_dtype_ == DT_BOOL) { - // XLA's ConvertElementType doesn't support casting to/from - // bools. So we need to handle those cases separately. - // Builds the equivalent of (input ? 1 : 0) - xla::ComputationBuilder l(builder->client(), "PredCast"); - xla::ComputationDataHandle x = - l.Parameter(0, xla::ShapeUtil::MakeShape(src_type_, {}), "x"); - l.Select(x, XlaHelpers::One(&l, dst_dtype_), - XlaHelpers::Zero(&l, dst_dtype_)); - xla::Computation computation = l.Build().ConsumeValueOrDie(); - output = builder->Map({input}, computation); } else if (dst_dtype_ == DT_BOOL) { output = builder->Ne(input, XlaHelpers::Zero(builder, src_dtype_)); } else { diff --git a/tensorflow/compiler/tf2xla/kernels/concat_op.cc b/tensorflow/compiler/tf2xla/kernels/concat_op.cc index e2eacb3839d39e6fa41192e8aa0f31d878d96aea..73a4740e29af7fa57e71ef42a342f46b0e24231d 100644 --- a/tensorflow/compiler/tf2xla/kernels/concat_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/concat_op.cc @@ -52,7 +52,7 @@ class ConcatBaseOp : public XlaOpKernel { 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 = xla::LiteralUtil::Get(literal, {}); + const int32 concat_dim = literal.Get({}); std::vector values; std::vector shapes; @@ -163,7 +163,7 @@ class ConcatOffsetOp : public XlaOpKernel { xla::Literal concat_dim_literal; OP_REQUIRES_OK(ctx, ctx->ConstantInput(0, &concat_dim_literal)); - const int64 cdim = xla::LiteralUtil::Get(concat_dim_literal, {}); + const int64 cdim = concat_dim_literal.Get({}); VLOG(1) << "ConcatOffset " << cdim << "," << dims; int32 axis = cdim < 0 ? cdim + dims : cdim; @@ -185,12 +185,10 @@ class ConcatOffsetOp : public XlaOpKernel { for (int64 j = 0; j < dims; ++j) { if (j == axis) { out_vec(j) = offset; - offset += xla::LiteralUtil::Get(inp_literal, {j}); + offset += inp_literal.Get({j}); } else { - const int32 inp0_element = - xla::LiteralUtil::Get(inp0_literal, {j}); - const int32 inp_element = - xla::LiteralUtil::Get(inp_literal, {j}); + const int32 inp0_element = inp0_literal.Get({j}); + const int32 inp_element = inp_literal.Get({j}); OP_REQUIRES( ctx, (inp0_element == inp_element), errors::InvalidArgument("input[", i, ",", j, "] mismatch: ", diff --git a/tensorflow/compiler/tf2xla/kernels/const_op.cc b/tensorflow/compiler/tf2xla/kernels/const_op.cc new file mode 100644 index 0000000000000000000000000000000000000000..9833323d851e00e7ca76d0b39cd2b216748a17fa --- /dev/null +++ b/tensorflow/compiler/tf2xla/kernels/const_op.cc @@ -0,0 +1,112 @@ +/* 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/type_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/core/framework/kernel_def_builder.h" +#include "tensorflow/core/framework/tensor.pb.h" + +namespace tensorflow { +namespace { + +class ConstOp : public XlaOpKernel { + public: + explicit ConstOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) { + const TensorProto* proto = nullptr; + OP_REQUIRES_OK(ctx, ctx->GetAttr("value", &proto)); + proto_ = *proto; + OP_REQUIRES( + ctx, ctx->output_type(0) == proto_.dtype(), + errors::InvalidArgument("Type mismatch between value (", + DataTypeString(proto_.dtype()), ") and dtype (", + DataTypeString(ctx->output_type(0)), ")")); + OP_REQUIRES_OK(ctx, TensorShape::IsValidShape(proto_.tensor_shape())); + } + + void Compile(XlaOpKernelContext* ctx) override { + TensorShape shape(proto_.tensor_shape()); + + xla::ComputationBuilder* b = ctx->builder(); + + // To avoid blowups for large constants filled with the same value, + // recognize that case and emit a scalar broadcast instead. + if (shape.num_elements() > 1) { + switch (proto_.dtype()) { + case DT_BOOL: + if (proto_.bool_val_size() == 1) { + ctx->SetOutput(0, + b->Broadcast(b->ConstantR0(proto_.bool_val(0)), + shape.dim_sizes())); + return; + } + break; + case DT_FLOAT: + if (proto_.float_val_size() == 1) { + ctx->SetOutput( + 0, b->Broadcast(b->ConstantR0(proto_.float_val(0)), + shape.dim_sizes())); + return; + } + break; + case DT_DOUBLE: + if (proto_.double_val_size() == 1) { + ctx->SetOutput( + 0, b->Broadcast(b->ConstantR0(proto_.double_val(0)), + shape.dim_sizes())); + return; + } + break; + case DT_INT32: + if (proto_.int_val_size() == 1) { + ctx->SetOutput(0, + b->Broadcast(b->ConstantR0(proto_.int_val(0)), + shape.dim_sizes())); + return; + } + break; + case DT_INT64: + if (proto_.int64_val_size() == 1) { + ctx->SetOutput( + 0, b->Broadcast(b->ConstantR0(proto_.int64_val(0)), + shape.dim_sizes())); + return; + } + break; + default: + break; + } + } + + // General case + Tensor tensor(proto_.dtype()); + OP_REQUIRES(ctx, tensor.FromProto(cpu_allocator(), proto_), + errors::InvalidArgument("Cannot parse tensor from proto: ", + proto_.DebugString())); + ctx->SetConstantOutput(0, tensor); + } + + private: + TensorProto proto_; + TF_DISALLOW_COPY_AND_ASSIGN(ConstOp); +}; + +// XLA_* devices also register a "real" Const operator so we suppress the +// dummy operator using CompilationOnly(). +REGISTER_XLA_OP(Name("Const").CompilationOnly(), ConstOp); + +} // namespace +} // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/kernels/conv_ops.cc b/tensorflow/compiler/tf2xla/kernels/conv_ops.cc index 67a0b803c5bb7add83d7a27f6f056eee53353881..0091b66d28ad62fcd5c0f3b09e90fed8347bb661 100644 --- a/tensorflow/compiler/tf2xla/kernels/conv_ops.cc +++ b/tensorflow/compiler/tf2xla/kernels/conv_ops.cc @@ -25,7 +25,6 @@ limitations under the License. #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/conv_grad_ops.h" #include "tensorflow/core/kernels/ops_util.h" #include "tensorflow/core/util/padding.h" @@ -35,12 +34,100 @@ namespace tensorflow { namespace { +// Returns the expanded size of a filter used for depthwise convolution. +// If `shape` is [H, W, ..., M, N] returns [H, W, ..., M, M*N]. +TensorShape ExpandedFilterShapeForDepthwiseConvolution( + const TensorShape& shape) { + int num_dims = shape.dims(); + CHECK_GE(num_dims, 2); + TensorShape expanded_shape = shape; + expanded_shape.set_dim(num_dims - 1, shape.dim_size(num_dims - 2) * + shape.dim_size(num_dims - 1)); + return expanded_shape; +} + +// Expands a filter of shape [H, W, ..., M, N] to [H, W, ..., M, M*N] by adding +// zeros for the cross-depth filters. Used to build a depthwise convolution. +xla::ComputationDataHandle ExpandFilterForDepthwiseConvolution( + const TensorShape& filter_shape, DataType dtype, + const xla::ComputationDataHandle& filter, + xla::ComputationBuilder* builder) { + // Filter has shape [H, W, ..., M, N] + // Dilate to [H, W, ..., M*M, N] using M inter-element padding, and then + // reshape to [H, W, ..., M, M*N]. + int num_spatial_dims = filter_shape.dims() - 2; + const int64 in_depth = filter_shape.dim_size(num_spatial_dims); + xla::PaddingConfig padding = xla::MakeNoPaddingConfig(filter_shape.dims()); + padding.mutable_dimensions(num_spatial_dims)->set_interior_padding(in_depth); + auto dilated_filter = + builder->Pad(filter, XlaHelpers::Zero(builder, dtype), padding); + + TensorShape expanded_filter_shape = + ExpandedFilterShapeForDepthwiseConvolution(filter_shape); + return builder->Reshape(dilated_filter, expanded_filter_shape.dim_sizes()); +} + +// Inverse of ExpandFilterForDepthwiseConvolution. +xla::ComputationDataHandle ContractFilterForDepthwiseBackprop( + const TensorShape& filter_shape, DataType dtype, + const xla::ComputationDataHandle& filter_backprop, + xla::ComputationBuilder* builder) { + int num_spatial_dims = filter_shape.dims() - 2; + + // Reshape to [H, W, ..., M*M, N] + TensorShape shape = filter_shape; + int64 in_depth = filter_shape.dim_size(num_spatial_dims); + shape.set_dim(num_spatial_dims, in_depth * in_depth); + auto reshaped = builder->Reshape(filter_backprop, shape.dim_sizes()); + + std::vector zeros(filter_shape.dims()); + std::vector strides(filter_shape.dims(), 1LL); + strides[num_spatial_dims] = in_depth + 1; + return builder->Slice(reshaped, zeros, shape.dim_sizes(), strides); + + // Alternate implementation for backends without strided Slice() support. + // TODO(phawkins): Remove when all backends support strided slice. + // // Pad [..., M * (M + 1), N] + // xla::PaddingConfig config = + // xla::MakeNoPaddingConfig(filter_shape.dims()); + // config.mutable_dimensions(num_spatial_dims) + // ->set_edge_padding_high(in_depth); + // auto zero = XlaHelpers::Zero(builder, dtype); + // auto padded = builder->Pad(reshaped, zero, config); + // + // // Reshape to [..., M, M + 1, N] + // shape = filter_shape; + // shape.set_dim(num_spatial_dims, in_depth); + // shape.set_dim(num_spatial_dims + 1, in_depth + 1); + // int64 out_depth = filter_shape.dim_size(num_spatial_dims + 1); + // shape.AddDim(out_depth); + // reshaped = builder->Reshape(padded, shape.dim_sizes()); + // + // // Slice to [..., M, 1, N] + // std::vector zeros(shape.dims()); + // std::vector strides(shape.dims(), 1LL); + // shape.set_dim(num_spatial_dims + 1, 1); + // auto sliced = builder->Slice(reshaped, zeros, shape.dim_sizes(), + // strides); + // + // // Reshape to [..., M, N] + // return builder->Reshape(sliced, filter_shape.dim_sizes()); +} + class ConvOp : public XlaOpKernel { public: - explicit ConvOp(OpKernelConstruction* ctx, int num_spatial_dims) - : XlaOpKernel(ctx), num_spatial_dims_(num_spatial_dims) { + explicit ConvOp(OpKernelConstruction* ctx, int num_spatial_dims, + bool depthwise) + : XlaOpKernel(ctx), + num_spatial_dims_(num_spatial_dims), + depthwise_(depthwise) { OP_REQUIRES_OK(ctx, ctx->GetAttr("strides", &strides_)); OP_REQUIRES_OK(ctx, ctx->GetAttr("padding", &padding_)); + + 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; } @@ -82,9 +169,16 @@ class ConvOp : public XlaOpKernel { "input and filter must have the same depth: ", in_depth, " vs ", input_shape.dim_size(feature_dim))); + xla::ComputationBuilder* b = ctx->builder(); + + xla::ComputationDataHandle filter = ctx->Input(1); + if (depthwise_) { + filter = ExpandFilterForDepthwiseConvolution( + filter_shape, ctx->input_type(0), filter, b); + } + xla::ConvolutionDimensionNumbers dims; std::vector window_strides; - dims.set_batch_dimension(GetTensorBatchDimIndex(num_dims(), data_format_)); dims.set_feature_dimension(feature_dim); for (int i = 0; i < num_spatial_dims_; ++i) { @@ -99,13 +193,14 @@ class ConvOp : public XlaOpKernel { xla::Padding xla_padding = (padding_ == VALID) ? xla::Padding::kValid : xla::Padding::kSame; - xla::ComputationDataHandle conv = ctx->builder()->ConvWithGeneralDimensions( - ctx->Input(0), ctx->Input(1), window_strides, xla_padding, dims); + xla::ComputationDataHandle conv = b->ConvWithGeneralDimensions( + ctx->Input(0), filter, window_strides, xla_padding, dims); ctx->SetOutput(0, conv); } protected: const int num_spatial_dims_; + const bool depthwise_; std::vector strides_; Padding padding_; TensorFormat data_format_ = FORMAT_NHWC; @@ -117,29 +212,38 @@ class ConvOp : public XlaOpKernel { class Conv2DOp : public ConvOp { public: explicit Conv2DOp(OpKernelConstruction* ctx) - : ConvOp(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")); - } + : ConvOp(ctx, /*num_spatial_dims=*/2, /*depthwise=*/false) {} }; REGISTER_XLA_OP(Name("Conv2D"), Conv2DOp); class Conv3DOp : public ConvOp { public: explicit Conv3DOp(OpKernelConstruction* ctx) - : ConvOp(ctx, /*num_spatial_dims=*/3) {} + : ConvOp(ctx, /*num_spatial_dims=*/3, /*depthwise=*/false) {} }; REGISTER_XLA_OP(Name("Conv3D"), Conv3DOp); +class DepthwiseConv2DOp : public ConvOp { + public: + explicit DepthwiseConv2DOp(OpKernelConstruction* ctx) + : ConvOp(ctx, /*num_spatial_dims=*/2, /*depthwise=*/true) {} +}; +REGISTER_XLA_OP(Name("DepthwiseConv2dNative"), DepthwiseConv2DOp); + // Backprop for input. class ConvBackpropInputOp : public XlaOpKernel { public: - explicit ConvBackpropInputOp(OpKernelConstruction* ctx, int num_spatial_dims) - : XlaOpKernel(ctx), num_spatial_dims_(num_spatial_dims) { + explicit ConvBackpropInputOp(OpKernelConstruction* ctx, int num_spatial_dims, + bool depthwise) + : XlaOpKernel(ctx), + num_spatial_dims_(num_spatial_dims), + depthwise_(depthwise) { OP_REQUIRES_OK(ctx, ctx->GetAttr("strides", &strides_)); OP_REQUIRES_OK(ctx, ctx->GetAttr("padding", &padding_)); + 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; } @@ -162,13 +266,17 @@ class ConvBackpropInputOp : public XlaOpKernel { const TensorShape filter_shape = ctx->InputShape(1); const TensorShape out_backprop_shape = ctx->InputShape(2); + const TensorShape expanded_filter_shape = + depthwise_ ? ExpandedFilterShapeForDepthwiseConvolution(filter_shape) + : filter_shape; // Reuse dimension computation logic from conv_grad_ops.cc. ConvBackpropDimensions dims; - OP_REQUIRES_OK( - ctx, ConvBackpropComputeDimensions( - type_string(), num_spatial_dims_, input_shape, filter_shape, - out_backprop_shape, strides_, padding_, data_format_, &dims)); + OP_REQUIRES_OK(ctx, ConvBackpropComputeDimensions( + type_string(), num_spatial_dims_, input_shape, + expanded_filter_shape, out_backprop_shape, strides_, + padding_, data_format_, &dims)); + xla::ComputationBuilder* b = ctx->builder(); auto filter = ctx->Input(1); auto out_backprop = ctx->Input(2); @@ -200,13 +308,19 @@ class ConvBackpropInputOp : public XlaOpKernel { lhs_dilation[i] = dims.spatial_dims[i].stride; } + // If this is a depthwise convolution, expand the filter. + if (depthwise_) { + filter = ExpandFilterForDepthwiseConvolution( + filter_shape, ctx->input_type(1), filter, b); + } + // Mirror the filter in the spatial dimensions. xla::ComputationDataHandle mirrored_weights = - ctx->builder()->Rev(filter, kernel_spatial_dims); + b->Rev(filter, kernel_spatial_dims); // activation gradients // = gradients (with padding and dilation) mirrored_weights - xla::ComputationDataHandle in_backprop = ctx->builder()->ConvGeneralDilated( + xla::ComputationDataHandle in_backprop = b->ConvGeneralDilated( out_backprop, mirrored_weights, /*window_strides=*/ones, padding, lhs_dilation, /*rhs_dilation=*/ones, dnums); @@ -215,6 +329,7 @@ class ConvBackpropInputOp : public XlaOpKernel { protected: const int num_spatial_dims_; + const bool depthwise_; std::vector strides_; Padding padding_; TensorFormat data_format_ = FORMAT_NHWC; @@ -226,28 +341,38 @@ class ConvBackpropInputOp : public XlaOpKernel { class Conv2DBackpropInputOp : public ConvBackpropInputOp { public: explicit Conv2DBackpropInputOp(OpKernelConstruction* ctx) - : ConvBackpropInputOp(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")); - } + : ConvBackpropInputOp(ctx, /*num_spatial_dims=*/2, /*depthwise=*/false) {} }; REGISTER_XLA_OP(Name("Conv2DBackpropInput"), Conv2DBackpropInputOp); class Conv3DBackpropInputOp : public ConvBackpropInputOp { public: explicit Conv3DBackpropInputOp(OpKernelConstruction* ctx) - : ConvBackpropInputOp(ctx, /*num_spatial_dims=*/3) {} + : ConvBackpropInputOp(ctx, /*num_spatial_dims=*/3, /*depthwise=*/false) {} }; REGISTER_XLA_OP(Name("Conv3DBackpropInputV2"), Conv3DBackpropInputOp); +class DepthwiseConv2DBackpropInputOp : public ConvBackpropInputOp { + public: + explicit DepthwiseConv2DBackpropInputOp(OpKernelConstruction* ctx) + : ConvBackpropInputOp(ctx, /*num_spatial_dims=*/2, /*depthwise=*/true) {} +}; +REGISTER_XLA_OP(Name("DepthwiseConv2dNativeBackpropInput"), + DepthwiseConv2DBackpropInputOp); + class ConvBackpropFilterOp : public XlaOpKernel { public: - explicit ConvBackpropFilterOp(OpKernelConstruction* ctx, int num_spatial_dims) - : XlaOpKernel(ctx), num_spatial_dims_(num_spatial_dims) { + explicit ConvBackpropFilterOp(OpKernelConstruction* ctx, int num_spatial_dims, + bool depthwise) + : XlaOpKernel(ctx), + num_spatial_dims_(num_spatial_dims), + depthwise_(depthwise) { OP_REQUIRES_OK(ctx, ctx->GetAttr("strides", &strides_)); OP_REQUIRES_OK(ctx, ctx->GetAttr("padding", &padding_)); + 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; } @@ -266,13 +391,18 @@ class ConvBackpropFilterOp : public XlaOpKernel { OP_REQUIRES_OK(ctx, ctx->ConstantInputAsShape(1, &filter_shape)); const TensorShape out_backprop_shape = ctx->InputShape(2); + const TensorShape expanded_filter_shape = + depthwise_ ? ExpandedFilterShapeForDepthwiseConvolution(filter_shape) + : filter_shape; + // Reuse dimension computation logic from conv_grad_ops.cc. ConvBackpropDimensions dims; OP_REQUIRES_OK(ctx, ConvBackpropComputeDimensions( type_string(), num_spatial_dims_, activations_shape, - filter_shape, out_backprop_shape, strides_, + expanded_filter_shape, out_backprop_shape, strides_, padding_, data_format_, &dims)); + xla::ComputationBuilder* b = ctx->builder(); xla::ComputationDataHandle activations = ctx->Input(0); xla::ComputationDataHandle gradients = ctx->Input(2); @@ -357,10 +487,10 @@ class ConvBackpropFilterOp : public XlaOpKernel { // // This is done by specifying the window dilation factors in the // convolution HLO below. - auto filter_backprop = ctx->builder()->ConvGeneralDilated( - activations, gradients, - /*window_strides=*/ones, padding, /*lhs_dilation=*/ones, rhs_dilation, - dnums); + auto filter_backprop = + b->ConvGeneralDilated(activations, gradients, + /*window_strides=*/ones, padding, + /*lhs_dilation=*/ones, rhs_dilation, dnums); // The layout of filter_backprop will match the layout of // padded_activations @@ -375,12 +505,18 @@ class ConvBackpropFilterOp : public XlaOpKernel { transpose_dims.push_back(c_dim); transpose_dims.push_back(n_dim); xla::ComputationDataHandle filter_backprop_reshaped = - ctx->builder()->Transpose(filter_backprop, transpose_dims); + b->Transpose(filter_backprop, transpose_dims); + + if (depthwise_) { + filter_backprop_reshaped = ContractFilterForDepthwiseBackprop( + filter_shape, ctx->input_type(0), filter_backprop_reshaped, b); + } ctx->SetOutput(0, filter_backprop_reshaped); } protected: - int num_spatial_dims_; + const int num_spatial_dims_; + const bool depthwise_; std::vector strides_; Padding padding_; TensorFormat data_format_ = FORMAT_NHWC; @@ -392,11 +528,7 @@ class ConvBackpropFilterOp : public XlaOpKernel { class Conv2DBackpropFilterOp : public ConvBackpropFilterOp { public: explicit Conv2DBackpropFilterOp(OpKernelConstruction* ctx) - : ConvBackpropFilterOp(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")); + : ConvBackpropFilterOp(ctx, /*num_spatial_dims=*/2, /*depthwise=*/false) { } }; REGISTER_XLA_OP(Name("Conv2DBackpropFilter"), Conv2DBackpropFilterOp); @@ -404,9 +536,18 @@ REGISTER_XLA_OP(Name("Conv2DBackpropFilter"), Conv2DBackpropFilterOp); class Conv3DBackpropFilterOp : public ConvBackpropFilterOp { public: explicit Conv3DBackpropFilterOp(OpKernelConstruction* ctx) - : ConvBackpropFilterOp(ctx, /*num_spatial_dims=*/3) {} + : ConvBackpropFilterOp(ctx, /*num_spatial_dims=*/3, /*depthwise=*/false) { + } }; REGISTER_XLA_OP(Name("Conv3DBackpropFilterV2"), Conv3DBackpropFilterOp); +class DepthwiseConv2DBackpropFilterOp : public ConvBackpropFilterOp { + public: + explicit DepthwiseConv2DBackpropFilterOp(OpKernelConstruction* ctx) + : ConvBackpropFilterOp(ctx, /*num_spatial_dims=*/2, /*depthwise=*/true) {} +}; +REGISTER_XLA_OP(Name("DepthwiseConv2dNativeBackpropFilter"), + DepthwiseConv2DBackpropFilterOp); + } // namespace } // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/kernels/cross_op.cc b/tensorflow/compiler/tf2xla/kernels/cross_op.cc new file mode 100644 index 0000000000000000000000000000000000000000..3df8c00f1b83556d7d954aedc8eeac0728251c3e --- /dev/null +++ b/tensorflow/compiler/tf2xla/kernels/cross_op.cc @@ -0,0 +1,87 @@ +/* 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_helpers.h" +#include "tensorflow/compiler/tf2xla/xla_op_kernel.h" +#include "tensorflow/compiler/tf2xla/xla_op_registry.h" + +namespace tensorflow { +namespace { + +class CrossOp : public XlaOpKernel { + public: + explicit CrossOp(OpKernelConstruction* context) : XlaOpKernel(context) {} + + void Compile(XlaOpKernelContext* ctx) override { + TensorShape in0_shape = ctx->InputShape(0); + TensorShape in1_shape = ctx->InputShape(1); + OP_REQUIRES(ctx, in0_shape == in1_shape, + errors::InvalidArgument("Both inputs must be of same shape: ", + in0_shape.DebugString(), " vs. ", + in1_shape.DebugString())); + OP_REQUIRES(ctx, in0_shape.dims() >= 1, + errors::InvalidArgument("Input must be at least 1D", + in0_shape.DebugString())); + + auto inner_dim = in0_shape.dim_size(in0_shape.dims() - 1); + OP_REQUIRES(ctx, inner_dim == 3, + errors::FailedPrecondition( + "Cross-products are only defined for 3-element vectors.")); + + // in0 is a [...,X,Y,Z,3] + // in1 is the same shape as in0 + // So slice 0 is: in0[...,:,:,:,0:1] + // So slice 1 is: in0[...,:,:,:,1:2] + // So slice 2 is: in0[...,:,:,:,2:3] + + std::vector starts(in0_shape.dims(), 0); + std::vector limits; + for (auto dim_size : in0_shape.dim_sizes()) { + limits.push_back(dim_size); + } + std::vector strides(in0_shape.dims(), 1); + + xla::ComputationBuilder* b = ctx->builder(); + auto in0 = ctx->Input(0); + auto in1 = ctx->Input(1); + starts.back() = 0; + limits.back() = 1; + auto u1 = b->Slice(in0, starts, limits, strides); + auto v1 = b->Slice(in1, starts, limits, strides); + starts.back() = 1; + limits.back() = 2; + auto u2 = b->Slice(in0, starts, limits, strides); + auto v2 = b->Slice(in1, starts, limits, strides); + starts.back() = 2; + limits.back() = 3; + auto u3 = b->Slice(in0, starts, limits, strides); + auto v3 = b->Slice(in1, starts, limits, strides); + + auto s1 = b->Sub(b->Mul(u2, v3), b->Mul(u3, v2)); + auto s2 = b->Sub(b->Mul(u3, v1), b->Mul(u1, v3)); + auto s3 = b->Sub(b->Mul(u1, v2), b->Mul(u2, v1)); + auto output = b->ConcatInDim({s1, s2, s3}, in0_shape.dims() - 1); + + ctx->SetOutput(0, output); + } + + private: + TF_DISALLOW_COPY_AND_ASSIGN(CrossOp); +}; + +REGISTER_XLA_OP(Name("Cross"), CrossOp); + +} // namespace +} // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/kernels/cwise_ops.cc b/tensorflow/compiler/tf2xla/kernels/cwise_ops.cc index de93a88f064a76fc8e75fe8f35b2afd69dde9faa..0cf03ceb948a5165a71e902eef5264eaddbd71e9 100644 --- a/tensorflow/compiler/tf2xla/kernels/cwise_ops.cc +++ b/tensorflow/compiler/tf2xla/kernels/cwise_ops.cc @@ -137,41 +137,4 @@ XlaBinaryOp::Broadcast(xla::ComputationBuilder* builder, return {lhs_output, rhs_output}; } -xla::ComputationDataHandle XlaBinaryMapOp::Computation( - XlaOpKernelContext* ctx, const xla::ComputationDataHandle& lhs, - const gtl::ArraySlice& lhs_shape, - const xla::ComputationDataHandle& rhs, - const gtl::ArraySlice& rhs_shape, const BCast& broadcast_helper, - const std::vector& extend_dimensions) { - xla::ComputationBuilder* builder = ctx->builder(); - - // Construct the builder for the lambda computation. - xla::ComputationBuilder l(builder->client(), ctx->op_kernel().name()); - xla::PrimitiveType type; - TF_CHECK_OK(DataTypeToPrimitiveType(input_type(0), &type)); - - // Make two scalar parameters of the desired type for the lambda. - xla::ComputationDataHandle x = - l.Parameter(0, xla::ShapeUtil::MakeShape(type, {}), "x"); - xla::ComputationDataHandle y = - l.Parameter(1, xla::ShapeUtil::MakeShape(type, {}), "y"); - - // Call virtual method to build the lambda. - BuildMapLambda(&l, x, y); - xla::Computation computation = l.Build().ConsumeValueOrDie(); - - xla::ComputationDataHandle lhs_broadcast = lhs; - xla::ComputationDataHandle rhs_broadcast = rhs; - if (lhs_shape == rhs_shape) { - // There's no broadcasting to do. - CHECK_EQ(0, extend_dimensions.size()); - return builder->Map({lhs, rhs}, computation); - } else { - std::tie(lhs_broadcast, rhs_broadcast) = - Broadcast(builder, lhs, rhs, broadcast_helper); - } - // Now the two sides are broadcast to the final shape we can do the map. - return builder->Map({lhs_broadcast, rhs_broadcast}, computation); -} - } // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/kernels/cwise_ops.h b/tensorflow/compiler/tf2xla/kernels/cwise_ops.h index ba38693325c84578619f9fdd7f3c255ddc3448bf..5bc1d5fb1f08fb576df654e1f4068b6be9114096 100644 --- a/tensorflow/compiler/tf2xla/kernels/cwise_ops.h +++ b/tensorflow/compiler/tf2xla/kernels/cwise_ops.h @@ -74,36 +74,6 @@ class XlaBinaryOp : public XlaOpKernel { const BCast& broadcast_helper); }; -// Coefficient-wise binary operations that map a scalar function. Each -// BinaryMap Op expects two inputs that can be broadcast to the same -// shape and maps a (scalar,scalar)->scalar function across the zipped -// elements of its (broadcast) inputs. The base class contains pure -// virtual methods to override: description is a textual description -// of the mapped function; and BuildMapLambda adds the -// implementation of the lambda to a xla::ComputationBuilder. -// Operations may have better performance if implemented as graphs of -// element-wise tensor operations. -class XlaBinaryMapOp : public XlaBinaryOp { - public: - explicit XlaBinaryMapOp(OpKernelConstruction* ctx) : XlaBinaryOp(ctx) {} - ~XlaBinaryMapOp() override {} - - // Implement the (scalar,scalar)->scalar lambda that should be - // applied to each pair of elements of the inputs. The desired - // computation should be added to 'builder' and - // '(scalar_lhs,scalar_rhs)' are the function's inputs. - virtual void BuildMapLambda(xla::ComputationBuilder* builder, - const xla::ComputationDataHandle& scalar_lhs, - const xla::ComputationDataHandle& scalar_rhs) = 0; - - xla::ComputationDataHandle Computation( - XlaOpKernelContext* ctx, const xla::ComputationDataHandle& lhs, - const gtl::ArraySlice& lhs_shape, - const xla::ComputationDataHandle& rhs, - const gtl::ArraySlice& rhs_shape, const BCast& broadcast_helper, - const std::vector& extend_dimensions) override; -}; - } // namespace tensorflow #endif // TENSORFLOW_COMPILER_TF2XLA_KERNELS_CWISE_OPS_H_ diff --git a/tensorflow/compiler/tf2xla/kernels/declaration_op.cc b/tensorflow/compiler/tf2xla/kernels/declaration_op.cc deleted file mode 100644 index be2ce038016e852e48c312e26bf959ca5b9215af..0000000000000000000000000000000000000000 --- a/tensorflow/compiler/tf2xla/kernels/declaration_op.cc +++ /dev/null @@ -1,129 +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/type_util.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" - -namespace tensorflow { -namespace { - -// This OpKernel implements the Constant Op for XLA JIT -// devices. It extracts the constant Tensor from the Proto at kernel -// construction time, and then every time the Constant Op is executed -// an expression containing the constant is compiled. -class ConstantDeclarationOp : public XlaOpKernel { - public: - explicit ConstantDeclarationOp(OpKernelConstruction* ctx) - : XlaOpKernel(ctx), tensor_(ctx->output_type(0)) { - const TensorProto* proto = nullptr; - OP_REQUIRES_OK(ctx, ctx->GetAttr("value", &proto)); - // MakeTensorFromProto uses the cpu_allocator, so tensor_ is a - // "real" tensor backed by CPU memory, holding the value of the - // constant. - OP_REQUIRES_OK(ctx, MakeTensorFromProto(*proto, &tensor_)); - OP_REQUIRES( - ctx, ctx->output_type(0) == tensor_.dtype(), - errors::InvalidArgument( - "Type mismatch between value (", DataTypeString(tensor_.dtype()), - ") and dtype (", DataTypeString(ctx->output_type(0)), ")")); - } - - void Compile(XlaOpKernelContext* ctx) override { - ctx->SetConstantOutput(0, tensor_); - } - - private: - // Extract the value of the constant from the Proto during Op kernel - // construction. The constant must be stored in a Tensor allocated - // using the cpu_allocator so that it is backed by real memory. The - // OpKernelConstruction's default allocator is the JITAllocator - // which only allocates enough space for metadata for each Tensor. - static Status MakeTensorFromProto(const TensorProto& tensor_proto, - Tensor* tensor) { - Tensor parsed(tensor_proto.dtype()); - if (!parsed.FromProto(cpu_allocator(), tensor_proto)) { - return errors::InvalidArgument("Cannot parse tensor from proto: ", - tensor_proto.DebugString()); - } - *tensor = parsed; - return Status::OK(); - } - - // This is a "real" tensor backed by CPU memory, containing the - // constant values. - Tensor tensor_; - TF_DISALLOW_COPY_AND_ASSIGN(ConstantDeclarationOp); -}; - -// XLA_* devices also register a "real" Identity operator so we suppress the -// dummy operator using CompilationOnly(). -REGISTER_XLA_OP(Name("Const").CompilationOnly(), ConstantDeclarationOp); - -// This OpKernel implements the _Arg Op for XLA JIT devices. It -// associates its output with one of the arguments to a -// subcomputation. -class ArgOp : public XlaOpKernel { - public: - explicit ArgOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) { - OP_REQUIRES_OK(ctx, ctx->GetAttr("T", &dtype_)); - OP_REQUIRES_OK(ctx, ctx->GetAttr("index", &index_)); - } - - void Compile(XlaOpKernelContext* ctx) override { - // If 'frame' is non-null, this is a function call inside an outer JIT - // compilation. Use the usual implementation of _Arg. - auto frame = ctx->call_frame(); - if (frame != nullptr) { - Tensor val; - OP_REQUIRES_OK(ctx, frame->GetArg(index_, &val)); - OP_REQUIRES(ctx, val.dtype() == dtype_, - errors::InvalidArgument( - "Type mismatch: actual ", DataTypeString(val.dtype()), - " vs. expect ", DataTypeString(dtype_))); - // Forwards the argument from the frame. - ctx->op_kernel_context()->set_output(0, val); - return; - } - - XlaContext& tc = XlaContext::Get(ctx); - const XlaContext::Argument& arg = tc.args()[index_]; - if (arg.is_variable) { - // We use the argument position of the variable input as a unique ID. - // TODO(phawkins): this code assumes that variables do not alias. - OP_REQUIRES_OK(ctx, tc.CreateVariable(index_, arg.name, arg.value.type, - arg.value.handle)); - ctx->SetVariableOutput(0, index_); - } else if (arg.value.is_constant) { - ctx->SetConstantOutput(0, arg.value.constant_value); - } else { - ctx->SetOutput(0, arg.value.handle); - } - } - - private: - int index_; - DataType dtype_; - - TF_DISALLOW_COPY_AND_ASSIGN(ArgOp); -}; - -REGISTER_XLA_OP(Name("_Arg").AllowResourceTypes(), ArgOp); - -} // namespace -} // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/kernels/depthwise_conv_ops.cc b/tensorflow/compiler/tf2xla/kernels/depthwise_conv_ops.cc deleted file mode 100644 index 92b371cc4ec0daf5137adf11f3e5c907c4bda153..0000000000000000000000000000000000000000 --- a/tensorflow/compiler/tf2xla/kernels/depthwise_conv_ops.cc +++ /dev/null @@ -1,236 +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. -==============================================================================*/ - -// XLA-specific Ops for 2D depthwise convolution. - -#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/literal_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/depthwise_conv_op.h" -#include "tensorflow/core/kernels/ops_util.h" -#include "tensorflow/core/util/padding.h" -#include "tensorflow/core/util/tensor_format.h" - -namespace tensorflow { -namespace { - -// Name of the function to use as the implementation for depthwise 2D -// convolution. Default is empty string; another possible value is -// "DummyDepthwiseConv2dKernel". -static const char kDepthwiseConv2dCustomFunc[] = ""; - -class DepthwiseConv2dNativeOp : public XlaOpKernel { - public: - explicit DepthwiseConv2dNativeOp(OpKernelConstruction* ctx) - : XlaOpKernel(ctx) { - // TODO(keveman): Refactor this (and other XLA OpKernel constructors) so - // that they use a common implementation shared with non-XLA kernels. - OP_REQUIRES_OK(ctx, ctx->GetAttr("strides", &strides_)); - OP_REQUIRES(ctx, strides_.size() == 4, - errors::InvalidArgument("Sliding window strides field must " - "specify 4 dimensions")); - OP_REQUIRES(ctx, strides_[1] == strides_[2], - errors::InvalidArgument( - "Current implementation only supports equal length " - "strides in the row and column dimensions.")); - OP_REQUIRES( - ctx, (strides_[0] == 1 && strides_[3] == 1), - errors::InvalidArgument("Current implementation does not yet support " - "strides in the batch and depth dimensions.")); - OP_REQUIRES_OK(ctx, ctx->GetAttr("padding", &padding_)); - } - - void Compile(XlaOpKernelContext* ctx) override { - // Input tensor is of the following dimensions: - // [ batch, in_rows, in_cols, in_depth ] - const TensorShape input_shape = ctx->InputShape(0); - - // Input filter is of the following dimensions: - // [ filter_rows, filter_cols, in_depth, depth_multiplier] - const TensorShape filter_shape = ctx->InputShape(1); - - // For 2D convolution, there should be 4 dimensions. - OP_REQUIRES(ctx, input_shape.dims() == 4, - errors::InvalidArgument("input must be 4-dimensional", - input_shape.DebugString())); - OP_REQUIRES(ctx, filter_shape.dims() == 4, - errors::InvalidArgument("filter must be 4-dimensional: ", - filter_shape.DebugString())); - - // The last dimension for input is in_depth. It must be the same as the - // filter's in_depth. - const int64 in_depth = input_shape.dim_size(3); - OP_REQUIRES( - ctx, in_depth == filter_shape.dim_size(2), - errors::InvalidArgument("input and filter must have the same depth: ", - in_depth, " vs ", filter_shape.dim_size(2))); - - // The last dimension for filter is depth multiplier. - const int64 depth_multiplier = filter_shape.dim_size(3); - - // The output depth is input depth x depth multiplier. - const int64 out_depth = in_depth * depth_multiplier; - - // The second dimension for input is rows/height. - // The first dimension for filter is rows/height. - const int64 input_rows = input_shape.dim_size(1); - const int64 filter_rows = filter_shape.dim_size(0); - - // The third dimension for input is columns/width. - // The second dimension for filter is columns/width. - const int64 input_cols = input_shape.dim_size(2); - const int64 filter_cols = filter_shape.dim_size(1); - - // The first dimension for input is batch. - const int64 batch = input_shape.dim_size(0); - - // For now we take the stride from the second dimension only (we - // assume row = col stride, and do not support striding on the - // batch or depth dimension). - const int32 stride = strides_[1]; - - int64 out_rows = 0, out_cols = 0, pad_rows = 0, pad_cols = 0; - OP_REQUIRES_OK(ctx, GetWindowedOutputSize(input_rows, filter_rows, stride, - padding_, &out_rows, &pad_rows)); - OP_REQUIRES_OK(ctx, GetWindowedOutputSize(input_cols, filter_cols, stride, - padding_, &out_cols, &pad_cols)); - TensorShape out_shape({batch, out_rows, out_cols, out_depth}); - OP_REQUIRES( - ctx, out_shape.num_elements() <= 2147483647, - errors::InvalidArgument("total number of outputs should be within the " - "range of int which is used in the GPU kernel", - in_depth, " vs ", filter_shape.dim_size(2))); - - // Output tensor is of the following dimensions: - // [ in_batch, out_rows, out_cols, out_depth ] - - VLOG(2) << "DepthwiseConv2dNative: " - << " Input: [" << batch << ", " << input_rows << ", " << input_cols - << ", " << in_depth << "]; Filter: [" << filter_rows << ", " - << filter_cols << ", " << in_depth << ", " << depth_multiplier - << "]; stride = " << stride << ", pad_rows = " << pad_rows - << ", pad_cols = " << pad_cols << ", output: [" << batch << ", " - << out_rows << ", " << out_cols << ", " << out_depth << "]"; - - xla::ComputationBuilder& b = *ctx->builder(); - xla::ComputationDataHandle input = ctx->Input(0); - xla::ComputationDataHandle filter = ctx->Input(1); - xla::ComputationDataHandle output; - - const string custom_function_name = kDepthwiseConv2dCustomFunc; - if (!custom_function_name.empty()) { - xla::Shape xla_out_shape; - OP_REQUIRES_OK( - ctx, TensorShapeToXLAShape(input_type(0), out_shape, &xla_out_shape)); - - // The custom function for depthwise should interpret its arguments - // as follows : - // func(T* output, - // const T* input, const T* filter, - // const int32* input_size, const int32* filter_size, - // const int32* output_size, - // int32 stride, int32 pad_rows, int32 pad_cols) - // - // where T is the type of Tensor that this kernel is registered for. - // Note that the custom call op passes uses the following calling - // convention: - // func(void* output, void** inputs) - // - // Therefore the custom function should first construct the above - // inputs by unparsing the second argument passed to it. - output = b.CustomCall( - custom_function_name, - {input, filter, - b.ConstantR1({batch, input_rows, input_cols, in_depth}), - b.ConstantR1( - {filter_rows, filter_cols, in_depth, depth_multiplier}), - b.ConstantR1({batch, out_rows, out_cols, out_depth}), - b.ConstantR0(stride), b.ConstantR0(pad_rows), - b.ConstantR0(pad_cols)}, - xla_out_shape); - } else { - // These will be used to define the bounds of each slice. - // Within the loop, the input_channel index will be modified. - gtl::InlinedVector filter_begin; - gtl::InlinedVector filter_limits; - gtl::InlinedVector input_begin; - gtl::InlinedVector input_limits; - for (int i = 0; i < 4; ++i) { - filter_begin.push_back(0); - filter_limits.push_back(filter_shape.dim_size(i)); - input_begin.push_back(0); - input_limits.push_back(input_shape.dim_size(i)); - } - - std::vector strides_for_tla{strides_[1], strides_[2]}; - - xla::Padding xla_padding = - (padding_ == VALID) ? xla::Padding::kValid : xla::Padding::kSame; - - xla::ConvolutionDimensionNumbers dims; - dims.set_batch_dimension(0); - dims.set_feature_dimension(3); - dims.add_spatial_dimensions(1); - dims.add_spatial_dimensions(2); - - // TF filter shape is [ H, W, inC, outC ] - dims.add_kernel_spatial_dimensions(0); - dims.add_kernel_spatial_dimensions(1); - dims.set_kernel_input_feature_dimension(2); - dims.set_kernel_output_feature_dimension(3); - - // Create one convolution for each input channel - std::vector convs; - for (int i = 0; i < in_depth; ++i) { - filter_begin[2] = i; - filter_limits[2] = i + 1; - input_begin[3] = i; - input_limits[3] = i + 1; - - xla::ComputationDataHandle filter_slice = - b.Slice(filter, filter_begin, filter_limits); - xla::ComputationDataHandle input_slice = - b.Slice(input, input_begin, input_limits); - convs.push_back(b.ConvWithGeneralDimensions( - input_slice, filter_slice, strides_for_tla, xla_padding, dims)); - } - // Concatenate the per-channel convolutions along the depth dimension. - output = b.ConcatInDim(convs, 3); - } - - ctx->SetOutput(0, output); - } - - private: - std::vector strides_; - Padding padding_; - - TF_DISALLOW_COPY_AND_ASSIGN(DepthwiseConv2dNativeOp); -}; - -REGISTER_XLA_OP(Name("DepthwiseConv2dNative").TypeConstraint("T", kFloatTypes), - DepthwiseConv2dNativeOp); - -} // namespace -} // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/kernels/diag_op.cc b/tensorflow/compiler/tf2xla/kernels/diag_op.cc index 74994d89619552fb0cb18342e7a2a7e7a939a3b9..ec5017f6ab96bd3fc273a746b77fbb7e74fd9f35 100644 --- a/tensorflow/compiler/tf2xla/kernels/diag_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/diag_op.cc @@ -125,7 +125,7 @@ class DiagPartOp : public XlaOpKernel { diag = builder->Reshape(diag, {new_size, new_size + 1}); // Slices out the first column and reshapes to the final shape. - diag = builder->Slice(diag, {0, 0}, {new_size, 1}); + diag = builder->Slice(diag, {0, 0}, {new_size, 1}, {1, 1}); diag = builder->Reshape(diag, new_dims); ctx->SetOutput(0, diag); @@ -224,8 +224,9 @@ class MatrixDiagPartOp : public XlaOpKernel { } else if (actual_size > target_size) { std::vector start(flattened_dims.size(), 0); std::vector limits(flattened_dims.begin(), flattened_dims.end()); + std::vector strides(flattened_dims.size(), 1); limits[flattened_dims.size() - 1] = target_size; - diag = builder->Slice(diag, start, limits); + diag = builder->Slice(diag, start, limits, strides); } // Reshape so the target values are in the first position of the last @@ -238,8 +239,9 @@ class MatrixDiagPartOp : public XlaOpKernel { // Slices out the first column and reshapes to the final shape. std::vector start(dims.size(), 0); std::vector limits(dims.begin(), dims.end()); + std::vector strides(dims.size(), 1); limits[last_dim] = 1; - diag = builder->Slice(diag, start, limits); + diag = builder->Slice(diag, start, limits, strides); // Collapses away the last dimension. dims.pop_back(); diff --git a/tensorflow/compiler/tf2xla/kernels/dynamic_stitch_op.cc b/tensorflow/compiler/tf2xla/kernels/dynamic_stitch_op.cc index 2d1a0567192600776f6f70fc1a4063180433ebb0..dde7898015e73190c96fa6effddfd3fc892264ea 100644 --- a/tensorflow/compiler/tf2xla/kernels/dynamic_stitch_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/dynamic_stitch_op.cc @@ -63,11 +63,14 @@ class DynamicStitchOp : public XlaOpKernel { std::vector indices(indices_input.size()); const TensorShape& data0_shape = data_shapes[0]; - const TensorShape indices0_shape = - XLAShapeToTensorShape(indices_input[0].shape()); + TensorShape indices0_shape; + OP_REQUIRES_OK( + ctx, XLAShapeToTensorShape(indices_input[0].shape(), &indices0_shape)); for (int input_num = 0; input_num < indices_input.size(); input_num++) { - const TensorShape indices_shape = - XLAShapeToTensorShape(indices_input[input_num].shape()); + TensorShape indices_shape; + OP_REQUIRES_OK(ctx, + XLAShapeToTensorShape(indices_input[input_num].shape(), + &indices_shape)); const TensorShape& data_shape = data_shapes[input_num]; OP_REQUIRES(ctx, TensorShapeUtils::StartsWith(data_shape, indices_shape), errors::InvalidArgument( @@ -103,8 +106,7 @@ class DynamicStitchOp : public XlaOpKernel { int max_index = -1; for (int input_num = 0; input_num < indices.size(); input_num++) { for (int i = 0; i < indices[input_num].shape().dimensions(0); ++i) { - max_index = std::max( - max_index, xla::LiteralUtil::Get(indices[input_num], {i})); + max_index = std::max(max_index, indices[input_num].Get({i})); } } int number_of_indices = max_index + 1; @@ -118,7 +120,7 @@ class DynamicStitchOp : public XlaOpKernel { int index_used_count = 0; for (int input_num = 0; input_num < indices.size(); input_num++) { for (int i = 0; i < indices[input_num].shape().dimensions(0); ++i) { - int index = xla::LiteralUtil::Get(indices[input_num], {i}); + int index = indices[input_num].Get({i}); src_input_vector[index] = input_num; src_slice_vector[index] = i; if (!src_index_used[index]) { @@ -157,6 +159,8 @@ class DynamicStitchOp : public XlaOpKernel { 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); for (int d = indices0_shape.dims(); d < data0_shape.dims(); d++) { slice_limit[1 + d - indices0_shape.dims()] = data0_shape.dim_size(d); } @@ -169,7 +173,7 @@ class DynamicStitchOp : public XlaOpKernel { // And place it in the concat list in the place indicated by // the index. to_concat[index_num] = - ctx->builder()->Slice(expression, slice_start, slice_limit); + ctx->builder()->Slice(expression, slice_start, slice_limit, stride); } ctx->SetOutput(0, ctx->builder()->ConcatInDim(to_concat, 0)); diff --git a/tensorflow/compiler/tf2xla/kernels/elu_op.cc b/tensorflow/compiler/tf2xla/kernels/elu_op.cc new file mode 100644 index 0000000000000000000000000000000000000000..2fd27c5ca7e87c8b387d9d0854b787d30e7f7b6f --- /dev/null +++ b/tensorflow/compiler/tf2xla/kernels/elu_op.cc @@ -0,0 +1,109 @@ +/* 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. +==============================================================================*/ + +// Native XLA implementations of XLA Elu Ops + +#include "tensorflow/compiler/tf2xla/kernels/cwise_ops.h" +#include "tensorflow/compiler/tf2xla/xla_helpers.h" +#include "tensorflow/compiler/tf2xla/xla_op_registry.h" +#include "tensorflow/compiler/xla/client/computation_builder.h" +#include "tensorflow/compiler/xla/literal_util.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 { + +class EluOp : public XlaOpKernel { + public: + explicit EluOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {} + // Computes the max of the scalar input x and 0. + void Compile(XlaOpKernelContext* ctx) override { + xla::ComputationBuilder* b = ctx->builder(); + const auto zero = XlaHelpers::Zero(b, input_type(0)); + const auto one = XlaHelpers::One(b, input_type(0)); + const auto pred = b->Gt(ctx->Input(0), zero); + const auto expm1 = b->Sub(b->Exp(ctx->Input(0)), one); + ctx->SetOutput(0, b->Select(pred, ctx->Input(0), expm1)); + } +}; + +class EluGradOp : public XlaOpKernel { + public: + explicit EluGradOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {} + // Return the lhs (incoming gradient) if the rhs (input feature) > 0, + // otherwise return lhs * (1 + rhs). + void Compile(XlaOpKernelContext* ctx) override { + xla::ComputationBuilder* b = ctx->builder(); + const auto zero = XlaHelpers::Zero(b, input_type(0)); + const auto one = XlaHelpers::One(b, input_type(0)); + const auto grad = ctx->Input(0); + const auto activation = ctx->Input(1); + const auto exp_grad = b->Mul(grad, b->Add(activation, one)); + const auto pred = b->Gt(activation, zero); + ctx->SetOutput(0, b->Select(pred, grad, exp_grad)); + } +}; + +REGISTER_XLA_OP(Name("Elu"), EluOp); +REGISTER_XLA_OP(Name("EluGrad"), EluGradOp); + +class SeluOp : public XlaOpKernel { + public: + explicit SeluOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {} + // Computes the max of the scalar input x and 0. + void Compile(XlaOpKernelContext* ctx) override { + xla::ComputationBuilder* b = ctx->builder(); + const auto zero = XlaHelpers::Zero(b, input_type(0)); + const auto one = XlaHelpers::One(b, input_type(0)); + const auto scale = XlaHelpers::FloatLiteral(b, input_type(0), + 1.0507009873554804934193349852946); + const auto scale_alpha = XlaHelpers::FloatLiteral(b, input_type(0), + 1.7580993408473768599402175208123); + const auto pred = b->Gt(ctx->Input(0), zero); + const auto expm1 = b->Sub(b->Exp(ctx->Input(0)), one); + ctx->SetOutput(0, b->Select(pred, b->Mul(scale, ctx->Input(0)), + b->Mul(scale_alpha, expm1))); + } +}; + +class SeluGradOp : public XlaOpKernel { + public: + explicit SeluGradOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {} + // Return the lhs (incoming gradient) if the rhs (input feature) > 0, + // otherwise return lhs * (1 + rhs). + void Compile(XlaOpKernelContext* ctx) override { + xla::ComputationBuilder* b = ctx->builder(); + const auto zero = XlaHelpers::Zero(b, input_type(0)); + const auto one = XlaHelpers::One(b, input_type(0)); + const auto scale = XlaHelpers::FloatLiteral(b, input_type(0), + 1.0507009873554804934193349852946); + const auto scale_alpha = XlaHelpers::FloatLiteral(b, input_type(0), + 1.7580993408473768599402175208123); + const auto grad = ctx->Input(0); + const auto activation = ctx->Input(1); + const auto lin_grad = b->Mul(grad, scale); + const auto exp_grad = b->Mul(grad, b->Add(activation, scale_alpha)); + const auto pred = b->Gt(activation, zero); + ctx->SetOutput(0, b->Select(pred, lin_grad, exp_grad)); + } +}; + +REGISTER_XLA_OP(Name("Selu"), SeluOp); +REGISTER_XLA_OP(Name("SeluGrad"), SeluGradOp); + +} // namespace +} // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/kernels/fill_op.cc b/tensorflow/compiler/tf2xla/kernels/fill_op.cc index 1b5f94d4e5f3816c52b7f32338a7fa12ce918419..9e090fe01cbfd4dab81b0de21e3a44e42c2ef18e 100644 --- a/tensorflow/compiler/tf2xla/kernels/fill_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/fill_op.cc @@ -50,8 +50,9 @@ class FillOp : public XlaOpKernel { // Convert the dims literal into a vector that we can pass to // ComputationBuilder. std::vector broadcast; + broadcast.reserve(dims_literal.shape().dimensions(0)); for (int i = 0; i < dims_literal.shape().dimensions(0); ++i) { - broadcast.push_back(xla::LiteralUtil::Get(dims_literal, {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). diff --git a/tensorflow/compiler/tf2xla/kernels/function_ops.cc b/tensorflow/compiler/tf2xla/kernels/function_ops.cc index d718f98545f66cb79a77d758a3fb7ee486d87b4b..af1085d5b35077b7ebd144bfb2473485e3b3de6b 100644 --- a/tensorflow/compiler/tf2xla/kernels/function_ops.cc +++ b/tensorflow/compiler/tf2xla/kernels/function_ops.cc @@ -18,6 +18,7 @@ limitations under the License. #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/framework/node_def.pb.h" namespace tensorflow { namespace { @@ -68,7 +69,8 @@ class SymbolicGradientOp : public AsyncOpKernel { done); OP_REQUIRES_OK_ASYNC( - ctx, lib->Instantiate(kGradientOp, def().attr(), &handle_), done); + ctx, lib->Instantiate(kGradientOp, AttrSlice(&def().attr()), &handle_), + done); FunctionLibraryRuntime::Options opts; opts.step_id = ctx->step_id(); diff --git a/tensorflow/compiler/tf2xla/kernels/gather_op.cc b/tensorflow/compiler/tf2xla/kernels/gather_op.cc index 49eadaf9d1f0ff1dbfa2321f20f9f833a0d4eb9a..ee59ac4d999ac6cdf54d20c80791db3121e0539b 100644 --- a/tensorflow/compiler/tf2xla/kernels/gather_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/gather_op.cc @@ -13,70 +13,229 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ +#include "tensorflow/compiler/tf2xla/kernels/gather_op.h" +#include "tensorflow/compiler/tf2xla/kernels/gather_op_helpers.h" #include "tensorflow/compiler/tf2xla/shape_util.h" +#include "tensorflow/compiler/tf2xla/type_util.h" #include "tensorflow/compiler/tf2xla/xla_context.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/framework/op_kernel.h" namespace tensorflow { + +xla::ComputationDataHandle XlaComputeGatherDynamicSlice( + XlaOpKernelContext* context, const xla::ComputationDataHandle& input, + const TensorShape& input_shape, const xla::ComputationDataHandle& indices, + const TensorShape& indices_shape, DataType dtype, + xla::ComputationBuilder* builder) { + // Although the indices Tensor is flattened into rank 1 during the lookup, + // and each scalar entry is used as an index into the first dimension of the + // input, the output is returned with shape indices.shape + input.shape[1:] + const int num_indices = indices_shape.num_elements(); + TensorShape input_shape_1(input_shape); + input_shape_1.RemoveDim(0); + + // Each slice of the input tensor is [1, ] + TensorShape slice_shape(input_shape); + slice_shape.set_dim(0, 1); + + // TODO(b/37575001) The tensor in which we construct the output during + // the loop must have rank >= 3 as a workaround for lowering issues. + int64 extra_dims = 0; + if (input_shape.dims() < 3) extra_dims = 3 - input_shape.dims(); + + TensorShape loop_out_shape; + for (int64 k = 0; k < extra_dims; ++k) loop_out_shape.AddDim(1); + loop_out_shape.AddDim(num_indices); + loop_out_shape.AppendShape(input_shape_1); + + // Slices are reshaped into the rank >= 3 shape of the loop carried output. + TensorShape loop_out_slice_shape; + for (int64 k = 0; k < extra_dims; ++k) loop_out_slice_shape.AddDim(1); + loop_out_slice_shape.AddDim(1); + loop_out_slice_shape.AppendShape(input_shape_1); + + // Finally, the loop-carried rank >= 3 output is reshaped to the op's + // specified result shape. + TensorShape out_shape(indices_shape); + out_shape.AppendShape(input_shape_1); + + // Degenerate case: empty indices. + if (num_indices == 0) { + return builder->Broadcast(XlaHelpers::Zero(builder, dtype), + out_shape.dim_sizes()); + } + + // Specify the shape of the loop-carried Tensor tuple. + xla::PrimitiveType ptype; + TF_CHECK_OK(DataTypeToPrimitiveType(dtype, &ptype)); + std::vector tuple_shapes( + {// The iteration counter i is a scalar, incremented each iteration. + xla::ShapeUtil::MakeShape(xla::S32, {}), + // The input array has shape input_shape. Loop invariant. + xla::ShapeUtil::MakeShape(ptype, input_shape.dim_sizes()), + // The gather indices are reshaped to rank 1. Loop invariant. + xla::ShapeUtil::MakeShape(xla::S32, {num_indices}), + // The output array is rank >= 3, and is updated on each loop iteration. + xla::ShapeUtil::MakeShape(ptype, loop_out_shape.dim_sizes())}); + xla::Shape tuple_shape = xla::ShapeUtil::MakeTupleShape(tuple_shapes); + + // Construct the initial values of the loop-carried Tensors. + auto init_i = builder->ConstantR0(0); + auto init_out = + builder->Broadcast(builder->ConstantLiteral(xla::Literal::Zero(ptype)), + loop_out_shape.dim_sizes()); + // Flatten the indices into 1-D for ease of iteration. + auto indices_1d = builder->Reshape(indices, {num_indices}); + auto init = builder->Tuple({init_i, input, indices_1d, init_out}); + + // Construct the while loop condition (i < num_indices) + xla::ComputationBuilder condb(context->builder()->client(), + "GatherWhileCond"); + condb.Lt(condb.GetTupleElement( + condb.Parameter(0, tuple_shape, "GatherWhileTuple"), 0), + condb.ConstantR0(num_indices)); + auto cond_status = condb.Build(); + auto cond = cond_status.ConsumeValueOrDie(); + + // Construct the while loop body's function. The implementation of gather is: + // for i in range(num_indices): + // index = dynamic-slice(indices, i) + // xi = dynamic-slice(input, index) + // output = dynamic-update-slice(output, xi, i) + xla::ComputationBuilder bodyb(context->builder()->client(), + "GatherWhileBody"); + { + // The four loop carried values. + auto loop_tuple = bodyb.Parameter(0, tuple_shape, "GatherWhileTuple"); + auto i = bodyb.GetTupleElement(loop_tuple, 0); + auto input = bodyb.GetTupleElement(loop_tuple, 1); + auto indices = bodyb.GetTupleElement(loop_tuple, 2); + auto output = bodyb.GetTupleElement(loop_tuple, 3); + + // Slice from the input array. + auto index = bodyb.DynamicSlice(indices, bodyb.Reshape(i, {1}), {1}); + auto start_indices = + XlaHelpers::PadWithZeros(&bodyb, index, input_shape.dims() - 1); + auto slice_i = bodyb.Reshape( + bodyb.DynamicSlice(input, start_indices, slice_shape.dim_sizes()), + loop_out_slice_shape.dim_sizes()); + + // Construct the index into the R3+ output Tensor 0, ..., , 0, ... + std::vector out_index_vals( + loop_out_shape.dims(), bodyb.ConstantR1({0})); + out_index_vals[extra_dims] = bodyb.Reshape(i, {1}); + auto out_index = bodyb.ConcatInDim(out_index_vals, 0); + + // Update the output Tensor + auto updated_output = bodyb.DynamicUpdateSlice(output, slice_i, out_index); + + bodyb.Tuple({bodyb.Add(i, bodyb.ConstantR0(1)), input, indices, + updated_output}); + } + auto body_status = bodyb.Build(); + auto body = body_status.ConsumeValueOrDie(); + + // Construct the While loop, extract and reshape the output. + auto gather_while = builder->While(cond, body, init); + auto gather_output = builder->GetTupleElement(gather_while, 3); + return builder->Reshape(gather_output, out_shape.dim_sizes()); +} + namespace { -class GatherOp : public XlaOpKernel { +class GatherOpCustomCall : public XlaOpKernel { public: - explicit GatherOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {} + explicit GatherOpCustomCall(OpKernelConstruction* context) + : XlaOpKernel(context) {} - void Compile(XlaOpKernelContext* ctx) override { - const TensorShape params_shape = ctx->InputShape(0); - const TensorShape indices_shape = ctx->InputShape(1); + void Compile(XlaOpKernelContext* context) override { + const TensorShape params_shape = context->InputShape(0); + const auto params_dims = params_shape.dims(); + const TensorShape indices_shape = context->InputShape(1); OP_REQUIRES( - ctx, TensorShapeUtils::IsVectorOrHigher(params_shape), + context, TensorShapeUtils::IsVectorOrHigher(params_shape), errors::InvalidArgument("params must be at least 1 dimensional")); DataType index_type = input_type(1); - OP_REQUIRES(ctx, index_type == DT_INT32 || index_type == DT_INT64, + OP_REQUIRES(context, index_type == DT_INT32 || index_type == DT_INT64, errors::InvalidArgument("index must be int32 or int64")); + // GatherV2 added an axis argument. We support both Gather and GatherV2 in + // this kernel by defaulting axis to 0 if there are 2 inputs. + int64 axis = 0; + if (context->num_inputs() == 3) { + const TensorShape axis_shape = context->InputShape(2); + OP_REQUIRES(context, TensorShapeUtils::IsScalar(axis_shape), + errors::InvalidArgument("axis must be scalar")); + DataType axis_type = input_type(2); + OP_REQUIRES(context, axis_type == DT_INT32 || axis_type == DT_INT64, + errors::InvalidArgument("axis must be int32 or int64")); + + xla::Literal literal; + OP_REQUIRES_OK(context, context->ConstantInput(2, &literal)); + int64 axis_input = axis_type == DT_INT32 ? literal.Get({}) + : literal.Get({}); + axis = axis_input < 0 ? axis_input + params_dims : axis_input; + OP_REQUIRES(context, 0 <= axis && axis < params_dims, + errors::InvalidArgument("Expected axis in the range [", + -params_dims, ", ", params_dims, + "), but got ", axis_input)); + } + // Check that we have enough index space. const int64 limit = index_type == DT_INT32 ? std::numeric_limits::max() : std::numeric_limits::max(); - OP_REQUIRES( - ctx, params_shape.dim_size(0) <= limit, - errors::InvalidArgument("params.shape[0] too large for ", - DataTypeString(index_type), " indexing: ", - params_shape.dim_size(0), " > ", limit)); - - // The result shape is indices.shape + params.shape[1:]. - TensorShape result_shape = indices_shape; - for (int i = 1; i < params_shape.dims(); i++) { + OP_REQUIRES(context, params_shape.dim_size(axis) <= limit, + errors::InvalidArgument( + "params.shape[", axis, "] too large for ", + DataTypeString(index_type), + " indexing: ", params_shape.dim_size(axis), " > ", limit)); + + // The result shape is params.shape[0:axis] + indices.shape + + // params.shape[axis + 1:]. + TensorShape result_shape; + int64 outer_size = 1; + int64 inner_size = 1; + for (int i = 0; i < axis; i++) { + result_shape.AddDim(params_shape.dim_size(i)); + outer_size *= params_shape.dim_size(i); + } + result_shape.AppendShape(indices_shape); + for (int i = axis + 1; i < params_dims; i++) { result_shape.AddDim(params_shape.dim_size(i)); + inner_size *= params_shape.dim_size(i); } - XlaContext& tc = XlaContext::Get(ctx); + XlaContext& tc = XlaContext::Get(context); OP_REQUIRES( - ctx, tc.allow_cpu_custom_calls(), + context, tc.allow_cpu_custom_calls(), errors::InvalidArgument("Gather op requires CustomCall on CPU")); - xla::ComputationBuilder& b = *ctx->builder(); + xla::ComputationBuilder& b = *context->builder(); // Call gather_xla_float_kernel (from gather_op_kernel_float.cc). // XLA passes to the function, so it is not included here. std::vector args; args.push_back(tc.GetOrCreateRuntimeContextParameter()); args.push_back(b.ConstantLiteral( - *xla::LiteralUtil::CreateR0(indices_shape.num_elements()))); + *xla::Literal::CreateR0(indices_shape.num_elements()))); + args.push_back( + b.ConstantLiteral(*xla::Literal::CreateR0(outer_size))); args.push_back(b.ConstantLiteral( - *xla::LiteralUtil::CreateR0(params_shape.dim_size(0)))); - args.push_back(b.ConstantLiteral(*xla::LiteralUtil::CreateR0( - params_shape.num_elements() / params_shape.dim_size(0)))); - args.push_back(ctx->Input(0)); - args.push_back(ctx->Input(1)); + *xla::Literal::CreateR0(params_shape.dim_size(axis)))); + args.push_back( + b.ConstantLiteral(*xla::Literal::CreateR0(inner_size))); + args.push_back(context->Input(0)); + args.push_back(context->Input(1)); xla::Shape xla_out_shape; OP_REQUIRES_OK( - ctx, TensorShapeToXLAShape(DT_FLOAT, result_shape, &xla_out_shape)); + context, TensorShapeToXLAShape(DT_FLOAT, result_shape, &xla_out_shape)); // Call the custom code with args: xla::ComputationDataHandle output; @@ -86,17 +245,41 @@ class GatherOp : public XlaOpKernel { output = b.CustomCall("gather_float_int64_xla_impl", args, xla_out_shape); } - ctx->SetOutput(0, output); + context->SetOutput(0, output); } private: - TF_DISALLOW_COPY_AND_ASSIGN(GatherOp); + TF_DISALLOW_COPY_AND_ASSIGN(GatherOpCustomCall); }; REGISTER_XLA_OP(Name("Gather") .TypeConstraint("Tparams", DT_FLOAT) .Device(DEVICE_CPU_XLA_JIT), - GatherOp); + GatherOpCustomCall); +REGISTER_XLA_OP(Name("GatherV2") + .TypeConstraint("Tparams", DT_FLOAT) + .Device(DEVICE_CPU_XLA_JIT), + GatherOpCustomCall); } // namespace + +GatherOpDynamicSlice::GatherOpDynamicSlice(OpKernelConstruction* context) + : XlaOpKernel(context) {} + +void GatherOpDynamicSlice::Compile(XlaOpKernelContext* context) { + xla::ComputationBuilder* builder = context->builder(); + auto input = context->Input(0); + auto input_shape = context->InputShape(0); + auto indices = context->Input(1); + auto indices_shape = context->InputShape(1); + xla::ComputationDataHandle gather = XlaComputeGatherDynamicSlice( + context, input, input_shape, indices, indices_shape, DT_FLOAT, builder); + context->SetOutput(0, gather); +} + +REGISTER_XLA_OP(Name("Gather") + .TypeConstraint("Tparams", DT_FLOAT) + .Device(DEVICE_GPU_XLA_JIT), + GatherOpDynamicSlice); + } // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/kernels/gather_op.h b/tensorflow/compiler/tf2xla/kernels/gather_op.h new file mode 100644 index 0000000000000000000000000000000000000000..df86e1fcdd1a4860ed7ee0c5017d25ccf9d227ea --- /dev/null +++ b/tensorflow/compiler/tf2xla/kernels/gather_op.h @@ -0,0 +1,41 @@ +/* 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. +==============================================================================*/ + +// Declaration of the Gather Op using the XLA dynamic slice implementation. + +#ifndef TENSORFLOW_COMPILER_TF2XLA_KERNELS_GATHER_OP_H_ +#define TENSORFLOW_COMPILER_TF2XLA_KERNELS_GATHER_OP_H_ + +#include "tensorflow/compiler/tf2xla/xla_op_kernel.h" +#include "tensorflow/compiler/xla/client/client_library.h" +#include "tensorflow/compiler/xla/client/computation_builder.h" +#include "tensorflow/core/framework/op_kernel.h" +#include "tensorflow/core/util/bcast.h" + +namespace tensorflow { + +class GatherOpDynamicSlice : public XlaOpKernel { + public: + explicit GatherOpDynamicSlice(OpKernelConstruction* context); + + void Compile(XlaOpKernelContext* context) override; + + private: + TF_DISALLOW_COPY_AND_ASSIGN(GatherOpDynamicSlice); +}; + +} // namespace tensorflow + +#endif // TENSORFLOW_COMPILER_TF2XLA_KERNELS_GATHER_OP_H_ diff --git a/tensorflow/compiler/tf2xla/kernels/gather_op_helpers.h b/tensorflow/compiler/tf2xla/kernels/gather_op_helpers.h new file mode 100644 index 0000000000000000000000000000000000000000..4e8d505e12ff7f377de44e1c077a34d6311fd662 --- /dev/null +++ b/tensorflow/compiler/tf2xla/kernels/gather_op_helpers.h @@ -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. +==============================================================================*/ + +// Helper methods for XLA Gather Ops. + +#ifndef TENSORFLOW_COMPILER_TF2XLA_KERNELS_GATHER_OP_HELPERS_H_ +#define TENSORFLOW_COMPILER_TF2XLA_KERNELS_GATHER_OP_HELPERS_H_ + +#include "tensorflow/compiler/tf2xla/xla_op_kernel.h" +#include "tensorflow/compiler/xla/client/client_library.h" +#include "tensorflow/compiler/xla/client/computation_builder.h" +#include "tensorflow/core/framework/op_kernel.h" +#include "tensorflow/core/util/bcast.h" + +namespace tensorflow { + +// Adds to builder an XLA computation that performs a gather on input (of +// shape input_shape) keyed on indices (of shape indices_shape). +xla::ComputationDataHandle XlaComputeGatherDynamicSlice( + XlaOpKernelContext* ctx, const xla::ComputationDataHandle& input, + const TensorShape& input_shape, const xla::ComputationDataHandle& indices, + const TensorShape& indices_shape, DataType dtype, + xla::ComputationBuilder* builder); + +} // namespace tensorflow + +#endif // TENSORFLOW_COMPILER_TF2XLA_KERNELS_GATHER_OP_HELPERS_H_ diff --git a/tensorflow/compiler/tf2xla/kernels/gather_op_kernel_float_int32.cc b/tensorflow/compiler/tf2xla/kernels/gather_op_kernel_float_int32.cc index eff23bd77d23afc882c67f8168270d1cb4413977..33b1b087d00d8263cd80f7d5d879401e4ed6c0fb 100644 --- a/tensorflow/compiler/tf2xla/kernels/gather_op_kernel_float_int32.cc +++ b/tensorflow/compiler/tf2xla/kernels/gather_op_kernel_float_int32.cc @@ -20,33 +20,37 @@ limitations under the License. #include "tensorflow/core/framework/tensor_types.h" #include "tensorflow/core/kernels/gather_functor.h" #include "tensorflow/core/platform/dynamic_annotations.h" +#include "tensorflow/core/platform/macros.h" namespace tensorflow { EIGEN_STRONG_INLINE void gather_float_int32_xla_impl(float* out, void** data) { // data is managed by the JIT code so msan can't tell it's initialized. - TF_ANNOTATE_MEMORY_IS_INITIALIZED(data, 6 * sizeof(void*)); + TF_ANNOTATE_MEMORY_IS_INITIALIZED(data, 7 * sizeof(void*)); int64 indices_size = *static_cast(data[1]); int64 params_x = *static_cast(data[2]); int64 params_y = *static_cast(data[3]); + int64 params_z = *static_cast(data[4]); - float* in = static_cast(data[4]); + float* in = static_cast(data[5]); - int32* indices = static_cast(data[5]); - Eigen::DSizes in_eig_sizes; + int32* indices = static_cast(data[6]); + Eigen::DSizes in_eig_sizes; in_eig_sizes[0] = params_x; in_eig_sizes[1] = params_y; - tensorflow::TTypes::ConstMatrix in_eig(in, in_eig_sizes); + in_eig_sizes[2] = params_z; + tensorflow::TTypes::ConstTensor in_eig(in, in_eig_sizes); Eigen::DSizes indices_eig_sizes; indices_eig_sizes[0] = indices_size; tensorflow::TTypes::ConstFlat indices_eig(indices, indices_eig_sizes); - Eigen::DSizes out_eig_sizes; - out_eig_sizes[0] = indices_size; - out_eig_sizes[1] = params_y; - tensorflow::TTypes::Matrix out_eig(out, out_eig_sizes); + Eigen::DSizes out_eig_sizes; + out_eig_sizes[0] = params_x; + out_eig_sizes[1] = indices_size; + out_eig_sizes[2] = params_z; + tensorflow::TTypes::Tensor out_eig(out, out_eig_sizes); tensorflow::functor::GatherFunctorCPU f; const int64 bad_i = f(in_eig, indices_eig, out_eig); @@ -63,7 +67,6 @@ EIGEN_STRONG_INLINE void gather_float_int32_xla_impl(float* out, void** data) { // Implements gather on CPU. This is called by an XLA custom call, set up by // gather_op.cc. -extern "C" void __attribute__((visibility("default"))) -gather_float_int32_xla_impl(float* out, void** data) { +extern "C" void TF_EXPORT gather_float_int32_xla_impl(float* out, void** data) { tensorflow::gather_float_int32_xla_impl(out, data); } diff --git a/tensorflow/compiler/tf2xla/kernels/gather_op_kernel_float_int64.cc b/tensorflow/compiler/tf2xla/kernels/gather_op_kernel_float_int64.cc index ae31f6f2006959c03941a1eb04b31aecf52424b0..5e2d872ce0b28ab479c73ed1fea5f32804c21e22 100644 --- a/tensorflow/compiler/tf2xla/kernels/gather_op_kernel_float_int64.cc +++ b/tensorflow/compiler/tf2xla/kernels/gather_op_kernel_float_int64.cc @@ -20,33 +20,37 @@ limitations under the License. #include "tensorflow/core/framework/tensor_types.h" #include "tensorflow/core/kernels/gather_functor.h" #include "tensorflow/core/platform/dynamic_annotations.h" +#include "tensorflow/core/platform/macros.h" namespace tensorflow { EIGEN_STRONG_INLINE void gather_float_int64_xla_impl(float* out, void** data) { // data is managed by the JIT code so msan can't tell it's initialized. - TF_ANNOTATE_MEMORY_IS_INITIALIZED(data, 6 * sizeof(void*)); + TF_ANNOTATE_MEMORY_IS_INITIALIZED(data, 7 * sizeof(void*)); int64 indices_size = *static_cast(data[1]); int64 params_x = *static_cast(data[2]); int64 params_y = *static_cast(data[3]); + int64 params_z = *static_cast(data[4]); - float* in = static_cast(data[4]); + float* in = static_cast(data[5]); - int64* indices = static_cast(data[5]); - Eigen::DSizes in_eig_sizes; + int64* indices = static_cast(data[6]); + Eigen::DSizes in_eig_sizes; in_eig_sizes[0] = params_x; in_eig_sizes[1] = params_y; - tensorflow::TTypes::ConstMatrix in_eig(in, in_eig_sizes); + in_eig_sizes[2] = params_z; + tensorflow::TTypes::ConstTensor in_eig(in, in_eig_sizes); Eigen::DSizes indices_eig_sizes; indices_eig_sizes[0] = indices_size; tensorflow::TTypes::ConstFlat indices_eig(indices, indices_eig_sizes); - Eigen::DSizes out_eig_sizes; - out_eig_sizes[0] = indices_size; - out_eig_sizes[1] = params_y; - tensorflow::TTypes::Matrix out_eig(out, out_eig_sizes); + Eigen::DSizes out_eig_sizes; + out_eig_sizes[0] = params_x; + out_eig_sizes[1] = indices_size; + out_eig_sizes[2] = params_z; + tensorflow::TTypes::Tensor out_eig(out, out_eig_sizes); tensorflow::functor::GatherFunctorCPU f; const int64 bad_i = f(in_eig, indices_eig, out_eig); @@ -63,7 +67,6 @@ EIGEN_STRONG_INLINE void gather_float_int64_xla_impl(float* out, void** data) { // Implements gather on CPU. This is called by an XLA custom call, set up by // gather_op.cc. -extern "C" void __attribute__((visibility("default"))) -gather_float_int64_xla_impl(float* out, void** data) { +extern "C" void TF_EXPORT gather_float_int64_xla_impl(float* out, void** data) { tensorflow::gather_float_int64_xla_impl(out, data); } diff --git a/tensorflow/compiler/tf2xla/kernels/index_ops.cc b/tensorflow/compiler/tf2xla/kernels/index_ops.cc index df002dddd043c6795481436586a31c74b20d33d1..6be66cf66ec19cad33858f36a3239048efce9de3 100644 --- a/tensorflow/compiler/tf2xla/kernels/index_ops.cc +++ b/tensorflow/compiler/tf2xla/kernels/index_ops.cc @@ -69,7 +69,7 @@ class ArgMaxOp : public XlaOpKernel { // XLA op would have the same requirement. xla::Literal literal; OP_REQUIRES_OK(ctx, ctx->ConstantInput(1, &literal)); - const int32 dim = xla::LiteralUtil::Get(literal, {}); + const int32 dim = literal.Get({}); OP_REQUIRES(ctx, dim >= 0, errors::InvalidArgument("dim must be >= 0")); OP_REQUIRES( ctx, dim < input_shape.dims(), @@ -97,14 +97,13 @@ class ArgMaxOp : public XlaOpKernel { std::vector args; args.push_back(ctx->Input(0)); args.push_back(b.ConstantLiteral( - *xla::LiteralUtil::CreateR1(input_shape.dim_sizes()))); + *xla::Literal::CreateR1(input_shape.dim_sizes()))); if (input_shape.dims() > 1) { // Don't bother passing the output shape and dim for the 1d case, since // the shape is always a scalar and the dim is always 0. args.push_back(b.ConstantLiteral( - *xla::LiteralUtil::CreateR1(output_shape.dim_sizes()))); - args.push_back( - b.ConstantLiteral(*xla::LiteralUtil::CreateR0(dim))); + *xla::Literal::CreateR1(output_shape.dim_sizes()))); + args.push_back(b.ConstantLiteral(*xla::Literal::CreateR0(dim))); } xla::Shape xla_shape = diff --git a/tensorflow/compiler/tf2xla/kernels/index_ops_kernel_argmax_float_1d.cc b/tensorflow/compiler/tf2xla/kernels/index_ops_kernel_argmax_float_1d.cc index 0033a949a372684caadce70bf46a996a942e9ec4..afbd64ca5038378d48744d6d773e0dfb1376e1f9 100644 --- a/tensorflow/compiler/tf2xla/kernels/index_ops_kernel_argmax_float_1d.cc +++ b/tensorflow/compiler/tf2xla/kernels/index_ops_kernel_argmax_float_1d.cc @@ -18,6 +18,7 @@ limitations under the License. #include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor" #include "tensorflow/core/framework/tensor_types.h" #include "tensorflow/core/platform/dynamic_annotations.h" +#include "tensorflow/core/platform/macros.h" #include "tensorflow/core/platform/types.h" namespace tensorflow { @@ -43,7 +44,6 @@ EIGEN_STRONG_INLINE void argmax_float_1d_xla_impl(void* out, void** data) { // Implements argmax on CPU. This is called by an XLA custom call, set up by // index_ops.cc. -extern "C" void __attribute__((visibility("default"))) -argmax_float_1d_xla_impl(void* out, void** data) { +extern "C" void TF_EXPORT argmax_float_1d_xla_impl(void* out, void** data) { tensorflow::argmax_float_1d_xla_impl(out, data); } diff --git a/tensorflow/compiler/tf2xla/kernels/index_ops_kernel_argmax_float_2d.cc b/tensorflow/compiler/tf2xla/kernels/index_ops_kernel_argmax_float_2d.cc index be8ad2317c9ba6a39f839c4a535440fb94365aa9..841ff2f4df79fdd790ee3aace9e38aaeb01a3080 100644 --- a/tensorflow/compiler/tf2xla/kernels/index_ops_kernel_argmax_float_2d.cc +++ b/tensorflow/compiler/tf2xla/kernels/index_ops_kernel_argmax_float_2d.cc @@ -18,6 +18,7 @@ limitations under the License. #include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor" #include "tensorflow/core/framework/tensor_types.h" #include "tensorflow/core/platform/dynamic_annotations.h" +#include "tensorflow/core/platform/macros.h" #include "tensorflow/core/platform/types.h" namespace tensorflow { @@ -45,7 +46,6 @@ EIGEN_STRONG_INLINE void argmax_float_2d_xla_impl(void* out, void** data) { // Implements argmax on CPU. This is called by an XLA custom call, set up by // index_ops.cc. -extern "C" void __attribute__((visibility("default"))) -argmax_float_2d_xla_impl(void* out, void** data) { +extern "C" void TF_EXPORT argmax_float_2d_xla_impl(void* out, void** data) { tensorflow::argmax_float_2d_xla_impl(out, data); } diff --git a/tensorflow/compiler/tf2xla/kernels/mirror_pad_op.cc b/tensorflow/compiler/tf2xla/kernels/mirror_pad_op.cc new file mode 100644 index 0000000000000000000000000000000000000000..bea1d1600b5b5fc0c44f0208d394f25061ecbb68 --- /dev/null +++ b/tensorflow/compiler/tf2xla/kernels/mirror_pad_op.cc @@ -0,0 +1,98 @@ +/* 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_helpers.h" +#include "tensorflow/compiler/tf2xla/xla_op_kernel.h" +#include "tensorflow/compiler/tf2xla/xla_op_registry.h" +#include "tensorflow/core/util/mirror_pad_mode.h" + +namespace tensorflow { +namespace { + +class MirrorPadOp : public XlaOpKernel { + public: + explicit MirrorPadOp(OpKernelConstruction* context) : XlaOpKernel(context) {} + + xla::StatusOr DoMirrorPad( + const xla::ComputationDataHandle& t, const xla::Shape& original_shape, + const xla::Literal& pad_literal, xla::ComputationBuilder* b) { + xla::ComputationDataHandle accum = t; + for (int64 dimno = xla::ShapeUtil::Rank(original_shape) - 1; dimno >= 0; + --dimno) { + auto t_rev = b->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 dim_size = original_shape.dimensions(dimno); + auto lhs_pad = b->SliceInDim(t_rev, dim_size - 1 - lhs_padding, + dim_size - 1, 1, dimno); + auto rhs_pad = b->SliceInDim(t_rev, 1, 1 + rhs_padding, 1, dimno); + accum = b->ConcatInDim({lhs_pad, accum, rhs_pad}, dimno); + } + return accum; + } + + void Compile(XlaOpKernelContext* ctx) override { + const TensorShape input_shape = ctx->InputShape(0); + const TensorShape pad_shape = ctx->InputShape(1); + + MirrorPadMode mode; + OP_REQUIRES_OK(ctx, GetNodeAttr(def(), "mode", &mode)); + OP_REQUIRES(ctx, mode == MirrorPadMode::REFLECT, + xla::Unimplemented( + "Only REFLECT MirrorPad mode is currently supported")); + + 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), + 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)); + + xla::ComputationBuilder* b = ctx->builder(); + auto in0 = ctx->Input(0); + xla::StatusOr> in0_shape = b->GetShape(in0); + OP_REQUIRES(ctx, in0_shape.ok(), in0_shape.status()); + xla::StatusOr accum_status = + DoMirrorPad(in0, *in0_shape.ValueOrDie(), pad_literal, b); + + OP_REQUIRES_OK(ctx, accum_status.status()); + + ctx->SetOutput(0, accum_status.ValueOrDie()); + } + + private: + TF_DISALLOW_COPY_AND_ASSIGN(MirrorPadOp); +}; + +REGISTER_XLA_OP(Name("MirrorPad"), MirrorPadOp); + +} // namespace +} // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/kernels/no_op.cc b/tensorflow/compiler/tf2xla/kernels/no_op.cc index b8f0c0b9fe6087a7719a689628ca4738cc13aab9..8c8a9bbe787f3224e7444b62dcf8ad99130cf37f 100644 --- a/tensorflow/compiler/tf2xla/kernels/no_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/no_op.cc @@ -23,4 +23,9 @@ namespace tensorflow { // dummy operator using CompilationOnly(). REGISTER_XLA_OP(Name("NoOp").CompilationOnly(), NoOp); +// We register ControlTrigger as a no-op. This is correct since nodes seen +// by the XLA compiler are never dead. This may need rethinking when we add +// support for conditionals to XLA. +REGISTER_XLA_OP(Name("ControlTrigger"), NoOp); + } // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/kernels/pad_op.cc b/tensorflow/compiler/tf2xla/kernels/pad_op.cc index 22476f4a0c51930cabf146313347e5e3bd2eaebe..d841bd37b33c31dbc156fa824ff62a58169a99cb 100644 --- a/tensorflow/compiler/tf2xla/kernels/pad_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/pad_op.cc @@ -60,8 +60,8 @@ class PadOp : public XlaOpKernel { xla::PaddingConfig config; for (int i = 0; i < fixed_dims; ++i) { auto* dim = config.add_dimensions(); - int before = xla::LiteralUtil::Get(pad_literal, {i, 0}); - int after = xla::LiteralUtil::Get(pad_literal, {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)); @@ -69,12 +69,22 @@ class PadOp : public XlaOpKernel { dim->set_edge_padding_high(after); } - auto zero = XlaHelpers::Zero(ctx->builder(), input_type(0)); - ctx->SetOutput(0, ctx->builder()->Pad(ctx->Input(0), zero, config)); + // PadV2 added a "constant_values" input that indicates the pad value. + xla::ComputationDataHandle 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, + ctx->builder()->Pad(ctx->Input(0), ctx->Input(2), config)); + } else { + auto zero = XlaHelpers::Zero(ctx->builder(), input_type(0)); + ctx->SetOutput(0, ctx->builder()->Pad(ctx->Input(0), zero, config)); + } } }; REGISTER_XLA_OP(Name("Pad"), PadOp); +REGISTER_XLA_OP(Name("PadV2"), PadOp); } // namespace } // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/kernels/reduction_ops.cc b/tensorflow/compiler/tf2xla/kernels/reduction_ops.cc index 518a9372c4fa3f195ff7c77e8ef0de1ba0a8807b..dae2eb9d2a92ef8d4eabb8d6f9a79758c42d446d 100644 --- a/tensorflow/compiler/tf2xla/kernels/reduction_ops.cc +++ b/tensorflow/compiler/tf2xla/kernels/reduction_ops.cc @@ -63,7 +63,7 @@ class MinOp : public XlaReductionOp { xla::ComputationBuilder* builder) override { xla::PrimitiveType type; TF_CHECK_OK(DataTypeToPrimitiveType(input_type(0), &type)); - return builder->ConstantLiteral(xla::LiteralUtil::MaxValue(type)); + return builder->ConstantLiteral(xla::Literal::MaxValue(type)); } void BuildReducer(xla::ComputationBuilder* builder, @@ -83,7 +83,7 @@ class MaxOp : public XlaReductionOp { xla::ComputationBuilder* builder) override { xla::PrimitiveType type; TF_CHECK_OK(DataTypeToPrimitiveType(input_type(0), &type)); - return builder->ConstantLiteral(xla::LiteralUtil::MinValue(type)); + return builder->ConstantLiteral(xla::Literal::MinValue(type)); } void BuildReducer(xla::ComputationBuilder* builder, diff --git a/tensorflow/compiler/tf2xla/kernels/reduction_ops_common.cc b/tensorflow/compiler/tf2xla/kernels/reduction_ops_common.cc index 8798c80ad5354c76a9b4061ad8913b76ae0629b0..4b5d09eb9fd4110cdc4221099ff55767e9132540 100644 --- a/tensorflow/compiler/tf2xla/kernels/reduction_ops_common.cc +++ b/tensorflow/compiler/tf2xla/kernels/reduction_ops_common.cc @@ -66,13 +66,13 @@ void XlaReductionOp::Compile(XlaOpKernelContext* ctx) { 1, {axes_tensor_shape.num_elements()}, &axes_literal)); VLOG(1) << "data shape: " << data_shape.DebugString(); - VLOG(1) << "axes : " << xla::LiteralUtil::ToString(axes_literal); + VLOG(1) << "axes : " << axes_literal.ToString(); gtl::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) { - int32 index = xla::LiteralUtil::Get(axes_literal, {i}); + int32 index = axes_literal.Get({i}); OP_REQUIRES(ctx, !(index < -data_shape.dims() || index >= data_shape.dims()), errors::InvalidArgument("Invalid reduction dimension (", index, diff --git a/tensorflow/compiler/tf2xla/kernels/relu_op.cc b/tensorflow/compiler/tf2xla/kernels/relu_op.cc index dc6c9c579d4df20bc77a00670abf91a1537a7b62..a137d28118e6b4c66c70253817be9b3f0b75088a 100644 --- a/tensorflow/compiler/tf2xla/kernels/relu_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/relu_op.cc @@ -31,7 +31,7 @@ class ReluOp : public XlaOpKernel { public: explicit ReluOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {} // Computes the max of the scalar input x and 0. - void Compile(XlaOpKernelContext* ctx) { + void Compile(XlaOpKernelContext* ctx) override { xla::ComputationBuilder* builder = ctx->builder(); auto zero = XlaHelpers::Zero(builder, input_type(0)); ctx->SetOutput(0, builder->Max(zero, ctx->Input(0))); @@ -42,7 +42,7 @@ class Relu6Op : public XlaOpKernel { public: explicit Relu6Op(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {} // Clamp the scalar input between 0 and 6. - void Compile(XlaOpKernelContext* ctx) { + void Compile(XlaOpKernelContext* ctx) override { xla::ComputationBuilder* builder = ctx->builder(); auto zero = XlaHelpers::Zero(builder, input_type(0)); auto six = XlaHelpers::IntegerLiteral(builder, input_type(0), 6); @@ -55,7 +55,7 @@ class ReluGradOp : public XlaOpKernel { explicit ReluGradOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {} // Return the lhs (incoming gradient) if the rhs (input feature) > 0, // otherwise return 0. - void Compile(XlaOpKernelContext* ctx) { + void Compile(XlaOpKernelContext* ctx) override { xla::ComputationBuilder* b = ctx->builder(); const TensorShape shape = ctx->InputShape(0); const auto zero = @@ -70,7 +70,7 @@ class Relu6GradOp : public XlaOpKernel { explicit Relu6GradOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {} // Return the lhs (incoming gradient) if the rhs (input feature) > 0, // otherwise return 0. - void Compile(XlaOpKernelContext* ctx) { + void Compile(XlaOpKernelContext* ctx) override { xla::ComputationBuilder* b = ctx->builder(); const TensorShape shape = ctx->InputShape(0); const auto zero = diff --git a/tensorflow/compiler/tf2xla/kernels/reshape_op.cc b/tensorflow/compiler/tf2xla/kernels/reshape_op.cc index df542350b443b765a1ab35be9632cf61a38be49c..5952e752724d1e6953dd4dbb6a8099b847c64d08 100644 --- a/tensorflow/compiler/tf2xla/kernels/reshape_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/reshape_op.cc @@ -50,7 +50,7 @@ class ReshapeOp : public XlaOpKernel { int64 product = 1; int unknown_index = -1; for (int d = 0; d < num_dims; ++d) { - const int32 size = xla::LiteralUtil::Get(literal, {d}); + const int32 size = literal.Get({d}); if (size == -1) { OP_REQUIRES( ctx, unknown_index == -1, diff --git a/tensorflow/compiler/tf2xla/kernels/scatter_op_helpers.h b/tensorflow/compiler/tf2xla/kernels/scatter_op_helpers.h new file mode 100644 index 0000000000000000000000000000000000000000..a5ab7de17adb734014fe2dcbd60ae5c219c8e486 --- /dev/null +++ b/tensorflow/compiler/tf2xla/kernels/scatter_op_helpers.h @@ -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. +==============================================================================*/ +// Helper methods for XLA Scatter Ops. +#ifndef TENSORFLOW_COMPILER_TF2XLA_KERNELS_SCATTER_OP_HELPERS_H_ +#define TENSORFLOW_COMPILER_TF2XLA_KERNELS_SCATTER_OP_HELPERS_H_ + +#include "tensorflow/compiler/tf2xla/xla_op_kernel.h" +#include "tensorflow/compiler/xla/client/client_library.h" +#include "tensorflow/compiler/xla/client/computation_builder.h" +#include "tensorflow/core/framework/op_kernel.h" +#include "tensorflow/core/util/bcast.h" + +namespace tensorflow { + +// Adds to builder an XLA computation that performs a scatter-add of input (of +// shape input_shape) keyed on indices (of shape indices_shape). The shape +// of the Tensor returned by this is num_segments input_shape[indices.dims():] +// +static xla::ComputationDataHandle XlaComputeScatterAddDynamicSlice( + XlaOpKernelContext* ctx, const xla::ComputationDataHandle& input, + const TensorShape& input_shape, const xla::ComputationDataHandle& indices, + const TensorShape& indices_shape, int64 num_segments, DataType dtype, + xla::ComputationBuilder* builder); + +} // namespace tensorflow + +#endif // TENSORFLOW_COMPILER_TF2XLA_KERNELS_SCATTER_OP_HELPERS_H_ diff --git a/tensorflow/compiler/tf2xla/kernels/segment_reduction_ops.cc b/tensorflow/compiler/tf2xla/kernels/segment_reduction_ops.cc new file mode 100644 index 0000000000000000000000000000000000000000..8a67c0b67fcd95f4841c5e011a4e51638eea5b0f --- /dev/null +++ b/tensorflow/compiler/tf2xla/kernels/segment_reduction_ops.cc @@ -0,0 +1,208 @@ +/* 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/compiler/tf2xla/kernels/cwise_ops.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_registry.h" +#include "tensorflow/compiler/xla/client/computation_builder.h" +#include "tensorflow/compiler/xla/literal_util.h" +#include "tensorflow/core/framework/kernel_def_builder.h" +#include "tensorflow/core/framework/types.h" + +namespace tensorflow { + +xla::ComputationDataHandle XlaComputeScatterAddDynamicSlice( + XlaOpKernelContext* ctx, const xla::ComputationDataHandle& input, + const TensorShape& input_shape, const xla::ComputationDataHandle& indices, + const TensorShape& indices_shape, int64 num_segments, DataType dtype, + xla::ComputationBuilder* builder) { + // Flatten data for dynamic indexing via indices_1d. + TensorShape input_shape_i(input_shape); + for (int64 d = 0; d < indices_shape.dims(); ++d) { + input_shape_i.RemoveDim(0); + } + TensorShape flat_shape({indices_shape.num_elements()}); + flat_shape.AppendShape(input_shape_i); + + // output is same as flattened input shape with dim_size(0) = num_segments. + TensorShape out_shape(flat_shape); + out_shape.set_dim(0, num_segments); + + // TODO(b/37575001) The tensor in which we construct the output during + // the loop must have rank >= 3 as a workaround for lowering issues. + int64 extra_dims = 0; + if (out_shape.dims() < 3) { + extra_dims = 3 - out_shape.dims(); + } + TensorShape loop_out_shape; + for (int64 k = 0; k < extra_dims; ++k) { + loop_out_shape.AddDim(1); + } + loop_out_shape.AppendShape(out_shape); + + // Slices from the input data are same shape as the input data, except dim 0. + TensorShape slice_shape(flat_shape); + slice_shape.set_dim(0, 1); + // slices are reshaped into the rank >= 3 shape of the loop-carried output + TensorShape loop_out_slice_shape(loop_out_shape); + loop_out_slice_shape.set_dim(extra_dims, 1); + + // Construct the initial values of the loop-carried variables + // Flatten the indices into 1-D for ease of iteration. + auto indices_1d = builder->Reshape(indices, {indices_shape.num_elements()}); + // Flatten the data for ease of indexing via values in indices_1d. + auto data_flat = builder->Reshape(input, flat_shape.dim_sizes()); + + auto init_i = builder->ConstantR0(0); + auto init_out = builder->Broadcast(XlaHelpers::Zero(builder, dtype), + loop_out_shape.dim_sizes()); + + xla::PrimitiveType ptype; + TF_CHECK_OK(DataTypeToPrimitiveType(dtype, &ptype)); + + std::vector tuple_shapes( + {// The loop iteration counter is a scalar, incremented each iteration. + xla::ShapeUtil::MakeShape(xla::S32, {}), + // The flattened input data is loop invariant. + xla::ShapeUtil::MakeShape(ptype, flat_shape.dim_sizes()), + // The scatter indices tensor is loop invariant. + xla::ShapeUtil::MakeShape(xla::S32, {indices_shape.num_elements()}), + // The output data array is updated each loop iteration. + xla::ShapeUtil::MakeShape(ptype, loop_out_shape.dim_sizes())}); + xla::Shape tuple_shape = xla::ShapeUtil::MakeTupleShape(tuple_shapes); + + auto init = builder->Tuple({init_i, data_flat, indices_1d, init_out}); + + // Construct the while loop condition (i < num_indices) + xla::ComputationBuilder condb(ctx->builder()->client(), + "ScatterAddWhileCond"); + condb.Lt(condb.GetTupleElement( + condb.Parameter(0, tuple_shape, "ScatterAddWhileTuple"), 0), + condb.ConstantR0(indices_shape.num_elements())); + auto cond_status = condb.Build(); + // TF_CHECK_OK(cond_status); + auto cond = cond_status.ConsumeValueOrDie(); + + // Construct the while loop body's function. The implementation of scatter is: + // for i in range(num_indices): + // index = dynamic-slice(indices, i) + // xi = dynamic-slice(input, i) + // output = dynamic-update-slice(output, xi, index) + xla::ComputationBuilder bodyb(ctx->builder()->client(), + "ScatterAddWhileBody"); + { + auto input_tuple = bodyb.Parameter(0, tuple_shape, "ScatterAddWhileTuple"); + auto i = bodyb.GetTupleElement(input_tuple, 0); + auto data = bodyb.GetTupleElement(input_tuple, 1); + auto idcs = bodyb.GetTupleElement(input_tuple, 2); + auto output = bodyb.GetTupleElement(input_tuple, 3); + + // Index into the data array at i. + auto zero = bodyb.ConstantR1({0}); + std::vector index_vals(flat_shape.dims(), zero); + index_vals[0] = bodyb.Reshape(i, {1}); + auto index = bodyb.ConcatInDim(index_vals, 0); + + auto data_slice = + bodyb.Reshape(bodyb.DynamicSlice(data, index, slice_shape.dim_sizes()), + loop_out_slice_shape.dim_sizes()); + + // Index into the output array. + // Construct the index into the R3+ output array 0, ..., , 0, ... + std::vector out_index_vals( + loop_out_shape.dims(), zero); + out_index_vals[extra_dims] = + bodyb.DynamicSlice(idcs, bodyb.Reshape(i, {1}), {1}); + auto out_index = bodyb.ConcatInDim(out_index_vals, 0); + + // Slice the output array, update value, and update the output slice. + auto updated_output = bodyb.DynamicUpdateSlice( + output, + bodyb.Add(data_slice, + bodyb.DynamicSlice(output, out_index, + loop_out_slice_shape.dim_sizes())), + out_index); + + auto ip1 = bodyb.Add(i, bodyb.ConstantR0(1)); + bodyb.Tuple({ip1, data, idcs, updated_output}); + } + auto body_status = bodyb.Build(); + // TF_CHECK_OK(body_status); + auto body = body_status.ConsumeValueOrDie(); + + auto gather_while = builder->While(cond, body, init); + auto updated_output = builder->GetTupleElement(gather_while, 3); + return builder->Reshape(updated_output, out_shape.dim_sizes()); +} + +namespace { + +class UnsortedSegmentSum : public XlaOpKernel { + public: + explicit UnsortedSegmentSum(OpKernelConstruction* ctx) : XlaOpKernel(ctx) { + OP_REQUIRES_OK(ctx, ctx->GetAttr("T", &dtype_)); + } + + void Compile(XlaOpKernelContext* ctx) override { + // output = unsorted_segment_sum(data, indices, num_segments) + // Compute a tensor such that: + // output[i] = sum over {j where indices[j] == i} of data[j] + // output[i] == 0 if i does not appear in indices + // + // Contrast with segment_sum(), which assumes indices are sorted and that + // max(indices)+1 is the desired size of the output. + // + // The returned output tensor has the same type as data, and the same shape + // as data with the first indices.rank dimensions are replaced + // by a single dimension with size num_segments. + auto data = ctx->Input(0); + auto data_shape = ctx->InputShape(0); + + auto indices = ctx->Input(1); + auto indices_shape = ctx->InputShape(1); + + int64 num_segments; + OP_REQUIRES_OK(ctx, ctx->ConstantInputAsIntScalar(2, &num_segments)); + + OP_REQUIRES(ctx, data_shape.dims() >= indices_shape.dims(), + errors::InvalidArgument( + "UnsortedSegmentSum requires that indices' rank be" + " less than or equal to data's rank.")); + // Validate that indices.shape is a prefix of data.shape. + for (int d = 0; d < indices_shape.dims(); ++d) { + OP_REQUIRES(ctx, (data_shape.dim_size(d) == indices_shape.dim_size(d)), + errors::InvalidArgument( + "UnsortedSegmentSum requires indices shape to be prefix" + " of data_shape, but dimension ", + d, " differs ", data_shape.dim_size(d), " vs. ", + indices_shape.dim_size(d))); + } + auto result = XlaComputeScatterAddDynamicSlice( + ctx, data, data_shape, indices, indices_shape, num_segments, dtype_, + ctx->builder()); + ctx->SetOutput(0, result); + } + + private: + DataType dtype_; +}; + +REGISTER_XLA_OP(Name("UnsortedSegmentSum"), UnsortedSegmentSum); + +} // namespace +} // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/kernels/sendrecv_ops.cc b/tensorflow/compiler/tf2xla/kernels/sendrecv_ops.cc new file mode 100644 index 0000000000000000000000000000000000000000..2a245298504874a664042fcc06c744db87b26ebd --- /dev/null +++ b/tensorflow/compiler/tf2xla/kernels/sendrecv_ops.cc @@ -0,0 +1,85 @@ +/* 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/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/computation_builder.h" +#include "tensorflow/compiler/xla/xla_data.pb.h" +#include "tensorflow/core/framework/kernel_def_builder.h" +#include "tensorflow/core/framework/types.h" + +namespace tensorflow { +namespace { + +class SendOp : public XlaOpKernel { + public: + explicit SendOp(OpKernelConstruction* ctx); + void Compile(XlaOpKernelContext* ctx) override; + + private: + string tensor_name_; + + TF_DISALLOW_COPY_AND_ASSIGN(SendOp); +}; + +SendOp::SendOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) { + OP_REQUIRES_OK(ctx, ctx->GetAttr("tensor_name", &tensor_name_)); +} + +void SendOp::Compile(XlaOpKernelContext* ctx) { + XlaCompiler* compiler = XlaContext::Get(ctx).compiler(); + xla::ChannelHandle channel; + OP_REQUIRES_OK(ctx, compiler->GetChannelHandle(tensor_name_, &channel)); + ctx->builder()->Send(ctx->Input(0), channel); + ctx->SetOpHasSideEffects(); +} + +REGISTER_XLA_OP(Name("_XLASend"), SendOp); + +class RecvOp : public XlaOpKernel { + public: + explicit RecvOp(OpKernelConstruction* ctx); + void Compile(XlaOpKernelContext* ctx) override; + + private: + string tensor_name_; + xla::Shape shape_; + + TF_DISALLOW_COPY_AND_ASSIGN(RecvOp); +}; + +RecvOp::RecvOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) { + OP_REQUIRES_OK(ctx, ctx->GetAttr("tensor_name", &tensor_name_)); + + TensorShape tensor_shape; + DataType dtype; + OP_REQUIRES_OK(ctx, ctx->GetAttr("shape", &tensor_shape)); + OP_REQUIRES_OK(ctx, ctx->GetAttr("T", &dtype)); + OP_REQUIRES_OK(ctx, TensorShapeToXLAShape(dtype, tensor_shape, &shape_)); +} + +void RecvOp::Compile(XlaOpKernelContext* ctx) { + XlaCompiler* compiler = XlaContext::Get(ctx).compiler(); + xla::ChannelHandle channel; + OP_REQUIRES_OK(ctx, compiler->GetChannelHandle(tensor_name_, &channel)); + ctx->SetOutput(0, ctx->builder()->Recv(shape_, channel)); +} + +REGISTER_XLA_OP(Name("_XLARecv"), RecvOp); + +} // namespace +} // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/kernels/sequence_ops.cc b/tensorflow/compiler/tf2xla/kernels/sequence_ops.cc index 5b6fa64fa825894b5d7bf938c5892d30f4fc11b0..c2b0e1bb4c1a141d0ab3f5b3ff5397d9da620bd8 100644 --- a/tensorflow/compiler/tf2xla/kernels/sequence_ops.cc +++ b/tensorflow/compiler/tf2xla/kernels/sequence_ops.cc @@ -32,7 +32,7 @@ template Status GetValue(int index, XlaOpKernelContext* ctx, T* value) { xla::Literal literal; TF_RETURN_IF_ERROR(ctx->ConstantInput(index, &literal)); - *value = xla::LiteralUtil::Get(literal, {}); + *value = literal.Get({}); return Status::OK(); } @@ -41,10 +41,10 @@ Status GetIntValue(int index, XlaOpKernelContext* ctx, int64* value) { TF_RETURN_IF_ERROR(ctx->ConstantInput(index, &literal)); switch (literal.shape().element_type()) { case xla::S32: - *value = xla::LiteralUtil::Get(literal, {}); + *value = literal.Get({}); break; case xla::S64: - *value = xla::LiteralUtil::Get(literal, {}); + *value = literal.Get({}); break; default: return errors::InvalidArgument("Invalid argument type for argument", @@ -58,9 +58,9 @@ template Status CreateRangeTensor(const xla::Literal& start_literal, const xla::Literal& limit_literal, const xla::Literal& delta_literal, Tensor* output) { - T start = xla::LiteralUtil::Get(start_literal, {}); - T limit = xla::LiteralUtil::Get(limit_literal, {}); - T delta = xla::LiteralUtil::Get(delta_literal, {}); + T start = start_literal.Get({}); + T limit = limit_literal.Get({}); + T delta = delta_literal.Get({}); if (delta == 0) { return errors::InvalidArgument("Requires delta != 0: ", delta); diff --git a/tensorflow/compiler/tf2xla/kernels/slice_op.cc b/tensorflow/compiler/tf2xla/kernels/slice_op.cc index 87cd266708b7b5f255cd3781f81d02be77d0d74b..482c54a40cfe2f600b36344dff091481a93417a0 100644 --- a/tensorflow/compiler/tf2xla/kernels/slice_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/slice_op.cc @@ -50,10 +50,13 @@ class SliceOp : public XlaOpKernel { // slice will be an empty handle if the output has no elements. CHECK_EQ(begin.size(), size.size()); std::vector limits; + limits.reserve(begin.size()); for (int i = 0; i < begin.size(); ++i) { limits.push_back(begin[i] + size[i]); } - ctx->SetOutput(0, ctx->builder()->Slice(ctx->Input(0), begin, limits)); + std::vector strides(begin.size(), 1); + ctx->SetOutput(0, ctx->builder()->Slice(ctx->Input(0), begin, limits, + strides)); } private: diff --git a/tensorflow/compiler/tf2xla/kernels/spacetobatch_op.cc b/tensorflow/compiler/tf2xla/kernels/spacetobatch_op.cc new file mode 100644 index 0000000000000000000000000000000000000000..83a87f19a718ce86a105e3c33ab9eaf0faff3a76 --- /dev/null +++ b/tensorflow/compiler/tf2xla/kernels/spacetobatch_op.cc @@ -0,0 +1,190 @@ +/* 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_helpers.h" +#include "tensorflow/compiler/tf2xla/xla_op_kernel.h" +#include "tensorflow/compiler/tf2xla/xla_op_registry.h" + +namespace tensorflow { +namespace { + +void SpaceToBatch(XlaOpKernelContext* ctx, + const xla::ComputationDataHandle& input, DataType input_dtype, + const TensorShape& input_tensor_shape, + gtl::ArraySlice block_shape, + const xla::Literal& paddings) { + const int input_rank = input_tensor_shape.dims(); + const gtl::InlinedVector input_shape = + input_tensor_shape.dim_sizes(); + const int block_rank = block_shape.size(); + + OP_REQUIRES( + ctx, input_rank >= 1 + block_rank, + errors::InvalidArgument("input rank should be >= ", 1 + block_rank, + " instead of ", input_rank)); + gtl::ArraySlice remainder_shape(input_shape); + remainder_shape.remove_prefix(1 + block_rank); + + OP_REQUIRES( + ctx, + xla::ShapeUtil::Rank(paddings.shape()) == 2 && + block_rank == xla::ShapeUtil::GetDimension(paddings.shape(), 0) && + 2 == xla::ShapeUtil::GetDimension(paddings.shape(), 1), + errors::InvalidArgument("paddings should have shape [", block_rank, + ", 2] instead of ", + xla::ShapeUtil::HumanString(paddings.shape()))); + + xla::ComputationBuilder* b = ctx->builder(); + + // 1. Zero-pad the start and end of dimensions `[1, ..., M]` of the + // input according to `paddings` to produce `padded` of shape `padded_shape`. + xla::PaddingConfig padding_config; + std::vector padded_shape(input_shape.begin(), input_shape.end()); + int64 block_num_elems = 1LL; + padding_config.add_dimensions(); // Don't pad the batch dimension. + for (int i = 0; i < block_rank; ++i) { + auto* dim = padding_config.add_dimensions(); + int64 pad_start = paddings.Get({i, 0}); + int64 pad_end = paddings.Get({i, 1}); + OP_REQUIRES(ctx, pad_start >= 0 && pad_end >= 0, + errors::InvalidArgument("Paddings must be non-negative")); + dim->set_edge_padding_low(pad_start); + dim->set_edge_padding_high(pad_end); + padded_shape[1 + i] += pad_start + pad_end; + block_num_elems *= block_shape[i]; + } + // Don't pad the remainder dimensions. + for (int i = 0; i < remainder_shape.size(); ++i) { + padding_config.add_dimensions(); + } + OP_REQUIRES(ctx, block_num_elems > 0, + errors::InvalidArgument( + "The product of the block dimensions must be positive")); + + xla::ComputationDataHandle padded = + b->Pad(input, XlaHelpers::Zero(b, input_dtype), padding_config); + + // 2. Reshape `padded` to `reshaped_padded` of shape: + // + // [batch] + + // [padded_shape[1] / block_shape[0], + // block_shape[0], + // ..., + // padded_shape[M] / block_shape[M-1], + // block_shape[M-1]] + + // remaining_shape + const int64 batch_size = input_shape[0]; + std::vector reshaped_padded_shape(input_rank + block_rank); + reshaped_padded_shape[0] = batch_size; + for (int i = 0; i < block_rank; ++i) { + OP_REQUIRES(ctx, padded_shape[1 + i] % block_shape[i] == 0, + errors::InvalidArgument("padded_shape[", 1 + i, + "]=", padded_shape[1 + i], + " is not divisible by block_shape[", i, + "]=", block_shape[i])); + + reshaped_padded_shape[1 + i * 2] = padded_shape[1 + i] / block_shape[i]; + reshaped_padded_shape[1 + i * 2 + 1] = block_shape[i]; + } + std::copy(remainder_shape.begin(), remainder_shape.end(), + reshaped_padded_shape.begin() + 1 + 2 * block_rank); + + xla::ComputationDataHandle reshaped_padded = + b->Reshape(padded, reshaped_padded_shape); + + // 3. Permute dimensions of `reshaped_padded` to produce + // `permuted_reshaped_padded` of shape: + // + // block_shape + + // [batch] + + // [padded_shape[1] / block_shape[0], + // ..., + // padded_shape[M] / block_shape[M-1]] + + // remaining_shape + std::vector permutation(reshaped_padded_shape.size()); + for (int i = 0; i < block_rank; ++i) { + permutation[i] = 1 + 2 * i + 1; + permutation[block_rank + 1 + i] = 1 + 2 * i; + } + permutation[block_rank] = 0; + std::iota(permutation.begin() + 1 + block_rank * 2, permutation.end(), + 1 + block_rank * 2); + xla::ComputationDataHandle permuted_reshaped_padded = + b->Transpose(reshaped_padded, permutation); + + // 4. Reshape `permuted_reshaped_padded` to flatten `block_shape` into the + // batch dimension, producing an output tensor of shape: + // + // [batch * prod(block_shape)] + + // [padded_shape[1] / block_shape[0], + // ..., + // padded_shape[M] / block_shape[M-1]] + + // remaining_shape + // Determine the length of the prefix of block dims that can be combined + // into the batch dimension due to having no padding and block_shape=1. + std::vector output_shape(input_rank); + output_shape[0] = batch_size * block_num_elems; + for (int i = 0; i < block_rank; ++i) { + output_shape[1 + i] = padded_shape[1 + i] / block_shape[i]; + } + std::copy(remainder_shape.begin(), remainder_shape.end(), + output_shape.begin() + 1 + block_rank); + + xla::ComputationDataHandle output = + b->Reshape(permuted_reshaped_padded, output_shape); + ctx->SetOutput(0, output); +} + +class SpaceToBatchNDOp : public XlaOpKernel { + public: + explicit SpaceToBatchNDOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {} + + void Compile(XlaOpKernelContext* ctx) override { + std::vector block_shape; + OP_REQUIRES_OK(ctx, ctx->ConstantInputAsIntVector(1, &block_shape)); + + xla::Literal paddings; + OP_REQUIRES_OK(ctx, ctx->ConstantInputAsInt64Literal(2, &paddings)); + + SpaceToBatch(ctx, ctx->Input(0), input_type(0), ctx->InputShape(0), + block_shape, paddings); + } +}; +REGISTER_XLA_OP(Name("SpaceToBatchND"), SpaceToBatchNDOp); + +class SpaceToBatchOp : public XlaOpKernel { + public: + explicit SpaceToBatchOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) { + OP_REQUIRES_OK(ctx, ctx->GetAttr("block_size", &block_size_)); + OP_REQUIRES( + ctx, block_size_ > 1, + errors::InvalidArgument("Block size should be > 1: ", block_size_)); + } + + void Compile(XlaOpKernelContext* ctx) override { + xla::Literal paddings; + OP_REQUIRES_OK(ctx, ctx->ConstantInputAsInt64Literal(1, &paddings)); + + SpaceToBatch(ctx, ctx->Input(0), input_type(0), ctx->InputShape(0), + {block_size_, block_size_}, paddings); + } + + private: + int block_size_; +}; +REGISTER_XLA_OP(Name("SpaceToBatch"), SpaceToBatchOp); + +} // namespace +} // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/kernels/split_op.cc b/tensorflow/compiler/tf2xla/kernels/split_op.cc index f3cec5c3ca26368d80293ea3b0f62abc8a76eeba..44ee81461e5b31f15594c0dfb86f7219f9875768 100644 --- a/tensorflow/compiler/tf2xla/kernels/split_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/split_op.cc @@ -39,7 +39,7 @@ class SplitOp : public XlaOpKernel { int32 split_dim; if (index_shape.dims() == 0) { - split_dim = xla::LiteralUtil::Get(literal_index, {}); + split_dim = literal_index.Get({}); } else { OP_REQUIRES( ctx, index_shape.dims() == 1, @@ -49,7 +49,7 @@ class SplitOp : public XlaOpKernel { 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 = xla::LiteralUtil::Get(literal_index, {0}); + split_dim = literal_index.Get({0}); } const int32 num_split = num_outputs(); const TensorShape input_shape = ctx->InputShape(1); @@ -77,14 +77,14 @@ class SplitOp : public XlaOpKernel { // The vectors we will use to define the slice. The entry for the // split dimensions varies for each output. - std::vector begin; - std::vector limits; + std::vector begin(input_shape.dims(), 0); + std::vector limits(input_shape.dims()); + std::vector strides(input_shape.dims(), 1); for (int i = 0; i < input_shape.dims(); ++i) { // Initially set up the limits to be the full size of the input: // the split dimension is filled in below. int64 dim = input_shape.dim_size(i); - begin.push_back(0); - limits.push_back(dim); + limits[i] = dim; } auto input = ctx->Input(1); @@ -94,7 +94,7 @@ class SplitOp : public XlaOpKernel { // Slice out the ith split from the split dimension. begin[split_dim] = i * slice_size; limits[split_dim] = (i + 1) * slice_size; - ctx->SetOutput(i, ctx->builder()->Slice(input, begin, limits)); + ctx->SetOutput(i, ctx->builder()->Slice(input, begin, limits, strides)); } } }; @@ -115,7 +115,7 @@ class SplitVOp : public XlaOpKernel { OP_REQUIRES(ctx, index_shape.dims() == 0, errors::InvalidArgument("split_dim input to Split Op must be a " "scalar")); - split_dim = xla::LiteralUtil::Get(literal_index, {}); + split_dim = literal_index.Get({}); xla::ComputationDataHandle input = ctx->Input(0); const TensorShape input_shape = ctx->InputShape(0); @@ -152,7 +152,7 @@ class SplitVOp : public XlaOpKernel { for (int i = 0; i < num_split; ++i) { int slice_size; - slice_size = xla::LiteralUtil::Get(split_size_literal, {i}); + slice_size = split_size_literal.Get({i}); if (slice_size == -1) { OP_REQUIRES( ctx, neg_one_dim == -1, @@ -188,7 +188,7 @@ class SplitVOp : public XlaOpKernel { std::vector begin(input_shape.dims(), 0); auto dim_sizes = input_shape.dim_sizes(); std::vector limits(dim_sizes.begin(), dim_sizes.end()); - + 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]; @@ -196,7 +196,7 @@ class SplitVOp : public XlaOpKernel { // Slice out the ith split from the split dimension. limits[split_dim] = begin[split_dim] + slice_size; - ctx->SetOutput(i, ctx->builder()->Slice(input, begin, limits)); + ctx->SetOutput(i, ctx->builder()->Slice(input, begin, limits, strides)); begin[split_dim] = limits[split_dim]; } } diff --git a/tensorflow/compiler/tf2xla/kernels/stack_ops.cc b/tensorflow/compiler/tf2xla/kernels/stack_ops.cc new file mode 100644 index 0000000000000000000000000000000000000000..5d1394c280383b7e9b9be39da4ed028e15a005fd --- /dev/null +++ b/tensorflow/compiler/tf2xla/kernels/stack_ops.cc @@ -0,0 +1,250 @@ +/* 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. +==============================================================================*/ + +// XLA Stack operators. + +#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/literal_util.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" + +namespace tensorflow { +namespace { + +Status GetStackShape(xla::ComputationBuilder* builder, XlaResource* resource, + TensorShape* stack_shape) { + auto shape_or_status = builder->GetShape(resource->value); + if (!shape_or_status.ok()) { + return shape_or_status.status(); + } + xla::Shape shape = *shape_or_status.ValueOrDie(); + TF_RET_CHECK(xla::ShapeUtil::IsTuple(shape)); + return XLAShapeToTensorShape(xla::ShapeUtil::GetTupleElementShape(shape, 0), + stack_shape); +} + +// Since the element shape is not provided to the Stack operator, +// we lazily initialize the Stack at the time of the first write. +// +// If a Stack `resource` has not been initialized, constructs storage for the +// Stack with elements of `elem_shape`. For both initialized and +// uninitialized Stacks, checks that the tensor has a type compatible with +// 'dtype' and shape compatible with 'elem_shape'. +// +// TODO(phawkins): consider changing the API of the stack operators to +// allow an optional element shape at stack construction time. +Status MaybeInitializeStack(xla::ComputationBuilder* builder, + XlaResource* resource, DataType dtype, + const TensorShape& elem_shape) { + if (resource->type != dtype) { + return errors::InvalidArgument( + "Stack dtype is ", DataTypeString(resource->type), " but op has dtype ", + DataTypeString(dtype), "."); + } + + TensorShape stack_shape; + stack_shape.AddDim(resource->tensor_array_size); + stack_shape.AppendShape(elem_shape); + + if (resource->value.handle() == 0) { + // Stack has not been initialized. + xla::ComputationDataHandle zero = XlaHelpers::Zero(builder, resource->type); + resource->value = + builder->Tuple({builder->Broadcast(zero, stack_shape.dim_sizes()), + builder->ConstantR0(0)}); + } else { + // Checks the expected shape matches the actual shape. + TensorShape actual_shape; + TF_RETURN_IF_ERROR(GetStackShape(builder, resource, &actual_shape)); + if (stack_shape != actual_shape) { + return errors::InvalidArgument( + "Mismatched Stack shapes: ", stack_shape.DebugString(), " vs ", + actual_shape.DebugString()); + } + } + return Status::OK(); +} + +class StackOp : public XlaOpKernel { + public: + explicit StackOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) { + OP_REQUIRES_OK(ctx, ctx->GetAttr("elem_type", &dtype_)); + OP_REQUIRES_OK(ctx, ctx->GetAttr("stack_name", &stack_name_)); + } + + void Compile(XlaOpKernelContext* ctx) override { + int64 size; + OP_REQUIRES_OK(ctx, ctx->ConstantInputAsIntScalar(0, &size)); + OP_REQUIRES( + ctx, size >= 0, + errors::InvalidArgument( + "XLA compilation requires a fixed stack size upper bound.")); + + // We defer initializing the Stack resource until we see the first push. + // Otherwise we do not know the shape of the stack elements. + xla::ComputationDataHandle value; + XlaContext& xc = XlaContext::Get(ctx); + XlaResource* resource; + string name = strings::StrCat("Stack: ", stack_name_); + OP_REQUIRES_OK( + ctx, xc.CreateResource(XlaResource::kStack, -1, std::move(name), dtype_, + value, &resource)); + resource->tensor_array_size = size; + ctx->SetResourceOutput(0, resource); + } + + private: + DataType dtype_; + string stack_name_; + + TF_DISALLOW_COPY_AND_ASSIGN(StackOp); +}; + +REGISTER_XLA_OP(Name("StackV2"), StackOp); + +class StackPushOp : public XlaOpKernel { + public: + explicit StackPushOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) { + OP_REQUIRES_OK(ctx, ctx->GetAttr("T", &dtype_)); + } + + void Compile(XlaOpKernelContext* ctx) override { + xla::ComputationBuilder* b = ctx->builder(); + TensorShape elem_shape = ctx->InputShape(1); + + XlaResource* resource; + OP_REQUIRES_OK(ctx, ctx->GetResourceInput(0, &resource)); + + // Initializes the Stack, if the element shape was not already known. + OP_REQUIRES_OK(ctx, MaybeInitializeStack(b, resource, dtype_, elem_shape)); + + xla::ComputationDataHandle ta = b->GetTupleElement(resource->value, 0); + xla::ComputationDataHandle index = b->GetTupleElement(resource->value, 1); + xla::ComputationDataHandle value = ctx->Input(1); + + // start_indices of the DynamicUpdateSlice are [index, 0, 0, ..., 0]. + auto start_indices = XlaHelpers::PadWithZeros(b, index, elem_shape.dims()); + + TensorShape slice_shape = elem_shape; + slice_shape.InsertDim(0, 1LL); + auto update = b->Reshape(value, slice_shape.dim_sizes()); + + // TODO(phawkins): We don't check the index is in bounds --- there is no + // error mechanism in XLA. + resource->value = + b->Tuple({b->DynamicUpdateSlice(ta, update, start_indices), + b->Add(index, b->ConstantR0(1))}); + + ctx->SetOutput(0, value); + } + + private: + DataType dtype_; + + TF_DISALLOW_COPY_AND_ASSIGN(StackPushOp); +}; + +REGISTER_XLA_OP(Name("StackPushV2"), StackPushOp); + +class StackPopOp : public XlaOpKernel { + public: + explicit StackPopOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) { + OP_REQUIRES_OK(ctx, ctx->GetAttr("elem_type", &dtype_)); + } + + void Compile(XlaOpKernelContext* ctx) override { + xla::ComputationBuilder* b = ctx->builder(); + + XlaResource* resource; + OP_REQUIRES_OK(ctx, ctx->GetResourceInput(0, &resource)); + + OP_REQUIRES(ctx, resource->type == dtype_, + errors::InvalidArgument( + "Stack dtype is ", DataTypeString(resource->type), + " but Op requested dtype ", DataTypeString(dtype_), ".")); + + // There is a somewhat subtle issue here: here "uninitialized" means we have + // not yet seen a pop in the order that we compile operators, not the order + // that we run them. However, in practice the two orders should be the same + // for the sole user of the stack operators (loop gradients). + OP_REQUIRES(ctx, resource->value.handle() != 0, + errors::InvalidArgument("Stack pop on uninitialized stack")); + + TensorShape stack_shape; + OP_REQUIRES_OK(ctx, GetStackShape(b, resource, &stack_shape)); + + xla::ComputationDataHandle state = resource->value; + xla::ComputationDataHandle ta = b->GetTupleElement(state, 0); + xla::ComputationDataHandle index = b->GetTupleElement(state, 1); + + index = b->Sub(index, b->ConstantR0(1)); + resource->value = b->Tuple({ta, index}); + + // start_indices of the DynamicSlice are [index, 0, 0, ..., 0]. + auto start_indices = + XlaHelpers::PadWithZeros(b, index, stack_shape.dims() - 1); + + auto slice_shape = stack_shape.dim_sizes(); + slice_shape[0] = 1LL; + + // TODO(phawkins): We don't check the index is in bounds --- there is no + // error mechanism in XLA. + xla::ComputationDataHandle read = + b->DynamicSlice(ta, start_indices, slice_shape); + + // Remove the leading '1' dimension. + std::vector value_shape(slice_shape.begin() + 1, slice_shape.end()); + ctx->SetOutput(0, b->Reshape(read, value_shape)); + } + + private: + DataType dtype_; + + TF_DISALLOW_COPY_AND_ASSIGN(StackPopOp); +}; + +REGISTER_XLA_OP(Name("StackPopV2"), StackPopOp); + +class StackCloseOp : public XlaOpKernel { + public: + explicit StackCloseOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {} + + void Compile(XlaOpKernelContext* ctx) override { + // Do nothing. + } + + private: + TF_DISALLOW_COPY_AND_ASSIGN(StackCloseOp); +}; + +REGISTER_XLA_OP(Name("StackCloseV2"), StackCloseOp); + +} // anonymous namespace +} // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/kernels/strided_slice_op.cc b/tensorflow/compiler/tf2xla/kernels/strided_slice_op.cc index 03e02299e33a4e2bf62e757b2092db35288b0bea..6af4bd0496e0da926726e3f74376281f539e925a 100644 --- a/tensorflow/compiler/tf2xla/kernels/strided_slice_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/strided_slice_op.cc @@ -63,42 +63,39 @@ class StridedSliceOp : public XlaOpKernel { &strides_tensor)); TensorShape dummy_processing_shape; - ShapeReadWriteFromTensorShape wrapped_final_shape(&final_shape); - ShapeReadWriteFromTensorShape wrapped_dummy_processing_shape( - &dummy_processing_shape); bool dummy = false; - OP_REQUIRES_OK( - ctx, ValidateStridedSliceOp( - &begin_tensor, &end_tensor, strides_tensor, - ShapeReadWriteFromTensorShape(&input_shape), begin_mask_, - end_mask_, ellipsis_mask_, new_axis_mask_, shrink_axis_mask_, - &wrapped_dummy_processing_shape, &wrapped_final_shape, &dummy, - &dummy, &dummy, &begin, &end, &strides)); + OP_REQUIRES_OK(ctx, + ValidateStridedSliceOp( + &begin_tensor, &end_tensor, strides_tensor, input_shape, + begin_mask_, end_mask_, ellipsis_mask_, new_axis_mask_, + shrink_axis_mask_, &dummy_processing_shape, &final_shape, + &dummy, &dummy, &dummy, &begin, &end, &strides)); gtl::InlinedVector dimensions_to_reverse; - gtl::InlinedVector slice_begin, slice_end; + gtl::InlinedVector slice_begin, slice_end, slice_strides; + for (int i = 0; i < begin.size(); ++i) { - // TODO(phawkins): implement strides != 1 when b/30878775 is fixed. - 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_end.push_back(end[i]); + slice_strides.push_back(strides[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_end.push_back(begin[i] + 1); + slice_begin.push_back(input_shape.dim_size(i) - begin[i] - 1); + slice_end.push_back(input_shape.dim_size(i) - end[i] - 1); + slice_strides.push_back(-strides[i]); dimensions_to_reverse.push_back(i); } } - xla::ComputationDataHandle slice = - ctx->builder()->Slice(ctx->Input(0), slice_begin, slice_end); + + xla::ComputationDataHandle slice = ctx->Input(0); if (!dimensions_to_reverse.empty()) { slice = ctx->builder()->Rev(slice, dimensions_to_reverse); } + slice = ctx->builder()->Slice(slice, slice_begin, slice_end, slice_strides); + slice = ctx->builder()->Reshape(slice, final_shape.dim_sizes()); ctx->SetOutput(0, slice); } @@ -145,14 +142,11 @@ class StridedSliceGradOp : public XlaOpKernel { &strides_tensor)); bool dummy = false; - ShapeReadWriteFromTensorShape wrapped_final_shape(&final_shape); - ShapeReadWriteFromTensorShape wrapped_processing_shape(&processing_shape); OP_REQUIRES_OK( ctx, ValidateStridedSliceOp( - &begin_tensor, &end_tensor, strides_tensor, - ShapeReadWriteFromTensorShape(&input_shape), begin_mask_, - end_mask_, ellipsis_mask_, new_axis_mask_, shrink_axis_mask_, - &wrapped_processing_shape, &wrapped_final_shape, &dummy, + &begin_tensor, &end_tensor, strides_tensor, input_shape, + begin_mask_, end_mask_, ellipsis_mask_, new_axis_mask_, + shrink_axis_mask_, &processing_shape, &final_shape, &dummy, &dummy, &dummy, &begin, &end, &strides)); // Check to make sure dy is consistent with the original slice @@ -219,5 +213,114 @@ class StridedSliceGradOp : public XlaOpKernel { REGISTER_XLA_OP(Name("StridedSliceGrad"), StridedSliceGradOp); +class StridedSliceAssignOp : public XlaOpKernel { + public: + explicit StridedSliceAssignOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) { + OP_REQUIRES_OK(ctx, ctx->GetAttr("begin_mask", &begin_mask_)); + OP_REQUIRES_OK(ctx, ctx->GetAttr("end_mask", &end_mask_)); + OP_REQUIRES_OK(ctx, ctx->GetAttr("ellipsis_mask", &ellipsis_mask_)); + OP_REQUIRES_OK(ctx, ctx->GetAttr("new_axis_mask", &new_axis_mask_)); + OP_REQUIRES_OK(ctx, ctx->GetAttr("shrink_axis_mask", &shrink_axis_mask_)); + OP_REQUIRES_OK(ctx, ctx->GetAttr("Index", &index_type_)); + } + + void Compile(XlaOpKernelContext* ctx) override { + TensorShape final_shape; + gtl::InlinedVector begin; + gtl::InlinedVector end; + gtl::InlinedVector strides; + + xla::Literal begin_literal, end_literal, strides_literal; + OP_REQUIRES_OK(ctx, ctx->ConstantInput(1, &begin_literal)); + OP_REQUIRES_OK(ctx, ctx->ConstantInput(2, &end_literal)); + OP_REQUIRES_OK(ctx, ctx->ConstantInput(3, &strides_literal)); + + Tensor begin_tensor, end_tensor, strides_tensor; + OP_REQUIRES_OK( + ctx, LiteralToHostTensor(begin_literal, index_type_, &begin_tensor)); + OP_REQUIRES_OK(ctx, + LiteralToHostTensor(end_literal, index_type_, &end_tensor)); + OP_REQUIRES_OK(ctx, LiteralToHostTensor(strides_literal, index_type_, + &strides_tensor)); + + DataType lhs_type; + TensorShape lhs_shape; + OP_REQUIRES_OK(ctx, ctx->GetVariableTypeAndShape(0, &lhs_type, &lhs_shape)); + + const TensorShape rhs_shape = ctx->InputShape(4); + + TensorShape dummy_processing_shape; + bool dummy = false; + OP_REQUIRES_OK(ctx, + ValidateStridedSliceOp( + &begin_tensor, &end_tensor, strides_tensor, lhs_shape, + begin_mask_, end_mask_, ellipsis_mask_, new_axis_mask_, + shrink_axis_mask_, &dummy_processing_shape, &final_shape, + &dummy, &dummy, &dummy, &begin, &end, &strides)); + + if (final_shape.num_elements() == 0 && rhs_shape.num_elements() == 0) { + // DynamicUpdateSlice does not allow 0-element updates. We should probably + // check that rhs_shape can be broadcast to final_shape, but that is + // probably better handled when implementing broadcasting more generally. + return; + } + + // TODO(aselle): This check is too strong, we only should need + // input_shape to be broadcastable to final_shape + OP_REQUIRES(ctx, final_shape == rhs_shape, + errors::Unimplemented( + "sliced l-value shape ", final_shape.DebugString(), + " does not match r-value shape ", rhs_shape.DebugString(), + ". Automatic broadcasting not yet implemented.")); + + xla::ComputationDataHandle lhs; + OP_REQUIRES_OK(ctx, ctx->ReadVariableInput(0, &lhs)); + + xla::ComputationDataHandle rhs = ctx->Input(4); + + gtl::InlinedVector dimensions_to_reverse; + gtl::InlinedVector slice_begin, 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_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_dims.push_back(begin[i] - end[i]); + dimensions_to_reverse.push_back(i); + } + } + + if (!dimensions_to_reverse.empty()) { + rhs = ctx->builder()->Rev(rhs, dimensions_to_reverse); + } + rhs = ctx->builder()->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 = ctx->builder()->DynamicUpdateSlice( + lhs, rhs, ctx->builder()->ConstantR1(slice_begin)); + } + + OP_REQUIRES_OK(ctx, ctx->AssignVariable(0, lhs_type, lhs)); + } + + private: + int32 begin_mask_, end_mask_; + int32 ellipsis_mask_, new_axis_mask_, shrink_axis_mask_; + DataType index_type_; +}; + +REGISTER_XLA_OP(Name("ResourceStridedSliceAssign"), StridedSliceAssignOp); + } // namespace } // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/kernels/tensor_array_ops.cc b/tensorflow/compiler/tf2xla/kernels/tensor_array_ops.cc new file mode 100644 index 0000000000000000000000000000000000000000..bdd7e73302083540ca5a6caaea2315a8abc538c1 --- /dev/null +++ b/tensorflow/compiler/tf2xla/kernels/tensor_array_ops.cc @@ -0,0 +1,572 @@ +/* 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. +==============================================================================*/ + +// XLA TensorArray operators. + +#include +#include + +#include "tensorflow/compiler/tf2xla/kernels/gather_op_helpers.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/literal_util.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" + +namespace tensorflow { +namespace { + +// Since the element shape is not always provided to the TensorArrayV3 operator, +// we must support lazily initialization of the TensorArray at the time of the +// first write. +// If a TensorArray `resource` has not been initialized, constructs storage for +// the TensorArray with elements of `elem_shape`. For both initialized and +// uninitialized TensorArrays, checks that the tensor has a type compatible with +// 'dtype' and shape compatible with 'elem_shape'. +Status MaybeInitializeTensorArray(xla::ComputationBuilder* builder, + XlaResource* resource, DataType dtype, + const TensorShape& elem_shape) { + if (resource->kind != XlaResource::kTensorArray) { + return errors::InvalidArgument("Unexpected non-TensorArray resource"); + } + + if (resource->type != dtype) { + return errors::InvalidArgument( + "TensorArray dtype is ", DataTypeString(resource->type), + " but op has dtype ", DataTypeString(dtype), "."); + } + + TF_RET_CHECK(resource->tensor_array_size >= 0) + << resource->name << " size " << resource->tensor_array_size; + TensorShape ta_shape; + ta_shape.AddDim(resource->tensor_array_size); + ta_shape.AppendShape(elem_shape); + + if (resource->value.handle() == 0) { + // TensorArray has not been initialized. + xla::ComputationDataHandle zero = XlaHelpers::Zero(builder, resource->type); + resource->value = builder->Broadcast(zero, ta_shape.dim_sizes()); + } else { + // Checks the elem_shape matches the TensorArray shape. + auto shape_or_status = builder->GetShape(resource->value); + if (!shape_or_status.ok()) { + return shape_or_status.status(); + } + TensorShape shape; + TF_RETURN_IF_ERROR( + XLAShapeToTensorShape(*shape_or_status.ValueOrDie(), &shape)); + if (ta_shape != shape) { + return errors::InvalidArgument( + "Mismatched TensorArray sizes: ", ta_shape.DebugString(), " vs ", + shape.DebugString()); + } + } + return Status::OK(); +} + +// Checks that the TensorArray 'resource' has been initialized, and has type +// 'dtype'. Sets 'shape' to the shape +Status CheckTensorArrayIsInitialized(const string& op_name, + const XlaResource* resource, + DataType dtype) { + if (resource->kind != XlaResource::kTensorArray) { + return errors::InvalidArgument( + "Unexpected non-TensorArray resource passed " + "to ", + op_name); + } + if (resource->value.handle() == 0) { + return errors::InvalidArgument("Uninitialized TensorArray passed to ", + op_name); + } + if (resource->type != dtype) { + return errors::InvalidArgument( + "TensorArray dtype is ", DataTypeString(resource->type), + " but op has dtype ", DataTypeString(dtype), "."); + } + + return Status::OK(); +} + +Status GetTensorArrayShape(const XlaResource* resource, + xla::ComputationBuilder* builder, + TensorShape* shape) { + auto shape_or_status = builder->GetShape(resource->value); + if (!shape_or_status.ok()) { + return shape_or_status.status(); + } + TF_RETURN_IF_ERROR( + XLAShapeToTensorShape(*shape_or_status.ValueOrDie(), shape)); + if (shape->dims() < 1) { + return errors::InvalidArgument("TensorArray rank must be >= 1"); + } + return Status::OK(); +} + +// Like ComputationBuilder::DynamicUpdateSlice, but adds 'update' to the +// relevant slice of 'operand'. +xla::ComputationDataHandle DynamicAddSlice( + xla::ComputationBuilder* builder, const xla::ComputationDataHandle& operand, + const xla::ComputationDataHandle& update, + const gtl::ArraySlice& update_dims, + const xla::ComputationDataHandle& start_indices) { + xla::ComputationDataHandle current = + builder->DynamicSlice(operand, start_indices, update_dims); + xla::ComputationDataHandle sum = builder->Add(current, update); + return builder->DynamicUpdateSlice(operand, sum, start_indices); +} + +class TensorArrayOp : public XlaOpKernel { + public: + explicit TensorArrayOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) { + OP_REQUIRES_OK(ctx, ctx->GetAttr("element_shape", &element_shape_)); + OP_REQUIRES_OK(ctx, ctx->GetAttr("dtype", &dtype_)); + bool dynamic_size; + OP_REQUIRES_OK(ctx, ctx->GetAttr("dynamic_size", &dynamic_size)); + OP_REQUIRES( + ctx, !dynamic_size, + errors::Unimplemented( + "TensorArrays with dynamic size are not supported by XLA.")); + + OP_REQUIRES_OK(ctx, ctx->GetAttr("tensor_array_name", &tensor_array_name_)); + } + + void Compile(XlaOpKernelContext* ctx) override { + int64 size; + OP_REQUIRES_OK(ctx, ctx->ConstantInputAsIntScalar(0, &size)); + OP_REQUIRES(ctx, size >= 0, + errors::InvalidArgument("TensorArray size must be >= 0")); + + xla::ComputationBuilder* b = ctx->builder(); + + // Initializes the TensorArray value if we know the element shape. + // Otherwise, defer initialization to the first write. + xla::ComputationDataHandle value; + if (element_shape_.IsFullyDefined()) { + TensorShape shape; + CHECK(element_shape_.AsTensorShape(&shape)); + TensorShape ta_shape; + ta_shape.AddDim(size); + ta_shape.AppendShape(shape); + xla::ComputationDataHandle zero = XlaHelpers::Zero(b, dtype_); + value = b->Broadcast(zero, ta_shape.dim_sizes()); + } + + XlaContext& xc = XlaContext::Get(ctx); + XlaResource* var; + string name = strings::StrCat("TensorArray: ", tensor_array_name_); + OP_REQUIRES_OK( + ctx, xc.CreateResource(XlaResource::kTensorArray, -1, std::move(name), + dtype_, value, &var)); + var->tensor_array_size = size; + ctx->SetResourceOutput(0, var); + + Tensor flow(DT_FLOAT, TensorShape({})); + flow.scalar()() = 0.0f; + ctx->SetConstantOutput(1, flow); + } + + private: + PartialTensorShape element_shape_; + DataType dtype_; + string tensor_array_name_; + + TF_DISALLOW_COPY_AND_ASSIGN(TensorArrayOp); +}; + +REGISTER_XLA_OP(Name("TensorArrayV3"), TensorArrayOp); + +class TensorArrayWriteOp : public XlaOpKernel { + public: + explicit TensorArrayWriteOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) { + OP_REQUIRES_OK(ctx, ctx->GetAttr("T", &dtype_)); + } + + void Compile(XlaOpKernelContext* ctx) override { + xla::ComputationBuilder* b = ctx->builder(); + + TensorShape elem_shape = ctx->InputShape(2); + + // Initializes the TensorArray, if the element shape was not known at + // construction time. + XlaResource* resource; + OP_REQUIRES_OK(ctx, ctx->GetResourceInput(0, &resource)); + OP_REQUIRES_OK(ctx, + MaybeInitializeTensorArray(b, resource, dtype_, elem_shape)); + + xla::ComputationDataHandle ta = resource->value; + xla::ComputationDataHandle index = ctx->Input(1); + xla::ComputationDataHandle value = ctx->Input(2); + xla::ComputationDataHandle flow = ctx->Input(3); + + // start_indices of the DynamicUpdateSlice are [index, 0, 0, ..., 0]. + auto start_indices = XlaHelpers::PadWithZeros(b, index, elem_shape.dims()); + + TensorShape slice_shape = elem_shape; + slice_shape.InsertDim(0, 1LL); + auto update = b->Reshape(value, slice_shape.dim_sizes()); + + xla::ComputationDataHandle written = + DynamicAddSlice(b, ta, update, slice_shape.dim_sizes(), start_indices); + + resource->value = written; + ctx->SetOutput(0, flow); + } + + private: + DataType dtype_; + + TF_DISALLOW_COPY_AND_ASSIGN(TensorArrayWriteOp); +}; + +REGISTER_XLA_OP(Name("TensorArrayWriteV3"), TensorArrayWriteOp); + +class TensorArrayReadOp : public XlaOpKernel { + public: + explicit TensorArrayReadOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) { + OP_REQUIRES_OK(ctx, ctx->GetAttr("dtype", &dtype_)); + } + + void Compile(XlaOpKernelContext* ctx) override { + xla::ComputationBuilder* b = ctx->builder(); + + XlaResource* resource; + OP_REQUIRES_OK(ctx, ctx->GetResourceInput(0, &resource)); + + OP_REQUIRES_OK(ctx, + CheckTensorArrayIsInitialized(name(), resource, dtype_)); + TensorShape ta_shape; + OP_REQUIRES_OK(ctx, GetTensorArrayShape(resource, b, &ta_shape)); + + xla::ComputationDataHandle ta = resource->value; + xla::ComputationDataHandle index = ctx->Input(1); + + // start_indices of the DynamicSlice are [index, 0, 0, ..., 0]. + auto start_indices = + XlaHelpers::PadWithZeros(b, index, ta_shape.dims() - 1); + + auto slice_shape = ta_shape.dim_sizes(); + slice_shape[0] = 1LL; + + xla::ComputationDataHandle read = + b->DynamicSlice(ta, start_indices, slice_shape); + + // Remove the leading '1' dimension. + std::vector value_shape(slice_shape.begin() + 1, slice_shape.end()); + ctx->SetOutput(0, b->Reshape(read, value_shape)); + } + + private: + DataType dtype_; + + TF_DISALLOW_COPY_AND_ASSIGN(TensorArrayReadOp); +}; + +REGISTER_XLA_OP(Name("TensorArrayReadV3"), TensorArrayReadOp); + +class TensorArrayGatherOp : public XlaOpKernel { + public: + explicit TensorArrayGatherOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) { + OP_REQUIRES_OK(ctx, ctx->GetAttr("dtype", &dtype_)); + } + + void Compile(XlaOpKernelContext* ctx) override { + xla::ComputationBuilder* b = ctx->builder(); + + XlaResource* resource; + OP_REQUIRES_OK(ctx, ctx->GetResourceInput(0, &resource)); + + OP_REQUIRES_OK(ctx, + CheckTensorArrayIsInitialized(name(), resource, dtype_)); + TensorShape ta_shape; + OP_REQUIRES_OK(ctx, GetTensorArrayShape(resource, b, &ta_shape)); + + const TensorShape indices_shape = ctx->InputShape(1); + OP_REQUIRES(ctx, indices_shape.dims() == 1, + errors::InvalidArgument("indices must be rank 1")); + auto indices = ctx->Input(1); + + xla::ComputationDataHandle ta = resource->value; + + xla::ComputationDataHandle gather = XlaComputeGatherDynamicSlice( + ctx, ta, ta_shape, indices, indices_shape, dtype_, b); + ctx->SetOutput(0, gather); + } + + private: + DataType dtype_; + + TF_DISALLOW_COPY_AND_ASSIGN(TensorArrayGatherOp); +}; + +REGISTER_XLA_OP(Name("TensorArrayGatherV3"), TensorArrayGatherOp); + +class TensorArrayScatterOp : public XlaOpKernel { + public: + explicit TensorArrayScatterOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) { + OP_REQUIRES_OK(ctx, ctx->GetAttr("T", &dtype_)); + } + + void Compile(XlaOpKernelContext* ctx) override { + xla::ComputationBuilder* b = ctx->builder(); + + const TensorShape value_shape = ctx->InputShape(2); + + XlaResource* resource; + OP_REQUIRES_OK(ctx, ctx->GetResourceInput(0, &resource)); + TensorShape elem_shape = value_shape; + elem_shape.RemoveDim(0); + OP_REQUIRES_OK(ctx, + MaybeInitializeTensorArray(b, resource, dtype_, elem_shape)); + + const TensorShape indices_shape = ctx->InputShape(1); + OP_REQUIRES(ctx, indices_shape.dims() >= 1, + errors::InvalidArgument("indices must be rank 1")); + const int num_indices = indices_shape.dim_size(0); + const xla::ComputationDataHandle indices = ctx->Input(1); + + xla::ComputationDataHandle ta = resource->value; + const xla::ComputationDataHandle value = ctx->Input(2); + const xla::ComputationDataHandle flow = ctx->Input(3); + + auto slice_dims = value_shape.dim_sizes(); + slice_dims[0] = 1LL; + + std::vector value_starts(value_shape.dims(), 0); + auto value_ends = value_shape.dim_sizes(); + + std::vector value_strides(value_shape.dims(), 1); + + // For every (index, value) pair, update the corresponding TensorArray + // storage. + for (int i = 0; i < num_indices; ++i) { + // Slice out part of the value. + value_starts[0] = i; + value_ends[0] = i + 1; + auto slice = b->Slice(value, value_starts, value_ends, value_strides); + + // start_indices of the DynamicUpdateSlice are [index, 0, 0, ..., 0]. + auto index = b->Slice(indices, {i}, {i + 1}, {1}); + auto start_indices = + XlaHelpers::PadWithZeros(b, index, elem_shape.dims()); + ta = DynamicAddSlice(b, ta, slice, slice_dims, start_indices); + } + + resource->value = ta; + ctx->SetOutput(0, flow); + } + + private: + DataType dtype_; + + TF_DISALLOW_COPY_AND_ASSIGN(TensorArrayScatterOp); +}; + +REGISTER_XLA_OP(Name("TensorArrayScatterV3"), TensorArrayScatterOp); + +class TensorArrayConcatOp : public XlaOpKernel { + public: + explicit TensorArrayConcatOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) { + OP_REQUIRES_OK(ctx, ctx->GetAttr("dtype", &dtype_)); + } + + void Compile(XlaOpKernelContext* ctx) override { + xla::ComputationBuilder* b = ctx->builder(); + + XlaResource* resource; + OP_REQUIRES_OK(ctx, ctx->GetResourceInput(0, &resource)); + + OP_REQUIRES_OK(ctx, + CheckTensorArrayIsInitialized(name(), resource, dtype_)); + TensorShape ta_shape; + OP_REQUIRES_OK(ctx, GetTensorArrayShape(resource, b, &ta_shape)); + + xla::ComputationDataHandle ta = resource->value; + + auto ta_dims = ta_shape.dim_sizes(); + std::vector shape(ta_dims.begin() + 1, ta_dims.end()); + shape[0] *= ta_shape.dim_size(0); + ctx->SetOutput(0, b->Reshape(ta, shape)); + + Tensor lengths(DT_INT64, {ta_dims[0]}); + auto lengths_vec = lengths.vec(); + for (int i = 0; i < ta_dims[0]; ++i) { + lengths_vec(i) = ta_dims[1]; + } + ctx->SetConstantOutput(1, lengths); + } + + private: + DataType dtype_; + + TF_DISALLOW_COPY_AND_ASSIGN(TensorArrayConcatOp); +}; + +REGISTER_XLA_OP(Name("TensorArrayConcatV3"), TensorArrayConcatOp); + +class TensorArraySplitOp : public XlaOpKernel { + public: + explicit TensorArraySplitOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) { + OP_REQUIRES_OK(ctx, ctx->GetAttr("T", &dtype_)); + } + + void Compile(XlaOpKernelContext* ctx) override { + std::vector lengths; + OP_REQUIRES_OK(ctx, ctx->ConstantInputAsIntVector(2, &lengths)); + + int64 length = 0; + if (!lengths.empty()) { + length = lengths[0]; + for (int i = 1; i < lengths.size(); ++i) { + OP_REQUIRES(ctx, lengths[i] == length, + errors::InvalidArgument("lengths must be equal: ", length, + " vs. ", lengths[i])); + } + } + + TensorShape value_shape = ctx->InputShape(1); + OP_REQUIRES(ctx, value_shape.dims() >= 1, + errors::InvalidArgument("value must have rank >= 1, got ", + value_shape.DebugString())); + TensorShape elem_shape = value_shape; + elem_shape.set_dim(0, length); + + xla::ComputationBuilder* b = ctx->builder(); + XlaResource* resource; + OP_REQUIRES_OK(ctx, ctx->GetResourceInput(0, &resource)); + OP_REQUIRES_OK(ctx, + MaybeInitializeTensorArray(b, resource, dtype_, elem_shape)); + xla::ComputationDataHandle ta = resource->value; + + TensorShape ta_shape; + ta_shape.AddDim(resource->tensor_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, ")")); + + const xla::ComputationDataHandle value = ctx->Input(1); + const xla::ComputationDataHandle flow = ctx->Input(3); + + OP_REQUIRES(ctx, value_shape.num_elements() == ta_shape.num_elements(), + errors::InvalidArgument("mismatched element count ", + value_shape.DebugString(), " vs. ", + ta_shape.DebugString())); + + resource->value = b->Add(ta, b->Reshape(value, ta_shape.dim_sizes())); + + ctx->SetOutput(0, flow); + } + + private: + DataType dtype_; + + TF_DISALLOW_COPY_AND_ASSIGN(TensorArraySplitOp); +}; + +REGISTER_XLA_OP(Name("TensorArraySplitV3"), TensorArraySplitOp); + +class TensorArraySizeOp : public XlaOpKernel { + public: + explicit TensorArraySizeOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {} + + void Compile(XlaOpKernelContext* ctx) override { + XlaResource* var; + OP_REQUIRES_OK(ctx, ctx->GetResourceInput(0, &var)); + Tensor size_tensor(DT_INT32, {}); + size_tensor.scalar()() = static_cast(var->tensor_array_size); + ctx->SetConstantOutput(0, size_tensor); + } + + private: + TF_DISALLOW_COPY_AND_ASSIGN(TensorArraySizeOp); +}; + +REGISTER_XLA_OP(Name("TensorArraySizeV3"), TensorArraySizeOp); + +class TensorArrayGradOp : public XlaOpKernel { + public: + explicit TensorArrayGradOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) { + OP_REQUIRES_OK(ctx, ctx->GetAttr("source", &source_)); + } + + void Compile(XlaOpKernelContext* ctx) override { + xla::ComputationBuilder* b = ctx->builder(); + + XlaResource* resource; + OP_REQUIRES_OK(ctx, ctx->GetResourceInput(0, &resource)); + + OP_REQUIRES_OK( + ctx, CheckTensorArrayIsInitialized(name(), resource, resource->type)); + TensorShape ta_shape; + OP_REQUIRES_OK(ctx, GetTensorArrayShape(resource, b, &ta_shape)); + + // Finds or looks up the corresponding gradient TensorArray, which stores + // gradients computed during backpropagation. + XlaResource*& gradient = resource->tensor_array_gradient[source_]; + if (!gradient) { + xla::ComputationDataHandle zero = XlaHelpers::Zero(b, resource->type); + xla::ComputationDataHandle value = + b->Broadcast(zero, ta_shape.dim_sizes()); + + XlaContext& xc = XlaContext::Get(ctx); + string name = strings::StrCat("TensorArrayGrad: ", resource->name); + OP_REQUIRES_OK( + ctx, xc.CreateResource(XlaResource::kTensorArray, -1, std::move(name), + resource->type, value, &gradient)); + gradient->tensor_array_size = resource->tensor_array_size; + } + + ctx->SetResourceOutput(0, gradient); + ctx->SetConstantOutput(1, Tensor(DT_FLOAT)); + } + + private: + string source_; + + TF_DISALLOW_COPY_AND_ASSIGN(TensorArrayGradOp); +}; + +REGISTER_XLA_OP(Name("TensorArrayGradV3"), TensorArrayGradOp); + +class TensorArrayCloseOp : public XlaOpKernel { + public: + explicit TensorArrayCloseOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {} + + void Compile(XlaOpKernelContext* ctx) override { + // Do nothing; XLA handles resource management. + } + + private: + TF_DISALLOW_COPY_AND_ASSIGN(TensorArrayCloseOp); +}; + +REGISTER_XLA_OP(Name("TensorArrayCloseV3"), TensorArrayCloseOp); + +} // anonymous namespace +} // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/kernels/tile_ops.cc b/tensorflow/compiler/tf2xla/kernels/tile_ops.cc index 4cc2eb8f877a873593f0460346e3379e851e8e08..9ee6bd892504e683a191484fb09259619759f36d 100644 --- a/tensorflow/compiler/tf2xla/kernels/tile_ops.cc +++ b/tensorflow/compiler/tf2xla/kernels/tile_ops.cc @@ -68,7 +68,7 @@ class TileOp : public XlaOpKernel { bool all_multiples_are_one = true; bool one_dimension_is_broadcasted_without_multiple = true; for (int i = 0; i < input_dims; ++i) { - int multiple = xla::LiteralUtil::Get(literal, {i}); + int multiple = literal.Get({i}); OP_REQUIRES(ctx, multiple, errors::InvalidArgument("Expected multiples[", i, "] >= 0, but got ", multiple)); diff --git a/tensorflow/compiler/tf2xla/kernels/training_ops.cc b/tensorflow/compiler/tf2xla/kernels/training_ops.cc index f1d81f871423b220c6859c1dedf79b1c36a43e65..82ae0df5cc501cf1b51c2b25b9330d582fbdc44c 100644 --- a/tensorflow/compiler/tf2xla/kernels/training_ops.cc +++ b/tensorflow/compiler/tf2xla/kernels/training_ops.cc @@ -165,6 +165,106 @@ class ResourceApplyAdagrad : public XlaOpKernel { }; REGISTER_XLA_OP(Name("ResourceApplyAdagrad"), ResourceApplyAdagrad); +class ResourceApplyAdam : public XlaOpKernel { + public: + explicit ResourceApplyAdam(OpKernelConstruction* ctx) : XlaOpKernel(ctx) { + OP_REQUIRES_OK(ctx, ctx->GetAttr("T", &dtype_)); + } + + void Compile(XlaOpKernelContext* ctx) override { + DataType var_type, m_type, v_type; + TensorShape var_shape, m_shape, v_shape; + OP_REQUIRES_OK(ctx, ctx->GetVariableTypeAndShape(0, &var_type, &var_shape)); + OP_REQUIRES_OK(ctx, ctx->GetVariableTypeAndShape(1, &m_type, &m_shape)); + OP_REQUIRES_OK(ctx, ctx->GetVariableTypeAndShape(2, &v_type, &v_shape)); + + OP_REQUIRES( + ctx, dtype_ == var_type && dtype_ == m_type && dtype_ == v_type, + errors::InvalidArgument( + "Types of variable arguments to ResourceApplyRMSProp must match: ", + DataTypeString(dtype_), " vs. ", DataTypeString(var_type), " vs. ", + DataTypeString(m_type), " vs. ", DataTypeString(v_type))); + + TensorShape beta1_power_shape = ctx->InputShape(3); + TensorShape beta2_power_shape = ctx->InputShape(4); + TensorShape lr_shape = ctx->InputShape(5); + TensorShape beta1_shape = ctx->InputShape(6); + TensorShape beta2_shape = ctx->InputShape(7); + TensorShape epsilon_shape = ctx->InputShape(8); + TensorShape grad_shape = ctx->InputShape(9); + + OP_REQUIRES(ctx, TensorShapeUtils::IsScalar(beta1_power_shape), + errors::InvalidArgument("beta1_power is not a scalar: ", + beta1_power_shape.DebugString())); + OP_REQUIRES(ctx, TensorShapeUtils::IsScalar(beta2_power_shape), + errors::InvalidArgument("beta2_power is not a scalar: ", + beta2_power_shape.DebugString())); + OP_REQUIRES(ctx, TensorShapeUtils::IsScalar(lr_shape), + errors::InvalidArgument("lr is not a scalar : ", + lr_shape.DebugString())); + OP_REQUIRES(ctx, TensorShapeUtils::IsScalar(beta1_shape), + errors::InvalidArgument("beta1 is not a scalar: ", + beta1_shape.DebugString())); + OP_REQUIRES(ctx, TensorShapeUtils::IsScalar(beta2_shape), + errors::InvalidArgument("beta2 is not a scalar: ", + beta2_shape.DebugString())); + OP_REQUIRES(ctx, TensorShapeUtils::IsScalar(epsilon_shape), + errors::InvalidArgument("epsilon is not a scalar: ", + epsilon_shape.DebugString())); + + OP_REQUIRES(ctx, var_shape.IsSameSize(m_shape), + errors::InvalidArgument("var and m do not have the same shape", + var_shape.DebugString(), " ", + m_shape.DebugString())); + OP_REQUIRES(ctx, var_shape.IsSameSize(v_shape), + errors::InvalidArgument("var and v do not have the same shape", + var_shape.DebugString(), " ", + v_shape.DebugString())); + 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())); + + xla::ComputationDataHandle var, m, v; + OP_REQUIRES_OK(ctx, ctx->ReadVariableInput(0, &var)); + OP_REQUIRES_OK(ctx, ctx->ReadVariableInput(1, &m)); + OP_REQUIRES_OK(ctx, ctx->ReadVariableInput(2, &v)); + xla::ComputationDataHandle beta1_power = ctx->Input(3); + xla::ComputationDataHandle beta2_power = ctx->Input(4); + xla::ComputationDataHandle lr = ctx->Input(5); + xla::ComputationDataHandle beta1 = ctx->Input(6); + xla::ComputationDataHandle beta2 = ctx->Input(7); + xla::ComputationDataHandle epsilon = ctx->Input(8); + xla::ComputationDataHandle grad = ctx->Input(9); + + // alpha <- learning_rate * sqrt(1 - beta2^t) / (1 - beta1^t) + // m_t <- beta1 * m_{t-1} + (1 - beta1) * g_t + // v_t <- beta2 * v_{t-1} + (1 - beta2) * g_t * g_t + // variable <- variable - alpha * m_t / (sqrt(v_t) + epsilon) + + xla::ComputationBuilder* b = ctx->builder(); + xla::ComputationDataHandle half = XlaHelpers::FloatLiteral(b, dtype_, 0.5); + xla::ComputationDataHandle one = XlaHelpers::FloatLiteral(b, dtype_, 1.0); + xla::ComputationDataHandle two = XlaHelpers::FloatLiteral(b, dtype_, 2.0); + + xla::ComputationDataHandle alpha = + b->Div(b->Mul(lr, b->Pow(b->Sub(one, beta2_power), half)), + b->Sub(one, beta1_power)); + m = b->Add(m, b->Mul(b->Sub(grad, m), b->Sub(one, beta1))); + v = b->Add(v, b->Mul(b->Sub(b->Pow(grad, two), v), b->Sub(one, beta2))); + var = + b->Sub(var, b->Div(b->Mul(m, alpha), b->Add(b->Pow(v, half), epsilon))); + + OP_REQUIRES_OK(ctx, ctx->AssignVariable(0, dtype_, var)); + OP_REQUIRES_OK(ctx, ctx->AssignVariable(1, dtype_, m)); + OP_REQUIRES_OK(ctx, ctx->AssignVariable(2, dtype_, v)); + } + + private: + DataType dtype_; +}; +REGISTER_XLA_OP(Name("ResourceApplyAdam"), ResourceApplyAdam); + class ResourceApplyRMSProp : public XlaOpKernel { public: explicit ResourceApplyRMSProp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {} @@ -252,9 +352,9 @@ class ResourceApplyRMSProp : public XlaOpKernel { b->Sub(XlaHelpers::FloatLiteral(b, type, 1.0), rho))); xla::ComputationDataHandle new_mom = b->Add(b->Mul(mom, momentum), - b->Div(b->Mul(grad, lr), + b->Mul(b->Mul(grad, lr), b->Pow(b->Add(new_ms, epsilon), - XlaHelpers::FloatLiteral(b, type, 0.5)))); + XlaHelpers::FloatLiteral(b, type, -0.5)))); xla::ComputationDataHandle new_var = b->Sub(var, new_mom); OP_REQUIRES_OK(ctx, ctx->AssignVariable(0, type, new_var)); @@ -264,5 +364,158 @@ class ResourceApplyRMSProp : public XlaOpKernel { }; REGISTER_XLA_OP(Name("ResourceApplyRMSProp"), ResourceApplyRMSProp); +void CompileFtrl(XlaOpKernelContext* ctx, DataType dtype, + bool has_l2_shrinkage) { + xla::ComputationBuilder* b = ctx->builder(); + + DataType var_type, accum_type, linear_type; + TensorShape var_shape, accum_shape, linear_shape; + OP_REQUIRES_OK(ctx, ctx->GetVariableTypeAndShape(0, &var_type, &var_shape)); + OP_REQUIRES_OK(ctx, + ctx->GetVariableTypeAndShape(1, &accum_type, &accum_shape)); + OP_REQUIRES_OK(ctx, + ctx->GetVariableTypeAndShape(2, &linear_type, &linear_shape)); + + OP_REQUIRES( + ctx, dtype == var_type && dtype == accum_type && dtype == linear_type, + errors::InvalidArgument( + "Types of variable arguments to ResourceApplyFtrlV2 must match: ", + DataTypeString(dtype), " vs. ", DataTypeString(var_type), " and ", + DataTypeString(accum_type), " and ", DataTypeString(linear_type))); + + 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())); + + OP_REQUIRES(ctx, var_shape.IsSameSize(linear_shape), + errors::InvalidArgument( + "var and linear do not have the same shape", + var_shape.DebugString(), " ", linear_shape.DebugString())); + + TensorShape grad_shape = ctx->InputShape(3); + TensorShape lr_shape = ctx->InputShape(4); + TensorShape l1_shape = ctx->InputShape(5); + TensorShape l2_shape = ctx->InputShape(6); + TensorShape l2_shrinkage_shape; + TensorShape lr_power_shape; + if (has_l2_shrinkage) { + l2_shrinkage_shape = ctx->InputShape(7); + lr_power_shape = ctx->InputShape(8); + } else { + lr_power_shape = ctx->InputShape(7); + } + + 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())); + + OP_REQUIRES( + ctx, TensorShapeUtils::IsScalar(lr_shape), + errors::InvalidArgument("lr is not a scalar: ", lr_shape.DebugString())); + + OP_REQUIRES( + ctx, TensorShapeUtils::IsScalar(l1_shape), + errors::InvalidArgument("l1 is not a scalar: ", l1_shape.DebugString())); + + OP_REQUIRES( + ctx, TensorShapeUtils::IsScalar(l2_shape), + errors::InvalidArgument("l2 is not a scalar: ", l2_shape.DebugString())); + + if (has_l2_shrinkage) { + OP_REQUIRES(ctx, TensorShapeUtils::IsScalar(l2_shrinkage_shape), + errors::InvalidArgument("l2_shrinkage is not a scalar: ", + l2_shrinkage_shape.DebugString())); + } + + OP_REQUIRES(ctx, TensorShapeUtils::IsScalar(lr_power_shape), + errors::InvalidArgument("lr_power is not a scalar: ", + lr_power_shape.DebugString())); + + xla::ComputationDataHandle var, accum, linear; + OP_REQUIRES_OK(ctx, ctx->ReadVariableInput(0, &var)); + OP_REQUIRES_OK(ctx, ctx->ReadVariableInput(1, &accum)); + OP_REQUIRES_OK(ctx, ctx->ReadVariableInput(2, &linear)); + xla::ComputationDataHandle grad = ctx->Input(3); + xla::ComputationDataHandle lr = ctx->Input(4); + xla::ComputationDataHandle l1 = ctx->Input(5); + xla::ComputationDataHandle l2 = ctx->Input(6); + xla::ComputationDataHandle l2_shrinkage; + xla::ComputationDataHandle lr_power; + if (has_l2_shrinkage) { + l2_shrinkage = ctx->Input(7); + lr_power = ctx->Input(8); + } else { + lr_power = ctx->Input(7); + } + + // grad_to_use = grad + 2 * l2_shrinkage * var + // new_accum = accum + grad_to_use * grad_to_use + // linear += grad_to_use - + // (new_accum^(-lr_power) - accum^(-lr_power)) / lr * var + // quadratic = (new_accum^(-lr_power) / lr) + 2 * l2 + // linear_clipped = clamp linear in [-l1, l1] + // var = (linear_clipped - linear) / quadratic + // accum = new_accum + + xla::ComputationDataHandle two = XlaHelpers::FloatLiteral(b, dtype, 2.0); + xla::ComputationDataHandle grad_to_use; + if (has_l2_shrinkage) { + grad_to_use = b->Add(grad, b->Mul(two, b->Mul(l2_shrinkage, var))); + } else { + grad_to_use = grad; + } + + xla::ComputationDataHandle new_accum = + b->Add(accum, b->Pow(grad_to_use, two)); + xla::ComputationDataHandle new_accum_lr_pow = + b->Pow(new_accum, b->Neg(lr_power)); + xla::ComputationDataHandle accum_lr_pow = b->Pow(accum, b->Neg(lr_power)); + linear = b->Add( + linear, + b->Sub(grad_to_use, + b->Mul(b->Div(b->Sub(new_accum_lr_pow, accum_lr_pow), lr), var))); + xla::ComputationDataHandle linear_clipped = b->Clamp(b->Neg(l1), linear, l1); + xla::ComputationDataHandle quadratic = + b->Add(b->Div(new_accum_lr_pow, lr), b->Mul(two, l2)); + var = b->Div(b->Sub(linear_clipped, linear), quadratic); + accum = new_accum; + + OP_REQUIRES_OK(ctx, ctx->AssignVariable(0, dtype, var)); + OP_REQUIRES_OK(ctx, ctx->AssignVariable(1, dtype, accum)); + OP_REQUIRES_OK(ctx, ctx->AssignVariable(2, dtype, linear)); +} + +class ResourceApplyFtrl : public XlaOpKernel { + public: + explicit ResourceApplyFtrl(OpKernelConstruction* ctx) : XlaOpKernel(ctx) { + OP_REQUIRES_OK(ctx, ctx->GetAttr("T", &dtype_)); + } + + void Compile(XlaOpKernelContext* ctx) override { + CompileFtrl(ctx, dtype_, /*has_l2_shrinkage=*/false); + } + + private: + DataType dtype_; +}; +REGISTER_XLA_OP(Name("ResourceApplyFtrl"), ResourceApplyFtrl); + +class ResourceApplyFtrlV2 : public XlaOpKernel { + public: + explicit ResourceApplyFtrlV2(OpKernelConstruction* ctx) : XlaOpKernel(ctx) { + OP_REQUIRES_OK(ctx, ctx->GetAttr("T", &dtype_)); + } + + void Compile(XlaOpKernelContext* ctx) override { + CompileFtrl(ctx, dtype_, /*has_l2_shrinkage=*/true); + } + + private: + DataType dtype_; +}; +REGISTER_XLA_OP(Name("ResourceApplyFtrlV2"), ResourceApplyFtrlV2); + } // namespace } // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/kernels/unary_ops.cc b/tensorflow/compiler/tf2xla/kernels/unary_ops.cc index abe4949f5dbc8034fa46828e3ff872cae7591d90..6b8f5ec7b33cd448a7b06c5dfe4aac288e53e9c9 100644 --- a/tensorflow/compiler/tf2xla/kernels/unary_ops.cc +++ b/tensorflow/compiler/tf2xla/kernels/unary_ops.cc @@ -43,11 +43,42 @@ namespace { // Return x if x>0, otherwise -x. XLAJIT_MAKE_UNARY(Abs, b->Abs(x)); + +// acosh(x) = log(x + sqrt(x^2 - 1)) +XLAJIT_MAKE_UNARY( + Acosh, + b->Log(b->Add(x, b->Pow(b->Sub(b->Mul(x, x), + XlaHelpers::One(b, input_type(0))), + XlaHelpers::FloatLiteral(b, input_type(0), 0.5))))); +// asinh(x) = log(x + sqrt(x^2 + 1)) +XLAJIT_MAKE_UNARY( + Asinh, + b->Log(b->Add(x, b->Pow(b->Add(b->Mul(x, x), + XlaHelpers::One(b, input_type(0))), + XlaHelpers::FloatLiteral(b, input_type(0), 0.5))))); +// atanh(x) = 0.5 * log((1 + x) / (1 - x)) +XLAJIT_MAKE_UNARY( + Atanh, b->Mul(b->Log(b->Div(b->Add(XlaHelpers::One(b, input_type(0)), x), + b->Sub(XlaHelpers::One(b, input_type(0)), x))), + XlaHelpers::FloatLiteral(b, input_type(0), 0.5))); XLAJIT_MAKE_UNARY(Ceil, b->Ceil(x)); +XLAJIT_MAKE_UNARY(Cos, b->Cos(x)); +XLAJIT_MAKE_UNARY(Cosh, + b->Mul(b->Add(b->Exp(x), b->Exp(b->Neg(x))), + XlaHelpers::FloatLiteral(b, input_type(0), 0.5))); +XLAJIT_MAKE_UNARY(Sin, b->Sin(x)); XLAJIT_MAKE_UNARY(Exp, b->Exp(x)); + +// TODO(b/34703906): use a more accurate implementation of expm1. +XLAJIT_MAKE_UNARY(Expm1, b->Sub(b->Exp(x), XlaHelpers::One(b, input_type(0)))); + XLAJIT_MAKE_UNARY(Floor, b->Floor(x)); -// Returns 0 if x is 0, -1 if x < 0 and 1 if x > 0. -XLAJIT_MAKE_UNARY(Sign, b->Sign(x)); +XLAJIT_MAKE_UNARY(IsFinite, b->IsFinite(x)); +XLAJIT_MAKE_UNARY(IsInf, b->Eq(b->Abs(x), + XlaHelpers::FloatLiteral( + b, input_type(0), + std::numeric_limits::infinity()))); +XLAJIT_MAKE_UNARY(IsNan, b->Ne(x, x)); // Return 1/x XLAJIT_MAKE_UNARY(Inv, b->Div(XlaHelpers::One(b, input_type(0)), x)); XLAJIT_MAKE_UNARY(Reciprocal, b->Div(XlaHelpers::One(b, input_type(0)), x)); @@ -77,17 +108,37 @@ static xla::ComputationDataHandle Round(xla::ComputationBuilder* b, b->LogicalAnd(b->Eq(fraction, half), is_odd)), b->Add(round_val, one), round_val); } + +XLAJIT_MAKE_UNARY(Rint, Round(b, input_type(0), x)); XLAJIT_MAKE_UNARY(Round, Round(b, input_type(0), x)); XLAJIT_MAKE_UNARY(Rsqrt, b->Pow(x, XlaHelpers::FloatLiteral(b, input_type(0), -0.5))); -XLAJIT_MAKE_UNARY(Sigmoid, - b->Map({x}, *ctx->GetOrCreateSigmoid(input_type(0)))); + +// Expresses sigmoid as a rescaled tanh: sigmoid(x) == (tanh(x/2) + 1) / 2. +static xla::ComputationDataHandle Sigmoid(xla::ComputationBuilder* b, + DataType dtype, + const xla::ComputationDataHandle& x) { + auto half = XlaHelpers::FloatLiteral(b, dtype, 0.5); + return b->Add(half, b->Mul(half, b->Tanh(b->Mul(half, x)))); +} +XLAJIT_MAKE_UNARY(Sigmoid, Sigmoid(b, input_type(0), x)); + +// Returns 0 if x is 0, -1 if x < 0 and 1 if x > 0. +XLAJIT_MAKE_UNARY(Sign, b->Sign(x)); +XLAJIT_MAKE_UNARY(Sinh, + b->Mul(b->Sub(b->Exp(x), b->Exp(b->Neg(x))), + XlaHelpers::FloatLiteral(b, input_type(0), 0.5))); XLAJIT_MAKE_UNARY(Softplus, b->Log(b->Add(b->Exp(x), XlaHelpers::One(b, input_type(0))))); +// softsign(x) = x / (abs(x) + 1) +XLAJIT_MAKE_UNARY(Softsign, + b->Div(x, + b->Add(b->Abs(x), XlaHelpers::One(b, input_type(0))))); XLAJIT_MAKE_UNARY(Sqrt, b->Pow(x, XlaHelpers::FloatLiteral(b, input_type(0), 0.5))); XLAJIT_MAKE_UNARY(Square, b->Mul(x, x)); +XLAJIT_MAKE_UNARY(Tan, b->Div(b->Sin(x), b->Cos(x))); XLAJIT_MAKE_UNARY(Tanh, b->Tanh(x)); #undef XLAJIT_MAKE_UNARY diff --git a/tensorflow/compiler/tf2xla/kernels/unpack_op.cc b/tensorflow/compiler/tf2xla/kernels/unpack_op.cc index a5ce78e52068d816cb8e3a2702024d419e25abc5..f87586ba578a6138e7fb921032e1a71f8c9ac80c 100644 --- a/tensorflow/compiler/tf2xla/kernels/unpack_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/unpack_op.cc @@ -66,6 +66,7 @@ class UnpackOp : public XlaOpKernel { std::vector start_indices(input_shape.dims(), 0); std::vector limit_indices(input_shape.dims()); + std::vector strides(input_shape.dims(), 1); for (int i = 0; i < input_shape.dims(); ++i) { limit_indices[i] = input_shape.dim_size(i); } @@ -73,7 +74,8 @@ class UnpackOp : public XlaOpKernel { for (int i = 0; i < num; ++i) { start_indices[axis] = i; limit_indices[axis] = i + 1; - auto slice = ctx->builder()->Slice(input, start_indices, limit_indices); + auto slice = ctx->builder()->Slice(input, start_indices, limit_indices, + strides); // Reshape to drop the 'axis' dimension. auto result = ctx->builder()->Reshape(slice, output_shape.dim_sizes()); ctx->SetOutput(i, result); diff --git a/tensorflow/compiler/tf2xla/kernels/variable_ops.cc b/tensorflow/compiler/tf2xla/kernels/variable_ops.cc index 1b04b8b802c5c6e9da337933a7c4cd99233ebe8d..b1fb656c7310cecb017f19e5844dc296042f4152 100644 --- a/tensorflow/compiler/tf2xla/kernels/variable_ops.cc +++ b/tensorflow/compiler/tf2xla/kernels/variable_ops.cc @@ -14,6 +14,8 @@ limitations under the License. ==============================================================================*/ #include "tensorflow/compiler/tf2xla/kernels/cwise_ops.h" +#include "tensorflow/compiler/tf2xla/kernels/gather_op_helpers.h" +#include "tensorflow/compiler/tf2xla/shape_util.h" #include "tensorflow/compiler/tf2xla/xla_helpers.h" #include "tensorflow/compiler/tf2xla/xla_op_registry.h" #include "tensorflow/compiler/xla/client/computation_builder.h" @@ -86,5 +88,38 @@ REGISTER_XLA_OP( Name("AssignSubVariableOp").TypeConstraint("dtype", kNumericTypes), AssignSubVariableOp); +class ResourceGatherOp : public XlaOpKernel { + public: + explicit ResourceGatherOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {} + void Compile(XlaOpKernelContext* ctx) override { + xla::ComputationBuilder* builder = ctx->builder(); + + // Get the shape of the resource tensor. + TensorShape resource_shape; + DataType resource_dtype; + OP_REQUIRES_OK( + ctx, ctx->GetVariableTypeAndShape(0, &resource_dtype, &resource_shape)); + + DataType expected_output_dtype = ctx->expected_output_dtype(0); + OP_REQUIRES(ctx, resource_dtype == expected_output_dtype, + errors::InvalidArgument( + "Variable dtype is ", DataTypeString(resource_dtype), + " but expected output dtype is ", + DataTypeString(expected_output_dtype), ".")); + + xla::ComputationDataHandle resource_handle; + OP_REQUIRES_OK(ctx, ctx->ReadVariableInput(0, &resource_handle)); + + auto indices = ctx->Input(1); + auto indices_shape = ctx->InputShape(1); + xla::ComputationDataHandle gather = XlaComputeGatherDynamicSlice( + ctx, resource_handle, resource_shape, indices, indices_shape, + resource_dtype, builder); + ctx->SetOutput(0, gather); + } +}; +REGISTER_XLA_OP(Name("ResourceGather").TypeConstraint("dtype", kNumericTypes), + ResourceGatherOp); + } // namespace } // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/kernels/while_op.cc b/tensorflow/compiler/tf2xla/kernels/while_op.cc new file mode 100644 index 0000000000000000000000000000000000000000..2c2031fc761e55ddb08a19dbc1b34a4d60e19562 --- /dev/null +++ b/tensorflow/compiler/tf2xla/kernels/while_op.cc @@ -0,0 +1,277 @@ +/* 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/kernels/while_op.h" + +#include "tensorflow/compiler/tf2xla/shape_util.h" +#include "tensorflow/compiler/tf2xla/type_util.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/compiler/xla/client/computation_builder.h" +#include "tensorflow/compiler/xla/literal_util.h" +#include "tensorflow/core/framework/function.h" +#include "tensorflow/core/framework/op_kernel.h" + +namespace tensorflow { + +namespace { + +// Builds XlaCompiler argument descriptions `args` from `ctx`. +Status MakeXlaCompilerArgumentsFromInputs( + XlaOpKernelContext* ctx, std::vector* args, + bool* has_uninitialized_vars) { + VLOG(2) << "Num inputs " << ctx->num_inputs(); + args->resize(ctx->num_inputs()); + *has_uninitialized_vars = false; + for (int i = 0; i < ctx->num_inputs(); ++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); + // When reading a resource input, use the type and shape of the resource's + // current value. + if (type == DT_RESOURCE) { + XlaResource* resource; + TF_RETURN_IF_ERROR(ctx->GetResourceInput(i, &resource)); + + arg.initialized = resource->value.handle() > 0; + switch (resource->kind) { + case XlaResource::kVariable: + arg.kind = XlaCompiler::Argument::kVariable; + break; + case XlaResource::kTensorArray: + arg.kind = XlaCompiler::Argument::kTensorArray; + break; + case XlaResource::kStack: + arg.kind = XlaCompiler::Argument::kStack; + break; + case XlaResource::kInvalid: + CHECK(false); + } + arg.type = resource->type; + if (arg.initialized) { + auto shape = ctx->builder()->GetShape(resource->value); + TF_RETURN_IF_ERROR(shape.status()); + arg.shape = *shape.ValueOrDie(); + } else { + *has_uninitialized_vars = true; + } + arg.tensor_array_size = resource->tensor_array_size; + arg.name = resource->name; + // TODO(phawkins): propagate TensorArray gradients into loops. + VLOG(2) << " resource " << resource->name + << " type: " << DataTypeString(arg.type) + << " shape: " << arg.shape.DebugString() + << " initialized: " << arg.initialized; + + } else { + arg.kind = XlaCompiler::Argument::kParameter; + arg.type = ctx->input_type(i); + TF_RETURN_IF_ERROR( + TensorShapeToXLAShape(arg.type, ctx->InputShape(i), &arg.shape)); + } + } + return Status::OK(); +} + +} // anonymous namespace + +XlaWhileOp::XlaWhileOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) { + const NameAttrList* name_attr; + OP_REQUIRES_OK(ctx, ctx->GetAttr("cond", &name_attr)); + cond_name_attr_ = *name_attr; + OP_REQUIRES_OK(ctx, ctx->GetAttr("body", &name_attr)); + body_name_attr_ = *name_attr; +} + +void XlaWhileOp::Compile(XlaOpKernelContext* ctx) { + VLOG(1) << "WhileOp::Compile"; + + std::vector arguments; + bool has_uninitialized_vars; + OP_REQUIRES_OK(ctx, MakeXlaCompilerArgumentsFromInputs( + ctx, &arguments, &has_uninitialized_vars)); + + const bool use_tuple_arg = (arguments.size() != 1); + + xla::ComputationBuilder* builder = ctx->builder(); + XlaCompiler* compiler = ctx->compiler(); + + VLOG(1) << "Compiling body"; + + // All resource that are inputs to the loop's body must also be + // present as loop body outputs; the signature of the loop's input and + // output must match. We ensure this by asking the compiler to include the + // current values of all resources, even if they haven't been updated by the + // computation. We must also ask the compiler to keep compile-time constant + // outputs as part of the generated computation, for the same reason. + // TODO(phawkins): consider adding loop-invariant inputs to XLA's While() + // operator. + XlaCompiler::CompileOptions body_options; + body_options.use_tuple_arg = use_tuple_arg; + body_options.return_updated_values_for_all_resources = true; + body_options.resolve_compile_time_constants = false; + XlaCompiler::CompilationResult body; + OP_REQUIRES_OK(ctx, compiler->CompileFunction(body_options, body_name_attr_, + arguments, &body)); + + // We must use a static shape for parameters to an XLA compilation. However, + // we may not know the shape of a TensorArray if it is first written inside + // the loop. Ideally we would require the user to provide a static shape, + // but this is not always easy. + // So if uninitialized resource are used by the loop body, we compile the + // body function twice: + // 1) once with uninitialized resource inputs. We discard the computation + // but we assume resource shapes reach a fixpoint after one iteration. + // So we can use the output shapes of the resource as the "true" shapes. + // 2) again with the "correct" input shapes determined by (1). + if (has_uninitialized_vars) { + // Initializes any uninitialized resource with zero values of the + // shape determined by the first compilation. + for (int i = 0; i < body.resource_updates.size(); ++i) { + const XlaCompiler::ResourceUpdate& update = body.resource_updates[i]; + XlaCompiler::Argument& arg = arguments[update.input_index]; + if (!arg.initialized) { + VLOG(2) << "Update shape for argument " << update.input_index << " " + << xla::ShapeUtil::HumanString(update.shape); + arg.initialized = true; + arg.shape = update.shape; + + XlaResource* resource; + OP_REQUIRES_OK(ctx, + ctx->GetResourceInput(update.input_index, &resource)); + + std::unique_ptr zero = + xla::Literal::CreateFromShape(update.shape); + resource->value = builder->ConstantLiteral(*zero); + } + } + // Recompile the body with the "correct" shapes. + VLOG(1) << "Recompiling body with non-placeholder shapes"; + body = {}; + OP_REQUIRES_OK(ctx, compiler->CompileFunction(body_options, body_name_attr_, + arguments, &body)); + } + + VLOG(1) << "Compiling condition"; + + XlaCompiler::CompileOptions cond_options; + cond_options.use_tuple_arg = use_tuple_arg; + cond_options.resolve_compile_time_constants = false; + XlaCompiler::CompilationResult cond; + OP_REQUIRES_OK(ctx, compiler->CompileFunction(cond_options, cond_name_attr_, + arguments, &cond)); + + xla::Shape body_input_shape, cond_input_shape; + if (use_tuple_arg) { + body_input_shape = xla::ShapeUtil::MakeTupleShape(body.xla_input_shapes); + cond_input_shape = xla::ShapeUtil::MakeTupleShape(cond.xla_input_shapes); + } else { + CHECK(!body.xla_input_shapes.empty()); + body_input_shape = body.xla_input_shapes[0]; + CHECK(!cond.xla_input_shapes.empty()); + cond_input_shape = cond.xla_input_shapes[0]; + } + + VLOG(2) << "Body shape: " << xla::ShapeUtil::HumanString(body_input_shape) + << " -> " << xla::ShapeUtil::HumanString(body.xla_output_shape); + VLOG(2) << "Cond shape: " << xla::ShapeUtil::HumanString(cond_input_shape) + << " -> " << xla::ShapeUtil::HumanString(cond.xla_output_shape); + + OP_REQUIRES(ctx, + xla::ShapeUtil::Compatible(body_input_shape, cond_input_shape), + errors::InvalidArgument( + "Input shapes of loop body and condition do not match: ", + xla::ShapeUtil::HumanString(body_input_shape), " vs. ", + xla::ShapeUtil::HumanString(cond_input_shape))); + OP_REQUIRES( + ctx, xla::ShapeUtil::Compatible(body_input_shape, body.xla_output_shape), + errors::InvalidArgument( + "Input and output shapes of loop body do not match: ", + xla::ShapeUtil::HumanString(body_input_shape), " vs. ", + xla::ShapeUtil::HumanString(body.xla_output_shape))); + + xla::ComputationDataHandle data; + + int num_inputs = body.input_mapping.size(); + + std::vector inputs(num_inputs); + for (int i = 0; i < num_inputs; ++i) { + int input_num = body.input_mapping[i]; + if (ctx->input_type(input_num) == DT_RESOURCE) { + XlaResource* resource; + OP_REQUIRES_OK(ctx, ctx->GetResourceInput(input_num, &resource)); + inputs[i] = resource->value; + } else { + inputs[i] = ctx->Input(i); + } + } + + xla::ComputationDataHandle init; + if (use_tuple_arg) { + init = builder->Tuple(inputs); + } else { + init = inputs[0]; + } + + VLOG(1) << "Building while loop"; + + xla::ComputationDataHandle while_result = + builder->While(*cond.computation, *body.computation, init); + + auto get_loop_output = [&](int i) { + if (use_tuple_arg) { + return builder->GetTupleElement(while_result, i); + } else { + return while_result; + } + }; + + // Sets non-variable outputs. + for (int i = 0; i < ctx->num_outputs(); ++i) { + if (ctx->input_type(i) != DT_RESOURCE) { + ctx->SetOutput(body.input_mapping[i], get_loop_output(i)); + } + } + + // Updates the values of any resource variables modified by the loop. + for (int i = 0; i < body.resource_updates.size(); ++i) { + const XlaCompiler::ResourceUpdate& update = body.resource_updates[i]; + XlaResource* resource; + OP_REQUIRES_OK(ctx, ctx->GetResourceInput(update.input_index, &resource)); + if (update.modified) { + int pos = body.outputs.size() + i; + resource->value = get_loop_output(pos); + } + VLOG(2) << "Loop-carried variable: pos: " << update.input_index + << " name: " << resource->name << " modified: " << update.modified + << " type: " << DataTypeString(update.type) + << " shape: " << update.shape.DebugString(); + // Copies the identity of the resource variable from input to output + // unchanged, even if the variable was not modified. + ctx->op_kernel_context()->set_output( + update.input_index, + ctx->op_kernel_context()->input(update.input_index)); + } + + VLOG(1) << "Done building while loop"; +} + +REGISTER_XLA_OP(Name("XlaWhile").AllowResourceTypes(), XlaWhileOp); + +} // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/kernels/while_op.h b/tensorflow/compiler/tf2xla/kernels/while_op.h new file mode 100644 index 0000000000000000000000000000000000000000..67edebabf9f643a919d0f06c228e2d224a49a2af --- /dev/null +++ b/tensorflow/compiler/tf2xla/kernels/while_op.h @@ -0,0 +1,65 @@ +/* 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_KERNELS_WHILE_OP_H_ +#define TENSORFLOW_COMPILER_TF2XLA_KERNELS_WHILE_OP_H_ + +#include "tensorflow/compiler/tf2xla/xla_op_kernel.h" +#include "tensorflow/core/framework/attr_value.pb.h" + +namespace tensorflow { + +// This TensorFlow op provides a functional iteration primitive. +// +// The inputs and outputs of the loop body must agree on the number, types, and +// shapes of the Tensors carried around the loop body. +// +// Computations in while loops may read from and write to resource variables. +// Resource variables may be passed as arguments to a function's body and +// condition functions. The XlaCompiler converts resource variable arguments +// into parameters to the XLA computation and moves them to the end of the +// parameter list, and by using the `return_updated_values_for_all_variables` +// we ensure that all variables that appear in the input also appear at the +// end of the body's output. This ensures the loop body's input and output +// signatures match. +// +// It is the user's responsibility to ensure that each non-variable _Arg matches +// the corresponding _Retval. +// +// For example, suppose we have a loop body with arguments: +// DT_INT32, DT_RESOURCE (pointing to a DT_BOOL var), DT_FLOAT +// and return values +// DT_INT32, DT_FLOAT +// It is an error for the body to return DT_RESOURCE values. +// +// The body will be lowered into an XLA computation that takes and returns a +// tuple with XLA type (I32, F32, PRED). Note the resource variable appears at +// the end of both the loop body's input and output argument lists. +class XlaWhileOp : public XlaOpKernel { + public: + explicit XlaWhileOp(OpKernelConstruction* ctx); + + void Compile(XlaOpKernelContext* ctx) override; + + private: + NameAttrList cond_name_attr_; + NameAttrList body_name_attr_; + + TF_DISALLOW_COPY_AND_ASSIGN(XlaWhileOp); +}; + +} // namespace tensorflow + +#endif // TENSORFLOW_COMPILER_TF2XLA_KERNELS_WHILE_OP_H_ diff --git a/tensorflow/compiler/tf2xla/literal_util.cc b/tensorflow/compiler/tf2xla/literal_util.cc index 1f2bc01cf4a48b37de585c55b781c239ee4b8f2a..576cd9bf9abb43e29d9eb8f706e0f42ac2d038e9 100644 --- a/tensorflow/compiler/tf2xla/literal_util.cc +++ b/tensorflow/compiler/tf2xla/literal_util.cc @@ -27,13 +27,13 @@ Status HostTensorToLiteral(const Tensor& host_tensor, xla::Literal* literal) { TF_RETURN_IF_ERROR(TensorShapeToXLAShape( host_tensor.dtype(), host_tensor.shape(), literal->mutable_shape())); - xla::LiteralUtil::Reserve(host_tensor.NumElements(), literal); + literal->Reserve(host_tensor.NumElements()); // memcpy over the payload ... // TODO(phawkins): handle string types. size_t total_bytes = host_tensor.TotalBytes(); if (total_bytes > 0) { - void* dst_ptr = xla::LiteralUtil::MutableInternalData(literal); + void* dst_ptr = literal->MutableInternalData(); const void* src_ptr = DMAHelper::base(&host_tensor); memcpy(dst_ptr, src_ptr, total_bytes); } @@ -51,11 +51,12 @@ Status LiteralToHostTensor(const xla::Literal& literal, DataType target_type, " to tensor of type ", DataTypeString(target_type)); } - TensorShape shape = XLAShapeToTensorShape(literal.shape()); + TensorShape shape; + TF_RETURN_IF_ERROR(XLAShapeToTensorShape(literal.shape(), &shape)); *host_tensor = Tensor(target_type, shape); size_t total_bytes = host_tensor->TotalBytes(); if (total_bytes > 0) { - const void* src_ptr = xla::LiteralUtil::InternalData(literal); + const void* src_ptr = literal.InternalData(); void* dst_ptr = DMAHelper::base(host_tensor); memcpy(dst_ptr, src_ptr, total_bytes); } diff --git a/tensorflow/compiler/tf2xla/literal_util.h b/tensorflow/compiler/tf2xla/literal_util.h index e8b2233853d2eb8eb06172bb6209ca537538ed5a..fe08e83c2391a8b24696961cacfd909d46e49e7d 100644 --- a/tensorflow/compiler/tf2xla/literal_util.h +++ b/tensorflow/compiler/tf2xla/literal_util.h @@ -18,6 +18,7 @@ limitations under the License. #ifndef TENSORFLOW_COMPILER_TF2XLA_LITERAL_UTIL_H_ #define TENSORFLOW_COMPILER_TF2XLA_LITERAL_UTIL_H_ +#include "tensorflow/compiler/xla/literal_util.h" #include "tensorflow/compiler/xla/xla_data.pb.h" #include "tensorflow/core/framework/tensor.h" #include "tensorflow/core/lib/core/status.h" diff --git a/tensorflow/compiler/tf2xla/literal_util_test.cc b/tensorflow/compiler/tf2xla/literal_util_test.cc index 56993bc58534d1225f9177719804a69f561b3a06..f3d6787daaa1165b28ce63dfd501533fa0963edd 100644 --- a/tensorflow/compiler/tf2xla/literal_util_test.cc +++ b/tensorflow/compiler/tf2xla/literal_util_test.cc @@ -27,7 +27,7 @@ TEST(LiteralUtil, LiteralToHostTensor) { { std::vector int64_values = {1, 2, 3}; std::unique_ptr int64_values_literal = - xla::LiteralUtil::CreateR1(gtl::ArraySlice(int64_values)); + xla::Literal::CreateR1(gtl::ArraySlice(int64_values)); Tensor host_tensor; EXPECT_EQ("Cannot convert literal of type S64 to tensor of type int32", LiteralToHostTensor(*int64_values_literal, DT_INT32, &host_tensor) @@ -48,7 +48,7 @@ TEST(LiteralUtil, LiteralToHostTensor) { Tensor host_tensor; std::vector int32_values = {10, 11}; std::unique_ptr int32_values_literal = - xla::LiteralUtil::CreateR1(gtl::ArraySlice(int32_values)); + xla::Literal::CreateR1(gtl::ArraySlice(int32_values)); EXPECT_TRUE( LiteralToHostTensor(*int32_values_literal, DT_INT32, &host_tensor) .ok()); diff --git a/tensorflow/compiler/tf2xla/ops/BUILD b/tensorflow/compiler/tf2xla/ops/BUILD new file mode 100644 index 0000000000000000000000000000000000000000..98f72b3792eb147f5a1847c5e1ecef18bccbca5f --- /dev/null +++ b/tensorflow/compiler/tf2xla/ops/BUILD @@ -0,0 +1,55 @@ +package( + default_visibility = ["//tensorflow/compiler/tf2xla:internal"], +) + +licenses(["notice"]) # Apache 2.0 + +load("//tensorflow:tensorflow.bzl", "tf_gen_op_wrapper_py") + +cc_library( + name = "functional_ops", + srcs = ["functional_ops.cc"], + deps = [ + "//tensorflow/core:framework", + ], + alwayslink = 1, +) + +cc_library( + name = "sendrecv_ops", + srcs = ["sendrecv_ops.cc"], + deps = [ + "//tensorflow/core:framework", + ], + alwayslink = 1, +) + +tf_gen_op_wrapper_py( + name = "gen_functional_ops", + out = "gen_functional_ops.py", + deps = [ + ":functional_ops", + ], +) + +tf_gen_op_wrapper_py( + name = "gen_sendrecv_ops", + out = "gen_sendrecv_ops.py", + deps = [ + ":sendrecv_ops", + ], +) + +# ----------------------------------------------------------------------------- + +filegroup( + name = "all_files", + srcs = glob( + ["**/*"], + exclude = [ + "**/METADATA", + "**/OWNERS", + ], + ), + visibility = ["//tensorflow:__subpackages__"], +) diff --git a/tensorflow/compiler/tf2xla/ops/functional_ops.cc b/tensorflow/compiler/tf2xla/ops/functional_ops.cc new file mode 100644 index 0000000000000000000000000000000000000000..c1005405f9a9b09e4a6480332861d0cce2c52291 --- /dev/null +++ b/tensorflow/compiler/tf2xla/ops/functional_ops.cc @@ -0,0 +1,47 @@ +/* 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" + +namespace tensorflow { + +// TODO(b/37549631) setting the While Op to always be stateful is too +// conservative. +REGISTER_OP("XlaWhile") + .Input("input: T") + .Output("output: T") + .Attr("T: list(type) >= 0") + .Attr("cond: func") + .Attr("body: func") + .SetIsStateful() + .SetShapeFn(shape_inference::UnknownShape) + .Doc(R"doc( +output = input; While (Cond(output)) { output = Body(output) } + +input: A list of input tensors whose types are T. +output: A list of output tensors whose types are T. +cond: A function takes 'input' and returns a tensor. If the tensor is + a scalar of non-boolean, the scalar is converted to a boolean + according to the following rule: if the scalar is a numerical + value, non-zero means True and zero means False; if the scalar is + a string, non-empty means True and empty means False. If the + tensor is not a scalar, non-emptiness means True and False + otherwise. +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"); + +} // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/ops/sendrecv_ops.cc b/tensorflow/compiler/tf2xla/ops/sendrecv_ops.cc new file mode 100644 index 0000000000000000000000000000000000000000..b6947bfe570c75dd0c7c6301b972e2012bae26bd --- /dev/null +++ b/tensorflow/compiler/tf2xla/ops/sendrecv_ops.cc @@ -0,0 +1,49 @@ +/* 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" + +namespace tensorflow { + +REGISTER_OP("_XLASend") + .Input("tensor: T") + .Attr("T: type") + .Attr("tensor_name: string") + .SetIsStateful() + .SetShapeFn(shape_inference::UnknownShape) + .Doc(R"doc( +Sends the named tensor to another XLA computation. + +tensor: The tensor to send. +tensor_name: The name of the tensor to send. +)doc"); + +REGISTER_OP("_XLARecv") + .Output("tensor: T") + .Attr("T: type") + .Attr("tensor_name: string") + .Attr("shape: shape") + .SetIsStateful() + .SetShapeFn(shape_inference::UnknownShape) + .Doc(R"doc( +Receives the named tensor from another XLA computation. + +tensor: The tensor to receive. +tensor_name: The name of the tensor to receive. +shape: The shape of the input tensor. +)doc"); + +} // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/shape_util.cc b/tensorflow/compiler/tf2xla/shape_util.cc index f5ecb51a5b77e36e606ed1c48b8e2dbe76de0074..9d1992205b02665b99b1bd15b7b65a1fb8c35a51 100644 --- a/tensorflow/compiler/tf2xla/shape_util.cc +++ b/tensorflow/compiler/tf2xla/shape_util.cc @@ -24,12 +24,18 @@ limitations under the License. namespace tensorflow { // Convert an XLA Shape into the equivalent TensorFlow shape. -TensorShape XLAShapeToTensorShape(const xla::Shape& shape) { - TensorShape tensor_shape; +Status XLAShapeToTensorShape(const xla::Shape& shape, + TensorShape* tensor_shape) { + if (xla::ShapeUtil::IsTuple(shape)) { + 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) { - tensor_shape.AddDim(shape.dimensions(i)); + tensor_shape->AddDim(shape.dimensions(i)); } - return tensor_shape; + return Status::OK(); } // Convert a TensorShape into the equivalent XLA Shape proto. diff --git a/tensorflow/compiler/tf2xla/shape_util.h b/tensorflow/compiler/tf2xla/shape_util.h index 516dd636a970f78fda363a0b13961b8244dc2cd9..58240b9c965a194b9380ac7cd477ce7344e5ebe3 100644 --- a/tensorflow/compiler/tf2xla/shape_util.h +++ b/tensorflow/compiler/tf2xla/shape_util.h @@ -24,8 +24,10 @@ limitations under the License. namespace tensorflow { -// Convert an XLA Shape into the equivalent TensorFlow shape. -TensorShape XLAShapeToTensorShape(const xla::Shape& shape); +// Convert an XLA Shape into the equivalent TensorFlow shape. May fail since +// not all XLA shapes can be represented as TensorShapes. +Status XLAShapeToTensorShape(const xla::Shape& shape, + TensorShape* tensor_shape); // Convert a TensorShape into the equivalent XLA Shape proto. Unlike Tensorflow, // XLA shapes include the type. Not all `dtype` values can be represented by diff --git a/tensorflow/compiler/tf2xla/str_util.cc b/tensorflow/compiler/tf2xla/str_util.cc index ce25d631271b54a36078cd0d3ac4d318d58db9fa..2b0834fe7b6c4d2199267dbe0ec1f7c2785aa9c7 100644 --- a/tensorflow/compiler/tf2xla/str_util.cc +++ b/tensorflow/compiler/tf2xla/str_util.cc @@ -22,7 +22,7 @@ limitations under the License. namespace tensorflow { namespace str_util { -void ReplaceAll(string* text, StringPiece from, StringPiece to) { +static void ReplaceAll(string* text, StringPiece from, StringPiece to) { size_t pos = 0; while ((pos = text->find(from.data(), pos, from.size())) != string::npos) { text->replace(pos, from.size(), to.data(), to.size()); diff --git a/tensorflow/compiler/tf2xla/str_util.h b/tensorflow/compiler/tf2xla/str_util.h index 4920b1a4d4875192d6f06988b810ad388bc6293b..51f25009d7003db0d72296619a469ecbbbb1808d 100644 --- a/tensorflow/compiler/tf2xla/str_util.h +++ b/tensorflow/compiler/tf2xla/str_util.h @@ -29,10 +29,6 @@ limitations under the License. namespace tensorflow { namespace str_util { -// Replace all non-overlapping occurrences of from with to in-place in text. If -// from is empty, it matches at the beginning of the text and after every byte. -void ReplaceAll(string* text, StringPiece from, StringPiece to); - // Replace all non-overlapping occurrences of the given (from,to) pairs in-place // in text. If from is empty, it matches at the beginning of the text and after // every byte. Each (from,to) replacement pair is processed in the order it is diff --git a/tensorflow/compiler/tf2xla/str_util_test.cc b/tensorflow/compiler/tf2xla/str_util_test.cc index f992007a34532157f86c90c717a5e24c3923f22d..8817f6902a8e58e796ca5240a9a24d7506d38793 100644 --- a/tensorflow/compiler/tf2xla/str_util_test.cc +++ b/tensorflow/compiler/tf2xla/str_util_test.cc @@ -25,36 +25,6 @@ limitations under the License. namespace tensorflow { namespace str_util { -class ReplaceAllTest : public ::testing::Test { - protected: - void ExpectReplaceAll(string text, StringPiece from, StringPiece to, - StringPiece want) { - ReplaceAll(&text, from, to); - EXPECT_EQ(text, want); - } -}; - -TEST_F(ReplaceAllTest, Simple) { - ExpectReplaceAll("", "", "", ""); - ExpectReplaceAll("", "", "X", "X"); - ExpectReplaceAll("", "", "XYZ", "XYZ"); - ExpectReplaceAll("banana", "", "", "banana"); - ExpectReplaceAll("banana", "", "_", "_b_a_n_a_n_a_"); - ExpectReplaceAll("banana", "", "__", "__b__a__n__a__n__a__"); - ExpectReplaceAll("banana", "a", "a", "banana"); - ExpectReplaceAll("banana", "a", "", "bnn"); - ExpectReplaceAll("banana", "a", "X", "bXnXnX"); - ExpectReplaceAll("banana", "a", "XX", "bXXnXXnXX"); - ExpectReplaceAll("banana", "an", "an", "banana"); - ExpectReplaceAll("banana", "an", "", "ba"); - ExpectReplaceAll("banana", "an", "X", "bXXa"); - ExpectReplaceAll("banana", "an", "XY", "bXYXYa"); - ExpectReplaceAll("banana", "an", "XYZ", "bXYZXYZa"); - ExpectReplaceAll("foo {{bar}} baz {{bar}}", "{{bar}}", "X", "foo X baz X"); - ExpectReplaceAll("foo {{bar}} baz {{bar}}", "{{bar}}", "ABCDEFGHIJKLMNOP", - "foo ABCDEFGHIJKLMNOP baz ABCDEFGHIJKLMNOP"); -} - class ReplaceAllPairsTest : public ::testing::Test { protected: void ExpectReplaceAllPairs( diff --git a/tensorflow/compiler/tf2xla/test_util.cc b/tensorflow/compiler/tf2xla/test_util.cc new file mode 100644 index 0000000000000000000000000000000000000000..3c6c9a91b6d2fb47f6dee1c347e9b852f1eea3ec --- /dev/null +++ b/tensorflow/compiler/tf2xla/test_util.cc @@ -0,0 +1,43 @@ +/* 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/test_util.h" + +#include "tensorflow/compiler/xla/status_macros.h" +#include "tensorflow/core/framework/node_def.pb.h" + +namespace tensorflow { + +Status InstantiateFunctionForTest(const string& name, + const FunctionLibraryDefinition& library, + InstantiationResultForTest* result) { + const FunctionDef* fdef = library.Find(name); + TF_RET_CHECK(fdef != nullptr); + + auto get_func_sig = [&library](const string& op, const OpDef** sig) { + return library.LookUpOpDef(op, sig); + }; + InstantiationResult inst; + TF_RETURN_IF_ERROR( + InstantiateFunction(*fdef, AttrSlice(), get_func_sig, &inst)); + result->arg_types = inst.arg_types; + result->ret_types = inst.ret_types; + for (NodeDef& n : inst.nodes) { + *result->gdef.add_node() = std::move(n); + } + return Status::OK(); +} + +} // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/test_util.h b/tensorflow/compiler/tf2xla/test_util.h new file mode 100644 index 0000000000000000000000000000000000000000..e6e4ae92ed23f3fca0f59b131dc73152e0947b72 --- /dev/null +++ b/tensorflow/compiler/tf2xla/test_util.h @@ -0,0 +1,47 @@ +/* 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. +==============================================================================*/ + +// Helper functions for tests. + +#ifndef TENSORFLOW_COMPILER_TF2XLA_TEST_UTIL_H_ +#define TENSORFLOW_COMPILER_TF2XLA_TEST_UTIL_H_ + +#include +#include +#include + +#include "tensorflow/core/framework/function.h" +#include "tensorflow/core/framework/graph.pb.h" +#include "tensorflow/core/graph/graph.h" +#include "tensorflow/core/lib/core/status.h" + +namespace tensorflow { + +// Same as InstantiationResult, but has a GraphDef instead of just nodes. +struct InstantiationResultForTest { + DataTypeVector arg_types; + DataTypeVector ret_types; + GraphDef gdef; +}; + +// Instantiates a function, producing a GraphDef to compare against the +// expected graph. +Status InstantiateFunctionForTest(const string& name, + const FunctionLibraryDefinition& library, + InstantiationResultForTest* result); + +} // namespace tensorflow + +#endif // TENSORFLOW_COMPILER_TF2XLA_TEST_UTIL_H_ diff --git a/tensorflow/compiler/tf2xla/xla_compilation_device.cc b/tensorflow/compiler/tf2xla/xla_compilation_device.cc index d86e741b69e08652bac2dd7b5295c8ab2d94433a..1d0098591e3b4834ff5447416e6233cf1d313a96 100644 --- a/tensorflow/compiler/tf2xla/xla_compilation_device.cc +++ b/tensorflow/compiler/tf2xla/xla_compilation_device.cc @@ -76,8 +76,7 @@ XlaCompilationDevice::XlaCompilationDevice(const SessionOptions& options, options, Device::BuildDeviceAttributes( "", type, Bytes(256 << 20), DeviceLocality(), - strings::StrCat("device: XLA compilation device ", type.type())), - cpu_allocator()), + strings::StrCat("device: XLA compilation device ", type.type()))), allocator_(new XlaCompilationAllocator()) {} XlaCompilationDevice::~XlaCompilationDevice() {} @@ -120,6 +119,4 @@ void XlaExpression::set_constant_value(Tensor value) { constant_value_ = std::move(value); } -void XlaExpression::set_variable_id(int id) { variable_id_ = id; } - } // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/xla_compilation_device.h b/tensorflow/compiler/tf2xla/xla_compilation_device.h index 1ee96e5e6c06783e7d9228968b7e71a3dc5be591..ec28bdccda47a326a0f60f2f73e8837b68e668cb 100644 --- a/tensorflow/compiler/tf2xla/xla_compilation_device.h +++ b/tensorflow/compiler/tf2xla/xla_compilation_device.h @@ -64,6 +64,49 @@ class XlaCompilationDevice : public LocalDevice { std::unique_ptr allocator_; }; +// Represents a resource, such as a Variable or TensorArray. +struct XlaResource { + enum Kind { + kInvalid, + kVariable, + kTensorArray, + kStack, + }; + + Kind kind = kInvalid; + + // If this resource is visible externally, what was its argument number? + int arg_num = -1; + + // A descriptive name for the resource, used in error messages. + string name; + + // Current type and value of the resource. Uninitialized resources are + // represented by a default (zero) handle and type DT_INVALID. + // While the type of a resource is notionally fixed during execution, when + // a resource is first initialized we do not yet know its type, so we keep + // track of its type dynamically. + DataType type = DT_INVALID; + xla::ComputationDataHandle value; + + // Value of the resource at computation entry. Used to detect which + // variables have new values that need to be written back. + xla::ComputationDataHandle initial_value; + + // TensorArray-specific fields + + // 'tensor_array_size' stores the expected size of the TensorArray. We need + // to store this since sometimes TensorArrays must be initialized lazily since + // we do not know the element shape at construction time. + int64 tensor_array_size = -1; + + // 'tensor_array_gradient' is a map from TensorArrayGradV3 'source' attributes + // to an XlaResource containing the gradient TensorArrays. We store a pointer + // here since there should only be one gradient TensorArray per 'source' + // string, irrespective of the number of calls to TensorArrayGrad. + std::unordered_map tensor_array_gradient; +}; + // 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 ComputationDataHandle. Each @@ -82,8 +125,8 @@ class XlaExpression { bool has_constant_value() const { return has_constant_value_; } const Tensor& constant_value() const { return constant_value_; } - void set_variable_id(int id); - int variable_id() const { return variable_id_; } + void set_resource(XlaResource* resource) { resource_ = resource; } + XlaResource* resource() const { return resource_; } private: // The XLA handle of the expression's computation. @@ -95,7 +138,7 @@ class XlaExpression { bool has_constant_value_ = false; Tensor constant_value_; - int variable_id_ = -1; + XlaResource* resource_ = nullptr; // Not owned. TF_DISALLOW_COPY_AND_ASSIGN(XlaExpression); }; diff --git a/tensorflow/compiler/tf2xla/xla_compiler.cc b/tensorflow/compiler/tf2xla/xla_compiler.cc index 33b4a43aa1544f883d4242148ce77eebb8a4c54c..ae13147a18e220dd320798c0250dffa7b22c37c3 100644 --- a/tensorflow/compiler/tf2xla/xla_compiler.cc +++ b/tensorflow/compiler/tf2xla/xla_compiler.cc @@ -18,6 +18,7 @@ limitations under the License. #include #include "tensorflow/compiler/tf2xla/dump_graph.h" +#include "tensorflow/compiler/tf2xla/functionalize_control_flow.h" #include "tensorflow/compiler/tf2xla/shape_util.h" #include "tensorflow/compiler/tf2xla/type_util.h" #include "tensorflow/compiler/tf2xla/xla_context.h" @@ -35,7 +36,6 @@ limitations under the License. #include "tensorflow/core/public/version.h" namespace tensorflow { - namespace { // Checks that arguments `args` match types `types`. @@ -57,81 +57,120 @@ Status CheckSignature(const DataTypeVector& types, } // namespace +bool XlaCompiler::Argument::operator==( + const XlaCompiler::Argument& other) const { + if (std::tie(kind, type, name, tensor_array_size) != + std::tie(other.kind, other.type, other.name, other.tensor_array_size)) { + return false; + } + if (!xla::ShapeUtil::Equal(shape, other.shape)) { + return false; + } + if (constant_value.shape() != other.constant_value.shape()) { + return false; + } + return constant_value.tensor_data() == other.constant_value.tensor_data(); +} + XlaCompiler::XlaCompiler(XlaCompiler::Options options) - : options_(std::move(options)), + : options_(options), initialization_status_(Status::OK()), next_step_id_(1), - device_(new XlaCompilationDevice(SessionOptions(), options_.device_type)), + device_( + new XlaCompilationDevice(SessionOptions(), *options_.device_type)), device_mgr_({device_}) { + // We no longer need the device_type. + options_.device_type = nullptr; + if (options_.populate_resource_manager) { initialization_status_ = (*options_.populate_resource_manager)(device_->resource_manager()); } + + local_flib_def_.reset(new FunctionLibraryDefinition(OpRegistry::Global(), + + FunctionDefLibrary{})); + local_pflr_.reset(new ProcessFunctionLibraryRuntime( + &device_mgr_, Env::Default(), options.graph_def_version, + local_flib_def_.get(), OptimizerOptions(), + nullptr /* custom_kernel_creator */)); + pflr_.reset(new ProcessFunctionLibraryRuntime( + &device_mgr_, Env::Default(), options.graph_def_version, options.flib_def, + OptimizerOptions(), nullptr /* custom_kernel_creator */)); + + local_flib_runtime_ = local_pflr_->GetFLR(device_->name()); + flib_runtime_ = pflr_->GetFLR(device_->name()); } XlaCompiler::~XlaCompiler() = default; int64 XlaCompiler::NextStepId() { - mutex_lock l(mu_); return next_step_id_++; } -// Prunes any nodes from a function that are not dependencies of the _Retval -// nodes. Used to prune stateful ops from within a function body, such as -// variable initializers, that should not be executed unless requested. -static void PruneUnreachableNodes(Graph* graph) { - std::unordered_set nodes; - for (Node* node : graph->nodes()) { - if (node->type_string() == "_Retval" || - StringPiece(node->type_string()).ends_with("Send")) { - nodes.insert(node); - } - } - PruneForReverseReachability(graph, nodes); +uint64 XlaCompiler::SignatureHash::operator()( + const std::pair>& signature) const { + return std::hash()(signature.first); +} + +static Status GetFunctionBody(const NameAttrList& function, + FunctionLibraryRuntime* flib_runtime, + const FunctionBody** fbody) { + FunctionLibraryRuntime::Handle handle; + TF_RETURN_IF_ERROR(flib_runtime->Instantiate( + function.name(), AttrSlice(&function.attr()), &handle)); + + *fbody = flib_runtime->GetFunctionBody(handle); + TF_RET_CHECK(*fbody); + return Status::OK(); } Status XlaCompiler::CompileFunction( - FunctionLibraryRuntime* flr, const NameAttrList& function, + const XlaCompiler::CompileOptions& options, const NameAttrList& function, const std::vector& args, XlaCompiler::CompilationResult* result) { - const string function_id = Canonicalize(function.name(), function.attr()); + const string function_id = + Canonicalize(function.name(), AttrSlice(&function.attr())); VLOG(1) << "XlaCompiler::CompileFunction " << function_id; - FunctionLibraryRuntime::Handle handle; - TF_RETURN_IF_ERROR( - flr->Instantiate(function.name(), function.attr(), &handle)); + auto it = cache_.find({function_id, args}); + if (it != cache_.end()) { + *result = it->second; + return Status::OK(); + } - const FunctionBody* fbody = flr->GetFunctionBody(handle); - CHECK(fbody); + const FunctionBody* fbody; + if (!GetFunctionBody(function, local_flib_runtime_, &fbody).ok()) { + TF_RETURN_IF_ERROR(GetFunctionBody(function, flib_runtime_, &fbody)); + } TF_RETURN_IF_ERROR(CheckSignature(fbody->arg_types, args)); - std::unique_ptr graph(new Graph(flr->GetFunctionLibraryDefinition())); + std::unique_ptr graph(new Graph(options_.flib_def)); CopyGraph(*fbody->graph, graph.get()); - if (VLOG_IS_ON(1)) { - dump_graph::DumpGraphToFile( - strings::StrCat("xla_compile_function_input_", function_id), *graph); + if (VLOG_IS_ON(2)) { + VLOG(2) << "XlaCompiler::CompileFunction: " + << dump_graph::DumpGraphToFile( + strings::StrCat("xla_compile_function_", function_id), + *graph); } // Optimize the graph before running the compiler. - // TODO(pbar): The constant folder currently does not simplify int32 - // operations for devices other than CPU. OptimizerOptions opts; + opts.set_do_common_subexpression_elimination(true); + opts.set_do_function_inlining(true); + opts.set_do_constant_folding(true); GraphOptimizer optimizer(opts); - OptimizeGraph(flr, &graph); - - if (VLOG_IS_ON(1)) { - dump_graph::DumpGraphToFile( - strings::StrCat("xla_compile_function_optimized_", function_id), - *graph); - } + optimizer.Optimize(flib_runtime_, flib_runtime_->env(), + /*device=*/nullptr, &graph, /*shape_map=*/nullptr); VLOG(1) << "===================================================="; TF_RETURN_IF_ERROR( - CompileGraph(function_id, std::move(graph), flr, args, result)); + CompileGraph(options, function_id, std::move(graph), args, result)); VLOG(1) << "===================================================="; + cache_[{function_id, args}] = *result; return Status::OK(); } @@ -158,7 +197,7 @@ Status XlaCompiler::BuildExecutable( build_options.set_has_hybrid_result( options_.local_executable_has_hybrid_result); - auto compile_result = local_client->Compile(result.computation, + auto compile_result = local_client->Compile(*result.computation, argument_layouts, build_options); if (!compile_result.ok()) { return compile_result.status(); @@ -226,33 +265,37 @@ Status BuildArguments(const std::vector& args, std::vector* input_shapes) { context_args->resize(args.size()); - // Argument numbers of arguments and variables that are to be passed to the + // Argument numbers of arguments and resources that are to be passed to the // XLA computation as runtime parameters. - std::vector parameters, variables; + std::vector parameters, resources; parameters.reserve(args.size()); - variables.reserve(args.size()); + resources.reserve(args.size()); for (std::vector::size_type i = 0; i < args.size(); ++i) { XlaContext::Argument& context_arg = (*context_args)[i]; + context_arg.kind = args[i].kind; context_arg.name = args[i].name; context_arg.value.constant_value = args[i].constant_value; context_arg.value.type = args[i].type; switch (args[i].kind) { case XlaCompiler::Argument::kVariable: - variables.push_back(i); - context_arg.value.is_constant = false; - context_arg.is_variable = true; + case XlaCompiler::Argument::kTensorArray: + case XlaCompiler::Argument::kStack: + context_arg.is_resource = true; + if (args[i].initialized) { + resources.push_back(i); + context_arg.value.is_constant = false; + } else { + context_arg.value.is_constant = true; + } + context_arg.tensor_array_size = args[i].tensor_array_size; break; case XlaCompiler::Argument::kParameter: parameters.push_back(i); context_arg.value.is_constant = false; break; - case XlaCompiler::Argument::kUninitializedVariable: - context_arg.is_variable = true; - context_arg.value.is_constant = true; - break; case XlaCompiler::Argument::kConstant: context_arg.value.is_constant = true; break; @@ -263,7 +306,7 @@ Status BuildArguments(const std::vector& args, // Append parameters containing variable values after the other runtime // parameters. - parameters.insert(parameters.end(), variables.begin(), variables.end()); + parameters.insert(parameters.end(), resources.begin(), resources.end()); if (parameters.empty()) { return Status::OK(); } @@ -273,10 +316,7 @@ Status BuildArguments(const std::vector& args, for (std::vector::size_type i = 0; i < input_shapes->size(); ++i) { const XlaCompiler::Argument& arg = args[parameters[i]]; // Computes the shapes of non-constant arguments. - xla::PrimitiveType type; - TF_RETURN_IF_ERROR(DataTypeToPrimitiveType(arg.type, &type)); - xla::ShapeUtil::PopulateShape(type, arg.shape.dim_sizes(), - &(*input_shapes)[i]); + (*input_shapes)[i] = arg.shape; (*input_mapping)[i] = parameters[i]; } @@ -304,22 +344,22 @@ Status BuildArguments(const std::vector& args, // variable states, generated by the symbolic evaluation. // If `has_side_effects` is true, the computation has side effects and should be // built even if it has no outputs. -// If `return_updated_values_for_all_variables` is true, all variables will be -// included in `variable_updates`, regardless of whether their value changed. +// 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 `*variable_updates` to a description of variables whose values are +// Sets `*resource_updates` to a description of resources whose values are // written by the computation; the variable writes are the last -// `variable_updates.size()` return values from the computation. Each entry in -// `variable_updates` is a (input_index, type) pair, where `input_index` is the +// `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& retvals, - const std::unordered_map& variable_map, - bool has_side_effects, bool return_updated_values_for_all_variables, + const std::vector>& resources, + bool has_side_effects, bool return_updated_values_for_all_resources, xla::ComputationBuilder* builder, xla::Computation* computation, - int* num_nonconst_outputs, - std::vector* variable_updates) { + int* num_computation_outputs, int* num_nonconst_outputs, + std::vector* resource_updates) { std::vector elems; elems.reserve(retvals.size()); for (const XlaContext::HandleOrConstant& retval : retvals) { @@ -329,31 +369,32 @@ Status BuildComputation( } *num_nonconst_outputs = elems.size(); - // Add return values for variables whose values have changed. - std::vector> variables; - variables.reserve(variable_map.size()); - for (const auto& entry : variable_map) { - variables.emplace_back(entry.first, &entry.second); + // Add return values for resources whose values have changed. + std::vector arg_vars; + arg_vars.reserve(resources.size()); + for (const auto& var : resources) { + if (var->arg_num >= 0) { + arg_vars.push_back(var.get()); + } } - std::sort(variables.begin(), variables.end(), - [](const std::pair& a, - const std::pair& b) { - return a.first < b.first; + std::sort(arg_vars.begin(), arg_vars.end(), + [](const XlaResource* a, const XlaResource* b) { + return a->arg_num < b->arg_num; }); - for (const auto& entry : variables) { - bool modified = - entry.second->value.handle() != entry.second->initial_value.handle(); - if (return_updated_values_for_all_variables || modified) { - variable_updates->emplace_back(); - XlaCompiler::VariableUpdate& update = variable_updates->back(); - update.input_index = entry.first; - update.type = entry.second->type; + for (const XlaResource* var : arg_vars) { + bool modified = var->value.handle() != var->initial_value.handle(); + if (return_updated_values_for_all_resources || modified) { + resource_updates->emplace_back(); + XlaCompiler::ResourceUpdate& update = resource_updates->back(); + update.input_index = var->arg_num; + update.type = var->type; update.modified = modified; - elems.push_back(entry.second->value); + elems.push_back(var->value); } } + *num_computation_outputs = elems.size(); if (!elems.empty() || has_side_effects) { // Builds a empty tuple return value for computations that have side effects // but have no return values. @@ -376,49 +417,71 @@ Status BuildComputation( return Status::OK(); } +void AssignMajorToMinorLayout(xla::Shape* shape) { + if (xla::ShapeUtil::IsTuple(*shape)) { + for (xla::Shape& elem_shape : *shape->mutable_tuple_shapes()) { + AssignMajorToMinorLayout(&elem_shape); + } + } else { + auto& minor_to_major = *shape->mutable_layout()->mutable_minor_to_major(); + minor_to_major.Resize(xla::ShapeUtil::Rank(*shape), 0); + std::iota(minor_to_major.rbegin(), minor_to_major.rend(), 0); + } +} + } // namespace -Status XlaCompiler::CompileGraph(string const& name, +Status XlaCompiler::CompileGraph(const XlaCompiler::CompileOptions& options, + string const& name, std::unique_ptr graph, - FunctionLibraryRuntime* flib, const std::vector& args, CompilationResult* result) { VLOG(1) << "Executing graph symbolically to populate ComputationBuilder."; + if (VLOG_IS_ON(2)) { + VLOG(2) << "XlaCompiler::CompileGraph: " + << dump_graph::DumpGraphToFile( + strings::StrCat("xla_compile_graph_", name), *graph); + } + // Report the error here if initialization failed. TF_RETURN_IF_ERROR(initialization_status_); + // Converts Tensorflow's graph control-flow constructs into functional + // control-flow that can be compiled into XLA code. + TF_RETURN_IF_ERROR( + FunctionalizeControlFlow(graph.get(), local_flib_def_.get())); + xla::ComputationBuilder builder(client(), name); XlaContext* context = new XlaContext(this, &builder, options_.allow_cpu_custom_calls, - options_.resolve_compile_time_constants); + options.resolve_compile_time_constants); core::ScopedUnref context_unref(context); - result->tuple_arg = options_.use_tuple_arg; + result->tuple_arg = options.use_tuple_arg; std::vector context_args; - TF_RETURN_IF_ERROR(BuildArguments(args, options_.use_tuple_arg, &builder, + TF_RETURN_IF_ERROR(BuildArguments(args, options.use_tuple_arg, &builder, &context_args, &result->input_mapping, &result->xla_input_shapes)); context->set_args(std::move(context_args)); - if (options_.prune_unreachable_nodes) { - PruneUnreachableNodes(graph.get()); - } - - TF_RETURN_IF_ERROR( - ExecuteGraph(context, std::move(graph), device_, flib, NextStepId())); + TF_RETURN_IF_ERROR(ExecuteGraph(context, std::move(graph), device_, + flib_runtime_, NextStepId())); int num_nonconst_outputs; + int num_computation_outputs; + result->computation = std::make_shared(); TF_RETURN_IF_ERROR(BuildComputation( - context->retvals(), context->variables(), context->has_side_effects(), - options_.return_updated_values_for_all_variables, &builder, - &result->computation, &num_nonconst_outputs, &result->variable_updates)); + context->retvals(), context->resources(), context->has_side_effects(), + options.return_updated_values_for_all_resources, &builder, + result->computation.get(), &num_computation_outputs, + &num_nonconst_outputs, &result->resource_updates)); result->requires_runtime_context = context->has_context_parameter(); // Tuple arguments and runtime context parameters are incompatible. - CHECK(!(options_.use_tuple_arg && result->requires_runtime_context)); + CHECK(!(options.use_tuple_arg && result->requires_runtime_context)); VLOG(2) << "Outputs: total: " << context->retvals().size() << " nonconstant: " << num_nonconst_outputs; @@ -434,37 +497,24 @@ Status XlaCompiler::CompileGraph(string const& name, } } - if (result->computation.IsNull()) { + if (result->computation->IsNull()) { return Status::OK(); } // Compute the output shapes, if there is a computation with non-constant // outputs. - auto computation_shape = client()->GetComputationShape(result->computation); + auto computation_shape = client()->GetComputationShape(*result->computation); if (!computation_shape.ok()) { return computation_shape.status(); } result->xla_output_shape.Swap( computation_shape.ValueOrDie()->mutable_result()); + VLOG(2) << "XLA output shape: " + << xla::ShapeUtil::HumanString(result->xla_output_shape); - auto num_computation_outputs = - (xla::ShapeUtil::IsTuple(result->xla_output_shape)) - ? xla::ShapeUtil::TupleElementCount(result->xla_output_shape) - : 1; // Tensorflow expects a major-to-minor order of results. - if (1 == num_computation_outputs) { - xla::Shape& s = result->xla_output_shape; - auto& minor_to_major = *s.mutable_layout()->mutable_minor_to_major(); - minor_to_major.Resize(xla::ShapeUtil::Rank(s), 0); - std::iota(minor_to_major.rbegin(), minor_to_major.rend(), 0); - } else { - for (xla::Shape& s : *result->xla_output_shape.mutable_tuple_shapes()) { - auto& minor_to_major = *s.mutable_layout()->mutable_minor_to_major(); - minor_to_major.Resize(xla::ShapeUtil::Rank(s), 0); - std::iota(minor_to_major.rbegin(), minor_to_major.rend(), 0); - } - } + AssignMajorToMinorLayout(&result->xla_output_shape); // Converts the output shapes to TensorShapes. int computation_output = 0; @@ -472,30 +522,30 @@ Status XlaCompiler::CompileGraph(string const& name, i < context->retvals().size(); ++i) { const XlaContext::HandleOrConstant& retval = context->retvals()[i]; if (!retval.is_constant) { - CHECK_LT(computation_output, num_nonconst_outputs); + CHECK_LT(computation_output, num_computation_outputs); OutputDescription& output = result->outputs[i]; output.is_constant = false; - if (num_nonconst_outputs > 1) { - output.shape = - XLAShapeToTensorShape(xla::ShapeUtil::GetTupleElementShape( - result->xla_output_shape, computation_output)); + if (num_computation_outputs > 1) { + TF_RETURN_IF_ERROR(XLAShapeToTensorShape( + xla::ShapeUtil::GetTupleElementShape(result->xla_output_shape, + computation_output), + &output.shape)); } else { - output.shape = XLAShapeToTensorShape(result->xla_output_shape); + TF_RETURN_IF_ERROR( + XLAShapeToTensorShape(result->xla_output_shape, &output.shape)); } ++computation_output; } } - for (std::vector::size_type i = 0; - i < result->variable_updates.size(); ++i) { + for (std::vector::size_type i = 0; + i < result->resource_updates.size(); ++i) { if (num_computation_outputs > 1) { - result->variable_updates[i].shape = - XLAShapeToTensorShape(xla::ShapeUtil::GetTupleElementShape( - result->xla_output_shape, computation_output)); + result->resource_updates[i].shape = xla::ShapeUtil::GetTupleElementShape( + result->xla_output_shape, computation_output); } else { CHECK_EQ(0, computation_output); - result->variable_updates[i].shape = - XLAShapeToTensorShape(result->xla_output_shape); + result->resource_updates[i].shape = result->xla_output_shape; } ++computation_output; } @@ -504,7 +554,6 @@ Status XlaCompiler::CompileGraph(string const& name, Status XlaCompiler::GetChannelHandle(const string& key, xla::ChannelHandle* channel) { - mutex_lock lock(mu_); auto result = channels_.emplace(key, xla::ChannelHandle()); if (result.second) { TF_ASSIGN_OR_RETURN(result.first->second, client()->CreateChannelHandle()); diff --git a/tensorflow/compiler/tf2xla/xla_compiler.h b/tensorflow/compiler/tf2xla/xla_compiler.h index 3d28ca374609df28647d243544dcbf8cbf33e706..b5987c8ac8bd056b8a9707666fa42a623066d22a 100644 --- a/tensorflow/compiler/tf2xla/xla_compiler.h +++ b/tensorflow/compiler/tf2xla/xla_compiler.h @@ -21,6 +21,7 @@ limitations under the License. #include "tensorflow/core/common_runtime/device.h" #include "tensorflow/core/common_runtime/device_mgr.h" #include "tensorflow/core/common_runtime/function.h" +#include "tensorflow/core/framework/function.h" #include "tensorflow/core/platform/env.h" #include "tensorflow/core/platform/mutex.h" #include "tensorflow/core/platform/notification.h" @@ -84,27 +85,31 @@ class XlaCompiler { // Argument is a compile-time constant. No associated runtime parameter. kConstant, - // Argument is a variable that has not been initialized yet. No associated - // runtime parameter. - kUninitializedVariable, - - // Argument is a variable that already has a value set. Expects a runtime - // parameter containing the current value. + // Argument is a Variable resource. Has an associated runtime parameter + // iff `initialized` is true. kVariable, + // Argument is a TensorArray resource. Has an associated runtime parameter + // iff `initialized` is true. + kTensorArray, + + // Argument is a Stack resource. Has an associated runtime parameter + // iff `initialized` is true. + kStack, + // Argument is a run-time parameter. kParameter, }; Kind kind = kInvalid; - // The type of the argument. If the argument is a resource variable, this + // 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; - // The shape of the argument. If the argument is a resource variable, this - // is the shape of the variable's value. - TensorShape shape; + // The shape of the argument. If the argument is a resource, this is the + // shape of the resource's value. + xla::Shape shape; // The value of the argument, if it is a compile-time constant. Must be a // host-memory tensor. @@ -112,6 +117,15 @@ class XlaCompiler { // The name of this argument, used for debugging. string name; + + // For a kVariable or kTensorArray, has this resource been initialized? + bool initialized = false; + + // For a kTensorArray, what is the array's declared size? (Used for lazy + // initialization.) + int64 tensor_array_size = -1; + + bool operator==(const Argument& other) const; }; struct OutputDescription { @@ -126,23 +140,23 @@ class XlaCompiler { }; // Describes a variable write side effect of the computation. - struct VariableUpdate { + struct ResourceUpdate { // Index of the input that contains the variable resource to write to. int input_index; // Type and shape of the tensor to be written back. DataType type; - TensorShape shape; + xla::Shape shape; // Was the value of the variable modified by the computation? - // (Always true, unless `return_updated_values_for_all_variables` is true.) + // (Always true, unless `return_updated_values_for_all_resources` is true.) bool modified; }; struct CompilationResult { // Vector that maps from the parameters of the XLA computation to their // original argument positions. To handle compile-time constant inputs and - // variables, the parameters to the XLA computation may be a subset of the + // resources, the parameters to the XLA computation may be a subset of the // original arguments, and are not necessarily in the same order.) std::vector input_mapping; @@ -165,22 +179,29 @@ class XlaCompiler { // containing both constant and non-constant results. std::vector outputs; - // Variables whose values were updated by the computation, ordered - // by return value position. Variable updates follow the non-constant + // Resources whose values were updated by the computation, ordered + // by return value position. Resource updates follow the non-constant // results in the outputs of XLA computation. - std::vector variable_updates; + std::vector resource_updates; // The XLA computation built from the tensorflow subgraph. May be null // if the output consists solely of compile-time constants. - xla::Computation computation; + std::shared_ptr computation; }; struct Options { - // Name of the compilation device to use. - DeviceType device_type = DeviceType(""); + // Name of the compilation device to use. Needs to be live only during + // XlaCompiler's constructor. + const DeviceType* device_type = nullptr; xla::Client* client = nullptr; + // Function library in which to find function definitions. Must be non-null. + const FunctionLibraryDefinition* flib_def = nullptr; + + // The graph def version to be compiled. + int graph_def_version = TF_GRAPH_DEF_VERSION; + // If 'allow_cpu_custom_calls' is true, kernels may make use of CustomCall() // for CPU; additionally, an optional XlaLocalRuntimeContext* may be passed // to the computation. @@ -192,29 +213,6 @@ class XlaCompiler { // stored in device memory. bool local_executable_has_hybrid_result = false; - // If 'resolve_compile_time_constants' is true, then outputs of a - // computation that are known to be compile-time constants will be returned - // as Tensors at compile-time, rather than as run-time outputs of the - // computation. - bool resolve_compile_time_constants = true; - - // If `use_tuple_arg` is true, a single tuple parameter will be used for all - // arguments; if false, each argument gets its own parameter. - bool use_tuple_arg = false; - - // If 'return_updated_values_for_all_variables' is true, then updated - // values of all resource variables arguments will be included in the - // 'variable_updates' of the computation, even if the variable was not - // modified by the computation. Used when compiling loop bodies to ensure - // the input and output signatures match. - bool return_updated_values_for_all_variables = false; - - // If 'prune_unreachable_nodes' is true, then nodes that are not - // dependencies of graph's _Retval nodes will be pruned before compilation. - // This is useful to prune stateful operators that should not be executed - // from a function body. - bool prune_unreachable_nodes = false; - // If not nullptr, populate_resource_manager is called with the // compilation device's resource manager when the compilation // device is created, and can be used to create metadata objects @@ -225,6 +223,27 @@ class XlaCompiler { explicit XlaCompiler(Options options); ~XlaCompiler(); + // Options pertaining to an individual call to CompileGraph() or + // CompileFunction(). + struct CompileOptions { + // If `use_tuple_arg` is true, a single tuple parameter will be used for all + // arguments; if false, each argument gets its own parameter. + bool use_tuple_arg = false; + + // If 'return_updated_values_for_all_resources' is true, then updated + // values of all resource resources arguments will be included in the + // 'resource_updates' of the computation, even if the resource was not + // modified by the computation. Used when compiling loop bodies to ensure + // the input and output signatures match. + bool return_updated_values_for_all_resources = false; + + // If 'resolve_compile_time_constants' is true, then outputs of a + // computation that are known to be compile-time constants will be returned + // as Tensors at compile-time, rather than as run-time outputs of the + // computation. + bool resolve_compile_time_constants = true; + }; + // Compiles a Tensorflow function `fn_name_attrs` into an XLA computation. // `args` describes the arguments to the function, each of which must either // be a runtime-parameter to the XLA computation, a compile-time constant, or @@ -235,7 +254,7 @@ class XlaCompiler { // arguments are returned as host memory tensors in the output list and are // not included in the XLA computation's outputs. The XLA computation is // null if there are no data-dependent outputs and no side effects. - Status CompileFunction(FunctionLibraryRuntime* flr, + Status CompileFunction(const CompileOptions& options, const NameAttrList& fn_name_attrs, const std::vector& args, CompilationResult* result); @@ -243,8 +262,8 @@ class XlaCompiler { // Compiles a tensorflow::Graph into an xla::Computation. // Similar to CompileFunction, but takes a Graph as input rather than a // function. - Status CompileGraph(string const& name, std::unique_ptr graph, - FunctionLibraryRuntime* flr, + Status CompileGraph(const CompileOptions& options, string const& name, + std::unique_ptr graph, const std::vector& args, CompilationResult* result); @@ -257,6 +276,7 @@ class XlaCompiler { xla::Client* client() const { return options_.client; } XlaCompilationDevice* device() const { return device_; } const DeviceMgr* device_mgr() const { return &device_mgr_; } + FunctionLibraryRuntime* flib_runtime() const { return flib_runtime_; } // Retrieves the channel handle associated with `key`. Allocates // a new channel handle if none exists. @@ -273,15 +293,32 @@ class XlaCompiler { // Returns the next step sequence number. int64 NextStepId(); - mutex mu_; - // Internal sequence number for steps executed on the compilation device. - int64 next_step_id_ GUARDED_BY(mu_); + int64 next_step_id_; XlaCompilationDevice* device_; // Owned by device_mgr_ DeviceMgr device_mgr_; - std::unordered_map channels_ GUARDED_BY(mu_); + // To avoid copying the client's function library, use a local function + // library and runtime for functions created as part of the functionalize + // control flow transformation. + std::unique_ptr local_flib_def_; + std::unique_ptr pflr_; + std::unique_ptr local_pflr_; + + FunctionLibraryRuntime* local_flib_runtime_; // owned by local_pflr_. + FunctionLibraryRuntime* flib_runtime_; // owned by pflr_. + + struct SignatureHash { + uint64 operator()( + const std::pair>& signature) const; + }; + + std::unordered_map>, + CompilationResult, SignatureHash> + cache_; + + std::unordered_map channels_; TF_DISALLOW_COPY_AND_ASSIGN(XlaCompiler); }; diff --git a/tensorflow/compiler/tf2xla/xla_compiler_test.cc b/tensorflow/compiler/tf2xla/xla_compiler_test.cc index 1cc7f4abd15798b29fe065c65c618b0166007b7e..a1e4dcb68478582950c50542e5e132f815d01114 100644 --- a/tensorflow/compiler/tf2xla/xla_compiler_test.cc +++ b/tensorflow/compiler/tf2xla/xla_compiler_test.cc @@ -22,8 +22,10 @@ limitations under the License. #include "tensorflow/compiler/xla/client/client_library.h" #include "tensorflow/compiler/xla/client/local_client.h" #include "tensorflow/compiler/xla/literal_util.h" +#include "tensorflow/compiler/xla/shape_util.h" #include "tensorflow/compiler/xla/tests/literal_test_util.h" #include "tensorflow/core/common_runtime/function.h" +#include "tensorflow/core/framework/common_shape_fns.h" #include "tensorflow/core/framework/resource_mgr.h" #include "tensorflow/core/framework/tensor_testutil.h" #include "tensorflow/core/graph/graph.h" @@ -75,6 +77,8 @@ class DummyReadResourceCC { scope.UpdateBuilder(&builder); scope.UpdateStatus(builder.Finalize(scope.graph(), &ret)); if (!scope.ok()) return; + scope.UpdateStatus(scope.DoShapeInference(ret)); + if (!scope.ok()) return; this->output_ = Output(ret, 0); } Node* node() const { return output_.node(); } @@ -85,6 +89,7 @@ class DummyReadResourceCC { REGISTER_OP("DummyReadResource") .Input("input: int32") .Output("output: int32") + .SetShapeFn(shape_inference::UnknownShape) .Doc(R"doc( A dummy Op. @@ -94,8 +99,35 @@ output: dummy output. REGISTER_XLA_OP(Name("DummyReadResource"), DummyReadResourceOp); +// DummyDuplicateOp is present purely to test multiple REGISTER_XLA_OP calls +// on the same Op name below. +class DummyDuplicateOp : public XlaOpKernel { + public: + explicit DummyDuplicateOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {} + void Compile(XlaOpKernelContext* ctx) override { + ctx->SetOutput(0, ctx->Input(0)); + } +}; + +REGISTER_OP("DummyDuplicateOp") + .Input("input: int32") + .Output("output: int32") + .Doc(R"doc( +A dummy Op. + +input: dummy input. +output: dummy output. +)doc"); + +REGISTER_XLA_OP(Name("DummyDuplicateOp").Device(DEVICE_CPU_XLA_JIT), + DummyDuplicateOp); +REGISTER_XLA_OP(Name("DummyDuplicateOp").Device(DEVICE_GPU_XLA_JIT), + DummyDuplicateOp); + class XlaCompilerTest : public ::testing::Test { protected: + XlaCompilerTest() : cpu_device_type_(DEVICE_CPU_XLA_JIT) {} + void SetUp() override { client_ = xla::ClientLibrary::LocalClientOrDie(); @@ -107,19 +139,13 @@ class XlaCompilerTest : public ::testing::Test { XlaCompiler::Options DefaultOptions() { XlaCompiler::Options options; - options.device_type = DeviceType(DEVICE_CPU_XLA_JIT); + options.device_type = &cpu_device_type_; options.client = client_; + options.flib_def = flib_def_.get(); return options; } - std::unique_ptr BuildFunctionLibraryRuntime( - const XlaCompiler& compiler) { - return std::unique_ptr(NewFunctionLibraryRuntime( - compiler.device_mgr(), /*env=*/nullptr, compiler.device(), - TF_GRAPH_DEF_VERSION, flib_def_.get(), OptimizerOptions(), - /*custom_kernel_creator=*/nullptr)); - } - + DeviceType cpu_device_type_; xla::Client* client_; std::unique_ptr flib_def_; }; @@ -127,15 +153,15 @@ class XlaCompilerTest : public ::testing::Test { // Tests compilation of an empty graph. TEST_F(XlaCompilerTest, EmptyReturnValues) { XlaCompiler compiler(DefaultOptions()); - auto flr = BuildFunctionLibraryRuntime(compiler); std::unique_ptr graph(new Graph(OpRegistry::Global())); XlaCompiler::CompilationResult result; - TF_ASSERT_OK(compiler.CompileGraph("add", std::move(graph), flr.get(), + TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), "add", + std::move(graph), /*args=*/{}, &result)); // No computation should be generated. - EXPECT_EQ(0, result.computation.handle().handle()); + EXPECT_EQ(0, result.computation->handle().handle()); } // Tests compilation and execution of a graph that adds two tensors. @@ -153,24 +179,23 @@ TEST_F(XlaCompilerTest, Simple) { std::vector args(2); args[0].kind = XlaCompiler::Argument::kParameter; args[0].type = DT_INT32; - args[0].shape = TensorShape({2}); + args[0].shape = xla::ShapeUtil::MakeShape(xla::S32, {2}); args[1].kind = XlaCompiler::Argument::kParameter; args[1].type = DT_INT32; - args[1].shape = TensorShape({2}); + args[1].shape = xla::ShapeUtil::MakeShape(xla::S32, {2}); // Compiles the graph. XlaCompiler compiler(DefaultOptions()); - auto flr = BuildFunctionLibraryRuntime(compiler); XlaCompiler::CompilationResult result; - TF_ASSERT_OK( - compiler.CompileGraph("add", std::move(graph), flr.get(), args, &result)); + TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), "add", + std::move(graph), args, &result)); // Tests that the generated computation works. std::unique_ptr param0_literal = - xla::LiteralUtil::CreateR1({7, 42}); + xla::Literal::CreateR1({7, 42}); std::unique_ptr param1_literal = - xla::LiteralUtil::CreateR1({-3, 101}); + xla::Literal::CreateR1({-3, 101}); std::unique_ptr param0_data = client_->TransferToServer(*param0_literal).ConsumeValueOrDie(); std::unique_ptr param1_data = @@ -178,13 +203,13 @@ TEST_F(XlaCompilerTest, Simple) { std::unique_ptr actual = client_ - ->Execute(result.computation, {param0_data.get(), param1_data.get()}) + ->Execute(*result.computation, {param0_data.get(), param1_data.get()}) .ConsumeValueOrDie(); std::unique_ptr actual_literal = client_->Transfer(*actual).ConsumeValueOrDie(); std::unique_ptr expected_literal = - xla::LiteralUtil::CreateR1({4, 143}); + xla::Literal::CreateR1({4, 143}); xla::LiteralTestUtil::ExpectEqual(*expected_literal, *actual_literal); } @@ -206,21 +231,21 @@ TEST_F(XlaCompilerTest, ConstantOutputs) { std::vector args(1); args[0].kind = XlaCompiler::Argument::kParameter; args[0].type = DT_INT32; - args[0].shape = TensorShape({2}); + args[0].shape = xla::ShapeUtil::MakeShape(xla::S32, {2}); + XlaCompiler::Options options = DefaultOptions(); + XlaCompiler compiler(options); { // Compiles the graph, with resolve_compile_time_constants enabled. - XlaCompiler::Options options = DefaultOptions(); - options.resolve_compile_time_constants = true; - XlaCompiler compiler(options); - auto flr = BuildFunctionLibraryRuntime(compiler); std::unique_ptr graph_copy(new Graph(OpRegistry::Global())); CopyGraph(*graph, graph_copy.get()); + XlaCompiler::CompileOptions compile_options; + compile_options.resolve_compile_time_constants = true; XlaCompiler::CompilationResult result; - TF_ASSERT_OK(compiler.CompileGraph("constants", std::move(graph_copy), - flr.get(), args, &result)); + TF_ASSERT_OK(compiler.CompileGraph(compile_options, "constants", + std::move(graph_copy), args, &result)); ASSERT_EQ(2, result.outputs.size()); EXPECT_TRUE(result.outputs[0].is_constant); @@ -230,34 +255,31 @@ TEST_F(XlaCompilerTest, ConstantOutputs) { // Tests that the generated computation works. std::unique_ptr param0_literal = - xla::LiteralUtil::CreateR1({7, 42}); + xla::Literal::CreateR1({7, 42}); std::unique_ptr param0_data = client_->TransferToServer(*param0_literal).ConsumeValueOrDie(); std::unique_ptr actual = - client_->Execute(result.computation, {param0_data.get()}) + client_->Execute(*result.computation, {param0_data.get()}) .ConsumeValueOrDie(); std::unique_ptr actual_literal = client_->Transfer(*actual).ConsumeValueOrDie(); std::unique_ptr expected_literal = - xla::LiteralUtil::CreateR1({-7, -42}); + xla::Literal::CreateR1({-7, -42}); xla::LiteralTestUtil::ExpectEqual(*expected_literal, *actual_literal); } { // Compiles the graph, with resolve_compile_time_constants disabled. - XlaCompiler::Options options = DefaultOptions(); - options.resolve_compile_time_constants = false; - XlaCompiler compiler(options); - auto flr = BuildFunctionLibraryRuntime(compiler); - std::unique_ptr graph_copy(new Graph(OpRegistry::Global())); CopyGraph(*graph, graph_copy.get()); + XlaCompiler::CompileOptions compile_options; + compile_options.resolve_compile_time_constants = false; XlaCompiler::CompilationResult result; - TF_ASSERT_OK(compiler.CompileGraph("constants", std::move(graph_copy), - flr.get(), args, &result)); + TF_ASSERT_OK(compiler.CompileGraph(compile_options, "constants", + std::move(graph_copy), args, &result)); ASSERT_EQ(2, result.outputs.size()); EXPECT_FALSE(result.outputs[0].is_constant); @@ -265,22 +287,21 @@ TEST_F(XlaCompilerTest, ConstantOutputs) { // Tests that the generated computation works. std::unique_ptr param0_literal = - xla::LiteralUtil::CreateR1({7, 42}); + xla::Literal::CreateR1({7, 42}); std::unique_ptr param0_data = client_->TransferToServer(*param0_literal).ConsumeValueOrDie(); std::unique_ptr actual = - client_->Execute(result.computation, {param0_data.get()}) + client_->Execute(*result.computation, {param0_data.get()}) .ConsumeValueOrDie(); std::unique_ptr actual_literal = client_->Transfer(*actual).ConsumeValueOrDie(); - std::unique_ptr expected0 = - xla::LiteralUtil::CreateR0(7); + std::unique_ptr expected0 = xla::Literal::CreateR0(7); std::unique_ptr expected1 = - xla::LiteralUtil::CreateR1({-7, -42}); + xla::Literal::CreateR1({-7, -42}); std::unique_ptr expected = - xla::LiteralUtil::MakeTuple({expected0.get(), expected1.get()}); + xla::Literal::MakeTuple({expected0.get(), expected1.get()}); xla::LiteralTestUtil::ExpectEqual(*expected, *actual_literal); } } @@ -299,7 +320,7 @@ TEST_F(XlaCompilerTest, ResourceManager) { std::vector args(1); args[0].kind = XlaCompiler::Argument::kParameter; args[0].type = DT_INT32; - args[0].shape = TensorShape({2}); + args[0].shape = xla::ShapeUtil::MakeShape(xla::S32, {2}); DummyResourceForTest* resource = new DummyResourceForTest(); @@ -312,13 +333,12 @@ TEST_F(XlaCompilerTest, ResourceManager) { }; options.populate_resource_manager = &populate_function; XlaCompiler compiler(options); - auto flr = BuildFunctionLibraryRuntime(compiler); EXPECT_EQ(0, resource->Get()); XlaCompiler::CompilationResult result; - TF_ASSERT_OK(compiler.CompileGraph("dummy", std::move(graph), flr.get(), args, - &result)); + TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), "dummy", + std::move(graph), args, &result)); EXPECT_EQ(1, resource->Get()); diff --git a/tensorflow/compiler/tf2xla/xla_context.cc b/tensorflow/compiler/tf2xla/xla_context.cc index 57d946509b65a6d5ebf013857cf52297559431ea..d4d493b456f668ecfbdd0164c573b9ae2aa810e9 100644 --- a/tensorflow/compiler/tf2xla/xla_context.cc +++ b/tensorflow/compiler/tf2xla/xla_context.cc @@ -22,6 +22,8 @@ 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_helpers.h" +#include "tensorflow/compiler/tf2xla/xla_op_kernel.h" #include "tensorflow/compiler/xla/client/client_library.h" #include "tensorflow/compiler/xla/client/computation_builder.h" #include "tensorflow/compiler/xla/layout_util.h" @@ -52,6 +54,10 @@ 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); } @@ -123,26 +129,18 @@ void XlaContext::AddSideEffects() { xla::ComputationBuilder* XlaContext::builder() { return builder_; } -Status XlaContext::CreateVariable(int variable_id, string name, DataType type, - const xla::ComputationDataHandle& handle) { - auto result = variables_.emplace(variable_id, Variable()); - if (!result.second) { - return errors::InvalidArgument("Duplicate ID ", variable_id, - " for variable ", name); - } - Variable& var = result.first->second; - var.name = std::move(name); - var.type = type; - var.initial_value = var.value = handle; - return Status::OK(); -} - -Status XlaContext::GetVariable(int variable_id, Variable** variable) { - auto it = variables_.find(variable_id); - if (it == variables_.end()) { - return errors::InvalidArgument("Unknown variable ID ", variable_id); - } - *variable = &it->second; +Status XlaContext::CreateResource(XlaResource::Kind kind, int arg_num, + string name, DataType type, + const xla::ComputationDataHandle& handle, + XlaResource** resource) { + resources_.emplace_back(new XlaResource); + *resource = resources_.back().get(); + XlaResource& r = **resource; + r.kind = kind; + r.arg_num = arg_num; + r.name = std::move(name); + r.type = type; + r.initial_value = r.value = handle; return Status::OK(); } @@ -174,22 +172,6 @@ const xla::Computation* XlaContext::GetOrCreateAdd(const DataType type) { }); } -const xla::Computation* XlaContext::GetOrCreateSigmoid(const DataType type) { - return LookupOrCreate(type, &sigmoid_func_, [this, type] { - const string type_string = DataTypeString(type); - VLOG(1) << "Building Sigmoid() for " << type_string; - xla::ComputationBuilder b(builder()->client(), - "sigmoid<" + type_string + ">"); - xla::PrimitiveType xla_type; - TF_CHECK_OK(DataTypeToPrimitiveType(type, &xla_type)); - auto x = b.Parameter(0, xla::ShapeUtil::MakeShape(xla_type, {}), "x"); - auto one = b.ConstantLiteral(xla::LiteralUtil::One(xla_type)); - auto minus_one = b.Neg(one); - b.Div(one, b.Add(b.Exp(b.Mul(x, minus_one)), one)); - return b.Build().ConsumeValueOrDie(); - }); -} - const xla::Computation* XlaContext::LookupOrCreate( DataType type, ComputationMap* out, const std::function& create) { diff --git a/tensorflow/compiler/tf2xla/xla_context.h b/tensorflow/compiler/tf2xla/xla_context.h index 657ead539124f8f4e21f306340c63322633a0c11..544921b9e38fb52e70b9f67ba10f7c79dc53c657 100644 --- a/tensorflow/compiler/tf2xla/xla_context.h +++ b/tensorflow/compiler/tf2xla/xla_context.h @@ -21,7 +21,6 @@ limitations under the License. #include #include "tensorflow/compiler/tf2xla/xla_compiler.h" -#include "tensorflow/compiler/tf2xla/xla_op_kernel.h" #include "tensorflow/compiler/xla/client/computation.h" #include "tensorflow/compiler/xla/client/computation_builder.h" #include "tensorflow/compiler/xla/xla_data.pb.h" @@ -31,6 +30,8 @@ limitations under the License. namespace tensorflow { +class XlaOpKernelContext; + // The XlaContext is the data structure that holds the state of an XLA // compilation, that is accessible from OpKernelContexts when compiling a // subgraph of Ops using XLA. @@ -51,20 +52,22 @@ class XlaContext : public ResourceBase { }; struct Argument { - // Descriptive name for the variable, for use in error messages. + XlaCompiler::Argument::Kind kind; + + // Descriptive name for the resource, for use in error messages. string name; - // Is this a variable? - bool is_variable; + // Is this a resource? + bool is_resource = false; HandleOrConstant value; + + int64 tensor_array_size = -1; }; // Retrieves the XlaContext of the current compilation. static XlaContext& Get(const OpKernelContext* ctx); - static XlaContext& Get(const XlaOpKernelContext* ctx) { - return Get(ctx->op_kernel_context()); - } + static XlaContext& Get(const XlaOpKernelContext* ctx); // Creates a new XlaContext. XlaContext(XlaCompiler* compiler, xla::ComputationBuilder* builder, @@ -105,33 +108,16 @@ class XlaContext : public ResourceBase { bool has_side_effects() const { return has_side_effects_; } - struct Variable { - // A descriptive name for the variable, used in error messages. - string name; - - // Current type and value of the variable. Uninitialized variables are - // represented by a default (zero) handle and type DT_INVALID. - // While the type of a variable is notionally fixed during execution, when - // a variable is first initialized we do not yet know its type, so we keep - // track of its type dynamically. - DataType type = DT_INVALID; - xla::ComputationDataHandle value; - - // Value of the variable at computation entry. Used to detect which - // variables have new values that need to be written back. - xla::ComputationDataHandle initial_value; - }; - - // Creates a variable with variable `variable_id` and initial type `type` and + // Creates a resource with resource `kind` and initial type `type` and // value `handle`. `name` is a descriptive name for use in error messages. - // Fails if the variable already exists. - Status CreateVariable(int variable_id, string name, DataType type, - const xla::ComputationDataHandle& handle); - - // Retrieves variable `variable_id`. Fails if the variable does not exist. - Status GetVariable(int variable_id, Variable** variable); + // Fails if the resource already exists. + Status CreateResource(XlaResource::Kind kind, int arg_num, string name, + DataType type, const xla::ComputationDataHandle& handle, + XlaResource** resource); - const std::unordered_map& variables() { return variables_; } + const std::vector>& resources() { + return resources_; + } // Get an XLA lambda to compute Max. This is cached in the // XlaContext since it may be used by multiple Ops. There is a @@ -143,11 +129,6 @@ class XlaContext : public ResourceBase { // separate specialization of the computation for each DataType. const xla::Computation* GetOrCreateAdd(const DataType type); - // Get an XLA lambda to compute Sigmoid. 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. - const xla::Computation* GetOrCreateSigmoid(const DataType type); - // The name of the XlaContext resource during symbolic graph execution. static const char kXlaContextResourceName[]; @@ -182,8 +163,8 @@ class XlaContext : public ResourceBase { // Does the computation have side effects, i.e., Send() calls? bool has_side_effects_ = false; - // Map from variable ID to the current value of each variable. - std::unordered_map variables_; + // Holds ownership of resources. The resources are not ordered. + std::vector> resources_; // Cache of prebuilt computations indexed by their type. using ComputationMap = std::map; diff --git a/tensorflow/compiler/tf2xla/xla_cpu_backend.cc b/tensorflow/compiler/tf2xla/xla_cpu_backend.cc new file mode 100644 index 0000000000000000000000000000000000000000..8286480e0ea07429adbe31ec4f16d043e321df0a --- /dev/null +++ b/tensorflow/compiler/tf2xla/xla_cpu_backend.cc @@ -0,0 +1,38 @@ +/* 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_op_registry.h" +#include "tensorflow/core/framework/kernel_def.pb.h" + +namespace tensorflow { + +bool CpuOpFilter(KernelDef* kdef) { + // TODO(b/34339814): implement inverse erf for double types and remove this + // workaround. + if (kdef->op() == "RandomStandardNormal") { + kdef->clear_constraint(); + // Change the type constraint to permit only DTD_FLOAT. + KernelDef::AttrConstraint* attr_constraint = kdef->add_constraint(); + attr_constraint->set_name("dtype"); + attr_constraint->mutable_allowed_values()->mutable_list()->add_type( + DT_FLOAT); + return true; + } + return true; +} + +REGISTER_XLA_BACKEND(DEVICE_CPU_XLA_JIT, kCpuAllTypes, CpuOpFilter); + +} // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/xla_gpu_backend.cc b/tensorflow/compiler/tf2xla/xla_gpu_backend.cc new file mode 100644 index 0000000000000000000000000000000000000000..d504613d232c779e47a506657d2825d052e726dc --- /dev/null +++ b/tensorflow/compiler/tf2xla/xla_gpu_backend.cc @@ -0,0 +1,35 @@ +/* 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_op_registry.h" +#include "tensorflow/core/framework/kernel_def.pb.h" + +namespace tensorflow { + +bool GpuOpFilter(KernelDef* kdef) { + // TODO(b/31361304): The GPU backend does not parallelize PRNG ops, leading to + // slow code. + // TODO(b/34969189) The implementation of TruncatedNormal generates illegal + // code on GPU. + if (kdef->op() == "RandomStandardNormal" || kdef->op() == "RandomUniform" || + kdef->op() == "RandomUniformInt" || kdef->op() == "TruncatedNormal") { + return false; + } + return true; +} + +REGISTER_XLA_BACKEND(DEVICE_GPU_XLA_JIT, kGpuAllTypes, GpuOpFilter); + +} // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/xla_helpers.cc b/tensorflow/compiler/tf2xla/xla_helpers.cc index 10d8b67bbd2d0e897e3ca55e584f575448a3a4fd..3af866f9be516beae7e6fa64b5a4cf1fef843f67 100644 --- a/tensorflow/compiler/tf2xla/xla_helpers.cc +++ b/tensorflow/compiler/tf2xla/xla_helpers.cc @@ -20,6 +20,7 @@ limitations under the License. #include "tensorflow/compiler/tf2xla/type_util.h" #include "tensorflow/compiler/tf2xla/xla_context.h" #include "tensorflow/compiler/xla/client/computation_builder.h" +#include "tensorflow/compiler/xla/types.h" #include "tensorflow/core/framework/tensor.h" #include "tensorflow/core/lib/gtl/array_slice.h" @@ -29,28 +30,28 @@ xla::ComputationDataHandle XlaHelpers::MinValue(xla::ComputationBuilder* b, DataType data_type) { xla::PrimitiveType type; TF_CHECK_OK(DataTypeToPrimitiveType(data_type, &type)); - return b->ConstantLiteral(xla::LiteralUtil::MinValue(type)); + return b->ConstantLiteral(xla::Literal::MinValue(type)); } xla::ComputationDataHandle XlaHelpers::MaxValue(xla::ComputationBuilder* b, DataType data_type) { xla::PrimitiveType type; TF_CHECK_OK(DataTypeToPrimitiveType(data_type, &type)); - return b->ConstantLiteral(xla::LiteralUtil::MaxValue(type)); + return b->ConstantLiteral(xla::Literal::MaxValue(type)); } xla::ComputationDataHandle XlaHelpers::Zero(xla::ComputationBuilder* b, DataType data_type) { xla::PrimitiveType type; TF_CHECK_OK(DataTypeToPrimitiveType(data_type, &type)); - return b->ConstantLiteral(xla::LiteralUtil::Zero(type)); + return b->ConstantLiteral(xla::Literal::Zero(type)); } xla::ComputationDataHandle XlaHelpers::One(xla::ComputationBuilder* b, DataType data_type) { xla::PrimitiveType type; TF_CHECK_OK(DataTypeToPrimitiveType(data_type, &type)); - return b->ConstantLiteral(xla::LiteralUtil::One(type)); + return b->ConstantLiteral(xla::Literal::One(type)); } xla::ComputationDataHandle XlaHelpers::IntegerLiteral( @@ -60,28 +61,28 @@ xla::ComputationDataHandle XlaHelpers::IntegerLiteral( TF_CHECK_OK(DataTypeToPrimitiveType(data_type, &type)); switch (type) { case xla::U8: - literal = *xla::LiteralUtil::CreateR0(value); + literal = *xla::Literal::CreateR0(value); break; case xla::U32: - literal = *xla::LiteralUtil::CreateR0(value); + literal = *xla::Literal::CreateR0(value); break; case xla::U64: - literal = *xla::LiteralUtil::CreateR0(value); + literal = *xla::Literal::CreateR0(value); break; case xla::S8: - literal = *xla::LiteralUtil::CreateR0(value); + literal = *xla::Literal::CreateR0(value); break; case xla::S32: - literal = *xla::LiteralUtil::CreateR0(value); + literal = *xla::Literal::CreateR0(value); break; case xla::S64: - literal = *xla::LiteralUtil::CreateR0(value); + literal = *xla::Literal::CreateR0(value); break; case xla::F32: - literal = *xla::LiteralUtil::CreateR0(value); + literal = *xla::Literal::CreateR0(value); break; case xla::F64: - literal = *xla::LiteralUtil::CreateR0(value); + literal = *xla::Literal::CreateR0(value); break; case xla::PRED: LOG(FATAL) << "pred element type is not integral"; @@ -89,7 +90,9 @@ xla::ComputationDataHandle XlaHelpers::IntegerLiteral( case xla::U16: LOG(FATAL) << "u16/s16 literals not yet implemented"; case xla::F16: - LOG(FATAL) << "f16 literals not yet implemented"; + literal = + *xla::Literal::CreateR0(static_cast(value)); + break; case xla::TUPLE: LOG(FATAL) << "tuple element type is not integral"; case xla::OPAQUE: @@ -107,6 +110,9 @@ xla::ComputationDataHandle XlaHelpers::FloatLiteral(xla::ComputationBuilder* b, xla::PrimitiveType type; TF_CHECK_OK(DataTypeToPrimitiveType(data_type, &type)); switch (type) { + case xla::F16: + return b->ConstantR0(static_cast(value)); + break; case xla::F32: return b->ConstantR0(static_cast(value)); break; @@ -199,4 +205,13 @@ Status XlaHelpers::OneHot(xla::ComputationBuilder* builder, int64 depth, return Status::OK(); } +xla::ComputationDataHandle XlaHelpers::PadWithZeros( + xla::ComputationBuilder* builder, const xla::ComputationDataHandle& x, + int count) { + xla::ComputationDataHandle zero = builder->ConstantR1({0}); + std::vector xs(count + 1, zero); + xs[0] = builder->Reshape(x, {1}); + return builder->ConcatInDim(xs, 0); +} + } // end namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/xla_helpers.h b/tensorflow/compiler/tf2xla/xla_helpers.h index a141ee05c13ed2e09fab69946ba400ab6cd628a9..3f92008465303fcbeba78b844da8376b7b544d4b 100644 --- a/tensorflow/compiler/tf2xla/xla_helpers.h +++ b/tensorflow/compiler/tf2xla/xla_helpers.h @@ -13,7 +13,7 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -// This file defines helper routines for the TLA device. +// This file defines helper routines for the XLA device. #ifndef TENSORFLOW_COMPILER_TF2XLA_XLA_HELPERS_H_ #define TENSORFLOW_COMPILER_TF2XLA_XLA_HELPERS_H_ @@ -77,6 +77,11 @@ class XlaHelpers { const xla::ComputationDataHandle& on_value, const xla::ComputationDataHandle& off_value, xla::ComputationDataHandle* one_hot); + + // Pads 'x' with 'count' zeros. 'x' must have 1 element. + static xla::ComputationDataHandle PadWithZeros( + xla::ComputationBuilder* builder, const xla::ComputationDataHandle& x, + int count); }; } // end namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/xla_local_runtime_context.h b/tensorflow/compiler/tf2xla/xla_local_runtime_context.h index cd773d64ed4154aa2a05ac2d15e9358614239b1f..dca420d6ee3fec45f88ac3b450ab0cb4fb83d38a 100644 --- a/tensorflow/compiler/tf2xla/xla_local_runtime_context.h +++ b/tensorflow/compiler/tf2xla/xla_local_runtime_context.h @@ -23,7 +23,7 @@ limitations under the License. // actually used. E.g. some ahead-of-time compiled computations don't need a // thread pool. namespace Eigen { -class ThreadPoolDevice; +struct ThreadPoolDevice; } namespace tensorflow { diff --git a/tensorflow/compiler/tf2xla/xla_op_kernel.cc b/tensorflow/compiler/tf2xla/xla_op_kernel.cc index 53dcdec7a25a2ffc2d5986fe1587ca4dbfd79287..e36eafa6e4bbce83bafb67b40378a65fc28e841f 100644 --- a/tensorflow/compiler/tf2xla/xla_op_kernel.cc +++ b/tensorflow/compiler/tf2xla/xla_op_kernel.cc @@ -38,7 +38,8 @@ xla::ComputationBuilder* XlaOpKernelContext::builder() const { static const XlaExpression* CastExpressionFromTensor(const Tensor& tensor) { const XlaExpression* expression = reinterpret_cast(tensor.tensor_data().data()); - CHECK(expression->handle().handle() != 0 || expression->variable_id() >= 0); + CHECK(expression->handle().handle() != 0 || + expression->resource() != nullptr); VLOG(1) << "Fetched T" << expression->handle().handle(); return expression; } @@ -118,22 +119,15 @@ Status XlaOpKernelContext::ConstantInputReshaped( xla::Layout layout = xla::LayoutUtil::MakeLayout(layout_indices); // Ask the XLA compiler to evaluate the data handle to a literal. - xla::StatusOr> computed = + xla::StatusOr> computed = builder()->ComputeConstant(handle, &layout); if (!computed.ok()) { return errors::InvalidArgument( "Error evaluating ", context_->op_kernel().name(), " input ", index, ": ", computed.status().error_message()); } - // Fetch the literal from the compiler service. - xla::StatusOr> constant = - builder()->client()->Transfer(*computed.ValueOrDie()); - if (!constant.ok()) { - return errors::InvalidArgument( - "Error evaluating ", context_->op_kernel().name(), " input ", index, - ": ", constant.status().error_message()); - } - constant_literal->Swap(constant.ValueOrDie().get()); + constant_literal->Swap(computed.ValueOrDie().get()); + return Status::OK(); } @@ -143,9 +137,9 @@ static Status LiteralToInt64Scalar(const xla::Literal& literal, int64* out) { return errors::InvalidArgument("value is not a scalar"); } if (literal.shape().element_type() == xla::S32) { - *out = xla::LiteralUtil::Get(literal, {}); + *out = literal.Get({}); } else if (literal.shape().element_type() == xla::S64) { - *out = xla::LiteralUtil::Get(literal, {}); + *out = literal.Get({}); } else { return errors::InvalidArgument("value must be either int32 or int64"); } @@ -167,11 +161,11 @@ static Status LiteralToInt64Vector(const xla::Literal& literal, int64 size = xla::ShapeUtil::ElementsIn(literal.shape()); if (literal.shape().element_type() == xla::S32) { for (int64 i = 0; i < size; ++i) { - out->push_back(xla::LiteralUtil::Get(literal, {i})); + out->push_back(literal.Get({i})); } } else if (literal.shape().element_type() == xla::S64) { for (int64 i = 0; i < size; ++i) { - out->push_back(xla::LiteralUtil::Get(literal, {i})); + out->push_back(literal.Get({i})); } } else { return errors::InvalidArgument("value must be either int32 or int64"); @@ -186,6 +180,31 @@ Status XlaOpKernelContext::ConstantInputAsIntVector(int index, return LiteralToInt64Vector(literal, out); } +Status XlaOpKernelContext::ConstantInputAsInt64Literal(int index, + xla::Literal* out) { + xla::Literal literal; + TF_RETURN_IF_ERROR(ConstantInput(index, &literal)); + switch (literal.shape().element_type()) { + case xla::S32: + out->Clear(); + *out->mutable_shape() = literal.shape(); + out->mutable_shape()->set_element_type(xla::S64); + for (int32 x : literal.s32s()) { + out->add_s64s(x); + } + return Status::OK(); + + case xla::S64: + out->Swap(&literal); + return Status::OK(); + + default: + return errors::InvalidArgument( + "Invalid argument to ConstantInputAsInt64Literal: ", + xla::ShapeUtil::HumanString(literal.shape())); + } +} + // TODO(phawkins): validate that the dimensions form a valid shape, fail // gracefully if they do not. Status XlaOpKernelContext::ConstantInputAsShape(int index, TensorShape* shape) { @@ -226,11 +245,9 @@ Status XlaOpKernelContext::ReadVariableInput( int index, xla::ComputationDataHandle* value) { const Tensor& tensor = context_->input(index); const XlaExpression* expression = CastExpressionFromTensor(tensor); - int variable_id = expression->variable_id(); - - XlaContext::Variable* variable; - XlaContext& context = XlaContext::Get(this); - TF_RETURN_IF_ERROR(context.GetVariable(variable_id, &variable)); + XlaResource* variable = expression->resource(); + TF_RET_CHECK(variable != nullptr); + TF_RET_CHECK(variable->kind == XlaResource::kVariable); if (variable->value.handle() == 0) { return errors::InvalidArgument("Read of uninitialized variable ", variable->name); @@ -239,28 +256,13 @@ Status XlaOpKernelContext::ReadVariableInput( return Status::OK(); } -string XlaOpKernelContext::VariableDebugString(int index) { - const Tensor& tensor = context_->input(index); - const XlaExpression* expression = CastExpressionFromTensor(tensor); - int variable_id = expression->variable_id(); - - XlaContext::Variable* variable; - XlaContext& context = XlaContext::Get(this); - if (!context.GetVariable(variable_id, &variable).ok()) { - return ""; - } - return variable->name; -} - Status XlaOpKernelContext::GetVariableTypeAndShape(int index, DataType* type, TensorShape* shape) const { const Tensor& tensor = context_->input(index); const XlaExpression* expression = CastExpressionFromTensor(tensor); - int variable_id = expression->variable_id(); - - XlaContext::Variable* variable; - XlaContext& context = XlaContext::Get(this); - TF_RETURN_IF_ERROR(context.GetVariable(variable_id, &variable)); + XlaResource* variable = expression->resource(); + TF_RET_CHECK(variable != nullptr); + TF_RET_CHECK(variable->kind == XlaResource::kVariable); if (variable->value.handle() == 0) { return errors::InvalidArgument("Read of uninitialized variable ", variable->name); @@ -270,7 +272,8 @@ Status XlaOpKernelContext::GetVariableTypeAndShape(int index, DataType* type, if (!shape_or_status.ok()) { return shape_or_status.status(); } - *shape = XLAShapeToTensorShape(*shape_or_status.ValueOrDie()); + TF_RETURN_IF_ERROR( + XLAShapeToTensorShape(*shape_or_status.ValueOrDie(), shape)); return Status::OK(); } @@ -287,10 +290,11 @@ void XlaOpKernelContext::SetOutput(int index, // 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, XLAShapeToTensorShape(*shape.ValueOrDie()), &output)); + TensorShape tensor_shape; + OP_REQUIRES_OK(context_, + XLAShapeToTensorShape(*shape.ValueOrDie(), &tensor_shape)); + OP_REQUIRES_OK(context_, + context_->allocate_output(index, tensor_shape, &output)); // The expression is stored in the tensor's data buffer. Fill in the // fields now. @@ -320,25 +324,34 @@ void XlaOpKernelContext::SetConstantOutput(int index, const Tensor& constant) { expression->set_constant_value(constant); } -void XlaOpKernelContext::SetVariableOutput(int index, int variable_id) { +void XlaOpKernelContext::SetResourceOutput(int index, XlaResource* resource) { Tensor* output = nullptr; - // The shape of the output tensor is the shape of the variable resource - // (i.e., a scalar), not the shape of the variable's value. + // 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_variable_id(variable_id); + expression->set_resource(resource); +} + +Status XlaOpKernelContext::GetResourceInput(int index, XlaResource** resource) { + const XlaExpression* expression = + CastExpressionFromTensor(context_->input(index)); + TF_RET_CHECK(expression->resource() != nullptr); + *resource = expression->resource(); + return Status::OK(); } Status XlaOpKernelContext::AssignVariable( - int index, DataType type, const xla::ComputationDataHandle& handle) { + int input_index, DataType type, const xla::ComputationDataHandle& handle) { + TF_RET_CHECK(handle.handle() != 0); SetOpHasSideEffects(); const XlaExpression* expression = - CastExpressionFromTensor(context_->input(index)); - XlaContext& context = XlaContext::Get(this); - XlaContext::Variable* variable; - TF_RETURN_IF_ERROR(context.GetVariable(expression->variable_id(), &variable)); + CastExpressionFromTensor(context_->input(input_index)); + XlaResource* variable = expression->resource(); + TF_RET_CHECK(variable != nullptr); + TF_RET_CHECK(variable->kind == XlaResource::kVariable); if (!((variable->type == DT_INVALID && type != DT_INVALID) || (variable->type == type))) { return errors::InvalidArgument( @@ -354,8 +367,8 @@ void XlaOpKernelContext::SetOpHasSideEffects() { XlaContext::Get(context_).AddSideEffects(); } -const XlaCompiler::Options& XlaOpKernelContext::GetCompilerOptions() const { - return XlaContext::Get(context_).compiler()->options(); +XlaCompiler* XlaOpKernelContext::compiler() const { + return XlaContext::Get(context_).compiler(); } void XlaOpKernelContext::CtxFailure(Status s) { context_->CtxFailure(s); } @@ -373,11 +386,6 @@ const xla::Computation* XlaOpKernelContext::GetOrCreateAdd( return XlaContext::Get(context_).GetOrCreateAdd(type); } -const xla::Computation* XlaOpKernelContext::GetOrCreateSigmoid( - const DataType type) { - return XlaContext::Get(context_).GetOrCreateSigmoid(type); -} - XlaOpKernel::XlaOpKernel(OpKernelConstruction* context) : OpKernel(context) {} void XlaOpKernel::Compute(OpKernelContext* context) { diff --git a/tensorflow/compiler/tf2xla/xla_op_kernel.h b/tensorflow/compiler/tf2xla/xla_op_kernel.h index 60e3b59d32a8eb64809011743184d0d4cd48968f..30b794c8c198cae6bf3b11794b35049b729063e1 100644 --- a/tensorflow/compiler/tf2xla/xla_op_kernel.h +++ b/tensorflow/compiler/tf2xla/xla_op_kernel.h @@ -110,6 +110,9 @@ class XlaOpKernelContext { // Converts a constant 1D int32 or int64 tensor into a vector of int64s. Status ConstantInputAsIntVector(int index, std::vector* out); + // Converts a constant int32 or int64 Tensor into an xla int64 Literal. + Status ConstantInputAsInt64Literal(int index, xla::Literal* out); + // Converts a constant 1D int32 or int64 tensor into a TensorShape. Status ConstantInputAsShape(int index, TensorShape* shape); @@ -145,6 +148,12 @@ class XlaOpKernelContext { // Variables + // Sets '*resource' to the resource associated with input `index`. + Status GetResourceInput(int index, XlaResource** resource); + + // Sets output 'index' to be a reference to resource 'resource'. + void SetResourceOutput(int index, XlaResource* resource); + // Sets `*type` and `*shape` to the current type and shape of a variable's // value. Status GetVariableTypeAndShape(int index, DataType* type, @@ -154,18 +163,11 @@ class XlaOpKernelContext { // 'index'. Status ReadVariableInput(int index, xla::ComputationDataHandle* value); - // Sets output 'index' to be a reference to variable 'variable_id'. Used - // to propagate resource variables through the compilation. - void SetVariableOutput(int index, int variable_id); - // Assigns the value `handle` to the variable referenced by input - // `variable_index`. Marks the operator as having side effects. - Status AssignVariable(int variable_index, DataType type, + // `input_index`. Marks the operator as having side effects. + Status AssignVariable(int input_index, DataType type, const xla::ComputationDataHandle& handle); - // Returns a human-readable debug string describing 'variable_index'. - string VariableDebugString(int variable_index); - // Helper routines for the OP_REQUIRES macros void CtxFailure(Status s); void CtxFailureWithWarning(Status s); @@ -183,10 +185,9 @@ class XlaOpKernelContext { // Returns the underlying OpKernelContext. Use rarely. OpKernelContext* op_kernel_context() const { return context_; } - // Returns the options passed to the XlaCompiler that is being - // run. Used for, e.g., While to inherit options needed for nested - // computation. - const XlaCompiler::Options& GetCompilerOptions() const; + // Returns the XlaCompiler that is performing the compilation. Used for, e.g., + // While to compile nested computations. + XlaCompiler* compiler() const; // TODO(phawkins): find a better home for these helpers. @@ -200,11 +201,6 @@ class XlaOpKernelContext { // separate specialization of the computation for each DataType. const xla::Computation* GetOrCreateAdd(const DataType type); - // Get an XLA lambda to compute Sigmoid. 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. - const xla::Computation* GetOrCreateSigmoid(const DataType type); - private: OpKernelContext* const context_; }; diff --git a/tensorflow/compiler/tf2xla/xla_op_registry.cc b/tensorflow/compiler/tf2xla/xla_op_registry.cc index 13fdfc3b0c82e3d0018c72eebaaf7fa313111648..2cf3d4c1f2563b995d5cd84dc380928552b20f00 100644 --- a/tensorflow/compiler/tf2xla/xla_op_registry.cc +++ b/tensorflow/compiler/tf2xla/xla_op_registry.cc @@ -24,6 +24,8 @@ limitations under the License. #include "tensorflow/core/common_runtime/device_factory.h" #include "tensorflow/core/common_runtime/local_device.h" #include "tensorflow/core/framework/device_base.h" +#include "tensorflow/core/framework/kernel_def.pb.h" +#include "tensorflow/core/framework/node_def.pb.h" #include "tensorflow/core/platform/mem.h" #include "tensorflow/core/platform/stream_executor_no_cuda.h" @@ -31,17 +33,58 @@ namespace tensorflow { const char* const DEVICE_CPU_XLA_JIT = "XLA_CPU_JIT"; const char* const DEVICE_GPU_XLA_JIT = "XLA_GPU_JIT"; - -// Is platform 'id' supported by XLA? -static bool IsPlatformSupported(perftools::gputools::Platform::Id id) { - auto platform = perftools::gputools::MultiPlatformManager::PlatformWithId(id); - if (!platform.ok()) return false; - return xla::ClientLibrary::GetOrCreateLocalClient(platform.ValueOrDie()).ok(); +const char* const DEVICE_XLA_CPU = "XLA_CPU"; +const char* const DEVICE_XLA_GPU = "XLA_GPU"; + +static Status LaunchOpHasKernelForDevice(const DeviceType& device_type) { + const OpDef* op_def; + TF_RETURN_IF_ERROR(OpRegistry::Global()->LookUpOpDef("_XlaLaunch", &op_def)); + NodeDef node_def; + node_def.set_name("_XlaLaunch-op"); + node_def.set_op("_XlaLaunch"); + string kernel_class_name; + TF_RETURN_IF_ERROR(FindKernelDef(device_type, node_def, /*KernelDef*/ nullptr, + &kernel_class_name)); + VLOG(1) << "LaunchOpHasKernelForDevice" + << " kernel_class_name: " << kernel_class_name; + return Status::OK(); } XlaOpRegistry::XlaOpRegistry() = default; XlaOpRegistry::~XlaOpRegistry() = default; +// TODO(b/64575122) consider adding more sophisticated definitions of +// compatibility if needed by future use cases. +/* static */ bool XlaOpRegistry::IsCompatible(const OpRegistration& x, + const OpRegistration& y) { + if (x.name != y.name) return true; + // The registrations refer to the same Op: ensures they are compatible and + // are restricted to different device whitelists. + if (x.compilation_only != y.compilation_only) { + LOG(WARNING) << "Registrations of " << x.name + << " have incompatible compilation_only settings."; + return false; + } + if (x.allow_resource_types != y.allow_resource_types) { + LOG(WARNING) << "Registrations of " << x.name + << " have incompatible allow_resource_types settings."; + return false; + } + if (!x.has_device_whitelist || !y.has_device_whitelist) { + LOG(WARNING) << "Registrations of " << x.name + << " do not both have device whitelists."; + return false; + } + for (const auto& device : x.device_whitelist) { + if (y.device_whitelist.count(device) != 0) { + LOG(WARNING) << "Multiple registrations of " << x.name << " on device " + << device; + return false; + } + } + return true; +} + /* static */ void XlaOpRegistry::RegisterCompilationDevice( const string& device_name, const DeviceRegistration& registration) { XlaOpRegistry& registry = Instance(); @@ -73,7 +116,7 @@ XlaOpRegistry::~XlaOpRegistry() = default; // GetCompilationDevice is called. static void* registration_init = [®istry]() { mutex_lock lock(registry.mutex_); - if (IsPlatformSupported(perftools::gputools::host::kHostPlatformId)) { + if (LaunchOpHasKernelForDevice(DeviceType(DEVICE_CPU)).ok()) { DeviceRegistration& registration = registry.compilation_devices_[DEVICE_CPU]; registration.compilation_device_name = DEVICE_CPU_XLA_JIT; @@ -81,7 +124,7 @@ XlaOpRegistry::~XlaOpRegistry() = default; registration.enable_jit_by_default = false; registration.compile_resource_ops = false; } - if (IsPlatformSupported(perftools::gputools::cuda::kCudaPlatformId)) { + if (LaunchOpHasKernelForDevice(DeviceType(DEVICE_GPU)).ok()) { DeviceRegistration& registration = registry.compilation_devices_[DEVICE_GPU]; registration.compilation_device_name = DEVICE_GPU_XLA_JIT; @@ -109,8 +152,10 @@ void XlaOpRegistry::RegisterCompilationKernels() { OpRegistryInterface* op_registry = OpRegistry::Global(); for (const auto& op : registry.ops_) { + const string& op_name = op.first; + const std::unique_ptr& op_registration = op.second; const OpDef* op_def; - TF_CHECK_OK(op_registry->LookUpOpDef(op.first, &op_def)); + TF_CHECK_OK(op_registry->LookUpOpDef(op_name, &op_def)); std::unordered_set type_attrs; for (const OpDef::AttrDef& attr_def : op_def->attr()) { @@ -120,24 +165,24 @@ void XlaOpRegistry::RegisterCompilationKernels() { } // Checks there are no type constraints referring to unknown attributes. - for (const auto& constraint : op.second->type_constraints) { + for (const auto& constraint : op_registration->type_constraints) { if (type_attrs.find(constraint.first) == type_attrs.end()) { LOG(FATAL) << "Unknown type attribute " << constraint.first - << " in XLA op registration for " << op.first; + << " in XLA op registration for " << op_name; } } for (auto& backend : registry.backends_) { // If the operator has a device whitelist, only register on whitelisted // devices. - if (op.second->has_device_whitelist && - op.second->device_whitelist.find(backend.first) == - op.second->device_whitelist.end()) { + if (op_registration->has_device_whitelist && + op_registration->device_whitelist.find(backend.first) == + op_registration->device_whitelist.end()) { continue; } std::unique_ptr kdef(new KernelDef); - kdef->set_op(op.second->name); + kdef->set_op(op_registration->name); kdef->set_device_type(backend.first); // Constrain each type attribute to the intersection of: @@ -151,15 +196,15 @@ void XlaOpRegistry::RegisterCompilationKernels() { auto* allowed_values = attr_constraint->mutable_allowed_values()->mutable_list(); - auto it = op.second->type_constraints.find(type_attr); + auto it = op_registration->type_constraints.find(type_attr); for (DataType dtype : backend.second.supported_types) { - if (it == op.second->type_constraints.end() || - (it != op.second->type_constraints.end() && + if (it == op_registration->type_constraints.end() || + (it != op_registration->type_constraints.end() && it->second.find(dtype) != it->second.end())) { allowed_values->add_type(dtype); } } - if (op.second->allow_resource_types) { + if (op_registration->allow_resource_types) { allowed_values->add_type(DT_RESOURCE); } } @@ -168,10 +213,10 @@ void XlaOpRegistry::RegisterCompilationKernels() { continue; } VLOG(2) << "XLA op registration: device: " << backend.first - << " op: " << op.first; + << " op: " << op_name; registry.kernel_registrars_.emplace_back( new kernel_factory::OpKernelRegistrar( - new KernelDef(*kdef), "XlaJitOp", op.second->factory)); + new KernelDef(*kdef), "XlaJitOp", op_registration->factory)); backend.second.kernel_defs.push_back(std::move(kdef)); } } @@ -186,7 +231,12 @@ std::vector XlaOpRegistry::DeviceKernels( CHECK(it != registry.backends_.end()) << "Unknown backend " << compilation_device_name; for (const std::unique_ptr& k : it->second.kernel_defs) { - if (!registry.ops_.at(k->op())->compilation_only) { + auto op_iter = registry.ops_.find(k->op()); + CHECK(op_iter != registry.ops_.end()); + // The test in IsCompatible ensures that if there are multiple matching + // registrations for this op name, they all have the same value of + // compilation_only, so only the first match needs to be tested. + if (!op_iter->second->compilation_only) { kernels.push_back(k.get()); } } @@ -261,11 +311,16 @@ XlaOpRegistrar::XlaOpRegistrar( std::unique_ptr registration) { XlaOpRegistry& registry = XlaOpRegistry::Instance(); mutex_lock lock(registry.mutex_); - auto result = registry.ops_.emplace(registration->name, nullptr); - if (!result.second) { - LOG(FATAL) << "Duplicate XLA op registration " << registration->name; + auto existing_ops = registry.ops_.equal_range(registration->name); + for (auto existing = existing_ops.first; existing != existing_ops.second; + ++existing) { + if (!XlaOpRegistry::IsCompatible(*existing->second, *registration)) { + LOG(FATAL) + << "XLA op registration " << registration->name + << " is incompatible with existing registration of the same name."; + } } - result.first->second = std::move(registration); + registry.ops_.emplace(registration->name, std::move(registration)); } XlaBackendRegistrar::XlaBackendRegistrar( @@ -275,35 +330,4 @@ XlaBackendRegistrar::XlaBackendRegistrar( registry.RegisterBackend(name.ToString(), types, op_filter); } -bool CpuOpFilter(KernelDef* kdef) { - // TODO(b/34339814): implement inverse erf for double types and remove this - // workaround. - if (kdef->op() == "RandomStandardNormal") { - kdef->clear_constraint(); - // Change the type constraint to permit only DTD_FLOAT. - KernelDef::AttrConstraint* attr_constraint = kdef->add_constraint(); - attr_constraint->set_name("dtype"); - attr_constraint->mutable_allowed_values()->mutable_list()->add_type( - DT_FLOAT); - return true; - } - return true; -} - -REGISTER_XLA_BACKEND(DEVICE_CPU_XLA_JIT, kCpuAllTypes, CpuOpFilter); - -bool GpuOpFilter(KernelDef* kdef) { - // TODO(b/31361304): The GPU backend does not parallelize PRNG ops, leading to - // slow code. - // TODO(b/34969189) The implementation of TruncatedNormal generates illegal - // code on GPU. - if (kdef->op() == "RandomStandardNormal" || kdef->op() == "RandomUniform" || - kdef->op() == "RandomUniformInt" || kdef->op() == "TruncatedNormal") { - return false; - } - return true; -} - -REGISTER_XLA_BACKEND(DEVICE_GPU_XLA_JIT, kGpuAllTypes, GpuOpFilter); - } // namespace tensorflow diff --git a/tensorflow/compiler/tf2xla/xla_op_registry.h b/tensorflow/compiler/tf2xla/xla_op_registry.h index d7fb69270b442c9f5a970df8231c97f78911aec5..d74203c82a1932f0b064fea5e2451a10bf222def 100644 --- a/tensorflow/compiler/tf2xla/xla_op_registry.h +++ b/tensorflow/compiler/tf2xla/xla_op_registry.h @@ -41,10 +41,14 @@ namespace tensorflow { extern const char* const DEVICE_CPU_XLA_JIT; // "CPU_XLA_JIT" extern const char* const DEVICE_GPU_XLA_JIT; // "GPU_XLA_JIT" +extern const char* const DEVICE_XLA_CPU; +extern const char* const DEVICE_XLA_GPU; + constexpr std::array kIntTypes = {{DT_INT32, DT_INT64}}; -constexpr std::array kFloatTypes = {{DT_FLOAT, DT_DOUBLE}}; -constexpr std::array kNumericTypes = { - {DT_INT32, DT_INT64, DT_FLOAT, DT_DOUBLE}}; +constexpr std::array kFloatTypes = { + {DT_HALF, DT_FLOAT, DT_DOUBLE}}; +constexpr std::array kNumericTypes = { + {DT_INT32, DT_INT64, DT_HALF, DT_FLOAT, DT_DOUBLE}}; constexpr std::array kCpuAllTypes = { {DT_INT32, DT_INT64, DT_FLOAT, DT_DOUBLE, DT_BOOL}}; @@ -171,8 +175,17 @@ class XlaOpRegistry { Factory factory; }; + // Returns true if registrations x and y can both be added to the registry. + // This is always the case if they refer to different ops. If they refer to + // the same op name, they must: have the same values for compilation_only and + // allow_resource_types; use a device_whitelist; and their + // whitelists must not intersect. + static bool IsCompatible(const OpRegistration& x, const OpRegistration& y); + // Map from operator name to OpRegistrations, populated by REGISTER_XLA_OP. - std::unordered_map> ops_ + // Registrations present under the same key must satisfy IsCompatible above, + // and this is checked during registration. + std::unordered_multimap> ops_ GUARDED_BY(mutex_); // Have we already registered the JIT kernels on the JIT devices? diff --git a/tensorflow/compiler/xla/BUILD b/tensorflow/compiler/xla/BUILD index 35c0efb8f0f81fcd94cc39c654a6484fd715d46f..49f7bcb3406d3daa55112ab87487917b9da74f01 100644 --- a/tensorflow/compiler/xla/BUILD +++ b/tensorflow/compiler/xla/BUILD @@ -6,6 +6,7 @@ package_group( name = "friends", packages = [ "//tensorflow/compiler/...", + "//tensorflow/contrib/tpu/...", "//tensorflow/contrib/xla_tf_graph/...", ], ) @@ -17,6 +18,7 @@ package_group( ], ) +load("//tensorflow:tensorflow.bzl", "cc_header_only_library") load("//tensorflow/compiler/xla:xla.bzl", "xla_proto_library") # Filegroup used to collect source files for dependency checking. @@ -44,11 +46,40 @@ xla_proto_library( ], ) +cc_library( + name = "execution_options_util", + srcs = [ + "execution_options_util.cc", + ], + hdrs = [ + "execution_options_util.h", + ], + visibility = [":friends"], + deps = [ + ":xla_proto", + "//tensorflow/compiler/xla/legacy_flags:debug_options_flags", + ], +) + +cc_library( + name = "test", + testonly = 1, + hdrs = ["test.h"], + visibility = [":friends"], + deps = [ + "//tensorflow/core:lib_internal", + "//tensorflow/core:test", + ], +) + cc_library( name = "types", hdrs = ["types.h"], visibility = [":friends"], - deps = ["//tensorflow/core:lib"], + deps = [ + "//tensorflow/core:lib", + "//third_party/eigen3", + ], ) cc_library( @@ -81,9 +112,9 @@ cc_test( deps = [ ":status_macros", ":statusor", + ":test", ":test_helpers", "//tensorflow/core:lib", - "//tensorflow/core:test", "//tensorflow/core:test_main", ], ) @@ -101,7 +132,10 @@ cc_library( cc_library( name = "statusor", srcs = ["statusor.cc"], - hdrs = ["statusor.h"], + hdrs = [ + "statusor.h", + "statusor_internals.h", + ], visibility = ["//visibility:public"], deps = [ ":status", @@ -116,6 +150,7 @@ cc_test( srcs = ["statusor_test.cc"], deps = [ ":statusor", + ":test", ":types", "//tensorflow/core:lib", "//tensorflow/core:test", @@ -136,7 +171,6 @@ cc_library( ":status", ":types", ":xla_data_proto", - "//tensorflow/compiler/xla/legacy_flags:util_flags", "//tensorflow/core:lib", ], ) @@ -149,18 +183,22 @@ cc_library( ], visibility = ["//visibility:public"], deps = [ + ":status_macros", + ":statusor", ":types", + ":util", "//tensorflow/core:lib", ], ) cc_test( name = "util_test", + size = "small", srcs = ["util_test.cc"], deps = [ + ":test", ":types", ":util", - "//tensorflow/core:test", "//tensorflow/core:test_main", ], ) @@ -187,7 +225,6 @@ cc_library( ":types", ":util", ":xla_data_proto", - "//tensorflow/compiler/xla/legacy_flags:layout_util_flags", "//tensorflow/core:framework_internal", "//tensorflow/core:lib", "//tensorflow/core:regexp_internal", @@ -196,37 +233,39 @@ cc_library( cc_test( name = "shape_util_test", + size = "small", srcs = ["shape_util_test.cc"], deps = [ ":shape_util", + ":test", ":test_helpers", ":types", ":util", - "//tensorflow/core:test", + ":xla_data_proto", "//tensorflow/core:test_main", ], ) cc_test( name = "layout_util_test", + size = "small", srcs = ["layout_util_test.cc"], deps = [ ":shape_util", + ":test", ":test_helpers", - "//tensorflow/compiler/xla/legacy_flags:layout_util_flags", - "//tensorflow/core:test", "//tensorflow/core:test_main", ], ) cc_test( name = "index_util_test", + size = "small", srcs = ["index_util_test.cc"], deps = [ ":shape_util", - ":test_helpers", + ":test", ":xla_data_proto", - "//tensorflow/core:test", "//tensorflow/core:test_main", ], ) @@ -241,6 +280,7 @@ cc_library( ":array3d", ":array4d", ":shape_util", + ":status_macros", ":types", ":util", ":xla_data_proto", @@ -250,13 +290,14 @@ cc_library( cc_test( name = "literal_util_test", + size = "small", srcs = ["literal_util_test.cc"], deps = [ ":array3d", ":array4d", ":literal_util", ":shape_util", - ":test_helpers", + ":test", ":types", "//tensorflow/core:lib", "//tensorflow/core:test", @@ -271,7 +312,6 @@ cc_library( visibility = ["//visibility:public"], deps = [ ":util", - ":xla_data_proto", "//tensorflow/core:lib", ], ) @@ -301,10 +341,11 @@ cc_library( cc_test( name = "array2d_test", + size = "small", srcs = ["array2d_test.cc"], deps = [ ":array2d", - "//tensorflow/core:test", + ":test", "//tensorflow/core:test_main", ], ) @@ -321,11 +362,12 @@ cc_library( cc_test( name = "array3d_test", + size = "small", srcs = ["array3d_test.cc"], deps = [ ":array3d", + ":test", ":types", - "//tensorflow/core:test", "//tensorflow/core:test_main", ], ) @@ -343,11 +385,12 @@ cc_library( cc_test( name = "array4d_test", + size = "small", srcs = ["array4d_test.cc"], deps = [ ":array4d", + ":test", "//tensorflow/core:lib", - "//tensorflow/core:test", "//tensorflow/core:test_main", ], ) @@ -379,7 +422,6 @@ cc_library( cc_library( name = "test_helpers", testonly = 1, - srcs = ["test_helpers.cc"], hdrs = ["test_helpers.h"], visibility = [":internal"], deps = [ @@ -411,15 +453,16 @@ cc_library( cc_test( name = "text_literal_reader_test", + size = "small", srcs = ["text_literal_reader_test.cc"], deps = [ ":literal_util", ":shape_util", + ":test", ":text_literal_reader", ":types", ":xla_data_proto", "//tensorflow/core:lib", - "//tensorflow/core:test", "//tensorflow/core:test_main", ], ) @@ -441,14 +484,15 @@ cc_library( cc_test( name = "text_literal_writer_test", + size = "small", srcs = ["text_literal_writer_test.cc"], deps = [ ":literal_util", + ":test", ":test_helpers", ":text_literal_writer", ":types", "//tensorflow/core:lib", - "//tensorflow/core:test", "//tensorflow/core:test_main", ], ) @@ -463,17 +507,19 @@ cc_library( ":util", ":xla_data_proto", "//tensorflow/core:lib", + "//tensorflow/core:lib_internal", ], ) cc_test( name = "shape_tree_test", + size = "small", srcs = ["shape_tree_test.cc"], deps = [ ":shape_tree", ":shape_util", + ":test", ":xla_data_proto", - "//tensorflow/core:test", "//tensorflow/core:test_main", ], ) @@ -518,6 +564,9 @@ cc_library( ":xla_data_proto", "//tensorflow/compiler/xla/client:computation_builder", "//tensorflow/compiler/xla/client:padding", + "//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", ], @@ -525,17 +574,19 @@ cc_library( cc_test( name = "reference_util_test", + size = "small", srcs = ["reference_util_test.cc"], deps = [ ":array2d", + ":array3d", ":array4d", ":literal_util", ":reference_util", + ":test", ":util", ":xla_data_proto", "//tensorflow/compiler/xla/client:padding", "//tensorflow/compiler/xla/tests:literal_test_util", - "//tensorflow/core:test", "//tensorflow/core:test_main", ], ) @@ -553,3 +604,17 @@ filegroup( ), visibility = ["//tensorflow:__subpackages__"], ) + +# 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. +cc_header_only_library( + name = "xla_headers_lib", + visibility = ["//visibility:public"], + deps = [ + ":xla_data_proto", + ":xla_proto", + "//tensorflow/compiler/xla/client:client_library", + "//tensorflow/compiler/xla/service:hlo", + "//tensorflow/core:framework_headers_lib", + "//tensorflow/core:stream_executor_headers_lib", + ], +) diff --git a/tensorflow/compiler/xla/array2d.h b/tensorflow/compiler/xla/array2d.h index f885821210eb68dfb599303830c814c309e0a24d..593084a0c111690d9e239ed5837f6f0c6c713048 100644 --- a/tensorflow/compiler/xla/array2d.h +++ b/tensorflow/compiler/xla/array2d.h @@ -45,11 +45,15 @@ class Array2D { // Creates an array of dimensions n1 x n2, uninitialized values. Array2D(const int64 n1, const int64 n2) - : n1_(n1), n2_(n2), values_(n1 * n2) {} + : n1_(n1), n2_(n2), values_(new T[n1 * n2]()) { + Fill(T()); + } // Creates an array of dimensions n1 x n2, initialized to value. Array2D(const int64 n1, const int64 n2, const T value) - : n1_(n1), n2_(n2), values_(n1 * n2, value) {} + : n1_(n1), n2_(n2), values_(new T[n1 * n2]()) { + Fill(value); + } // Creates an array from the given nested initializer list. The outer // initializer list is the first dimension; the inner is the second dimension. @@ -65,16 +69,30 @@ class Array2D { } } - T& operator()(const int64 n1, const int64 n2) { - CHECK_LT(n1, n1_); - CHECK_LT(n2, n2_); - return values_[n1 * n2_ + n2]; + Array2D(const Array2D& other) : Array2D(other.n1(), other.n2()) { + std::copy(&other.values_[0], &other.values_[0] + num_elements(), + &values_[0]); + } + + Array2D& operator=(const Array2D& other) { + n1_ = other.n1(); + n2_ = other.n2(); + values_.reset(new T[num_elements()]); + std::copy(&other.values_[0], &other.values_[0] + num_elements(), + &values_[0]); + return *this; + } + + T& operator()(const int64 i1, const int64 i2) { + CHECK_LT(i1, n1_); + CHECK_LT(i2, n2_); + return values_[i1 * n2_ + i2]; } - const T& operator()(const int64 n1, const int64 n2) const { - CHECK_LT(n1, n1_); - CHECK_LT(n2, n2_); - return values_[n1 * n2_ + n2]; + const T& operator()(const int64 i1, const int64 i2) const { + CHECK_LT(i1, n1_); + CHECK_LT(i2, n2_); + return values_[i1 * n2_ + i2]; } // Access to the array's dimensions. height() and width() provide the @@ -84,15 +102,15 @@ class Array2D { int64 n2() const { return n2_; } int64 height() const { return n1_; } int64 width() const { return n2_; } - int64 num_elements() const { return values_.size(); } + int64 num_elements() const { return n1_ * n2_; } // Low-level accessor for stuff like memcmp, handle with care. Returns pointer // to the underlying storage of the array (similarly to std::vector::data()). - T* data() const { return const_cast(this)->values_.data(); } + T* data() const { return const_cast(this)->values_.get(); } // Fills the array with the given value. void Fill(const T& value) { - std::fill(values_.begin(), values_.end(), value); + std::fill(&values_[0], &values_[0] + num_elements(), value); } // Applies f to all cells in this array, in row-major order. @@ -124,8 +142,8 @@ class Array2D { std::mt19937 g(seed); std::normal_distribution distribution(mean, static_cast(value)); - for (auto& v : values_) { - v = static_cast(distribution(g)); + for (int64 i = 0; i < num_elements(); ++i) { + values_[i] = static_cast(distribution(g)); } } @@ -150,7 +168,7 @@ class Array2D { private: int64 n1_; int64 n2_; - std::vector values_; + std::unique_ptr values_; }; // Returns a linspace-populated Array2D in the range [from, to] (inclusive) diff --git a/tensorflow/compiler/xla/array2d_test.cc b/tensorflow/compiler/xla/array2d_test.cc index ac107b1c0d426c676629762dbc8191c74e2e1c7e..795d50ca5b56a60c34279a33e65aa635a65fa5ec 100644 --- a/tensorflow/compiler/xla/array2d_test.cc +++ b/tensorflow/compiler/xla/array2d_test.cc @@ -17,7 +17,7 @@ limitations under the License. #include -#include "tensorflow/core/platform/test.h" +#include "tensorflow/compiler/xla/test.h" namespace xla { namespace { @@ -84,6 +84,17 @@ TEST(Array2dTest, IndexingReadWrite) { EXPECT_EQ(arr(1, 2), 61); } +TEST(Array2dTest, IndexingReadWriteBool) { + Array2D arr = {{false, true, false}, {true, true, false}}; + + EXPECT_EQ(arr(1, 1), true); + EXPECT_EQ(arr(1, 2), false); + arr(1, 1) = false; + arr(1, 2) = true; + EXPECT_EQ(arr(1, 1), false); + EXPECT_EQ(arr(1, 2), true); +} + TEST(Array2dTest, Fill) { Array2D fullof7(2, 3, 7); for (int64 n1 = 0; n1 < fullof7.n1(); ++n1) { diff --git a/tensorflow/compiler/xla/array3d.h b/tensorflow/compiler/xla/array3d.h index 654af8f03074f30dd1561db412ad36f43a33aab9..124ccd1975b3a9ab047e9bbbfb38921fe7386fe4 100644 --- a/tensorflow/compiler/xla/array3d.h +++ b/tensorflow/compiler/xla/array3d.h @@ -20,9 +20,9 @@ limitations under the License. #include #include #include +#include #include #include -#include #include "tensorflow/compiler/xla/types.h" #include "tensorflow/core/platform/logging.h" @@ -39,11 +39,15 @@ class Array3D { public: // Creates an array of dimensions n1 x n2 x n3, uninitialized values. Array3D(const int64 n1, const int64 n2, const int64 n3) - : n1_(n1), n2_(n2), n3_(n3), values_(n1 * n2 * n3) {} + : n1_(n1), n2_(n2), n3_(n3), values_(new T[n1 * n2 * n3]) { + Fill(T()); + } // Creates an array of dimensions n1 x n2 x n3, initialized to value. Array3D(const int64 n1, const int64 n2, const int64 n3, const T value) - : n1_(n1), n2_(n2), n3_(n3), values_(n1 * n2 * n3, value) {} + : n1_(n1), n2_(n2), n3_(n3), values_(new T[n1 * n2 * n3]) { + Fill(value); + } // Creates an array from the given nested initializer list. The outer // initializer list is the first dimension, and so on. @@ -69,34 +73,50 @@ class Array3D { } } - T& operator()(const int64 n1, const int64 n2, const int64 n3) { - CHECK_LT(n1, n1_); - CHECK_LT(n2, n2_); - CHECK_LT(n3, n3_); - return values_[n1 * n2_ * n3_ + n2 * n3_ + n3]; + Array3D(const Array3D& other) + : Array3D(other.n1(), other.n2(), other.n3()) { + std::copy(&other.values_[0], &other.values_[0] + num_elements(), + &values_[0]); + } + + Array3D& operator=(const Array3D& other) { + n1_ = other.n1(); + n2_ = other.n2(); + n3_ = other.n3(); + values_.reset(new T[num_elements()]); + std::copy(&other.values_[0], &other.values_[0] + num_elements(), + &values_[0]); + return *this; + } + + T& operator()(const int64 i1, const int64 i2, const int64 i3) { + CHECK_LT(i1, n1_); + CHECK_LT(i2, n2_); + CHECK_LT(i3, n3_); + return values_[i1 * n2_ * n3_ + i2 * n3_ + i3]; } - const T& operator()(const int64 n1, const int64 n2, const int64 n3) const { - CHECK_LT(n1, n1_); - CHECK_LT(n2, n2_); - CHECK_LT(n3, n3_); - return values_[n1 * n2_ * n3_ + n2 * n3_ + n3]; + const T& operator()(const int64 i1, const int64 i2, const int64 i3) const { + CHECK_LT(i1, n1_); + CHECK_LT(i2, n2_); + CHECK_LT(i3, n3_); + return values_[i1 * n2_ * n3_ + i2 * n3_ + i3]; } // Access to the array's dimensions. int64 n1() const { return n1_; } int64 n2() const { return n2_; } int64 n3() const { return n3_; } - int64 num_elements() const { return values_.size(); } + int64 num_elements() const { return n1_ * n2_ * n3_; } // Fills the array with the given value. void Fill(const T& value) { - std::fill(values_.begin(), values_.end(), value); + std::fill(&values_[0], &values_[0] + num_elements(), value); } // Fills the array with sequentially increasing values. void FillIota(const T& value) { - std::iota(values_.begin(), values_.end(), value); + std::iota(&values_[0], &values_[0] + num_elements(), value); } // Fills the array with random normal values with a mean of 0 and standard @@ -106,8 +126,8 @@ class Array3D { std::mt19937 g(seed); std::normal_distribution distribution(mean, static_cast(value)); - for (auto& v : values_) { - v = static_cast(distribution(g)); + for (int64 i = 0; i < num_elements(); ++i) { + values_[i] = static_cast(distribution(g)); } } @@ -115,7 +135,7 @@ class Array3D { int64 n1_; int64 n2_; int64 n3_; - std::vector values_; + std::unique_ptr values_; }; } // namespace xla diff --git a/tensorflow/compiler/xla/array3d_test.cc b/tensorflow/compiler/xla/array3d_test.cc index fa4435dfc48edcd5b88230e7d2de21e29e269b7e..6b5f4b343b2113652758bbd5ce0fc803239c1266 100644 --- a/tensorflow/compiler/xla/array3d_test.cc +++ b/tensorflow/compiler/xla/array3d_test.cc @@ -17,8 +17,8 @@ limitations under the License. #include +#include "tensorflow/compiler/xla/test.h" #include "tensorflow/compiler/xla/types.h" -#include "tensorflow/core/platform/test.h" namespace xla { namespace { diff --git a/tensorflow/compiler/xla/array4d.h b/tensorflow/compiler/xla/array4d.h index 199ad2baaeb7999349fd6bb201a476706bb12ce7..4c7fce1aaf1faf4bd08bca38bc8eb2b47303b575 100644 --- a/tensorflow/compiler/xla/array4d.h +++ b/tensorflow/compiler/xla/array4d.h @@ -20,6 +20,7 @@ limitations under the License. #include #include #include +#include #include #include #include @@ -54,21 +55,21 @@ namespace xla { template class Array4D { public: - // Creates a 4D array, unitialized values. + // Creates a 4D array, uninitialized values. Array4D(int64 planes, int64 depth, int64 height, int64 width) : planes_(planes), depth_(depth), height_(height), width_(width), - values_(planes * depth * height * width) {} + values_(new T[planes * depth * height * width]) { + Fill(T()); + } - // Creates a 4D array, initalized to value. + // Creates a 4D array, initialized to value. Array4D(int64 planes, int64 depth, int64 height, int64 width, T value) - : planes_(planes), - depth_(depth), - height_(height), - width_(width), - values_(planes * depth * height * width, value) {} + : Array4D(planes, depth, height, width) { + Fill(value); + } // Creates a 4D array, filled with values. // @@ -111,6 +112,23 @@ class Array4D { } } + Array4D(const Array4D& other) + : Array4D(other.planes(), other.depth(), other.height(), other.width()) { + std::copy(&other.values_[0], &other.values_[0] + num_elements(), + &values_[0]); + } + + Array4D& operator=(const Array4D& other) { + planes_ = other.planes(); + depth_ = other.depth(); + height_ = other.height(); + width_ = other.width(); + values_.reset(new T[num_elements()]); + std::copy(&other.values_[0], &other.values_[0] + num_elements(), + &values_[0]); + return *this; + } + T& operator()(int64 plane, int64 depth, int64 height, int64 width) { CHECK_LT(plane, planes_); CHECK_LT(depth, depth_); @@ -135,24 +153,24 @@ class Array4D { int64 n3() const { return height_; } int64 n2() const { return depth_; } int64 n1() const { return planes_; } - int64 num_elements() const { return values_.size(); } + int64 num_elements() const { return width_ * height_ * depth_ * planes_; } // Sets all the values in the array to values. template > void SetValues(const Container& container) { CHECK_EQ(std::distance(std::begin(container), std::end(container)), num_elements()); - values_.assign(std::begin(container), std::end(container)); + std::copy(std::begin(container), std::end(container), &values_[0]); } // Fills the array with the given value. void Fill(const T& value) { - std::fill(values_.begin(), values_.end(), value); + std::fill(&values_[0], &values_[0] + num_elements(), value); } // Fills the array with iota. void FillIota(const T& value) { - std::iota(values_.begin(), values_.end(), value); + std::iota(&values_[0], &values_[0] + num_elements(), value); } // Fills the array with random variable with a deviation of value and a mean @@ -162,8 +180,8 @@ class Array4D { std::mt19937 g(seed); std::normal_distribution distribution(mean, static_cast(value)); - for (auto& v : values_) { - v = static_cast(distribution(g)); + for (int64 i = 0; i < num_elements(); ++i) { + values_[i] = static_cast(distribution(g)); } } @@ -189,6 +207,18 @@ class Array4D { } } + // Invokes a callback with the (indices, value) for each cell in the 4D array. + void Each( + std::function, T)> f) const { + // We const_cast to be able to use the common non-const implementation, + // but prevent modification of the data by passing it by-value to the + // caller. + const_cast(this)->Each( + [&f](tensorflow::gtl::ArraySlice indices, T* value) { + f(indices, *value); + }); + } + // Fills all of the {p,z} with the array provided, which specifies {y,x}. void FillWithYX(const Array2D& value) { CHECK_EQ(value.height(), height()); @@ -268,7 +298,7 @@ class Array4D { int64 depth_; int64 height_; int64 width_; - std::vector values_; + std::unique_ptr values_; }; } // namespace xla diff --git a/tensorflow/compiler/xla/array4d_test.cc b/tensorflow/compiler/xla/array4d_test.cc index 72ada467e515eff98a2e5845dc6a3714a770650e..3bc8148c911df0aeade364e4ac2e2ee828bacb53 100644 --- a/tensorflow/compiler/xla/array4d_test.cc +++ b/tensorflow/compiler/xla/array4d_test.cc @@ -18,8 +18,8 @@ limitations under the License. #include #include +#include "tensorflow/compiler/xla/test.h" #include "tensorflow/core/lib/gtl/array_slice.h" -#include "tensorflow/core/platform/test.h" namespace xla { namespace { diff --git a/tensorflow/compiler/xla/client/BUILD b/tensorflow/compiler/xla/client/BUILD index 3e9dfe2a922c913c528d586413c11e2da8cbdc39..a998b91c89d79ac5e354d2a3edf5fb78695d73cb 100644 --- a/tensorflow/compiler/xla/client/BUILD +++ b/tensorflow/compiler/xla/client/BUILD @@ -46,6 +46,7 @@ cc_library( cc_test( name = "padding_test", + size = "small", srcs = ["padding_test.cc"], deps = [ ":padding", @@ -61,6 +62,7 @@ cc_library( deps = [ ":computation", ":global_data", + "//tensorflow/compiler/xla:execution_options_util", "//tensorflow/compiler/xla:literal_util", "//tensorflow/compiler/xla:service_interface", "//tensorflow/compiler/xla:status_macros", @@ -69,6 +71,7 @@ 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/compiler/xla/service:session_proto", "//tensorflow/core:lib", ], @@ -99,6 +102,25 @@ cc_library( ], ) +cc_library( + name = "compile_only_client", + srcs = ["compile_only_client.cc"], + hdrs = ["compile_only_client.h"], + deps = [ + ":client", + ":computation", + "//tensorflow/compiler/xla:status_macros", + "//tensorflow/compiler/xla:statusor", + "//tensorflow/compiler/xla:util", + "//tensorflow/compiler/xla:xla_data_proto", + "//tensorflow/compiler/xla/service:compile_only_service", + "//tensorflow/compiler/xla/service:compiler", + "//tensorflow/compiler/xla/service/llvm_ir:llvm_util", + "//tensorflow/core:stream_executor_no_cuda", + "@llvm//:support", + ], +) + # This target is used to instantiate the XLA service in-process and create # a client for it. cc_library( @@ -106,12 +128,14 @@ cc_library( srcs = ["client_library.cc"], hdrs = ["client_library.h"], deps = [ + ":compile_only_client", ":local_client", "//tensorflow/compiler/xla:status_macros", "//tensorflow/compiler/xla:statusor", "//tensorflow/compiler/xla:types", "//tensorflow/compiler/xla:util", "//tensorflow/compiler/xla/service:backend", + "//tensorflow/compiler/xla/service:compile_only_service", "//tensorflow/compiler/xla/service:device_memory_allocator", "//tensorflow/compiler/xla/service:local_service", "//tensorflow/compiler/xla/service:platform_util", diff --git a/tensorflow/compiler/xla/client/client.cc b/tensorflow/compiler/xla/client/client.cc index 973204eae2d4f8e0ef209c360173982e038e882f..387253617e4f37a1561d4659eb796a181f0b5bee 100644 --- a/tensorflow/compiler/xla/client/client.cc +++ b/tensorflow/compiler/xla/client/client.cc @@ -18,6 +18,8 @@ limitations under the License. #include #include +#include "tensorflow/compiler/xla/execution_options_util.h" +#include "tensorflow/compiler/xla/legacy_flags/debug_options_flags.h" #include "tensorflow/compiler/xla/literal_util.h" #include "tensorflow/compiler/xla/ptr_util.h" #include "tensorflow/compiler/xla/status_macros.h" @@ -58,34 +60,13 @@ StatusOr> Client::Transfer( "server provided response without a literal in " "TransferToClient request"); } - - return WrapUnique(response.release_literal()); -} - -Status Client::TransferInProcess(const GlobalData& data, void* destination) { - TransferToClientInProcessRequest request; - *request.mutable_data() = data.handle(); - request.set_buffer(reinterpret_cast(destination)); - TransferToClientInProcessResponse response; - - VLOG(1) << "making transfer in-process request"; - VLOG(3) << "TransferToClientInProcessRequest: {" << request.DebugString() - << "}"; - Status s = stub_->TransferToClientInProcess(&request, &response); - VLOG(1) << "done with request"; - - if (!s.ok()) { - return s; - } - VLOG(3) << "TransferToClientInProcessResponse: {" << response.DebugString() - << "}"; - return Status::OK(); + return MakeUnique(response.literal()); } StatusOr> Client::TransferToServer( const Literal& literal, const DeviceHandle* device_handle) { TransferToServerRequest request; - *request.mutable_literal() = literal; + *request.mutable_literal() = literal.ToProto(); if (device_handle) { *request.mutable_device_handle() = *device_handle; } @@ -113,7 +94,7 @@ StatusOr> Client::TransferToServer( Status Client::TransferToInfeed(const Literal& literal, int64 replica_id, const DeviceHandle* device_handle) { TransferToInfeedRequest request; - *request.mutable_literal() = literal; + *request.mutable_literal() = literal.ToProto(); if (device_handle) { *request.mutable_device_handle() = *device_handle; } @@ -161,7 +142,8 @@ StatusOr> Client::TransferFromOutfeed( "TransferToClient request"); } - return WrapUnique(response.release_literal()); + Literal literal(response.literal()); + return MakeUnique(literal); } Status Client::ResetDevice() { @@ -196,34 +178,6 @@ StatusOr> Client::ExecuteAndTransfer( return Transfer(*data, shape_with_output_layout); } -StatusOr> Client::TransferToServerInProcess( - const Shape& shape, const void* buffer) { - TransferToServerInProcessRequest request; - request.set_buffer(reinterpret_cast(buffer)); - *request.mutable_shape() = shape; - TransferToServerInProcessResponse response; - - VLOG(1) << "making transfer to server in-process request"; - VLOG(3) << "TransferToServerInProcessRequest: {" << request.DebugString() - << "}"; - Status s = stub_->TransferToServerInProcess(&request, &response); - VLOG(1) << "done with request"; - - if (!s.ok()) { - return s; - } - VLOG(3) << "TransferToServerInProcessResponse: {" << response.DebugString() - << "}"; - - if (!response.has_data()) { - return FailedPrecondition( - "server provided response without a data handle in " - "TransferToServerInProcess request"); - } - - return MakeUnique(stub_, response.data()); -} - StatusOr Client::LoadSnapshot(const SessionModule& module) { LoadComputationSnapshotRequest request; *request.mutable_module() = module; @@ -245,7 +199,10 @@ StatusOr> Client::Execute( ExecutionProfile* execution_profile) { ExecuteRequest request; *request.mutable_computation() = computation.handle(); - if (execution_options != nullptr) { + + if (execution_options == nullptr) { + *request.mutable_execution_options() = CreateDefaultExecutionOptions(); + } else { *request.mutable_execution_options() = *execution_options; } for (GlobalData* argument : arguments) { @@ -337,59 +294,6 @@ StatusOr> Client::GetDeviceHandles( return device_handles; } -StatusOr Client::ExecuteAsync( - const Computation& computation, - tensorflow::gtl::ArraySlice arguments, - const ExecutionOptions* execution_options) { - ExecuteAsyncRequest request; - *request.mutable_computation() = computation.handle(); - for (GlobalData* argument : arguments) { - *request.add_arguments() = argument->handle(); - } - if (execution_options != nullptr) { - *request.mutable_execution_options() = *execution_options; - } - - ExecuteAsyncResponse response; - VLOG(1) << "making execute async request: " << request.ShortDebugString(); - Status s = stub_->ExecuteAsync(&request, &response); - VLOG(1) << "done with request"; - - if (!s.ok()) { - return s; - } - - return response.execution(); -} - -StatusOr> Client::WaitForExecution( - const Computation& computation, const ExecutionHandle& execution, - ExecutionProfile* execution_profile) { - WaitForExecutionRequest request; - *request.mutable_execution() = execution; - - WaitForExecutionResponse response; - VLOG(1) << "making wait-for-execute request: " << request.ShortDebugString(); - Status s = stub_->WaitForExecution(&request, &response); - VLOG(1) << "done with request"; - - if (!s.ok()) { - return s; - } - - 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 MakeUnique(stub_, response.output()); -} - Status Client::Unregister(const GlobalData& data) { UnregisterRequest request; *request.mutable_data() = data.handle(); @@ -424,9 +328,10 @@ StatusOr>> Client::DeconstructTuple( } StatusOr Client::GetComputationStats( - const Computation& computation) const { + const Computation& computation, const DebugOptions& debug_options) const { ComputationStatsRequest request; *request.mutable_computation() = computation.handle(); + *request.mutable_debug_options() = debug_options; ComputationStatsResponse response; VLOG(1) << "making computation stats request"; @@ -475,7 +380,10 @@ StatusOr Client::GetShape(const GlobalData& data) { StatusOr Client::ExecutionStatsAsString( const Computation& computation, const ExecutionProfile& profile) { - TF_ASSIGN_OR_RETURN(auto computation_stats, GetComputationStats(computation)); + TF_ASSIGN_OR_RETURN( + auto computation_stats, + GetComputationStats(computation, + legacy_flags::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 ea166acc91a24508f7c9e2e7df256c42144baf98..e72816a6217afd6a827642bbe3aa205409ef5718 100644 --- a/tensorflow/compiler/xla/client/client.h +++ b/tensorflow/compiler/xla/client/client.h @@ -21,6 +21,7 @@ limitations under the License. #include "tensorflow/compiler/xla/client/computation.h" #include "tensorflow/compiler/xla/client/global_data.h" +#include "tensorflow/compiler/xla/literal_util.h" #include "tensorflow/compiler/xla/service/session.pb.h" #include "tensorflow/compiler/xla/service_interface.h" #include "tensorflow/compiler/xla/statusor.h" @@ -74,24 +75,6 @@ class Client { // TransferToInfeed). StatusOr> GetDeviceHandles(int64 device_count); - // Executes the given computation as above Execute(), but launches the - // computation asynchronously and returns before the execution is complete. - // Returns an ExecutionHandle that represents the launched execution, which is - // used to call WaitForExecution() to wait for the execution's completion. - StatusOr ExecuteAsync( - const Computation& computation, - tensorflow::gtl::ArraySlice arguments, - const ExecutionOptions* execution_options = nullptr); - - // Waits until the given asynchronously launched execution of the computation - // is complete and returns the execution result. Once this is called, the - // given execution handle is no longer valid. If execution_profile is not - // nullptr then the pointed-to ExecutionProfile will be filled with profile - // data from the execution. - StatusOr> WaitForExecution( - const Computation& computation, const ExecutionHandle& execution, - ExecutionProfile* execution_profile = nullptr); - // Transfer the global data provided to this client process, which is // returned in the provided literal. Use sparingly to avoid transfer // overheads. @@ -149,11 +132,10 @@ class Client { // Retrieves the statistics of the given computation. StatusOr GetComputationStats( - const Computation& computation) const; + const Computation& computation, const DebugOptions& debug_options) const; // Returns the Shape of the given array specified by 'data'. The shape - // includes the Layout of the array as it is stored on the service. The layout - // information is useful for calling TransferInProcess. + // includes the Layout of the array as it is stored on the service. StatusOr GetShape(const GlobalData& data); // As above, but returns the shape of the provided computation (parameter @@ -165,24 +147,6 @@ class Client { // two computations via a pair of Send and Recv instructions. StatusOr CreateChannelHandle(); - // If the service is running in the same process as the client then the - // following "InProcess" transfer methods may be used. These methods enable - // more efficient transfer of arrays to and from the service. - - // Transfer array from the service into the given buffer. The buffer must be - // large enough to hold the array. The array is copied verbatim (memcpy) from - // the service. The method GetShape should be called ahead of time - // to get the shape and layout of the array as it is stored in the - // service. The shape and layout can be used to determine how large the buffer - // needs to be. - Status TransferInProcess(const GlobalData& data, void* destination); - - // Transfer array to the service from the given buffer with the given shape - // and layout. The service creates an internal copy of the data so the client - // can free the buffer when this method returns. - StatusOr> TransferToServerInProcess( - const Shape& shape, const void* buffer); - StatusOr LoadSnapshot(const SessionModule& module); ServiceInterface* stub() { return stub_; } diff --git a/tensorflow/compiler/xla/client/client_library.cc b/tensorflow/compiler/xla/client/client_library.cc index 93437023bc8956e449f828f5bf6dea7a6bff8610..b1663bc815719c3da75b37593ac665b1f3493db8 100644 --- a/tensorflow/compiler/xla/client/client_library.cc +++ b/tensorflow/compiler/xla/client/client_library.cc @@ -23,6 +23,13 @@ limitations under the License. namespace xla { +LocalClientOptions::LocalClientOptions(perftools::gputools::Platform* platform, + int number_of_replicas, + int intra_op_parallelism_threads) + : platform_(platform), + number_of_replicas_(number_of_replicas), + intra_op_parallelism_threads_(intra_op_parallelism_threads) {} + LocalClientOptions& LocalClientOptions::set_platform( perftools::gputools::Platform* platform) { platform_ = platform; @@ -43,6 +50,16 @@ int LocalClientOptions::number_of_replicas() const { return number_of_replicas_; } +LocalClientOptions& LocalClientOptions::set_intra_op_parallelism_threads( + int num_threads) { + intra_op_parallelism_threads_ = num_threads; + return *this; +} + +int LocalClientOptions::intra_op_parallelism_threads() const { + return intra_op_parallelism_threads_; +} + /* static */ ClientLibrary& ClientLibrary::Singleton() { static ClientLibrary* c = new ClientLibrary; return *c; @@ -69,22 +86,24 @@ ClientLibrary::~ClientLibrary() = default; TF_ASSIGN_OR_RETURN(platform, PlatformUtil::GetDefaultPlatform()); } - auto it = client_library.instances_.find(platform->id()); - if (it != client_library.instances_.end()) { + auto it = client_library.local_instances_.find(platform->id()); + if (it != client_library.local_instances_.end()) { return it->second->client.get(); } ServiceOptions service_options; service_options.set_platform(platform); service_options.set_number_of_replicas(replica_count); + service_options.set_intra_op_parallelism_threads( + options.intra_op_parallelism_threads()); - std::unique_ptr instance = MakeUnique(); + auto instance = MakeUnique(); TF_ASSIGN_OR_RETURN(instance->service, LocalService::NewService(service_options)); instance->client = MakeUnique(instance->service.get()); LocalClient* cl = instance->client.get(); - client_library.instances_.insert( + client_library.local_instances_.insert( std::make_pair(platform->id(), std::move(instance))); return cl; } @@ -99,9 +118,43 @@ ClientLibrary::~ClientLibrary() = default; perftools::gputools::Platform* platform) { ClientLibrary& client_library = Singleton(); tensorflow::mutex_lock lock(client_library.service_mutex_); - auto it = client_library.instances_.find(platform->id()); - CHECK(it != client_library.instances_.end()); + auto it = client_library.local_instances_.find(platform->id()); + CHECK(it != client_library.local_instances_.end()); return it->second->service.get(); } +/* static */ StatusOr +ClientLibrary::GetOrCreateCompileOnlyClient( + perftools::gputools::Platform* platform) { + ClientLibrary& client_library = Singleton(); + tensorflow::mutex_lock lock(client_library.service_mutex_); + + if (platform == nullptr) { + TF_ASSIGN_OR_RETURN(platform, PlatformUtil::GetDefaultPlatform()); + } + + auto it = client_library.compile_only_instances_.find(platform->id()); + if (it != client_library.compile_only_instances_.end()) { + return it->second->client.get(); + } + + auto instance = MakeUnique(); + TF_ASSIGN_OR_RETURN(instance->service, + CompileOnlyService::NewService(platform)); + instance->client = MakeUnique(instance->service.get()); + CompileOnlyClient* cl = instance->client.get(); + + client_library.compile_only_instances_.insert( + std::make_pair(platform->id(), std::move(instance))); + return cl; +} + +/* static */ void ClientLibrary::DestroyLocalInstances() { + ClientLibrary& client_library = Singleton(); + tensorflow::mutex_lock lock(client_library.service_mutex_); + + client_library.local_instances_.clear(); + client_library.compile_only_instances_.clear(); +} + } // namespace xla diff --git a/tensorflow/compiler/xla/client/client_library.h b/tensorflow/compiler/xla/client/client_library.h index 2bc319f9333368635690add017ad3d89947e2551..a6f30d82e43587135697e76e8bc7d122edc0f602 100644 --- a/tensorflow/compiler/xla/client/client_library.h +++ b/tensorflow/compiler/xla/client/client_library.h @@ -26,7 +26,9 @@ limitations under the License. #include #include +#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" #include "tensorflow/compiler/xla/service/device_memory_allocator.h" #include "tensorflow/compiler/xla/service/local_service.h" #include "tensorflow/compiler/xla/statusor.h" @@ -41,19 +43,27 @@ namespace xla { // Options to configure the local client when it is created. class LocalClientOptions { public: + LocalClientOptions(perftools::gputools::Platform* platform = nullptr, + int number_of_replicas = 1, + int intra_op_parallelism_threads = -1); + // Set the platform backing the service, or nullptr for the default platform. LocalClientOptions& set_platform(perftools::gputools::Platform* platform); perftools::gputools::Platform* platform() const; // Set the number of replicas to use when compiling replicated - // programs. The default is -1 meaning that the value is read from - // the xla_replicas flag. + // programs. LocalClientOptions& set_number_of_replicas(int number_of_replicas); int number_of_replicas() const; + // Sets the thread pool size for parallel execution of an individual operator. + LocalClientOptions& set_intra_op_parallelism_threads(int num_threads); + int intra_op_parallelism_threads() const; + private: - perftools::gputools::Platform* platform_ = nullptr; - int number_of_replicas_ = -1; + perftools::gputools::Platform* platform_; + int number_of_replicas_; + int intra_op_parallelism_threads_; }; class ClientLibrary { @@ -76,6 +86,18 @@ class ClientLibrary { // access user computations from client. static LocalService* GetXlaService(perftools::gputools::Platform* platform); + // Singleton constructor-or-accessor for compile-only clients. Arguments: + // + // platform : The platform the underlying XLA service should target. If + // null then default platform is used. + static StatusOr GetOrCreateCompileOnlyClient( + perftools::gputools::Platform* platform = nullptr); + + // Clears the local instance and compile only instance caches. The client + // pointers returned by the previous GetOrCreateLocalClient() or + // GetOrCreateCompileOnlyClient() invocations are not valid anymore. + static void DestroyLocalInstances(); + private: // Returns the singleton instance of ClientLibrary. static ClientLibrary& Singleton(); @@ -90,10 +112,21 @@ class ClientLibrary { std::unique_ptr client; }; + struct CompileOnlyInstance { + // Service that is wrapped by the singleton client object. + std::unique_ptr service; + // Singleton client object. + std::unique_ptr client; + }; + tensorflow::mutex service_mutex_; // Guards the singleton creation state. std::unordered_map> - instances_ GUARDED_BY(service_mutex_); + local_instances_ GUARDED_BY(service_mutex_); + + std::unordered_map> + compile_only_instances_ GUARDED_BY(service_mutex_); TF_DISALLOW_COPY_AND_ASSIGN(ClientLibrary); }; diff --git a/tensorflow/compiler/xla/client/compile_only_client.cc b/tensorflow/compiler/xla/client/compile_only_client.cc new file mode 100644 index 0000000000000000000000000000000000000000..999572ad37af73e5058c2064e5464938082bc2e6 --- /dev/null +++ b/tensorflow/compiler/xla/client/compile_only_client.cc @@ -0,0 +1,57 @@ +/* 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/compile_only_client.h" + +#include "llvm/ADT/Triple.h" +#include "tensorflow/compiler/xla/ptr_util.h" +#include "tensorflow/compiler/xla/service/llvm_ir/llvm_util.h" +#include "tensorflow/compiler/xla/status_macros.h" + +namespace xla { + +StatusOr>> +CompileOnlyClient::CompileAheadOfTime( + const tensorflow::gtl::ArraySlice computations, + const AotCompilationOptions& options) { + std::vector service_instances; + service_instances.reserve(computations.size()); + for (const AotComputationInstance& instance : computations) { + service_instances.push_back({}); + CompileOnlyService::AotComputationInstance& service_instance = + service_instances.back(); + TF_RET_CHECK(instance.computation != nullptr); + service_instance.computation = instance.computation->handle(); + service_instance.argument_layouts = instance.argument_layouts; + service_instance.result_layout = instance.result_layout; + } + return compiler_service_->CompileAheadOfTime(service_instances, options); +} + +int64 CompileOnlyClient::PointerSizeForTriple( + tensorflow::StringPiece target_triple) { + llvm::Triple triple( + llvm::Triple::normalize(llvm_ir::AsStringRef(target_triple))); + if (triple.isArch64Bit()) { + return 8; + } else if (triple.isArch32Bit()) { + return 4; + } else { + CHECK(triple.isArch16Bit()); + return 2; + } +} + +} // namespace xla diff --git a/tensorflow/compiler/xla/client/compile_only_client.h b/tensorflow/compiler/xla/client/compile_only_client.h new file mode 100644 index 0000000000000000000000000000000000000000..5900048711384e0240a3cd502260eb388eb40f51 --- /dev/null +++ b/tensorflow/compiler/xla/client/compile_only_client.h @@ -0,0 +1,66 @@ +/* Copyright 2017 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#ifndef TENSORFLOW_COMPILER_XLA_CLIENT_COMPILE_ONLY_CLIENT_H_ +#define TENSORFLOW_COMPILER_XLA_CLIENT_COMPILE_ONLY_CLIENT_H_ + +#include "tensorflow/compiler/xla/client/client.h" +#include "tensorflow/compiler/xla/client/computation.h" +#include "tensorflow/compiler/xla/service/compile_only_service.h" +#include "tensorflow/compiler/xla/service/compiler.h" +#include "tensorflow/compiler/xla/statusor.h" +#include "tensorflow/compiler/xla/xla_data.pb.h" +#include "tensorflow/core/platform/stream_executor_no_cuda.h" + +namespace xla { + +// An XLA Client specialization for doing ahead-of-time compilation. This does +// not require (or attempt to instantiate) an execution-capable backend for the +// relevant platform. +class CompileOnlyClient : public Client { + public: + explicit CompileOnlyClient(CompileOnlyService* service) + : Client(service), compiler_service_(service) {} + + CompileOnlyClient(const CompileOnlyClient&) = delete; + void operator=(const CompileOnlyClient&) = delete; + + // A description of a computation to compile using CompileAheadOfTime. + struct AotComputationInstance { + const Computation* computation; + // Inform the compiler of the expected layout for arguments. + std::vector argument_layouts; + // Specifies the expected result layout. + const Shape* result_layout; + }; + + // Compiles a list of computations for ahead-of-time execution. This is + // intended for use in static compilation. The |options| parameter describes + // the target for which the compiler should emit code. + StatusOr>> + CompileAheadOfTime( + const tensorflow::gtl::ArraySlice computations, + const AotCompilationOptions& options); + + // Returns the size of a pointer in bytes for a given triple. + static int64 PointerSizeForTriple(tensorflow::StringPiece triple); + + private: + CompileOnlyService* compiler_service_; +}; + +} // namespace xla + +#endif // TENSORFLOW_COMPILER_XLA_CLIENT_COMPILE_ONLY_CLIENT_H_ diff --git a/tensorflow/compiler/xla/client/computation.cc b/tensorflow/compiler/xla/client/computation.cc index 0f9ca7b4fe40215601f28c2f49327f570e1e8551..4baea8df6e3331200ee52f500fb7b961428e56be 100644 --- a/tensorflow/compiler/xla/client/computation.cc +++ b/tensorflow/compiler/xla/client/computation.cc @@ -28,7 +28,7 @@ Computation::Computation(ServiceInterface* parent, : handle_(handle), parent_(parent) {} Computation::Computation(Computation&& computation) - : handle_(computation.handle_), parent_(computation.parent_) { + : handle_(std::move(computation.handle_)), parent_(computation.parent_) { computation.ResetWithoutFreeing(); } diff --git a/tensorflow/compiler/xla/client/computation_builder.cc b/tensorflow/compiler/xla/client/computation_builder.cc index 22a70681468f16b12793274bf5ce72613534df42..30afaed732348b51c7a629e49199408aa53a4fef 100644 --- a/tensorflow/compiler/xla/client/computation_builder.cc +++ b/tensorflow/compiler/xla/client/computation_builder.cc @@ -111,13 +111,12 @@ bool ComputationBuilder::MakeWindow( return true; } else { NoteError(InvalidArgument( - "%s", - tensorflow::strings::StrCat( - "Window has different number of window dimensions than of ", - x_name, "\nNumber of window dimensions: ", - window_dimensions.size(), "\nNumber of ", x_name, ": ", x, - "\n") - .c_str())); // + "%s", tensorflow::strings::StrCat( + "Window has different number of window dimensions than of ", + x_name, "\nNumber of window dimensions: ", + window_dimensions.size(), "\nNumber of ", x_name, ": ", x, + "\n") + .c_str())); // return false; } }; @@ -165,9 +164,10 @@ ComputationDataHandle ComputationBuilder::ConstantOp( } ConstantRequest request; - Literal* literal = request.mutable_literal(); - populate(literal); - VLOG(3) << "created constant: " << literal->ShortDebugString(); + Literal literal; + populate(&literal); + *request.mutable_literal() = literal.ToProto(); + VLOG(3) << "created constant: " << request.literal().ShortDebugString(); OpRequest op_request; *op_request.mutable_constant_request() = request; *op_request.mutable_computation() = computation_.handle(); @@ -255,7 +255,8 @@ void ComputationBuilder::CheckSameShape(const ComputationDataHandle& lhs, ComputationDataHandle ComputationBuilder::Slice( const ComputationDataHandle& operand, tensorflow::gtl::ArraySlice start_indices, - tensorflow::gtl::ArraySlice limit_indices) { + tensorflow::gtl::ArraySlice limit_indices, + tensorflow::gtl::ArraySlice strides) { if (!first_error_.ok() || !PrepareComputation().ok()) { return ComputationDataHandle(); } @@ -268,6 +269,9 @@ ComputationDataHandle ComputationBuilder::Slice( for (int64 index : limit_indices) { request.add_limit_indices(index); } + for (int64 index : strides) { + request.add_strides(index); + } OpRequest op_request; *op_request.mutable_computation() = computation_.handle(); *op_request.mutable_slice_request() = request; @@ -279,6 +283,25 @@ ComputationDataHandle ComputationBuilder::Slice( return ParseOpResponse(s, &response); } +ComputationDataHandle ComputationBuilder::SliceInDim( + const ComputationDataHandle& operand, int64 start_index, int64 limit_index, + int64 stride, int64 dimno) { + StatusOr> shape_status = GetShape(operand); + if (!shape_status.ok()) { + NoteError(shape_status.status()); + return ComputationDataHandle{}; + } + const Shape& shape = *shape_status.ValueOrDie(); + std::vector starts(ShapeUtil::Rank(shape), 0); + std::vector limits(shape.dimensions().begin(), + shape.dimensions().end()); + std::vector strides(ShapeUtil::Rank(shape), 1); + starts[dimno] = start_index; + limits[dimno] = limit_index; + strides[dimno] = stride; + return Slice(operand, starts, limits, strides); +} + ComputationDataHandle ComputationBuilder::DynamicSlice( const ComputationDataHandle& operand, const ComputationDataHandle& start_indices, @@ -639,24 +662,26 @@ bool ComputationBuilder::VerifyConvolution( } int num_spatial_dims = num_dims - 2; - const auto check_spatial_dimensions = [&]( - const char* const field_name, - const tensorflow::protobuf::RepeatedField& - numbers) { - if (numbers.size() != num_spatial_dims) { - NoteError(InvalidArgument("Expected %d elements for %s, but got %d.", - num_spatial_dims, field_name, numbers.size())); - return false; - } - for (int i = 0; i < numbers.size(); ++i) { - if (numbers.Get(i) < 0 || numbers.Get(i) >= num_dims) { - NoteError(InvalidArgument("Convolution %s[%d] is out of bounds: %lld", - field_name, i, numbers.Get(i))); - return false; - } - } - return true; - }; + const auto check_spatial_dimensions = + [&](const char* const field_name, + const tensorflow::protobuf::RepeatedField& + numbers) { + if (numbers.size() != num_spatial_dims) { + NoteError(InvalidArgument("Expected %d elements for %s, but got %d.", + num_spatial_dims, field_name, + numbers.size())); + return false; + } + for (int i = 0; i < numbers.size(); ++i) { + if (numbers.Get(i) < 0 || numbers.Get(i) >= num_dims) { + NoteError( + InvalidArgument("Convolution %s[%d] is out of bounds: %lld", + field_name, i, numbers.Get(i))); + return false; + } + } + return true; + }; return check_spatial_dimensions("spatial_dimensions", dimension_numbers.spatial_dimensions()) && check_spatial_dimensions( @@ -966,6 +991,16 @@ ComputationDataHandle ComputationBuilder::Sign( return UnaryOp(UNOP_SIGN, operand); } +ComputationDataHandle ComputationBuilder::Cos( + const ComputationDataHandle& operand) { + return UnaryOp(UNOP_COS, operand); +} + +ComputationDataHandle ComputationBuilder::Sin( + const ComputationDataHandle& operand) { + return UnaryOp(UNOP_SIN, operand); +} + ComputationDataHandle ComputationBuilder::Tanh( const ComputationDataHandle& operand) { return UnaryOp(UNOP_TANH, operand); @@ -1234,7 +1269,7 @@ StatusOr ComputationBuilder::IsConstant( return response.is_constant(); } -StatusOr> ComputationBuilder::ComputeConstant( +StatusOr> ComputationBuilder::ComputeConstant( const ComputationDataHandle& operand, const Layout* output_layout) { if (!first_error_.ok()) { return first_error_; @@ -1257,8 +1292,14 @@ StatusOr> ComputationBuilder::ComputeConstant( return s; } - TF_RET_CHECK(response.output().handle() != 0); - return MakeUnique(client_->stub(), response.output()); + VLOG(3) << "ComputeConstant: {" << response.DebugString() << "}"; + + if (!response.has_literal()) { + return InternalError( + "no computed literal in the provided response in ComputeConstant " + "request"); + } + return MakeUnique(response.literal()); } ComputationDataHandle ComputationBuilder::Map( @@ -1406,6 +1447,92 @@ ComputationDataHandle ComputationBuilder::ReduceWindowWithGeneralPadding( return ParseOpResponse(s, &response); } +ComputationDataHandle ComputationBuilder::BatchNormTraining( + const ComputationDataHandle& operand, const ComputationDataHandle& scale, + const ComputationDataHandle& offset, float epsilon, int64 feature_index) { + if (!first_error_.ok() || !PrepareComputation().ok()) { + return ComputationDataHandle(); + } + BatchNormTrainingRequest request; + *request.mutable_operand() = operand; + *request.mutable_scale() = scale; + *request.mutable_offset() = offset; + request.set_epsilon(epsilon); + request.set_feature_index(feature_index); + + OpRequest op_request; + *op_request.mutable_batch_norm_training_request() = request; + *op_request.mutable_computation() = computation_.handle(); + AddOpMetadata(&op_request); + + OpResponse response; + + VLOG(2) << "making BatchNormTraining request"; + + Status s = client_->stub()->Op(&op_request, &response); + return ParseOpResponse(s, &response); +} + +ComputationDataHandle ComputationBuilder::BatchNormInference( + const ComputationDataHandle& operand, const ComputationDataHandle& scale, + const ComputationDataHandle& offset, const ComputationDataHandle& mean, + const ComputationDataHandle& variance, float epsilon, int64 feature_index) { + if (!first_error_.ok() || !PrepareComputation().ok()) { + return ComputationDataHandle(); + } + BatchNormInferenceRequest request; + *request.mutable_operand() = operand; + *request.mutable_scale() = scale; + *request.mutable_offset() = offset; + *request.mutable_mean() = mean; + *request.mutable_variance() = variance; + request.set_epsilon(epsilon); + request.set_feature_index(feature_index); + + OpRequest op_request; + *op_request.mutable_batch_norm_inference_request() = request; + *op_request.mutable_computation() = computation_.handle(); + AddOpMetadata(&op_request); + + OpResponse response; + + VLOG(2) << "making BatchNormInference request"; + + Status s = client_->stub()->Op(&op_request, &response); + return ParseOpResponse(s, &response); +} + +ComputationDataHandle ComputationBuilder::BatchNormGrad( + const ComputationDataHandle& operand, const ComputationDataHandle& scale, + const ComputationDataHandle& mean, const ComputationDataHandle& var, + const ComputationDataHandle& grad_output, float epsilon, + int64 feature_index) { + if (!first_error_.ok() || !PrepareComputation().ok()) { + return ComputationDataHandle(); + } + BatchNormGradRequest request; + *request.mutable_operand() = operand; + *request.mutable_scale() = scale; + *request.mutable_mean() = mean; + *request.mutable_variance() = var; + *request.mutable_grad_output() = grad_output; + request.set_epsilon(epsilon); + request.set_feature_index(feature_index); + + OpRequest op_request; + *op_request.mutable_batch_norm_grad_request() = request; + *op_request.mutable_computation() = computation_.handle(); + AddOpMetadata(&op_request); + + OpResponse response; + + VLOG(2) << "making BatchNormGrad request"; + + Status s = client_->stub()->Op(&op_request, &response); + + return ParseOpResponse(s, &response); +} + ComputationDataHandle ComputationBuilder::CrossReplicaSum( const ComputationDataHandle& operand) { if (!first_error_.ok() || !PrepareComputation().ok()) { @@ -1482,6 +1609,28 @@ ComputationDataHandle ComputationBuilder::SelectAndScatterWithGeneralPadding( return ParseOpResponse(s, &response); } +ComputationDataHandle ComputationBuilder::ReducePrecision( + const ComputationDataHandle& operand, const int exponent_bits, + const int mantissa_bits) { + if (!first_error_.ok() || !PrepareComputation().ok()) { + return ComputationDataHandle(); + } + + ReducePrecisionRequest request; + *request.mutable_operand() = operand; + request.set_exponent_bits(exponent_bits); + request.set_mantissa_bits(mantissa_bits); + OpRequest op_request; + *op_request.mutable_computation() = computation_.handle(); + *op_request.mutable_reduce_precision_request() = request; + AddOpMetadata(&op_request); + OpResponse response; + + VLOG(2) << "making reduce-precision request"; + Status s = client_->stub()->Op(&op_request, &response); + return ParseOpResponse(s, &response); +} + void ComputationBuilder::Send(const ComputationDataHandle& operand, const ChannelHandle& handle) { if (!first_error_.ok() || !PrepareComputation().ok()) { @@ -1498,8 +1647,8 @@ void ComputationBuilder::Send(const ComputationDataHandle& operand, OpResponse response; VLOG(2) << "making send request"; - tensorflow::Status s = client_->stub()->Op(&op_request, &response); - VLOG(2) << "done with request"; + Status s = client_->stub()->Op(&op_request, &response); + VLOG(2) << "done with op request"; if (!s.ok()) { NoteError(s); @@ -1523,9 +1672,7 @@ ComputationDataHandle ComputationBuilder::Recv(const Shape& shape, OpResponse response; VLOG(2) << "making recv request"; - tensorflow::Status s = client_->stub()->Op(&op_request, &response); - VLOG(2) << "done with request"; - + Status s = client_->stub()->Op(&op_request, &response); return ParseOpResponse(s, &response); } diff --git a/tensorflow/compiler/xla/client/computation_builder.h b/tensorflow/compiler/xla/client/computation_builder.h index 87ceb43d1fe6650e1d160f3099b883ea208d8aac..cf1f3b074ed007921e39da4444147a314bfc707f 100644 --- a/tensorflow/compiler/xla/client/computation_builder.h +++ b/tensorflow/compiler/xla/client/computation_builder.h @@ -211,9 +211,21 @@ class ComputationBuilder { // // 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 ComputationDataHandle Slice(const ComputationDataHandle& operand, tensorflow::gtl::ArraySlice start_indices, - tensorflow::gtl::ArraySlice limit_indices); + tensorflow::gtl::ArraySlice limit_indices, + tensorflow::gtl::ArraySlice 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, :] + ComputationDataHandle SliceInDim(const ComputationDataHandle& 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'. @@ -508,6 +520,12 @@ class ComputationBuilder { // Enqueues a sign instruction onto the computation. ComputationDataHandle Sign(const ComputationDataHandle& operand); + // Enqueues a cosine instruction onto the computation. + ComputationDataHandle Cos(const ComputationDataHandle& operand); + + // Enqueues a sine instruction onto the computation. + ComputationDataHandle Sin(const ComputationDataHandle& operand); + // Enqueues a tanh instruction onto the computation. ComputationDataHandle Tanh(const ComputationDataHandle& operand); @@ -595,6 +613,11 @@ class ComputationBuilder { const Computation& body, const ComputationDataHandle& init); + // Enqueues a ReducePrecision node onto the computation. + ComputationDataHandle ReducePrecision(const ComputationDataHandle& operand, + const int exponent_bits, + const int mantissa_bits); + // Enqueues a Send node onto the computation, to send the given operand to // a Recv instruction that shares the same channel handle. void Send(const ComputationDataHandle& operand, const ChannelHandle& handle); @@ -611,15 +634,57 @@ class ComputationBuilder { // computation. StatusOr IsConstant(const ComputationDataHandle& operand); + // 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. + ComputationDataHandle BatchNormTraining(const ComputationDataHandle& operand, + const ComputationDataHandle& scale, + const ComputationDataHandle& 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. + ComputationDataHandle BatchNormInference( + const ComputationDataHandle& operand, const ComputationDataHandle& scale, + const ComputationDataHandle& offset, const ComputationDataHandle& mean, + const ComputationDataHandle& 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` + ComputationDataHandle BatchNormGrad(const ComputationDataHandle& operand, + const ComputationDataHandle& scale, + const ComputationDataHandle& batch_mean, + const ComputationDataHandle& batch_var, + const ComputationDataHandle& grad_output, + float epsilon, int64 feature_index); + // Computes the value of a constant indicated by a // ComputationDataHandle. // - // The handle must be from the computation currently being built - + // The operand must be from the computation currently being built - // i.e., returned from this builder with no intervening call to // Build(). This happens to currently work regardless of that, but // that may stop working at any time. // - // The handle must represent a constant value, which in this case + // The operand must represent a constant value, which in this case // means that it must not statically depend on a parameter to the // computation that is being built. // @@ -637,8 +702,8 @@ class ComputationBuilder { // // If output_layout is non-null, then the output of the computation // will be stored using that layout. - StatusOr> ComputeConstant( - const ComputationDataHandle& handle, + StatusOr> ComputeConstant( + const ComputationDataHandle& operand, const Layout* output_layout = nullptr); // Returns a new ComputationBuilder whose resultant Computation is used only @@ -668,6 +733,14 @@ class ComputationBuilder { // then Build() should be used instead. Computation BuildAndNoteError(); + // Returns the first error that was encountered while building the + // computation. When an error is encountered, by default we return a vacuous + // ComputationDataHandle and inform the user of the error that occurred while + // building the computation when they make a final call to Build(). + // + // See also set_die_immediately_on_error(). + Status first_error() const { return first_error_; } + private: using PopulateLiteral = std::function; @@ -768,87 +841,80 @@ class ComputationBuilder { template ComputationDataHandle ComputationBuilder::ConstantR0(NativeT value) { - return ConstantOp( - [value](Literal* literal) { LiteralUtil::PopulateR0(value, literal); }); + return ConstantOp([value](Literal* literal) { literal->PopulateR0(value); }); } template ComputationDataHandle ComputationBuilder::ConstantR1( tensorflow::gtl::ArraySlice values) { - return ConstantOp([&values](Literal* literal) { - LiteralUtil::PopulateR1(values, literal); - }); + return ConstantOp( + [&values](Literal* literal) { literal->PopulateR1(values); }); } template ComputationDataHandle ComputationBuilder::ConstantR1(int64 length, NativeT value) { return ConstantOp([length, value](Literal* literal) { - LiteralUtil::PopulateWithValue(value, {length}, literal); + literal->PopulateWithValue(value, {length}); }); } inline ComputationDataHandle ComputationBuilder::ConstantR1( const tensorflow::core::Bitmap& values) { - return ConstantOp([&values](Literal* literal) { - LiteralUtil::PopulateR1(values, literal); - }); + return ConstantOp( + [&values](Literal* literal) { literal->PopulateR1(values); }); } template ComputationDataHandle ComputationBuilder::ConstantR2( std::initializer_list> values) { - return ConstantOp([&values](Literal* literal) { - LiteralUtil::PopulateR2(values, literal); - }); + return ConstantOp( + [&values](Literal* literal) { literal->PopulateR2(values); }); } template ComputationDataHandle ComputationBuilder::ConstantR2FromArray2DWithLayout( const Array2D& values, const Layout& layout) { return ConstantOp([&values, &layout](Literal* literal) { - LiteralUtil::PopulateR2FromArray2DWithLayout(values, layout, literal); + literal->PopulateR2FromArray2DWithLayout(values, layout); }); } template ComputationDataHandle ComputationBuilder::ConstantR2FromArray2D( const Array2D& values) { - return ConstantOp([&values](Literal* literal) { - LiteralUtil::PopulateR2FromArray2D(values, literal); - }); + return ConstantOp( + [&values](Literal* literal) { literal->PopulateR2FromArray2D(values); }); } template ComputationDataHandle ComputationBuilder::ConstantR3FromArray3DWithLayout( const Array3D& values, const Layout& layout) { return ConstantOp([&values, &layout](Literal* literal) { - LiteralUtil::PopulateR3FromArray3DWithLayout(values, layout, literal); + literal->PopulateR3FromArray3DWithLayout(values, layout); }); } template ComputationDataHandle ComputationBuilder::ConstantR3FromArray3D( const Array3D& values) { - return ConstantOp([&values](Literal* literal) { - LiteralUtil::PopulateR3FromArray3D(values, literal); - }); + return ConstantOp( + [&values](Literal* literal) { literal->PopulateR3FromArray3D(values); }); } template ComputationDataHandle ComputationBuilder::ConstantR4FromArray4DWithLayout( const Array4D& values, const Layout& layout) { return ConstantOp([&values, &layout](Literal* literal) { - LiteralUtil::PopulateR4FromArray4DWithLayout(values, layout, literal); + literal->PopulateR4FromArray4DWithLayout(values, layout); }); } template ComputationDataHandle ComputationBuilder::ConstantR4FromArray4D( const Array4D& values) { - return ConstantOp([&values](Literal* literal) { - LiteralUtil::PopulateR4FromArray4D(values, literal); - }); + return ConstantOp( + [&values](Literal* literal) { literal->PopulateR4FromArray4D(values); }); } } // namespace xla diff --git a/tensorflow/compiler/xla/client/global_data.cc b/tensorflow/compiler/xla/client/global_data.cc index be706f7d23250b14b75ba4208072a72c07dc374f..40f59eaa68ebeb47edbd2afbeabad0cd2623ebc6 100644 --- a/tensorflow/compiler/xla/client/global_data.cc +++ b/tensorflow/compiler/xla/client/global_data.cc @@ -16,6 +16,7 @@ limitations under the License. #include "tensorflow/compiler/xla/client/global_data.h" #include +#include #include "tensorflow/compiler/xla/types.h" #include "tensorflow/core/platform/logging.h" @@ -23,7 +24,7 @@ limitations under the License. namespace xla { GlobalData::GlobalData(ServiceInterface* parent, GlobalDataHandle handle) - : handle_(handle), parent_(parent) {} + : handle_(std::move(handle)), parent_(parent) {} GlobalData::~GlobalData() { UnregisterRequest request; diff --git a/tensorflow/compiler/xla/client/global_data.h b/tensorflow/compiler/xla/client/global_data.h index eb11d91034ba524f093ff80fa7cd0473e04eac2c..b7929357d06032b55c04bf0391f7fa703ee15f17 100644 --- a/tensorflow/compiler/xla/client/global_data.h +++ b/tensorflow/compiler/xla/client/global_data.h @@ -23,13 +23,15 @@ limitations under the License. namespace xla { -// Wraps a GlobalDataHandle with a lifetime. +// A GlobalData object represents a globally-accessible allocation of +// data in the associated XLA service. class GlobalData { public: // Gives ownership of the global data handle to this object. GlobalData(ServiceInterface* parent, GlobalDataHandle handle); - // Unregisters the wrapped handle. + // Unregisters the wrapped handle, which causes the service to + // deallocate the associated data. ~GlobalData(); const GlobalDataHandle& handle() const { return handle_; } diff --git a/tensorflow/compiler/xla/client/lib/BUILD b/tensorflow/compiler/xla/client/lib/BUILD index 86b16be62f041ae3e96591627501592b34203e16..ee3468208792879c3fe4ff5860e434ef5a0c0155 100644 --- a/tensorflow/compiler/xla/client/lib/BUILD +++ b/tensorflow/compiler/xla/client/lib/BUILD @@ -24,6 +24,7 @@ cc_library( "//tensorflow/compiler/xla:xla_data_proto", "//tensorflow/compiler/xla/client:computation", "//tensorflow/compiler/xla/client:computation_builder", + "//tensorflow/core:lib", ], ) @@ -32,6 +33,7 @@ cc_library( srcs = ["testing.cc"], hdrs = ["testing.h"], deps = [ + "//tensorflow/compiler/xla:execution_options_util", "//tensorflow/compiler/xla:literal_util", "//tensorflow/compiler/xla:shape_util", "//tensorflow/compiler/xla:statusor", diff --git a/tensorflow/compiler/xla/client/lib/arithmetic.cc b/tensorflow/compiler/xla/client/lib/arithmetic.cc index a45974b86b67c14868fcfe9c5f8a43445a35807e..969b0eee1d195a36728f16a598add4b3b850ed60 100644 --- a/tensorflow/compiler/xla/client/lib/arithmetic.cc +++ b/tensorflow/compiler/xla/client/lib/arithmetic.cc @@ -22,65 +22,85 @@ limitations under the License. #include "tensorflow/compiler/xla/status_macros.h" #include "tensorflow/compiler/xla/types.h" #include "tensorflow/compiler/xla/xla_data.pb.h" +#include "tensorflow/core/lib/strings/strcat.h" namespace xla { +namespace { +using InstructionGenerator = + ComputationDataHandle (*)(ComputationBuilder*, const ComputationDataHandle&, + const ComputationDataHandle&); + +Computation CreateScalarComputation(const string& name, PrimitiveType type, + ComputationBuilder* builder, + InstructionGenerator generator) { + std::unique_ptr b; + if (type == PRED) { + b = builder->CreateSubBuilder(name); + } else { + b = builder->CreateSubBuilder( + tensorflow::strings::StrCat(name, "_", PrimitiveType_Name(type))); + } -Computation CreateScalarAddComputation(PrimitiveType type, - ComputationBuilder* builder) { const Shape scalar = ShapeUtil::MakeShape(type, {}); - auto b = builder->CreateSubBuilder("add_" + PrimitiveType_Name(type)); auto lhs = b->Parameter(0, scalar, "lhs"); auto rhs = b->Parameter(1, scalar, "rhs"); - b->Add(lhs, rhs); + generator(b.get(), lhs, rhs); return b->BuildAndNoteError(); } +} // namespace + +Computation CreateScalarAddComputation(PrimitiveType type, + ComputationBuilder* builder) { + return CreateScalarComputation( + "add", type, builder, + [](ComputationBuilder* b, const ComputationDataHandle& lhs, + const ComputationDataHandle& rhs) { return b->Add(lhs, rhs); }); +} + +Computation CreateScalarMultiplyComputation(PrimitiveType type, + ComputationBuilder* builder) { + return CreateScalarComputation( + "add", type, builder, + [](ComputationBuilder* b, const ComputationDataHandle& lhs, + const ComputationDataHandle& rhs) { return b->Mul(lhs, rhs); }); +} Computation CreateScalarGeComputation(PrimitiveType type, ComputationBuilder* builder) { - const Shape scalar = ShapeUtil::MakeShape(type, {}); - auto b = builder->CreateSubBuilder("ge_" + PrimitiveType_Name(type)); - auto lhs = b->Parameter(0, scalar, "lhs"); - auto rhs = b->Parameter(1, scalar, "rhs"); - b->Ge(lhs, rhs); - return b->BuildAndNoteError(); + return CreateScalarComputation( + "ge", type, builder, + [](ComputationBuilder* b, const ComputationDataHandle& lhs, + const ComputationDataHandle& rhs) { return b->Ge(lhs, rhs); }); } Computation CreateScalarMaxComputation(PrimitiveType type, ComputationBuilder* builder) { - const Shape scalar = ShapeUtil::MakeShape(type, {}); - auto b = builder->CreateSubBuilder("max_" + PrimitiveType_Name(type)); - auto lhs = b->Parameter(0, scalar, "lhs"); - auto rhs = b->Parameter(1, scalar, "rhs"); - b->Max(lhs, rhs); - return b->BuildAndNoteError(); + return CreateScalarComputation( + "max", type, builder, + [](ComputationBuilder* b, const ComputationDataHandle& lhs, + const ComputationDataHandle& rhs) { return b->Max(lhs, rhs); }); } Computation CreateScalarMinComputation(PrimitiveType type, ComputationBuilder* builder) { - const Shape scalar = ShapeUtil::MakeShape(type, {}); - auto b = builder->CreateSubBuilder("min_" + PrimitiveType_Name(type)); - auto lhs = b->Parameter(0, scalar, "lhs"); - auto rhs = b->Parameter(1, scalar, "rhs"); - b->Min(lhs, rhs); - return b->BuildAndNoteError(); + return CreateScalarComputation( + "min", type, builder, + [](ComputationBuilder* b, const ComputationDataHandle& lhs, + const ComputationDataHandle& rhs) { return b->Min(lhs, rhs); }); } Computation CreateScalarLogicalAndComputation(ComputationBuilder* builder) { - const Shape scalar = ShapeUtil::MakeShape(PRED, {}); - auto b = builder->CreateSubBuilder("logical_and"); - auto lhs = b->Parameter(0, scalar, "lhs"); - auto rhs = b->Parameter(1, scalar, "rhs"); - b->LogicalAnd(lhs, rhs); - return b->BuildAndNoteError(); + return CreateScalarComputation( + "logical_and", PRED, builder, + [](ComputationBuilder* b, const ComputationDataHandle& lhs, + const ComputationDataHandle& rhs) { return b->LogicalAnd(lhs, rhs); }); } Computation CreateScalarLogicalOrComputation(ComputationBuilder* builder) { - const Shape scalar = ShapeUtil::MakeShape(PRED, {}); - auto b = builder->CreateSubBuilder("logical_or"); - auto lhs = b->Parameter(0, scalar, "lhs"); - auto rhs = b->Parameter(1, scalar, "rhs"); - b->LogicalOr(lhs, rhs); - return b->BuildAndNoteError(); + return CreateScalarComputation( + "logical_or", PRED, builder, + [](ComputationBuilder* b, const ComputationDataHandle& lhs, + const ComputationDataHandle& rhs) { return b->LogicalOr(lhs, rhs); }); } StatusOr Any(const ComputationDataHandle& predicates, diff --git a/tensorflow/compiler/xla/client/lib/arithmetic.h b/tensorflow/compiler/xla/client/lib/arithmetic.h index 633086a2e7e4609543c465c9f52dc452ce3fabb3..f43d35fe4a52016d4054af28835d6b66a35217d4 100644 --- a/tensorflow/compiler/xla/client/lib/arithmetic.h +++ b/tensorflow/compiler/xla/client/lib/arithmetic.h @@ -28,6 +28,10 @@ namespace xla { Computation CreateScalarAddComputation(PrimitiveType type, ComputationBuilder* builder); +// Creates a scalar multiply computation and returns it. +Computation CreateScalarMultiplyComputation(PrimitiveType type, + ComputationBuilder* builder); + // Creates a scalar ge computation and returns it. Computation CreateScalarGeComputation(PrimitiveType type, ComputationBuilder* builder); diff --git a/tensorflow/compiler/xla/client/lib/testing.cc b/tensorflow/compiler/xla/client/lib/testing.cc index daa1557df0b97ee20679f45b8d54164ca93555fa..482d53cf330f152f496b77233714f93991fef6f0 100644 --- a/tensorflow/compiler/xla/client/lib/testing.cc +++ b/tensorflow/compiler/xla/client/lib/testing.cc @@ -17,6 +17,7 @@ limitations under the License. #include "tensorflow/compiler/xla/client/computation.h" #include "tensorflow/compiler/xla/client/computation_builder.h" +#include "tensorflow/compiler/xla/execution_options_util.h" #include "tensorflow/compiler/xla/literal_util.h" #include "tensorflow/compiler/xla/shape_util.h" #include "tensorflow/compiler/xla/statusor.h" @@ -27,23 +28,82 @@ limitations under the License. #include "tensorflow/core/platform/types.h" namespace xla { +namespace { -std::unique_ptr MakeFakeDataOrDie(const Shape& shape, - Client* client) { +std::unique_ptr MakeFakeDataViaDeviceOrDie(const Shape& shape, + Client* client) { ComputationBuilder b( client, tensorflow::strings::StrCat("make_fake_", ShapeUtil::HumanString(shape))); // TODO(b/26811613): Replace this when RNG is supported on all backends. - b.Broadcast(b.ConstantLiteral(LiteralUtil::One(shape.element_type())), + b.Broadcast(b.ConstantLiteral(Literal::One(shape.element_type())), AsInt64Slice(shape.dimensions())); Computation computation = b.Build().ConsumeValueOrDie(); - ExecutionOptions execution_options; + auto execution_options = CreateDefaultExecutionOptions(); *execution_options.mutable_shape_with_output_layout() = shape; return client->Execute(computation, /*arguments=*/{}, &execution_options) .ConsumeValueOrDie(); } +} // namespace + +StatusOr> MakeFakeLiteral(const Shape& shape) { + if (ShapeUtil::IsTuple(shape)) { + std::vector> elements; + for (const Shape& element_shape : shape.tuple_shapes()) { + TF_ASSIGN_OR_RETURN(std::unique_ptr element, + MakeFakeLiteral(element_shape)); + elements.push_back(std::move(element)); + } + return Literal::MakeTupleOwned(std::move(elements)); + } + std::unique_ptr literal = Literal::CreateFromShape(shape); + std::minstd_rand0 engine; + switch (shape.element_type()) { + case F32: { + std::uniform_real_distribution generator(0.0f, 1.0f); + TF_CHECK_OK(literal->Populate( + [&](tensorflow::gtl::ArraySlice /*indices*/) { + return generator(engine); + })); + break; + } + case S32: { + std::uniform_int_distribution generator( + std::numeric_limits::lowest(), + std::numeric_limits::max()); + TF_CHECK_OK(literal->Populate( + [&](tensorflow::gtl::ArraySlice /*indices*/) { + return generator(engine); + })); + break; + } + default: + return Unimplemented("Unsupported type for fake literal generation: %s", + ShapeUtil::HumanString(shape).c_str()); + } + return std::move(literal); +} + +std::unique_ptr MakeFakeDataOrDie(const Shape& shape, + Client* client) { + if (ShapeUtil::ByteSizeOf(shape) < (1LL << 30)) { + StatusOr> literal_status = MakeFakeLiteral(shape); + if (!literal_status.ok()) { + // If we got an Unimplemented error, fall back to making the fake data via + // an on-device computation. + CHECK_EQ(literal_status.status().code(), + tensorflow::error::UNIMPLEMENTED); + return MakeFakeDataViaDeviceOrDie(shape, client); + } + return client->TransferToServer(*literal_status.ValueOrDie()).ValueOrDie(); + } + + // If the data is large, generate it on-device. + return MakeFakeDataViaDeviceOrDie(shape, client); +} + std::vector> MakeFakeArgumentsOrDie( const Computation& computation, Client* client) { auto program_shape = diff --git a/tensorflow/compiler/xla/client/lib/testing.h b/tensorflow/compiler/xla/client/lib/testing.h index 7e640d1307edcc3e2c021f4391c456f578a015ee..b5c4393dcc3e37c03a5b0e1a806b0f8b07a132ed 100644 --- a/tensorflow/compiler/xla/client/lib/testing.h +++ b/tensorflow/compiler/xla/client/lib/testing.h @@ -26,6 +26,10 @@ limitations under the License. namespace xla { +// Generates fake data in a literal of the given shape, or returns an error +// status if the element type is currently unhandled for fake data generation. +StatusOr> MakeFakeLiteral(const Shape& shape); + // 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. diff --git a/tensorflow/compiler/xla/client/local_client.cc b/tensorflow/compiler/xla/client/local_client.cc index bfd14bc1c010353e3e473f10dd6c030cb0438648..3bee24e03e6f4179df61f68a32858ac00e3579ad 100644 --- a/tensorflow/compiler/xla/client/local_client.cc +++ b/tensorflow/compiler/xla/client/local_client.cc @@ -17,7 +17,7 @@ limitations under the License. #include -#include "external/llvm/include/llvm/ADT/Triple.h" +#include "llvm/ADT/Triple.h" #include "tensorflow/compiler/xla/ptr_util.h" #include "tensorflow/compiler/xla/service/backend.h" #include "tensorflow/compiler/xla/service/llvm_ir/llvm_util.h" @@ -176,31 +176,30 @@ StatusOr> LocalExecutable::Run( TF_RETURN_IF_ERROR(ValidateExecutionOptions(arguments, options, *backend_)); ExecutableRunOptions actual_options = options; - Backend::StreamPtr stream; if (options.stream() == nullptr) { TF_ASSIGN_OR_RETURN( - stream, BorrowStreamForDevice(options.device_ordinal(), backend_)); + Backend::StreamPtr stream, + BorrowStreamForDevice(options.device_ordinal(), backend_)); actual_options.set_stream(stream.get()); } if (options.allocator() == nullptr) { actual_options.set_allocator(backend_->memory_allocator()); } - ServiceExecutableRunOptions service_options(actual_options, - backend_->StreamBorrower()); + + // For local client execution on CPU backends: + // *) The thread pool used for eigen CPU ops is from + // 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( + actual_options, backend_->StreamBorrower(), + backend_->eigen_intra_op_thread_pool()); if (executable_->dumping()) { return ExecuteAndDump(&service_options, arguments); } - return Service::ExecuteOnStreamWrapper< - StatusOr>>( - executable_.get(), &service_options, options.execution_profile(), - backend_, - [&arguments](Executable* executable, - const ServiceExecutableRunOptions* run_options, - HloExecutionProfile* hlo_execution_profile) { - return executable->ExecuteOnStream(run_options, arguments, - hlo_execution_profile); - }); + return executable_->ExecuteOnStreamWrapper>( + &service_options, options.execution_profile(), arguments); } StatusOr> LocalExecutable::ExecuteAndDump( @@ -223,8 +222,9 @@ tensorflow::Status LocalExecutable::RecordArguments( SessionModule* session_module) { session_module->clear_arguments(); for (const ShapedBuffer* argument : arguments) { - TF_RETURN_IF_ERROR( - LiteralFromShapedBuffer(*argument, session_module->add_arguments())); + Literal literal; + TF_RETURN_IF_ERROR(LiteralFromShapedBuffer(*argument, &literal)); + *session_module->add_arguments() = literal.ToProto(); } return tensorflow::Status::OK(); } @@ -232,9 +232,13 @@ tensorflow::Status LocalExecutable::RecordArguments( tensorflow::Status LocalExecutable::RecordResult( const ShapedBuffer* result, SessionModule* session_module) { session_module->clear_result(); - return LiteralFromShapedBuffer(*result, session_module->mutable_result()); + Literal literal(session_module->result()); + TF_RETURN_IF_ERROR(LiteralFromShapedBuffer(*result, &literal)); + *session_module->mutable_result() = literal.ToProto(); + return tensorflow::Status::OK(); } +// TODO(dnovillo) Change signature to return StatusOr. tensorflow::Status LocalExecutable::LiteralFromShapedBuffer( const ShapedBuffer& shaped_buffer, Literal* literal) { TF_ASSIGN_OR_RETURN( @@ -253,46 +257,6 @@ StatusOr> LocalClient::AllocateBufferOnDevice( return std::unique_ptr(new GlobalData(local_service_, handle)); } -tensorflow::Status LocalClient::ResolveArguments( - const tensorflow::gtl::ArraySlice arguments, - int device_ordinal, - std::vector* argument_ptrs) { - return local_service_->ResolveArguments(arguments, device_ordinal, - argument_ptrs); -} - -StatusOr>> -LocalClient::CompileAheadOfTime( - const tensorflow::gtl::ArraySlice - computations, - const AotCompilationOptions& options) { - std::vector service_instances; - service_instances.reserve(computations.size()); - for (const AheadOfTimeComputationInstance& instance : computations) { - service_instances.push_back({}); - LocalService::AheadOfTimeComputationInstance& service_instance = - service_instances.back(); - TF_RET_CHECK(instance.computation != nullptr); - service_instance.computation = instance.computation->handle(); - service_instance.argument_layouts = instance.argument_layouts; - service_instance.result_layout = instance.result_layout; - } - return local_service_->CompileAheadOfTime(service_instances, options); -} - -int64 LocalClient::PointerSizeForTriple(tensorflow::StringPiece target_triple) { - llvm::Triple triple( - llvm::Triple::normalize(llvm_ir::AsStringRef(target_triple))); - if (triple.isArch64Bit()) { - return 8; - } else if (triple.isArch32Bit()) { - return 4; - } else { - CHECK(triple.isArch16Bit()); - return 2; - } -} - se::Platform* LocalClient::platform() const { return local_service_->backend().platform(); } diff --git a/tensorflow/compiler/xla/client/local_client.h b/tensorflow/compiler/xla/client/local_client.h index 2c467efcea119b66ad08e0636eca0f1acec3a3b8..c903cd271125b44677f7bb191f100f6604f40bbc 100644 --- a/tensorflow/compiler/xla/client/local_client.h +++ b/tensorflow/compiler/xla/client/local_client.h @@ -56,7 +56,7 @@ class ExecutableBuildOptions { // If set, this specifies the layout of the result of the computation. If not // set, the service will chose the layout of the result. A Shape is used to - // store the layout to accomodate tuple result shapes. A value of nullptr + // store the layout to accommodate tuple result shapes. A value of nullptr // indicates the option has not been set. ExecutableBuildOptions& set_result_layout(const Shape& shape_with_layout); const Shape* result_layout() const; @@ -148,7 +148,7 @@ class LocalExecutable { const ExecutableBuildOptions& build_options_; }; -// An XLA service client object for use when the client and service run in +// An XLA Client specialization for use when the client and service run in // the same process. class LocalClient : public Client { public: @@ -158,14 +158,6 @@ class LocalClient : public Client { LocalClient(const LocalClient&) = delete; void operator=(const LocalClient&) = delete; - // For an array of arguments held on the local service, validate - // that each is placed on the specified device_ordinal, and return - // the DeviceMemoryBase corresponding to each argument. - tensorflow::Status ResolveArguments( - const tensorflow::gtl::ArraySlice arguments, - int device_ordinal, - std::vector* argument_ptrs); - // Return a handle to a buffer large enough to hold shape, allocated // on device_ordinal on the local service. If // allocate_space_for_deep_copy, the buffer is large enough to hold @@ -182,30 +174,6 @@ class LocalClient : public Client { const tensorflow::gtl::ArraySlice argument_layouts, const ExecutableBuildOptions& options); - // A description of a computation to compile using CompileAheadOfTime. - struct AheadOfTimeComputationInstance { - const Computation* computation; - // Inform the compiler of the expected layout for arguments. - std::vector argument_layouts; - // Specifies the expected result layout. - const Shape* result_layout; - }; - - // Compiles a list of computations for ahead-of-time execution. This is - // intended for use in static compilation. The |options| parameter describes - // the target for which the compiler should emit code. - // - // TODO(b/31222190): This doesn't really belong in LocalClient. Move it to its - // own library. - StatusOr>> - CompileAheadOfTime( - const tensorflow::gtl::ArraySlice - computations, - const AotCompilationOptions& options); - - // Returns the size of a pointer in bytes for a given triple. - static int64 PointerSizeForTriple(tensorflow::StringPiece triple); - // Returns the platform that the underlying service targets. perftools::gputools::Platform* platform() const; diff --git a/tensorflow/compiler/xla/executable_run_options.cc b/tensorflow/compiler/xla/executable_run_options.cc index 67f3a6c1df4d74e5ef714dcaa56bae1e81f8276a..33d5b6f1d4d15d5143a3421c87eab9b7a7d11345 100644 --- a/tensorflow/compiler/xla/executable_run_options.cc +++ b/tensorflow/compiler/xla/executable_run_options.cc @@ -77,4 +77,14 @@ ExecutionProfile* ExecutableRunOptions::execution_profile() const { return execution_profile_; } +ExecutableRunOptions& ExecutableRunOptions::set_device_assignment( + DeviceAssignment* device_assignment) { + device_assignment_ = device_assignment; + return *this; +} + +DeviceAssignment* ExecutableRunOptions::device_assignment() const { + return device_assignment_; +} + } // namespace xla diff --git a/tensorflow/compiler/xla/executable_run_options.h b/tensorflow/compiler/xla/executable_run_options.h index 03f2d016ad07b63e6b7d9681c86885ce947f5319..deb3ddb203d263d25bef0499a8a53a6098d0de0c 100644 --- a/tensorflow/compiler/xla/executable_run_options.h +++ b/tensorflow/compiler/xla/executable_run_options.h @@ -40,6 +40,7 @@ struct ThreadPoolDevice; namespace xla { class DeviceMemoryAllocator; +class DeviceAssignment; class ExecutionProfile; // Class containing options for running a LocalExecutable. @@ -79,9 +80,14 @@ class ExecutableRunOptions { ExecutionProfile* execution_profile() const; ExecutableRunOptions& set_execution_profile(ExecutionProfile* profile); + ExecutableRunOptions& set_device_assignment( + DeviceAssignment* device_assignment); + DeviceAssignment* device_assignment() const; + private: DeviceMemoryAllocator* allocator_ = nullptr; int device_ordinal_ = -1; + DeviceAssignment* device_assignment_ = nullptr; perftools::gputools::Stream* stream_ = nullptr; tensorflow::thread::ThreadPool* inter_op_thread_pool_ = nullptr; const Eigen::ThreadPoolDevice* intra_op_thread_pool_ = nullptr; diff --git a/tensorflow/compiler/xla/execution_options_util.cc b/tensorflow/compiler/xla/execution_options_util.cc new file mode 100644 index 0000000000000000000000000000000000000000..e83ff7cddd675197c7f6d7018257edb4c25b6228 --- /dev/null +++ b/tensorflow/compiler/xla/execution_options_util.cc @@ -0,0 +1,27 @@ +/* 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/execution_options_util.h" +#include "tensorflow/compiler/xla/legacy_flags/debug_options_flags.h" + +namespace xla { + +ExecutionOptions CreateDefaultExecutionOptions() { + ExecutionOptions execution_options; + *(execution_options.mutable_debug_options()) = + legacy_flags::GetDebugOptionsFromFlags(); + return execution_options; +} + +} // namespace xla diff --git a/tensorflow/compiler/xla/execution_options_util.h b/tensorflow/compiler/xla/execution_options_util.h new file mode 100644 index 0000000000000000000000000000000000000000..562da78e837ea6c4a01f0d1170797340fd421ad8 --- /dev/null +++ b/tensorflow/compiler/xla/execution_options_util.h @@ -0,0 +1,29 @@ +/* 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 THIRD_PARTY_TENSORFLOW_COMPILER_XLA_EXECUTION_OPTIONS_UTIL_H_ +#define THIRD_PARTY_TENSORFLOW_COMPILER_XLA_EXECUTION_OPTIONS_UTIL_H_ + +#include "tensorflow/compiler/xla/xla.pb.h" + +namespace xla { + +// Create a default ExecutionOptions proto; this proto has its debug options +// popupated to the default values taken from flags. +ExecutionOptions CreateDefaultExecutionOptions(); + +} // namespace xla + +#endif // THIRD_PARTY_TENSORFLOW_COMPILER_XLA_EXECUTION_OPTIONS_UTIL_H_ diff --git a/tensorflow/compiler/xla/index_util.cc b/tensorflow/compiler/xla/index_util.cc index e3248d8e908b60c7e6f7224d25b963601c92f24a..76c0168f370ff1f0749759705b7ecff359a80341 100644 --- a/tensorflow/compiler/xla/index_util.cc +++ b/tensorflow/compiler/xla/index_util.cc @@ -118,17 +118,36 @@ namespace xla { return multi_index; } -/* static */ bool IndexUtil::BumpIndices(const Shape& shape, - std::vector* indices) { - for (int64 dimno = indices->size() - 1; dimno >= 0; --dimno) { +/* static */ bool IndexUtil::BumpIndices( + const Shape& shape, tensorflow::gtl::MutableArraySlice indices) { + for (int64 dimno = indices.size() - 1; dimno >= 0; --dimno) { int64 limit = shape.dimensions(dimno); - if ((*indices)[dimno] + 1 < limit) { - (*indices)[dimno]++; - std::fill(indices->begin() + dimno + 1, indices->end(), 0); + if (indices[dimno] + 1 < limit) { + indices[dimno]++; + std::fill(indices.begin() + dimno + 1, indices.end(), 0); return true; } } return false; } +/* static */ int64 IndexUtil::GetDimensionStride(const Shape& shape, + int64 dimension) { + const Layout& layout = shape.layout(); + int64 pdim_size = layout.padded_dimensions_size(); + int64 stride = 1; + DCHECK(pdim_size == 0 || pdim_size == shape.dimensions_size()); + for (auto dim : layout.minor_to_major()) { + if (dim == dimension) { + break; + } + if (pdim_size == 0) { + stride *= shape.dimensions(dim); + } else { + stride *= layout.padded_dimensions(dim); + } + } + return stride; +} + } // namespace xla diff --git a/tensorflow/compiler/xla/index_util.h b/tensorflow/compiler/xla/index_util.h index 2d8753c3fe8fc05bdcdeaa18360ac5fe4a5e587b..c9838966a5b67397eb5fc4afe3ab9d98e82eb2b1 100644 --- a/tensorflow/compiler/xla/index_util.h +++ b/tensorflow/compiler/xla/index_util.h @@ -58,7 +58,16 @@ class IndexUtil { // // Returns true iff the indices were successfully bumped; false if we've hit // the limit where it can no longer be bumped in-bounds. - static bool BumpIndices(const Shape& shape, std::vector* indices); + static bool BumpIndices(const Shape& shape, + tensorflow::gtl::MutableArraySlice 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. + // Example: + // GetDimensionStride(F32[5,8,10,4]{3,2,1,0}, 1) == + // sizeof(dimension(3)) * sizeof(dimension(2)) == 4 * 10 + static int64 GetDimensionStride(const Shape& shape, int64 dimension); private: TF_DISALLOW_COPY_AND_ASSIGN(IndexUtil); diff --git a/tensorflow/compiler/xla/index_util_test.cc b/tensorflow/compiler/xla/index_util_test.cc index 85259b33f0beea4b508c0d5c1f3a6294dda76813..7c4efdee484d9530a69b31cbe3a0d69a8a3cffa7 100644 --- a/tensorflow/compiler/xla/index_util_test.cc +++ b/tensorflow/compiler/xla/index_util_test.cc @@ -18,9 +18,8 @@ limitations under the License. #include #include "tensorflow/compiler/xla/shape_util.h" -#include "tensorflow/compiler/xla/test_helpers.h" +#include "tensorflow/compiler/xla/test.h" #include "tensorflow/compiler/xla/xla_data.pb.h" -#include "tensorflow/core/platform/test.h" namespace xla { namespace { @@ -144,14 +143,11 @@ TEST(IndexUtilTest, BumpIndices2x2) { auto shape = ShapeUtil::MakeShape(S32, {2, 2}); std::vector indices = {0, 0}; EXPECT_TRUE(IndexUtil::BumpIndices(shape, &indices)); - EXPECT_MATCH(indices, - testing::VectorMatcher(std::vector{0, 1})); + EXPECT_THAT(indices, ::testing::ElementsAre(0, 1)); EXPECT_TRUE(IndexUtil::BumpIndices(shape, &indices)); - EXPECT_MATCH(indices, - testing::VectorMatcher(std::vector{1, 0})); + EXPECT_THAT(indices, ::testing::ElementsAre(1, 0)); EXPECT_TRUE(IndexUtil::BumpIndices(shape, &indices)); - EXPECT_MATCH(indices, - testing::VectorMatcher(std::vector{1, 1})); + EXPECT_THAT(indices, ::testing::ElementsAre(1, 1)); EXPECT_FALSE(IndexUtil::BumpIndices(shape, &indices)); } diff --git a/tensorflow/compiler/xla/layout_util.cc b/tensorflow/compiler/xla/layout_util.cc index 119c4e373f7c52993f6dbbdfe1554d818746ed1d..6271b59a5bfbac05e1bd6f1610dfc11e1cef102d 100644 --- a/tensorflow/compiler/xla/layout_util.cc +++ b/tensorflow/compiler/xla/layout_util.cc @@ -23,7 +23,6 @@ limitations under the License. #include #include -#include "tensorflow/compiler/xla/legacy_flags/layout_util_flags.h" #include "tensorflow/compiler/xla/protobuf_util.h" #include "tensorflow/compiler/xla/shape_util.h" #include "tensorflow/compiler/xla/status_macros.h" @@ -39,35 +38,17 @@ limitations under the License. namespace xla { namespace { -using DimensionOrder = legacy_flags::DefaultLayout::DimensionOrder; - // 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) { + // 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(); - legacy_flags::LayoutUtilFlags* flags = legacy_flags::GetLayoutUtilFlags(); - auto default_layout = flags->xla_default_layout; - switch (default_layout.dimension_order) { - case DimensionOrder::kMajorToMinor: - for (int64 i = 0; i < size; ++i) { - minor_to_major->Set(i, size - 1 - i); - } - break; - case DimensionOrder::kMinorToMajor: - for (int64 i = 0; i < size; ++i) { - minor_to_major->Set(i, i); - } - break; - case DimensionOrder::kRandom: - for (int64 i = 0; i < size; ++i) { - minor_to_major->Set(i, i); - } - std::shuffle( - minor_to_major->begin(), minor_to_major->end(), - std::mt19937(default_layout.seed != 0 ? default_layout.seed - : std::random_device()())); + for (int64 i = 0; i < size; ++i) { + minor_to_major->Set(i, size - 1 - i); } } @@ -152,7 +133,8 @@ Layout CreateDefaultLayoutForRank(int64 rank) { } else { // Array shape. if (!shape.has_layout()) { - return InvalidArgument("shape does not have a layout"); + return InvalidArgument("shape %s does not have a layout", + ShapeUtil::HumanString(shape).c_str()); } return ValidateLayoutForShape(shape.layout(), shape); } diff --git a/tensorflow/compiler/xla/layout_util_test.cc b/tensorflow/compiler/xla/layout_util_test.cc index 531a6e03dad4759416f56465a6c582a06e440a5a..331bb9afa94e9e7c97d9c880dbac31c60ac0da18 100644 --- a/tensorflow/compiler/xla/layout_util_test.cc +++ b/tensorflow/compiler/xla/layout_util_test.cc @@ -15,10 +15,8 @@ limitations under the License. #include "tensorflow/compiler/xla/layout_util.h" #include "tensorflow/compiler/xla/shape_util.h" - -#include "tensorflow/compiler/xla/legacy_flags/layout_util_flags.h" +#include "tensorflow/compiler/xla/test.h" #include "tensorflow/compiler/xla/test_helpers.h" -#include "tensorflow/core/platform/test.h" namespace xla { namespace { @@ -114,8 +112,8 @@ TEST_F(LayoutUtilTest, CopyLayoutNotCompatibleDifferentRank) { Shape dst = MakeShapeWithLayout(F32, {2, 3}, {1, 0}); auto status = LayoutUtil::CopyLayoutBetweenShapes(src, &dst); EXPECT_FALSE(status.ok()); - EXPECT_MATCH(status.error_message(), - testing::ContainsRegex("cannot copy layout from shape")); + EXPECT_THAT(status.error_message(), + ::testing::ContainsRegex("cannot copy layout from shape")); } TEST_F(LayoutUtilTest, CopyLayoutNotCompatibleTuple) { @@ -133,8 +131,8 @@ TEST_F(LayoutUtilTest, CopyLayoutNotCompatibleTuple) { auto status = LayoutUtil::CopyLayoutBetweenShapes(src, &dst); EXPECT_FALSE(status.ok()); - EXPECT_MATCH(status.error_message(), - testing::ContainsRegex("cannot copy layout from shape")); + EXPECT_THAT(status.error_message(), + ::testing::ContainsRegex("cannot copy layout from shape")); } TEST_F(LayoutUtilTest, CopyLayoutBogusLayout) { @@ -145,9 +143,10 @@ TEST_F(LayoutUtilTest, CopyLayoutBogusLayout) { auto status = LayoutUtil::CopyLayoutBetweenShapes(src, &dst); EXPECT_FALSE(status.ok()); - EXPECT_MATCH(status.error_message(), - testing::ContainsRegex("layout minor_to_major field contains .* " - "elements, but shape is rank")); + EXPECT_THAT( + status.error_message(), + ::testing::ContainsRegex("layout minor_to_major field contains .* " + "elements, but shape is rank")); } TEST_F(LayoutUtilTest, ClearLayoutTuple) { @@ -210,13 +209,6 @@ TEST_F(LayoutUtilTest, IsPadded) { } TEST_F(LayoutUtilTest, DefaultLayoutGettersMajorToMinor) { - // Test that LayoutUtil returns expected layouts when the xla_default_layout - // flag is set to kMajorToMinor. - legacy_flags::LayoutUtilFlags* flags = legacy_flags::GetLayoutUtilFlags(); - flags->xla_default_layout = xla::legacy_flags::DefaultLayout{ - .dimension_order = - legacy_flags::DefaultLayout::DimensionOrder::kMajorToMinor}; - EXPECT_TRUE(LayoutUtil::Equal(LayoutUtil::MakeLayout({1, 0}), LayoutUtil::GetDefaultLayoutForR2())); EXPECT_TRUE(LayoutUtil::Equal(LayoutUtil::MakeLayout({2, 1, 0}), @@ -229,25 +221,5 @@ TEST_F(LayoutUtilTest, DefaultLayoutGettersMajorToMinor) { ShapeUtil::MakeShape(F32, {10, 20, 30, 15, 25})))); } -TEST_F(LayoutUtilTest, DefaultLayoutGettersMinorToMajor) { - // Test that LayoutUtil returns expected layouts when the xla_default_layout - // flag is set to kMinorToMajor. - legacy_flags::LayoutUtilFlags* flags = legacy_flags::GetLayoutUtilFlags(); - flags->xla_default_layout = xla::legacy_flags::DefaultLayout{ - .dimension_order = - legacy_flags::DefaultLayout::DimensionOrder::kMinorToMajor}; - - EXPECT_TRUE(LayoutUtil::Equal(LayoutUtil::MakeLayout({0, 1}), - LayoutUtil::GetDefaultLayoutForR2())); - EXPECT_TRUE(LayoutUtil::Equal(LayoutUtil::MakeLayout({0, 1, 2}), - LayoutUtil::GetDefaultLayoutForR3())); - EXPECT_TRUE(LayoutUtil::Equal(LayoutUtil::MakeLayout({0, 1, 2, 3}), - LayoutUtil::GetDefaultLayoutForR4())); - EXPECT_TRUE( - LayoutUtil::Equal(LayoutUtil::MakeLayout({0, 1, 2, 3, 4}), - LayoutUtil::GetDefaultLayoutForShape( - ShapeUtil::MakeShape(F32, {10, 20, 30, 15, 25})))); -} - } // namespace } // namespace xla diff --git a/tensorflow/compiler/xla/legacy_flags/BUILD b/tensorflow/compiler/xla/legacy_flags/BUILD index 79ff81262e9eef56feeb6457e1415cd4580a577c..7977df81f5c93ebc3a53556844584ccd758e0f95 100644 --- a/tensorflow/compiler/xla/legacy_flags/BUILD +++ b/tensorflow/compiler/xla/legacy_flags/BUILD @@ -1,11 +1,11 @@ -# Legacy command line flags for the XLA libraries. +# 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 +# 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 commnd line flags in -# tests and binaries, and serious problems in servers, where one might wish +# 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 @@ -29,6 +29,7 @@ cc_library( cc_test( name = "parse_flags_from_env_test", + size = "small", srcs = ["parse_flags_from_env_test.cc"], deps = [ @@ -41,205 +42,39 @@ cc_test( ) cc_library( - name = "layout_util_flags", - srcs = ["layout_util_flags.cc"], - hdrs = ["layout_util_flags.h"], - deps = - [ - ":parse_flags_from_env", - "//tensorflow/compiler/xla:types", - "//tensorflow/core:framework_internal", - "//tensorflow/core:lib", - ], -) - -cc_library( - name = "util_flags", - srcs = ["util_flags.cc"], - hdrs = ["util_flags.h"], + 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", ], ) -cc_library( - name = "cpu_compiler_flags", - srcs = ["cpu_compiler_flags.cc"], - hdrs = ["cpu_compiler_flags.h"], - deps = - [ - ":parse_flags_from_env", - "//tensorflow/compiler/xla:types", - "//tensorflow/core:framework_internal", - "//tensorflow/core:lib", - ], -) - -cc_library( - name = "cpu_runtime_flags", - srcs = ["cpu_runtime_flags.cc"], - hdrs = ["cpu_runtime_flags.h"], +cc_test( + name = "debug_options_parsers_test", + size = "small", + srcs = [ + "debug_options_parsers.h", + "debug_options_parsers_test.cc", + ], deps = [ - ":parse_flags_from_env", + "//tensorflow/compiler/xla:xla_proto", + "//tensorflow/compiler/xla/service:hlo", "//tensorflow/core:framework_internal", "//tensorflow/core:lib", + "//tensorflow/core:test", ], ) -cc_library( - name = "compiler_functor_flags", - srcs = ["compiler_functor_flags.cc"], - hdrs = ["compiler_functor_flags.h"], - deps = [ - ":parse_flags_from_env", - "//tensorflow/compiler/xla:types", - "//tensorflow/core:framework_internal", - "//tensorflow/core:lib", - ], -) - -cc_library( - name = "convolution_thunk_flags", - srcs = ["convolution_thunk_flags.cc"], - hdrs = ["convolution_thunk_flags.h"], - deps = [ - ":parse_flags_from_env", - "//tensorflow/core:framework_internal", - "//tensorflow/core:lib", - ], -) - -cc_library( - name = "gpu_compiler_flags", - srcs = ["gpu_compiler_flags.cc"], - hdrs = ["gpu_compiler_flags.h"], - deps = [ - ":parse_flags_from_env", - "//tensorflow/compiler/xla:types", - "//tensorflow/core:framework_internal", - "//tensorflow/core:lib", - ], -) - -cc_library( - name = "gpu_backend_lib_flags", - srcs = ["gpu_backend_lib_flags.cc"], - hdrs = ["gpu_backend_lib_flags.h"], - deps = [ - ":parse_flags_from_env", - "//tensorflow/compiler/xla:types", - "//tensorflow/core:framework_internal", - "//tensorflow/core:lib", - ], -) - -cc_library( - name = "stream_assignment_flags", - srcs = ["stream_assignment_flags.cc"], - hdrs = ["stream_assignment_flags.h"], - deps = [ - ":parse_flags_from_env", - "//tensorflow/core:framework_internal", - "//tensorflow/core:lib", - ], -) - -cc_library( - name = "hlo_graph_dumper_flags", - srcs = ["hlo_graph_dumper_flags.cc"], - hdrs = ["hlo_graph_dumper_flags.h"], - deps = [ - ":parse_flags_from_env", - "//tensorflow/compiler/xla:types", - "//tensorflow/core:framework_internal", - "//tensorflow/core:lib", - ], -) - -cc_library( - name = "hlo_pass_pipeline_flags", - srcs = ["hlo_pass_pipeline_flags.cc"], - hdrs = ["hlo_pass_pipeline_flags.h"], - deps = [ - ":parse_flags_from_env", - "//tensorflow/compiler/xla:types", - "//tensorflow/core:framework_internal", - "//tensorflow/core:lib", - ], -) - -cc_library( - name = "alias_analysis_flags", - srcs = ["alias_analysis_flags.cc"], - hdrs = ["alias_analysis_flags.h"], - deps = [ - ":parse_flags_from_env", - "//tensorflow/core:framework_internal", - "//tensorflow/core:lib", - ], -) - -cc_library( - name = "llvm_util_flags", - srcs = ["llvm_util_flags.cc"], - hdrs = ["llvm_util_flags.h"], - deps = [ - ":parse_flags_from_env", - "//tensorflow/core:framework_internal", - "//tensorflow/core:lib", - ], -) - -cc_library( - name = "service_flags", - srcs = ["service_flags.cc"], - hdrs = ["service_flags.h"], - deps = [ - ":parse_flags_from_env", - "//tensorflow/compiler/xla:types", - "//tensorflow/core:framework_internal", - "//tensorflow/core:lib", - ], -) - -cc_library( - name = "buffer_assignment_flags", - srcs = ["buffer_assignment_flags.cc"], - hdrs = ["buffer_assignment_flags.h"], - deps = [ - ":parse_flags_from_env", - "//tensorflow/core:framework_internal", - "//tensorflow/core:lib", - ], -) - -cc_library( - name = "hlo_test_base_flags", - srcs = ["hlo_test_base_flags.cc"], - hdrs = ["hlo_test_base_flags.h"], - deps = [ - ":parse_flags_from_env", - "//tensorflow/core:framework_internal", - "//tensorflow/core:lib", - ], -) - -cc_library( - name = "backend_flags", - srcs = ["backend_flags.cc"], - hdrs = ["backend_flags.h"], - deps = [ - ":parse_flags_from_env", - "//tensorflow/compiler/xla:types", - "//tensorflow/core:framework_internal", - "//tensorflow/core:lib", - ], -) - # ----------------------------------------------------------------------------- filegroup( diff --git a/tensorflow/compiler/xla/legacy_flags/alias_analysis_flags.cc b/tensorflow/compiler/xla/legacy_flags/alias_analysis_flags.cc deleted file mode 100644 index 474753c10ad7ed5eb4a9a446c3f877280c5ad302..0000000000000000000000000000000000000000 --- a/tensorflow/compiler/xla/legacy_flags/alias_analysis_flags.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. -==============================================================================*/ - -// Legacy flags for XLA's alias_analysis module. - -#include // NOLINT(build/c++11): only using std::call_once, not mutex. -#include - -#include "tensorflow/compiler/xla/legacy_flags/alias_analysis_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 xla { -namespace legacy_flags { - -// Pointers to the parsed value of the flags and flag descriptors, initialized -// via flags_init. -static AliasAnalysisFlags* 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 AliasAnalysisFlags; - flags->xla_emit_alias_scope = true; - flag_list = new std::vector({ - tensorflow::Flag("xla_emit_alias_scope", &flags->xla_emit_alias_scope, - "Use buffer analysis to refine alias-analysis."), - }); - ParseFlagsFromEnv(*flag_list); -} - -// Append to *append_to flag definitions associated with XLA's alias_analysis -// module. -void AppendAliasAnalysisFlags(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 AliasAnalysisFlags struct; -// repeated calls return the same pointer. -// This should be called only after Flags::Parse() has returned. -AliasAnalysisFlags* GetAliasAnalysisFlags() { - std::call_once(flags_init, &AllocateFlags); - return flags; -} - -} // namespace legacy_flags -} // namespace xla diff --git a/tensorflow/compiler/xla/legacy_flags/alias_analysis_flags.h b/tensorflow/compiler/xla/legacy_flags/alias_analysis_flags.h deleted file mode 100644 index 369f8cd7caa6f42273cd405ca5f43d325e457128..0000000000000000000000000000000000000000 --- a/tensorflow/compiler/xla/legacy_flags/alias_analysis_flags.h +++ /dev/null @@ -1,46 +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_ALIAS_ANALYSIS_FLAGS_H_ -#define TENSORFLOW_COMPILER_XLA_LEGACY_FLAGS_ALIAS_ANALYSIS_FLAGS_H_ - -// Legacy flags for XLA's alias_analysis module. - -#include - -#include "tensorflow/core/platform/types.h" -#include "tensorflow/core/util/command_line_flags.h" - -namespace xla { -namespace legacy_flags { - -// Append to *flag_list flag definitions associated with XLA's alias_analysis -// module. -void AppendAliasAnalysisFlags(std::vector* flag_list); - -// The values of flags associated with XLA's alias_analysis module. -typedef struct { - bool xla_emit_alias_scope; // Use buffer analysis to refine alias-analysis. -} AliasAnalysisFlags; - -// Return a pointer to the AliasAnalysisFlags struct; -// repeated calls return the same pointer. -// This should be called only after Flags::Parse() has returned. -AliasAnalysisFlags* GetAliasAnalysisFlags(); - -} // namespace legacy_flags -} // namespace xla - -#endif // TENSORFLOW_COMPILER_XLA_LEGACY_FLAGS_ALIAS_ANALYSIS_FLAGS_H_ diff --git a/tensorflow/compiler/xla/legacy_flags/backend_flags.cc b/tensorflow/compiler/xla/legacy_flags/backend_flags.cc deleted file mode 100644 index 7c007f4435c088b35bffce40372f88f37af6ed5b..0000000000000000000000000000000000000000 --- a/tensorflow/compiler/xla/legacy_flags/backend_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 XLA's backend module. - -#include // NOLINT(build/c++11): only using std::call_once, not mutex. -#include - -#include "tensorflow/compiler/xla/legacy_flags/backend_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 xla { -namespace legacy_flags { - -// Pointers to the parsed value of the flags and flag descriptors, initialized -// via flags_init. -static BackendFlags* 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 BackendFlags; - // TODO(b/32648682): Decide if this should continue to be a flag longer term. - flags->xla_replicas = 1; - flag_list = new std::vector({ - tensorflow::Flag( - "xla_replicas", &flags->xla_replicas, - "The number of replicas to use. 1 means no replication."), - }); - ParseFlagsFromEnv(*flag_list); -} - -// Append to *append_to flag definitions associated with XLA's backend module. -void AppendBackendFlags(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 BackendFlags struct; -// repeated calls return the same pointer. -// This should be called only after Flags::Parse() has returned. -BackendFlags* GetBackendFlags() { - std::call_once(flags_init, &AllocateFlags); - return flags; -} - -} // namespace legacy_flags -} // namespace xla diff --git a/tensorflow/compiler/xla/legacy_flags/backend_flags.h b/tensorflow/compiler/xla/legacy_flags/backend_flags.h deleted file mode 100644 index 061238b7e690257f4eb681558dcd59b1f8ba2653..0000000000000000000000000000000000000000 --- a/tensorflow/compiler/xla/legacy_flags/backend_flags.h +++ /dev/null @@ -1,46 +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_BACKEND_FLAGS_H_ -#define TENSORFLOW_COMPILER_XLA_LEGACY_FLAGS_BACKEND_FLAGS_H_ - -// Legacy flags for XLA's backend module. - -#include - -#include "tensorflow/compiler/xla/types.h" -#include "tensorflow/core/util/command_line_flags.h" - -namespace xla { -namespace legacy_flags { - -// Append to *flag_list flag definitions associated with XLA's backend module. -void AppendBackendFlags(std::vector* flag_list); - -// The values of flags associated with XLA's backend module. -typedef struct { - int64 xla_replicas; // The number of replicas to use. 1 means no - // replication. -} BackendFlags; - -// Return a pointer to the BackendFlags struct; -// repeated calls return the same pointer. -// This should be called only after Flags::Parse() has returned. -BackendFlags* GetBackendFlags(); - -} // namespace legacy_flags -} // namespace xla - -#endif // TENSORFLOW_COMPILER_XLA_LEGACY_FLAGS_BACKEND_FLAGS_H_ diff --git a/tensorflow/compiler/xla/legacy_flags/buffer_assignment_flags.cc b/tensorflow/compiler/xla/legacy_flags/buffer_assignment_flags.cc deleted file mode 100644 index 71873f73afd5bb8c59832a4c82f87f4e51c31180..0000000000000000000000000000000000000000 --- a/tensorflow/compiler/xla/legacy_flags/buffer_assignment_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 XLA's buffer_assignment module. - -#include // NOLINT(build/c++11): only using std::call_once, not mutex. -#include - -#include "tensorflow/compiler/xla/legacy_flags/buffer_assignment_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 xla { -namespace legacy_flags { - -// Pointers to the parsed value of the flags and flag descriptors, initialized -// via flags_init. -static BufferAssignmentFlags* 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 BufferAssignmentFlags; - flags->xla_enable_buffer_reuse = true; - flag_list = new std::vector({ - tensorflow::Flag("xla_enable_buffer_reuse", - &flags->xla_enable_buffer_reuse, - "Enable reuse of buffers."), - }); - ParseFlagsFromEnv(*flag_list); -} - -// Append to *append_to flag definitions associated with XLA's buffer_assignment -// module. -void AppendBufferAssignmentFlags(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 BufferAssignmentFlags struct; -// repeated calls return the same pointer. -// This should be called only after Flags::Parse() has returned. -BufferAssignmentFlags* GetBufferAssignmentFlags() { - std::call_once(flags_init, &AllocateFlags); - return flags; -} - -} // namespace legacy_flags -} // namespace xla diff --git a/tensorflow/compiler/xla/legacy_flags/buffer_assignment_flags.h b/tensorflow/compiler/xla/legacy_flags/buffer_assignment_flags.h deleted file mode 100644 index 5f098c2663f638940aead45b74332edcf3fcc37f..0000000000000000000000000000000000000000 --- a/tensorflow/compiler/xla/legacy_flags/buffer_assignment_flags.h +++ /dev/null @@ -1,46 +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_BUFFER_ASSIGNMENT_FLAGS_H_ -#define TENSORFLOW_COMPILER_XLA_LEGACY_FLAGS_BUFFER_ASSIGNMENT_FLAGS_H_ - -// Legacy flags for XLA's buffer_assignment module. - -#include - -#include "tensorflow/core/platform/types.h" -#include "tensorflow/core/util/command_line_flags.h" - -namespace xla { -namespace legacy_flags { - -// Append to *flag_list flag definitions associated with XLA's buffer_assignment -// module. -void AppendBufferAssignmentFlags(std::vector* flag_list); - -// The values of flags associated with XLA's buffer_assignment module. -typedef struct { - bool xla_enable_buffer_reuse; // Enable reuse of buffers. -} BufferAssignmentFlags; - -// Return a pointer to the BufferAssignmentFlags struct; -// repeated calls return the same pointer. -// This should be called only after Flags::Parse() has returned. -BufferAssignmentFlags* GetBufferAssignmentFlags(); - -} // namespace legacy_flags -} // namespace xla - -#endif // TENSORFLOW_COMPILER_XLA_LEGACY_FLAGS_BUFFER_ASSIGNMENT_FLAGS_H_ diff --git a/tensorflow/compiler/xla/legacy_flags/compiler_functor_flags.cc b/tensorflow/compiler/xla/legacy_flags/compiler_functor_flags.cc deleted file mode 100644 index 617a9b712ed99d343dc28b6e6c0de4b54e271096..0000000000000000000000000000000000000000 --- a/tensorflow/compiler/xla/legacy_flags/compiler_functor_flags.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. -==============================================================================*/ - -// Legacy flags for XLA's compiler_functor module. - -#include // NOLINT(build/c++11): only using std::call_once, not mutex. -#include - -#include "tensorflow/compiler/xla/legacy_flags/compiler_functor_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 xla { -namespace legacy_flags { - -// Pointers to the parsed value of the flags and flag descriptors, initialized -// via flags_init. -static CompilerFunctorFlags* 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 CompilerFunctorFlags; - flag_list = new std::vector({ - tensorflow::Flag("xla_debug_cpu_dump_ir", &flags->xla_debug_cpu_dump_ir, - "Dump IR, before optimizations to a path"), - }); - ParseFlagsFromEnv(*flag_list); -} - -// Append to *append_to flag definitions associated with XLA's compiler_functor -// module. -void AppendCompilerFunctorFlags(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 CompilerFunctorFlags struct; -// repeated calls return the same pointer. -// This should be called only after Flags::Parse() has returned. -CompilerFunctorFlags* GetCompilerFunctorFlags() { - std::call_once(flags_init, &AllocateFlags); - return flags; -} - -} // namespace legacy_flags -} // namespace xla diff --git a/tensorflow/compiler/xla/legacy_flags/compiler_functor_flags.h b/tensorflow/compiler/xla/legacy_flags/compiler_functor_flags.h deleted file mode 100644 index 28b505ec5eac2d74879a22779137c6982a7c9ce8..0000000000000000000000000000000000000000 --- a/tensorflow/compiler/xla/legacy_flags/compiler_functor_flags.h +++ /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. -==============================================================================*/ - -#ifndef TENSORFLOW_COMPILER_XLA_LEGACY_FLAGS_COMPILER_FUNCTOR_FLAGS_H_ -#define TENSORFLOW_COMPILER_XLA_LEGACY_FLAGS_COMPILER_FUNCTOR_FLAGS_H_ - -// Legacy flags for the XLA's compiler_functor module. - -#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 { - -// Append to *flag_list flag definitions associated with XLA's compiler_functor -// module. -void AppendCompilerFunctorFlags(std::vector* flag_list); - -// The values of flags associated with XLA's compiler_functor module. -typedef struct { - string xla_debug_cpu_dump_ir; // Dump IR, before optimizations to a path -} CompilerFunctorFlags; - -// Return a pointer to the CompilerFunctorFlags struct; -// repeated calls return the same pointer. -// This should be called only after Flags::Parse() has returned. -CompilerFunctorFlags* GetCompilerFunctorFlags(); - -} // namespace legacy_flags -} // namespace xla - -#endif // TENSORFLOW_COMPILER_XLA_LEGACY_FLAGS_COMPILER_FUNCTOR_FLAGS_H_ diff --git a/tensorflow/compiler/xla/legacy_flags/convolution_thunk_flags.cc b/tensorflow/compiler/xla/legacy_flags/convolution_thunk_flags.cc deleted file mode 100644 index fe5d19147f09557817fee5c670f52058f21f5cdc..0000000000000000000000000000000000000000 --- a/tensorflow/compiler/xla/legacy_flags/convolution_thunk_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 XLA's convolution_thunk module. - -#include // NOLINT(build/c++11): only using std::call_once, not mutex. -#include - -#include "tensorflow/compiler/xla/legacy_flags/convolution_thunk_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 xla { -namespace legacy_flags { - -// Pointers to the parsed value of the flags and flag descriptors, initialized -// via flags_init. -static ConvolutionThunkFlags* 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 ConvolutionThunkFlags; - flags->xla_gpu_autotune_convolution_algorithm = true; - flag_list = new std::vector({ - tensorflow::Flag("xla_gpu_autotune_convolution_algorithm", - &flags->xla_gpu_autotune_convolution_algorithm, - "Auto-tune the algorithm used by convolution"), - }); - ParseFlagsFromEnv(*flag_list); -} - -// Append to *append_to flag definitions associated with XLA's convolution_thunk -// module. -void AppendConvolutionThunkFlags(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 ConvolutionThunkFlags struct; -// repeated calls return the same pointer. -// This should be called only after Flags::Parse() has returned. -ConvolutionThunkFlags* GetConvolutionThunkFlags() { - std::call_once(flags_init, &AllocateFlags); - return flags; -} - -} // namespace legacy_flags -} // namespace xla diff --git a/tensorflow/compiler/xla/legacy_flags/convolution_thunk_flags.h b/tensorflow/compiler/xla/legacy_flags/convolution_thunk_flags.h deleted file mode 100644 index 53d6806a71902d1227728f74bd45f12f9d11421d..0000000000000000000000000000000000000000 --- a/tensorflow/compiler/xla/legacy_flags/convolution_thunk_flags.h +++ /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. -==============================================================================*/ - -#ifndef TENSORFLOW_COMPILER_XLA_LEGACY_FLAGS_CONVOLUTION_THUNK_FLAGS_H_ -#define TENSORFLOW_COMPILER_XLA_LEGACY_FLAGS_CONVOLUTION_THUNK_FLAGS_H_ - -// Legacy flags for XLA's convolution_thunk module. - -#include - -#include "tensorflow/core/platform/types.h" -#include "tensorflow/core/util/command_line_flags.h" - -namespace xla { -namespace legacy_flags { - -// Append to *flag_list flag definitions associated with XLA's convolution_thunk -// module. -void AppendConvolutionThunkFlags(std::vector* flag_list); - -// The values of flags associated with XLA's convolution_thunk module. -typedef struct { - // Auto-tune the algorithm used by convolution - bool xla_gpu_autotune_convolution_algorithm; -} ConvolutionThunkFlags; - -// Return a pointer to the ConvolutionThunkFlags struct; -// repeated calls return the same pointer. -// This should be called only after Flags::Parse() has returned. -ConvolutionThunkFlags* GetConvolutionThunkFlags(); - -} // namespace legacy_flags -} // namespace xla - -#endif // TENSORFLOW_COMPILER_XLA_LEGACY_FLAGS_CONVOLUTION_THUNK_FLAGS_H_ diff --git a/tensorflow/compiler/xla/legacy_flags/cpu_compiler_flags.cc b/tensorflow/compiler/xla/legacy_flags/cpu_compiler_flags.cc deleted file mode 100644 index f8ae25552d4ed5ffce85b297e6cbc998ee28fabb..0000000000000000000000000000000000000000 --- a/tensorflow/compiler/xla/legacy_flags/cpu_compiler_flags.cc +++ /dev/null @@ -1,76 +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 XLA's cpu_compiler module. - -#include // NOLINT(build/c++11): only using std::call_once, not mutex. -#include - -#include "tensorflow/compiler/xla/legacy_flags/cpu_compiler_flags.h" -#include "tensorflow/compiler/xla/legacy_flags/parse_flags_from_env.h" -#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 { - -// Pointers to the parsed value of the flags and flag descriptors, initialized -// via flags_init. -static CpuCompilerFlags* 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 CpuCompilerFlags; - flags->xla_cpu_llvm_opt_level = 2; - flags->xla_cpu_llvm_cl_opts = ""; - flags->xla_cpu_embed_ir = false; - flags->xla_cpu_parallel = false; - flag_list = new std::vector({ - tensorflow::Flag( - "xla_cpu_llvm_opt_level", &flags->xla_cpu_llvm_opt_level, - "The LLVM optimization level for the CPU XLA backend. " - "Valid range is from 0 to 3 where 0 means no optimizations."), - tensorflow::Flag( - "xla_cpu_llvm_cl_opts", &flags->xla_cpu_llvm_cl_opts, - "Comma-separated list of command line options to pass to LLVM."), - tensorflow::Flag( - "xla_cpu_embed_ir", &flags->xla_cpu_embed_ir, - "Embed the LLVM IR module string in the resultant CpuExecutable."), - tensorflow::Flag("xla_cpu_parallel", &flags->xla_cpu_parallel, - "Use the multi-threaded CPU backend."), - }); - ParseFlagsFromEnv(*flag_list); -} - -// Append to *append_to flag definitions associated with XLA's cpu_compiler -// module. -void AppendCpuCompilerFlags(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 CpuCompilerFlags struct; -// repeated calls return the same pointer. -// This should be called only after Flags::Parse() has returned. -CpuCompilerFlags* GetCpuCompilerFlags() { - std::call_once(flags_init, &AllocateFlags); - return flags; -} - -} // namespace legacy_flags -} // namespace xla diff --git a/tensorflow/compiler/xla/legacy_flags/cpu_compiler_flags.h b/tensorflow/compiler/xla/legacy_flags/cpu_compiler_flags.h deleted file mode 100644 index 16a7b687116fbf30a3d7829270126aa9e85518ca..0000000000000000000000000000000000000000 --- a/tensorflow/compiler/xla/legacy_flags/cpu_compiler_flags.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_XLA_LEGACY_FLAGS_CPU_COMPILER_FLAGS_H_ -#define TENSORFLOW_COMPILER_XLA_LEGACY_FLAGS_CPU_COMPILER_FLAGS_H_ - -// Legacy flags for the XLA's cpu_compiler module. - -#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 { - -// Append to *flag_list flag definitions associated with XLA's cpu_compiler -// module. -void AppendCpuCompilerFlags(std::vector* flag_list); - -// The values of flags associated with XLA's cpu_compiler module. -typedef struct { - // The LLVM optimization level for the CPU XLA backend. - // Valid range is from 0 to 3 where 0 means no optimizations. - int32 xla_cpu_llvm_opt_level; - string xla_cpu_llvm_cl_opts; // Comma-separated list of command line options - // to pass to LLVM. - bool xla_cpu_embed_ir; // Embed the LLVM IR module string in the resultant - // CpuExecutable - bool xla_cpu_parallel; // Use the multi-threaded CPU backend. -} CpuCompilerFlags; - -// Return a pointer to the CpuCompilerFlags struct; -// repeated calls return the same pointer. -// This should be called only after Flags::Parse() has returned. -CpuCompilerFlags* GetCpuCompilerFlags(); - -} // namespace legacy_flags -} // namespace xla - -#endif // TENSORFLOW_COMPILER_XLA_LEGACY_FLAGS_CPU_COMPILER_FLAGS_H_ diff --git a/tensorflow/compiler/xla/legacy_flags/cpu_runtime_flags.cc b/tensorflow/compiler/xla/legacy_flags/cpu_runtime_flags.cc deleted file mode 100644 index d7817c5d54a047b1987a19dfbde9f48081ae6413..0000000000000000000000000000000000000000 --- a/tensorflow/compiler/xla/legacy_flags/cpu_runtime_flags.cc +++ /dev/null @@ -1,71 +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 XLA's cpu_runtime module. - -#include // NOLINT(build/c++11): only using std::call_once, not mutex. -#include - -#include "tensorflow/compiler/xla/legacy_flags/cpu_runtime_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 xla { -namespace legacy_flags { - -// Pointers to the parsed value of the flags and flag descriptors, initialized -// via flags_init. -static CpuRuntimeFlags* 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 CpuRuntimeFlags; - flags->xla_cpu_use_eigen = true; - flags->xla_cpu_multi_thread_eigen = true; - flag_list = new std::vector({ - tensorflow::Flag( - "xla_cpu_use_eigen", &flags->xla_cpu_use_eigen, - "Use Eigen for matrix multiply on the CPU platform. This " - "is a useful hack for performance comparisons against " - "XLA's implementation."), - tensorflow::Flag( - "xla_cpu_multi_thread_eigen", &flags->xla_cpu_multi_thread_eigen, - "When generating calls to Eigen for matmul and conv, should " - "single or multi-threaded eigen be used? " - "Only used when --xla_cpu_use_eigen is true."), - }); - ParseFlagsFromEnv(*flag_list); -} - -// Append to *append_to flag definitions associated with XLA's cpu_runtime -// module. -void AppendCpuRuntimeFlags(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 CpuRuntimeFlags struct; -// repeated calls return the same pointer. -// This should be called only after Flags::Parse() has returned. -CpuRuntimeFlags* GetCpuRuntimeFlags() { - std::call_once(flags_init, &AllocateFlags); - return flags; -} - -} // namespace legacy_flags -} // namespace xla diff --git a/tensorflow/compiler/xla/legacy_flags/cpu_runtime_flags.h b/tensorflow/compiler/xla/legacy_flags/cpu_runtime_flags.h deleted file mode 100644 index e3ff30da36a5fabd7d7798fd636cb3955a91b09f..0000000000000000000000000000000000000000 --- a/tensorflow/compiler/xla/legacy_flags/cpu_runtime_flags.h +++ /dev/null @@ -1,51 +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_CPU_RUNTIME_FLAGS_H_ -#define TENSORFLOW_COMPILER_XLA_LEGACY_FLAGS_CPU_RUNTIME_FLAGS_H_ - -// Legacy flags for the XLA's cpu_runtime module. - -#include - -#include "tensorflow/core/platform/types.h" -#include "tensorflow/core/util/command_line_flags.h" - -namespace xla { -namespace legacy_flags { - -// Append to *flag_list flag definitions associated with XLA's cpu_runtime -// module. -void AppendCpuRuntimeFlags(std::vector* flag_list); - -// The values of flags associated with XLA's cpu_runtime module. -typedef struct { - // Use Eigen for matrix multiply on the CPU platform. This is a useful hack - // for performance comparisons against XLA's implementation. - bool xla_cpu_use_eigen; - // When generating calls to Eigen for matmul and conv, should single or - // multi-threaded eigen be used? Only used when --xla_cpu_use_eigen is true. - bool xla_cpu_multi_thread_eigen; -} CpuRuntimeFlags; - -// Return a pointer to the CpuRuntimeFlags struct; -// repeated calls return the same pointer. -// This should be called only after Flags::Parse() has returned. -CpuRuntimeFlags* GetCpuRuntimeFlags(); - -} // namespace legacy_flags -} // namespace xla - -#endif // TENSORFLOW_COMPILER_XLA_LEGACY_FLAGS_CPU_RUNTIME_FLAGS_H_ diff --git a/tensorflow/compiler/xla/legacy_flags/debug_options_flags.cc b/tensorflow/compiler/xla/legacy_flags/debug_options_flags.cc new file mode 100644 index 0000000000000000000000000000000000000000..8892bfbe929d168c602af24cfbb507256dc05328 --- /dev/null +++ b/tensorflow/compiler/xla/legacy_flags/debug_options_flags.cc @@ -0,0 +1,277 @@ +/* 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/legacy_flags/debug_options_flags.h" + +#include // NOLINT(build/c++11): only using std::call_once, not mutex. +#include +#include "tensorflow/compiler/xla/legacy_flags/debug_options_parsers.h" +#include "tensorflow/compiler/xla/legacy_flags/parse_flags_from_env.h" +#include "tensorflow/core/lib/strings/str_util.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_hlo_graph_path("/tmp/"); + flags->set_xla_enable_fast_math(true); + 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); +} + +// 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); + + // Returns a lambda that calls "member_setter" on "flag_values" with the + // argument passed in to the lambda. + auto bool_setter_for = [](void (DebugOptions::*member_setter)(bool)) { + return [member_setter](bool value) { + (flag_values->*member_setter)(value); + return true; + }; + }; + + // Returns a lambda that calls "member_setter" on "flag_values" with the + // argument passed in to the lambda. + auto int32_setter_for = [](void (DebugOptions::*member_setter)(int32)) { + return [member_setter](int32 value) { + (flag_values->*member_setter)(value); + return true; + }; + }; + + // Custom "sub-parser" lambda for xla_disable_hlo_passes. + auto setter_for_xla_disable_hlo_passes = [](string comma_separated_values) { + std::vector disabled_passes = + tensorflow::str_util::Split(comma_separated_values, ','); + for (const auto& passname : disabled_passes) { + flag_values->add_xla_disable_hlo_passes(passname); + } + return true; + }; + + // Custom "sub-parser" lambda for xla_backend_extra_options. + auto setter_for_xla_backend_extra_options = + [](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); + return true; + }; + + // Custom "sub-parser" lambda for xla_reduce_precision. + auto setter_for_xla_reduce_precision = + [](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); + }; + + flag_objects = new std::vector( + {tensorflow::Flag( + "xla_generate_hlo_graph", + flag_values->mutable_xla_generate_hlo_graph(), + "HLO modules matching this regex will be dumped to a .dot file " + "throughout various stages in compilation."), + tensorflow::Flag( + "xla_hlo_graph_addresses", + bool_setter_for(&DebugOptions::set_xla_hlo_graph_addresses), + flag_values->xla_hlo_graph_addresses(), + "With xla_generate_hlo_graph, show addresses of HLO ops in " + "graph dump."), + tensorflow::Flag( + "xla_hlo_graph_path", flag_values->mutable_xla_hlo_graph_path(), + "With xla_generate_hlo_graph, dump the graphs into this path."), + tensorflow::Flag( + "xla_hlo_dump_as_graphdef", + 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_log_hlo_text", flag_values->mutable_xla_log_hlo_text(), + "HLO modules matching this regex will be dumped to LOG(INFO). "), + tensorflow::Flag( + "xla_generate_hlo_text_to", + flag_values->mutable_xla_generate_hlo_text_to(), + "Dump all HLO modules as text into the provided directory path."), + tensorflow::Flag( + "xla_enable_fast_math", + bool_setter_for(&DebugOptions::set_xla_enable_fast_math), + flag_values->xla_enable_fast_math(), + "Enable unsafe fast-math optimizations in the compiler; " + "this may produce faster code at the expense of some accuracy."), + tensorflow::Flag( + "xla_llvm_enable_alias_scope_metadata", + bool_setter_for( + &DebugOptions::set_xla_llvm_enable_alias_scope_metadata), + flag_values->xla_llvm_enable_alias_scope_metadata(), + "In LLVM-based backends, enable the emission of " + "!alias.scope metadata in the generated IR."), + tensorflow::Flag( + "xla_llvm_enable_noalias_metadata", + bool_setter_for(&DebugOptions::set_xla_llvm_enable_noalias_metadata), + flag_values->xla_llvm_enable_noalias_metadata(), + "In LLVM-based backends, enable the emission of " + "!noalias metadata in the generated IR."), + tensorflow::Flag( + "xla_llvm_enable_invariant_load_metadata", + bool_setter_for( + &DebugOptions::set_xla_llvm_enable_invariant_load_metadata), + flag_values->xla_llvm_enable_invariant_load_metadata(), + "In LLVM-based backends, enable the emission of " + "!invariant.load metadata in " + "the generated IR."), + tensorflow::Flag( + "xla_llvm_disable_expensive_passes", + bool_setter_for( + &DebugOptions::set_xla_llvm_disable_expensive_passes), + flag_values->xla_llvm_disable_expensive_passes(), + "In LLVM-based backends, disable a custom set of " + "expensive optimization passes."), + tensorflow::Flag( + "xla_backend_optimization_level", + int32_setter_for(&DebugOptions::set_xla_backend_optimization_level), + flag_values->xla_backend_optimization_level(), + "Numerical optimization level for the XLA compiler backend."), + tensorflow::Flag( + "xla_disable_hlo_passes", setter_for_xla_disable_hlo_passes, "", + "Comma-separated list of hlo passes to be disabled. These names " + "must exactly match the passes' names; no whitespace around " + "commas."), + tensorflow::Flag( + "xla_embed_ir_in_executable", + bool_setter_for(&DebugOptions::set_xla_embed_ir_in_executable), + flag_values->xla_embed_ir_in_executable(), + "Embed the compiler IR as a string in the executable."), + tensorflow::Flag( + "xla_dump_ir_to", flag_values->mutable_xla_dump_ir_to(), + "Dump the compiler IR into this directory as individual files."), + tensorflow::Flag( + "xla_eliminate_hlo_implicit_broadcast", + bool_setter_for( + &DebugOptions::set_xla_eliminate_hlo_implicit_broadcast), + flag_values->xla_eliminate_hlo_implicit_broadcast(), + "Eliminate implicit broadcasts when lowering user " + "computations to HLO instructions; use explicit " + "broadcast instead."), + tensorflow::Flag( + "xla_cpu_multi_thread_eigen", + bool_setter_for(&DebugOptions::set_xla_cpu_multi_thread_eigen), + flag_values->xla_cpu_multi_thread_eigen(), + "When generating calls to Eigen in the CPU backend, " + "use multi-threaded Eigen mode."), + tensorflow::Flag("xla_gpu_cuda_data_dir", + flag_values->mutable_xla_gpu_cuda_data_dir(), + "If non-empty, speficies a local directory containing " + "ptxas and nvvm libdevice files; otherwise we use " + "those from runfile directories."), + tensorflow::Flag("xla_gpu_ftz", + bool_setter_for(&DebugOptions::set_xla_gpu_ftz), + flag_values->xla_gpu_ftz(), + "If true, flush-to-zero semantics are enabled in the " + "code generated for GPUs."), + tensorflow::Flag( + "xla_gpu_disable_multi_streaming", + bool_setter_for(&DebugOptions::set_xla_gpu_disable_multi_streaming), + flag_values->xla_gpu_disable_multi_streaming(), + "If true, multi-streaming in the GPU backend is disabled."), + tensorflow::Flag( + "xla_dump_debug_json_to", + flag_values->mutable_xla_dump_debug_json_to(), + "Dump compilation artifacts as JSON into this directory."), + tensorflow::Flag( + "xla_test_all_output_layouts", + bool_setter_for(&DebugOptions::set_xla_test_all_output_layouts), + flag_values->xla_test_all_output_layouts(), + "Let ClientLibraryTestBase::ComputeAndCompare* test " + "all permutations of output layouts. For example, with " + "a 3D shape, all permutations of the set {0, 1, 2} are " + "tried."), + tensorflow::Flag( + "xla_test_all_input_layouts", + bool_setter_for(&DebugOptions::set_xla_test_all_input_layouts), + flag_values->xla_test_all_input_layouts(), + "Let ClientLibraryTestBase::ComputeAndCompare* test " + "all permutations of *input* layouts. For example, for " + "2 input arguments with 2D shape and 4D shape, the " + "computation will run 2! * 4! times for every possible " + "layouts"), + tensorflow::Flag( + "xla_hlo_profile", + bool_setter_for(&DebugOptions::set_xla_hlo_profile), + flag_values->xla_hlo_profile(), + "Instrument the computation to collect per-HLO cycle counts"), + tensorflow::Flag("xla_dump_computations_to", + flag_values->mutable_xla_dump_computations_to(), + "Dump computations that XLA executes into the provided " + "directory path"), + tensorflow::Flag("xla_dump_executions_to", + flag_values->mutable_xla_dump_executions_to(), + "Dump parameters and results of computations that XLA " + "executes into the provided directory path"), + tensorflow::Flag("xla_backend_extra_options", + setter_for_xla_backend_extra_options, "", + "Extra options to pass to a backend; " + "comma-separated list of 'key=val' strings (=val " + "may be omitted); no whitespace around commas."), + tensorflow::Flag("xla_reduce_precision", setter_for_xla_reduce_precision, + "", + "Directions for adding reduce-precision operations. " + "Format is 'LOCATION=E,M:OPS;NAMES' where LOCATION is " + "the class of locations in which to insert the " + "operations (e.g., 'OP_OUTPUTS'), E and M are the " + "exponent and matissa bit counts respectively, and " + "OPS and NAMES are comma-separated (no spaces) lists " + "of the operation types and names to which to attach " + "the reduce-precision operations. The NAMES string " + "and its preceding ';' may be omitted. This option " + "may be repeated to define multiple sets of added " + "reduce-precision operations.")}); + ParseFlagsFromEnv(*flag_objects); +} + +} // namespace + +void AppendDebugOptionsFlags(std::vector* flag_list) { + std::call_once(flags_init, &AllocateFlags); + flag_list->insert(flag_list->end(), flag_objects->begin(), + flag_objects->end()); +} + +xla::DebugOptions GetDebugOptionsFromFlags() { + std::call_once(flags_init, &AllocateFlags); + return *flag_values; +} + +} // namespace legacy_flags +} // namespace xla diff --git a/tensorflow/compiler/xla/legacy_flags/debug_options_flags.h b/tensorflow/compiler/xla/legacy_flags/debug_options_flags.h new file mode 100644 index 0000000000000000000000000000000000000000..d0ef8e66ab0bcbf88035ae31fe32eb161e32e998 --- /dev/null +++ b/tensorflow/compiler/xla/legacy_flags/debug_options_flags.h @@ -0,0 +1,38 @@ +/* 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 THIRD_PARTY_TENSORFLOW_COMPILER_XLA_LEGACY_FLAGS_DEBUG_OPTIONS_FLAGS_H_ +#define THIRD_PARTY_TENSORFLOW_COMPILER_XLA_LEGACY_FLAGS_DEBUG_OPTIONS_FLAGS_H_ + +#include + +#include "tensorflow/compiler/xla/xla.pb.h" +#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); + +// 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(); + +} // namespace legacy_flags +} // namespace xla + +#endif // THIRD_PARTY_TENSORFLOW_COMPILER_XLA_LEGACY_FLAGS_DEBUG_OPTIONS_FLAGS_H_ diff --git a/tensorflow/compiler/xla/legacy_flags/debug_options_parsers.h b/tensorflow/compiler/xla/legacy_flags/debug_options_parsers.h new file mode 100644 index 0000000000000000000000000000000000000000..0c238e6a5decffb0339f428e4ea676944479cf1b --- /dev/null +++ b/tensorflow/compiler/xla/legacy_flags/debug_options_parsers.h @@ -0,0 +1,151 @@ +/* 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 THIRD_PARTY_TENSORFLOW_COMPILER_XLA_LEGACY_FLAGS_DEBUG_OPTIONS_PARSERS_H_ +#define THIRD_PARTY_TENSORFLOW_COMPILER_XLA_LEGACY_FLAGS_DEBUG_OPTIONS_PARSERS_H_ + +#include +#include "tensorflow/compiler/xla/service/hlo_opcode.h" +#include "tensorflow/compiler/xla/xla.pb.h" +#include "tensorflow/core/lib/strings/str_util.h" +#include "tensorflow/core/lib/strings/stringprintf.h" + +namespace xla { +namespace legacy_flags { +namespace impl { + +template +void parse_xla_backend_extra_options(T* extra_options_map, + string comma_separated_values) { + std::vector extra_options_parts = + tensorflow::str_util::Split(comma_separated_values, ','); + + // The flag contains a comma-separated list of options; some options + // have arguments following "=", some don't. + for (const auto& part : extra_options_parts) { + size_t eq_pos = part.find_first_of('='); + if (eq_pos == string::npos) { + (*extra_options_map)[part] = ""; + } else { + string value = ""; + if (eq_pos + 1 < part.size()) { + value = part.substr(eq_pos + 1); + } + (*extra_options_map)[part.substr(0, eq_pos)] = value; + } + } +} + +// The --xla_reduce_precision option has the format "LOCATION=E,M:OPS;NAME", +// where LOCATION is an HloReducePrecisionOptions::location, E and M are +// integers for the exponent and matissa bit counts respectively, and OPS and +// NAMES are comma-separated of the operation types and names to which to +// attach the reduce-precision operations. The OPS values are matches to the +// strings produced by HloOpcodeString, while the NAME values are arbitrary +// strings subject to the requirements that they not contain any of "=,:;". +// The NAME string (with its preceding semicolon) is optional. +inline bool parse_xla_reduce_precision_option( + HloReducePrecisionOptions* options, string option_string) { + // Split off "LOCATION" from remainder of string. + std::vector eq_split = + tensorflow::str_util::Split(option_string, '='); + if (eq_split.size() != 2) { + return false; + } + string& location = eq_split[0]; + if (location == "OP_INPUTS") { + options->set_location(HloReducePrecisionOptions::OP_INPUTS); + } else if (location == "OP_OUTPUTS") { + options->set_location(HloReducePrecisionOptions::OP_OUTPUTS); + } else if (location == "UNFUSED_OP_OUTPUTS") { + options->set_location(HloReducePrecisionOptions::UNFUSED_OP_OUTPUTS); + } else if (location == "FUSION_INPUTS_BY_CONTENT") { + options->set_location(HloReducePrecisionOptions::FUSION_INPUTS_BY_CONTENT); + } else if (location == "FUSION_OUTPUTS_BY_CONTENT") { + options->set_location(HloReducePrecisionOptions::FUSION_OUTPUTS_BY_CONTENT); + } else { + return false; + } + + // Split off "E,M" from remainder of string. + std::vector colon_split = + tensorflow::str_util::Split(eq_split[1], ':'); + if (colon_split.size() != 2) { + return false; + } + + // Split E and M, and parse. + std::vector bitsizes; + if (!tensorflow::str_util::SplitAndParseAsInts(colon_split[0], ',', + &bitsizes) || + bitsizes.size() != 2) { + return false; + } + options->set_exponent_bits(bitsizes[0]); + options->set_mantissa_bits(bitsizes[1]); + + // Split off OPS comma-separated list from remainder of string, if the + // remainder exists. + std::vector semicolon_split = + tensorflow::str_util::Split(colon_split[1], ';'); + if (semicolon_split.size() > 2) { + return false; + } + // The opcode values are either 'all' (meaning all opcodes), or matches to + // the strings returned by HloOpcodeString. An empty string is also + // interpreted as 'all', for convenience. Note that 'all' may not be part + // of a comma-separated list; it must stand alone. + string& opcode_string = semicolon_split[0]; + if (opcode_string == "" || opcode_string == "all") { + for (int i = 0; i < HloOpcodeCount(); i++) { + options->add_opcodes_to_suffix(i); + } + } else { + std::vector opcodes = + tensorflow::str_util::Split(opcode_string, ','); + for (const string& opcode : opcodes) { + bool found = false; + for (int i = 0; i < HloOpcodeCount(); i++) { + if (opcode == HloOpcodeString(static_cast(i))) { + options->add_opcodes_to_suffix(i); + found = true; + break; + } + } + if (!found) { + return false; + } + } + } + + // Process the NAMES string, if it exists. + if (semicolon_split.size() == 2) { + std::vector opnames = + tensorflow::str_util::Split(semicolon_split[1], ','); + for (const string& opname : opnames) { + if (opname.length() > 0) { + options->add_opname_substrings_to_suffix(opname); + } + } + } + + return true; +} + +} // namespace impl +} // namespace legacy_flags +} // namespace xla + +#endif // THIRD_PARTY_TENSORFLOW_COMPILER_XLA_LEGACY_FLAGS_DEBUG_OPTIONS_PARSERS_H_ diff --git a/tensorflow/compiler/xla/legacy_flags/debug_options_parsers_test.cc b/tensorflow/compiler/xla/legacy_flags/debug_options_parsers_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..0ed788a9676fe9b1bd06fb3ceabf627c108a2c70 --- /dev/null +++ b/tensorflow/compiler/xla/legacy_flags/debug_options_parsers_test.cc @@ -0,0 +1,106 @@ +/* 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 for parse_flags_from_env.cc + +#include "tensorflow/compiler/xla/legacy_flags/debug_options_parsers.h" + +#include +#include + +#include "tensorflow/core/lib/strings/str_util.h" +#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); + EXPECT_EQ(test_map.size(), 4); + EXPECT_EQ(test_map.at("aa"), "bb"); + EXPECT_EQ(test_map.at("cc"), ""); + EXPECT_EQ(test_map.at("dd"), ""); + EXPECT_EQ(test_map.at("ee"), "ff=gg"); +} + +// Test that the xla_reduce_precision flag is parsed correctly. +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_EQ(proto.location(), HloReducePrecisionOptions::OP_OUTPUTS); + EXPECT_EQ(proto.exponent_bits(), 5); + EXPECT_EQ(proto.mantissa_bits(), 10); + EXPECT_EQ(proto.opcodes_to_suffix_size(), 2); + EXPECT_EQ(static_cast(proto.opcodes_to_suffix(0)), + HloOpcode::kAdd); + EXPECT_EQ(static_cast(proto.opcodes_to_suffix(1)), + HloOpcode::kDot); + EXPECT_EQ(proto.opname_substrings_to_suffix_size(), 0); +} + +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_EQ(proto.location(), HloReducePrecisionOptions::OP_OUTPUTS); + EXPECT_EQ(proto.exponent_bits(), 5); + EXPECT_EQ(proto.mantissa_bits(), 10); + EXPECT_EQ(proto.opcodes_to_suffix_size(), 2); + EXPECT_EQ(static_cast(proto.opcodes_to_suffix(0)), + HloOpcode::kAdd); + EXPECT_EQ(static_cast(proto.opcodes_to_suffix(1)), + HloOpcode::kDot); + EXPECT_EQ(proto.opname_substrings_to_suffix_size(), 0); +} + +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_EQ(proto.location(), HloReducePrecisionOptions::UNFUSED_OP_OUTPUTS); + EXPECT_EQ(proto.exponent_bits(), 5); + EXPECT_EQ(proto.mantissa_bits(), 10); + EXPECT_EQ(proto.opcodes_to_suffix_size(), HloOpcodeCount()); + EXPECT_EQ(proto.opname_substrings_to_suffix_size(), 2); + EXPECT_EQ(proto.opname_substrings_to_suffix(0), "foo"); + EXPECT_EQ(proto.opname_substrings_to_suffix(1), "bar/baz"); +} + +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_EQ(proto.location(), HloReducePrecisionOptions::UNFUSED_OP_OUTPUTS); + EXPECT_EQ(proto.exponent_bits(), 5); + EXPECT_EQ(proto.mantissa_bits(), 10); + EXPECT_EQ(proto.opcodes_to_suffix_size(), 1); + EXPECT_EQ(static_cast(proto.opcodes_to_suffix(0)), + HloOpcode::kSubtract); + EXPECT_EQ(proto.opname_substrings_to_suffix_size(), 2); + EXPECT_EQ(proto.opname_substrings_to_suffix(0), "foo"); + EXPECT_EQ(proto.opname_substrings_to_suffix(1), "bar/baz"); +} + +} // namespace legacy_flags +} // namespace xla + +int main(int argc, char* argv[]) { + testing::InitGoogleTest(&argc, argv); + return RUN_ALL_TESTS(); +} diff --git a/tensorflow/compiler/xla/legacy_flags/gpu_backend_lib_flags.cc b/tensorflow/compiler/xla/legacy_flags/gpu_backend_lib_flags.cc deleted file mode 100644 index f8f6ea26b1d0df67b934616fe60aa29199fc2eb9..0000000000000000000000000000000000000000 --- a/tensorflow/compiler/xla/legacy_flags/gpu_backend_lib_flags.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. -==============================================================================*/ - -// Legacy flags for XLA's gpu_backend_lib module. - -#include // NOLINT(build/c++11): only using std::call_once, not mutex. -#include - -#include "tensorflow/compiler/xla/legacy_flags/gpu_backend_lib_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 xla { -namespace legacy_flags { - -// Pointers to the parsed value of the flags and flag descriptors, initialized -// via flags_init. -static GpuBackendLibFlags* 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 GpuBackendLibFlags; - flags->dump_temp_products_to = ""; - flags->ftz = false; - flags->fma = true; - flags->verbose_ptx_asm = false; - flags->kernel = ""; - flags->llvm_dump_passes = false; - flags->llvm_cl_opts = ""; - flags->dump_ir_before_passes = false; - flags->opt_level = 3; - flag_list = new std::vector({ - tensorflow::Flag("dump_temp_products_to", &flags->dump_temp_products_to, - "dump temporary compilation products to this directory. " - "If empty, no dump is produced"), - tensorflow::Flag("ftz", &flags->ftz, "flush to zero semantics"), - tensorflow::Flag("fma", &flags->fma, "use FMA synthesis"), - tensorflow::Flag("verbose_ptx_asm", &flags->verbose_ptx_asm, - "emit PTX assembly with extra comments"), - tensorflow::Flag("kernel", &flags->kernel, - "only emit the IR and PTX for this kernel"), - tensorflow::Flag("llvm_dump_passes", &flags->llvm_dump_passes, - "dump the passes LLVM runs to stderr"), - tensorflow::Flag( - "llvm_cl_opts", &flags->llvm_cl_opts, - "comma-separated list of command line options to pass to " - "LLVM. For example, --llvm_cl_opts=--print-before=loop-unroll"), - tensorflow::Flag("dump_ir_before_passes", &flags->dump_ir_before_passes, - "dump the IR before each optimization pass in " - "sequentially-named files."), - tensorflow::Flag("opt_level", &flags->opt_level, - "optimization level (default to 3)"), - }); - ParseFlagsFromEnv(*flag_list); -} - -// Append to *append_to flag definitions associated with XLA's gpu_backend_lib -// module. -void AppendGpuBackendLibFlags(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 GpuBackendLibFlags struct; -// repeated calls return the same pointer. -// This should be called only after Flags::Parse() has returned. -GpuBackendLibFlags* GetGpuBackendLibFlags() { - std::call_once(flags_init, &AllocateFlags); - return flags; -} - -} // namespace legacy_flags -} // namespace xla diff --git a/tensorflow/compiler/xla/legacy_flags/gpu_backend_lib_flags.h b/tensorflow/compiler/xla/legacy_flags/gpu_backend_lib_flags.h deleted file mode 100644 index 31cb50e9da986b5bad3e71439a4976ec84e17be7..0000000000000000000000000000000000000000 --- a/tensorflow/compiler/xla/legacy_flags/gpu_backend_lib_flags.h +++ /dev/null @@ -1,55 +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_GPU_BACKEND_LIB_FLAGS_H_ -#define TENSORFLOW_COMPILER_XLA_LEGACY_FLAGS_GPU_BACKEND_LIB_FLAGS_H_ - -// Legacy flags for XLA's gpu_backend_lib module. - -#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 { - -// Append to *flag_list flag definitions associated with XLA's gpu_backend_lib -// module. -void AppendGpuBackendLibFlags(std::vector* flag_list); - -// The values of flags associated with XLA's gpu_backend_lib module. -typedef struct { - string dump_temp_products_to; // temporary compilation products dir - bool ftz; // flush to zero semantics - bool fma; // use FMA synthesis - bool verbose_ptx_asm; // emit PTX assembly with extra comments - string kernel; // only emit the IR and PTX for this kernel - bool llvm_dump_passes; // dump the passes LLVM runs to stderr - string llvm_cl_opts; // comma-separated list of LLVM options - bool dump_ir_before_passes; // dump IR before each pass - int32 opt_level; // optimization level -} GpuBackendLibFlags; - -// Return a pointer to the GpuBackendLibFlags struct; -// repeated calls return the same pointer. -// This should be called only after Flags::Parse() has returned. -GpuBackendLibFlags* GetGpuBackendLibFlags(); - -} // namespace legacy_flags -} // namespace xla - -#endif // TENSORFLOW_COMPILER_XLA_LEGACY_FLAGS_GPU_BACKEND_LIB_FLAGS_H_ diff --git a/tensorflow/compiler/xla/legacy_flags/gpu_compiler_flags.cc b/tensorflow/compiler/xla/legacy_flags/gpu_compiler_flags.cc deleted file mode 100644 index e79d3635095a0aacf20b37e586d2c9ac799cbe07..0000000000000000000000000000000000000000 --- a/tensorflow/compiler/xla/legacy_flags/gpu_compiler_flags.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. -==============================================================================*/ - -// Legacy flags for XLA's gpu_compiler module. - -#include // NOLINT(build/c++11): only using std::call_once, not mutex. -#include - -#include "tensorflow/compiler/xla/legacy_flags/gpu_compiler_flags.h" -#include "tensorflow/compiler/xla/legacy_flags/parse_flags_from_env.h" -#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 { - -// Pointers to the parsed value of the flags and flag descriptors, initialized -// via flags_init. -static GpuCompilerFlags* 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 GpuCompilerFlags; - flags->xla_gpu_embed_ir = false; - flags->xla_cuda_data_dir = "./cuda_sdk_lib"; - flags->xla_ptxas_path = "/usr/local/cuda/bin/ptxas"; - flag_list = new std::vector({ - tensorflow::Flag( - "xla_gpu_embed_ir", &flags->xla_gpu_embed_ir, - "Embed the LLVM IR module string in the resultant GpuExecutable."), - tensorflow::Flag( - "xla_cuda_data_dir", &flags->xla_cuda_data_dir, - "If non-empty, specifies a local directory containing ptxas and " - "nvvm libdevice files. Otherwise, by default, we use those from " - "runfile directories."), - tensorflow::Flag("xla_ptxas_path", &flags->xla_ptxas_path, - "The path to ptxas. Required to log stats of the ptx."), - }); - ParseFlagsFromEnv(*flag_list); -} - -// Append to *append_to flag definitions associated with XLA's gpu_compiler -// module. -void AppendGpuCompilerFlags(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 GpuCompilerFlags struct; -// repeated calls return the same pointer. -// This should be called only after Flags::Parse() has returned. -GpuCompilerFlags* GetGpuCompilerFlags() { - std::call_once(flags_init, &AllocateFlags); - return flags; -} - -} // namespace legacy_flags -} // namespace xla diff --git a/tensorflow/compiler/xla/legacy_flags/gpu_compiler_flags.h b/tensorflow/compiler/xla/legacy_flags/gpu_compiler_flags.h deleted file mode 100644 index 04ddedab7325ef728b3fa8ee6a99079fd382c4e7..0000000000000000000000000000000000000000 --- a/tensorflow/compiler/xla/legacy_flags/gpu_compiler_flags.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_XLA_LEGACY_FLAGS_GPU_COMPILER_FLAGS_H_ -#define TENSORFLOW_COMPILER_XLA_LEGACY_FLAGS_GPU_COMPILER_FLAGS_H_ - -// Legacy flags for XLA's gpu_compiler module. - -#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 { - -// Append to *flag_list flag definitions associated with XLA's gpu_compiler -// module. -void AppendGpuCompilerFlags(std::vector* flag_list); - -// The values of flags associated with XLA's gpu_compiler module. -typedef struct { - bool xla_gpu_embed_ir; // Embed the LLVM IR module string in the resultant - // GpuExecutable. - string xla_cuda_data_dir; // If non-empty, specifies a local directory - // containing ptxas and nvvm libdevice files. - // Otherwise, by default, we use those from runfile - // directories. - string xla_ptxas_path; // The path to ptxas. Required to log stats of - // the ptx. -} GpuCompilerFlags; - -// Return a pointer to the GpuCompilerFlags struct; -// repeated calls return the same pointer. -// This should be called only after Flags::Parse() has returned. -GpuCompilerFlags* GetGpuCompilerFlags(); - -} // namespace legacy_flags -} // namespace xla - -#endif // TENSORFLOW_COMPILER_XLA_LEGACY_FLAGS_GPU_COMPILER_FLAGS_H_ diff --git a/tensorflow/compiler/xla/legacy_flags/hlo_graph_dumper_flags.cc b/tensorflow/compiler/xla/legacy_flags/hlo_graph_dumper_flags.cc deleted file mode 100644 index 8822f6f6107d3d9ff121c04e5904a7367c604be7..0000000000000000000000000000000000000000 --- a/tensorflow/compiler/xla/legacy_flags/hlo_graph_dumper_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 XLA's hlo_graph_dumper module. - -#include // NOLINT(build/c++11): only using std::call_once, not mutex. -#include - -#include "tensorflow/compiler/xla/legacy_flags/hlo_graph_dumper_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 xla { -namespace legacy_flags { - -// Pointers to the parsed value of the flags and flag descriptors, initialized -// via flags_init. -static HloGraphDumperFlags* 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 HloGraphDumperFlags; - flags->xla_hlo_dump_graph_path = "/tmp/"; - flag_list = new std::vector({ - tensorflow::Flag("xla_hlo_dump_graph_path", - &flags->xla_hlo_dump_graph_path, - "Path to write dumped HLO graphs to"), - }); - ParseFlagsFromEnv(*flag_list); -} - -// Append to *append_to flag definitions associated with XLA's hlo_graph_dumper -// module. -void AppendHloGraphDumperFlags(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 HloGraphDumperFlags struct; -// repeated calls return the same pointer. -// This should be called only after Flags::Parse() has returned. -HloGraphDumperFlags* GetHloGraphDumperFlags() { - std::call_once(flags_init, &AllocateFlags); - return flags; -} - -} // namespace legacy_flags -} // namespace xla diff --git a/tensorflow/compiler/xla/legacy_flags/hlo_graph_dumper_flags.h b/tensorflow/compiler/xla/legacy_flags/hlo_graph_dumper_flags.h deleted file mode 100644 index b6dfced87cae90c67bd46975a8e36eaef10b19e7..0000000000000000000000000000000000000000 --- a/tensorflow/compiler/xla/legacy_flags/hlo_graph_dumper_flags.h +++ /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. -==============================================================================*/ - -#ifndef TENSORFLOW_COMPILER_XLA_LEGACY_FLAGS_HLO_GRAPH_DUMPER_FLAGS_H_ -#define TENSORFLOW_COMPILER_XLA_LEGACY_FLAGS_HLO_GRAPH_DUMPER_FLAGS_H_ - -// Legacy flags for XLA's hlo_graph_dumper module. - -#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 { - -// Append to *flag_list flag definitions associated with XLA's hlo_graph_dumper -// module. -void AppendHloGraphDumperFlags(std::vector* flag_list); - -// The values of flags associated with XLA's hlo_graph_dumper module. -typedef struct { - string xla_hlo_dump_graph_path; // Path to write dumped HLO graphs to -} HloGraphDumperFlags; - -// Return a pointer to the HloGraphDumperFlags struct; -// repeated calls return the same pointer. -// This should be called only after Flags::Parse() has returned. -HloGraphDumperFlags* GetHloGraphDumperFlags(); - -} // namespace legacy_flags -} // namespace xla - -#endif // TENSORFLOW_COMPILER_XLA_LEGACY_FLAGS_HLO_GRAPH_DUMPER_FLAGS_H_ diff --git a/tensorflow/compiler/xla/legacy_flags/hlo_pass_pipeline_flags.cc b/tensorflow/compiler/xla/legacy_flags/hlo_pass_pipeline_flags.cc deleted file mode 100644 index edc04d51a70f2746086273152833d8db36c46d28..0000000000000000000000000000000000000000 --- a/tensorflow/compiler/xla/legacy_flags/hlo_pass_pipeline_flags.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. -==============================================================================*/ - -// Legacy flags for XLA's hlo_pass_pipeline module. - -#include // NOLINT(build/c++11): only using std::call_once, not mutex. -#include - -#include "tensorflow/compiler/xla/legacy_flags/hlo_pass_pipeline_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 xla { -namespace legacy_flags { - -// Pointers to the parsed value of the flags and flag descriptors, initialized -// via flags_init. -static HloPassPipelineFlags* 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 HloPassPipelineFlags; - flags->xla_disable_hlo_passes = ""; - flag_list = new std::vector({ - tensorflow::Flag("xla_disable_hlo_passes", &flags->xla_disable_hlo_passes, - "Comma-separated list of HLO passes to disable."), - }); - ParseFlagsFromEnv(*flag_list); -} - -// Append to *append_to flag definitions associated with XLA's hlo_pass_pipeline -// module. -void AppendHloPassPipelineFlags(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 HloPassPipelineFlags struct; -// repeated calls return the same pointer. -// This should be called only after Flags::Parse() has returned. -HloPassPipelineFlags* GetHloPassPipelineFlags() { - std::call_once(flags_init, &AllocateFlags); - return flags; -} - -} // namespace legacy_flags -} // namespace xla diff --git a/tensorflow/compiler/xla/legacy_flags/hlo_pass_pipeline_flags.h b/tensorflow/compiler/xla/legacy_flags/hlo_pass_pipeline_flags.h deleted file mode 100644 index 520759bbf0d2f43c763ed915bcd63520d4b9cb58..0000000000000000000000000000000000000000 --- a/tensorflow/compiler/xla/legacy_flags/hlo_pass_pipeline_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_XLA_LEGACY_FLAGS_HLO_PASS_PIPELINE_FLAGS_H_ -#define TENSORFLOW_COMPILER_XLA_LEGACY_FLAGS_HLO_PASS_PIPELINE_FLAGS_H_ - -// Legacy flags for XLA's hlo_pass_pipeline module. - -#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 { - -// Append to *flag_list flag definitions associated with XLA's hlo_pass_pipeline -// module. -void AppendHloPassPipelineFlags(std::vector* flag_list); - -// The values of flags associated with XLA's hlo_pass_pipeline module. -typedef struct { - // Comma-separated list of HLO passes to disable. - string xla_disable_hlo_passes; -} HloPassPipelineFlags; - -// Return a pointer to the HloPassPipelineFlags struct; -// repeated calls return the same pointer. -// This should be called only after Flags::Parse() has returned. -HloPassPipelineFlags* GetHloPassPipelineFlags(); - -} // namespace legacy_flags -} // namespace xla - -#endif // TENSORFLOW_COMPILER_XLA_LEGACY_FLAGS_HLO_PASS_PIPELINE_FLAGS_H_ diff --git a/tensorflow/compiler/xla/legacy_flags/hlo_test_base_flags.cc b/tensorflow/compiler/xla/legacy_flags/hlo_test_base_flags.cc deleted file mode 100644 index c7893c138596b034dbb83df9fda2d4c5edd8e32b..0000000000000000000000000000000000000000 --- a/tensorflow/compiler/xla/legacy_flags/hlo_test_base_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 XLA's hlo_test_base module. - -#include // NOLINT(build/c++11): only using std::call_once, not mutex. -#include - -#include "tensorflow/compiler/xla/legacy_flags/hlo_test_base_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 xla { -namespace legacy_flags { - -// Pointers to the parsed value of the flags and flag descriptors, initialized -// via flags_init. -static HloTestBaseFlags* 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 HloTestBaseFlags; - flags->xla_hlo_test_generate_hlo_graph = false; - flag_list = new std::vector({ - tensorflow::Flag("xla_hlo_test_generate_hlo_graph", - &flags->xla_hlo_test_generate_hlo_graph, - "Generate graph output of HLO instructions"), - }); - ParseFlagsFromEnv(*flag_list); -} - -// Append to *append_to flag definitions associated with XLA's hlo_test_base -// module. -void AppendHloTestBaseFlags(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 HloTestBaseFlags struct; -// repeated calls return the same pointer. -// This should be called only after Flags::Parse() has returned. -HloTestBaseFlags* GetHloTestBaseFlags() { - std::call_once(flags_init, &AllocateFlags); - return flags; -} - -} // namespace legacy_flags -} // namespace xla diff --git a/tensorflow/compiler/xla/legacy_flags/hlo_test_base_flags.h b/tensorflow/compiler/xla/legacy_flags/hlo_test_base_flags.h deleted file mode 100644 index 23b808cecb7e5eaf480292f5207a4b87ebd4a2d5..0000000000000000000000000000000000000000 --- a/tensorflow/compiler/xla/legacy_flags/hlo_test_base_flags.h +++ /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. -==============================================================================*/ - -#ifndef TENSORFLOW_COMPILER_XLA_LEGACY_FLAGS_HLO_TEST_BASE_FLAGS_H_ -#define TENSORFLOW_COMPILER_XLA_LEGACY_FLAGS_HLO_TEST_BASE_FLAGS_H_ - -// Legacy flags for XLA's hlo_test_base module. - -#include - -#include "tensorflow/core/platform/types.h" -#include "tensorflow/core/util/command_line_flags.h" - -namespace xla { -namespace legacy_flags { - -// Append to *flag_list flag definitions associated with XLA's hlo_test_base -// module. -void AppendHloTestBaseFlags(std::vector* flag_list); - -// The values of flags associated with XLA's hlo_test_base module. -typedef struct { - bool xla_hlo_test_generate_hlo_graph; // Generate graph output of HLO - // instructions -} HloTestBaseFlags; - -// Return a pointer to the HloTestBaseFlags struct; -// repeated calls return the same pointer. -// This should be called only after Flags::Parse() has returned. -HloTestBaseFlags* GetHloTestBaseFlags(); - -} // namespace legacy_flags -} // namespace xla - -#endif // TENSORFLOW_COMPILER_XLA_LEGACY_FLAGS_HLO_TEST_BASE_FLAGS_H_ diff --git a/tensorflow/compiler/xla/legacy_flags/layout_util_flags.cc b/tensorflow/compiler/xla/legacy_flags/layout_util_flags.cc deleted file mode 100644 index 4242b501d41379286b14b984aebd2de648873159..0000000000000000000000000000000000000000 --- a/tensorflow/compiler/xla/legacy_flags/layout_util_flags.cc +++ /dev/null @@ -1,107 +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 XLA's layout_util module. - -#include // NOLINT(build/c++11): only using std::call_once, not mutex. -#include - -#include "tensorflow/compiler/xla/legacy_flags/layout_util_flags.h" -#include "tensorflow/compiler/xla/legacy_flags/parse_flags_from_env.h" -#include "tensorflow/compiler/xla/types.h" -#include "tensorflow/core/lib/strings/numbers.h" -#include "tensorflow/core/lib/strings/str_util.h" -#include "tensorflow/core/platform/types.h" -#include "tensorflow/core/util/command_line_flags.h" - -namespace xla { -namespace legacy_flags { - -// Pointers to the string value of the xla_default_layout flag and the flag -// descriptor, initialized via raw_flags_init. -static string* raw_flag; -static std::vector* flag_list; -static std::once_flag raw_flags_init; - -// Allocate *raw_flag. Called via call_once(&raw_flags_init,...). -static void AllocateRawFlag() { - raw_flag = new string; - flag_list = new std::vector({ - tensorflow::Flag( - "xla_default_layout", raw_flag, - "Default layout for Shapes in XLA. Valid values are: " - "'minor2major', 'major2minor', 'random', 'random:'. " - "For debugging purposes. If no seed (or 0) is given, a seed from " - "random_device is used."), - }); - ParseFlagsFromEnv(*flag_list); -} - -// Parse text into *layout. -static bool ParseDefaultLayout(const string& text, DefaultLayout* layout) { - bool result = true; - std::vector field = tensorflow::str_util::Split(text, ':'); - if (field.size() > 0) { - if (field[0] == "random") { - layout->dimension_order = DefaultLayout::DimensionOrder::kRandom; - if (field.size() > 1) { - uint64 seed = 0; - result = tensorflow::strings::safe_strtou64(field[1], &seed); - layout->seed = seed; - } - } else if (field[0] == "minor2major") { - layout->dimension_order = DefaultLayout::DimensionOrder::kMinorToMajor; - } else if (field[0] == "major2minor") { - layout->dimension_order = DefaultLayout::DimensionOrder::kMajorToMinor; - } else { - result = false; - } - } - return result; -} - -// Pointer to the parsed value of the flags, initialized via flags_init. -static LayoutUtilFlags* flags; -static std::once_flag flags_init; - -// Allocate *flags. Called via call_once(&flags_init,...). -static void AllocateFlags() { - std::call_once(raw_flags_init, &AllocateRawFlag); - flags = new LayoutUtilFlags; - flags->xla_default_layout.dimension_order = - DefaultLayout::DimensionOrder::kMajorToMinor; - flags->xla_default_layout.seed = 0; - if (!ParseDefaultLayout(*raw_flag, &flags->xla_default_layout)) { - flags = nullptr; - } -} - -// Append to *append_to the flag definitions associated with XLA's layout_util -// module. -void AppendLayoutUtilFlags(std::vector* append_to) { - std::call_once(raw_flags_init, &AllocateRawFlag); - append_to->insert(append_to->end(), flag_list->begin(), flag_list->end()); -} - -// Return a pointer to the LayoutUtilFlags struct; -// repeated calls return the same pointer. -// This should be called only after Flags::Parse() has returned. -LayoutUtilFlags* GetLayoutUtilFlags() { - std::call_once(flags_init, &AllocateFlags); - return flags; -} - -} // namespace legacy_flags -} // namespace xla diff --git a/tensorflow/compiler/xla/legacy_flags/layout_util_flags.h b/tensorflow/compiler/xla/legacy_flags/layout_util_flags.h deleted file mode 100644 index 177f428b734dcdf703472f3e240aef9792f988d7..0000000000000000000000000000000000000000 --- a/tensorflow/compiler/xla/legacy_flags/layout_util_flags.h +++ /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. -==============================================================================*/ - -#ifndef TENSORFLOW_COMPILER_XLA_LEGACY_FLAGS_LAYOUT_UTIL_FLAGS_H_ -#define TENSORFLOW_COMPILER_XLA_LEGACY_FLAGS_LAYOUT_UTIL_FLAGS_H_ - -// Legacy flags for the XLA's layout_util module. - -#include - -#include "tensorflow/core/platform/types.h" -#include "tensorflow/core/util/command_line_flags.h" - -namespace xla { -namespace legacy_flags { - -// The default layout for all newly created shapes. Specified by the flag -// --xla_default_layout. -struct DefaultLayout { - enum class DimensionOrder { - kRandom, - kMinorToMajor, - kMajorToMinor, - }; - - DimensionOrder dimension_order; - size_t seed; -}; - -// Append to *flag_list the flag definitions associated with XLA's layout_util -// module. -void AppendLayoutUtilFlags(std::vector* flag_list); - -// The values of flags associated with XLA's layout_util module. -typedef struct { - // Default layout for Shapes in XLA. Valid values are: 'minor2major', - // 'major2minor', 'random', 'random:'. For debugging purposes. If no - // seed (or 0) is given, a seed from random_device is used. - DefaultLayout xla_default_layout; -} LayoutUtilFlags; - -// Return a pointer to the LayoutFlags struct; -// repeated calls return the same pointer. -// This should be called only after Flags::Parse() has returned. -LayoutUtilFlags* GetLayoutUtilFlags(); - -} // namespace legacy_flags -} // namespace xla - -#endif // TENSORFLOW_COMPILER_XLA_LEGACY_FLAGS_LAYOUT_UTIL_FLAGS_H_ diff --git a/tensorflow/compiler/xla/legacy_flags/llvm_util_flags.cc b/tensorflow/compiler/xla/legacy_flags/llvm_util_flags.cc deleted file mode 100644 index 3c53729a67049fdac6b358149e06f39858ebd98f..0000000000000000000000000000000000000000 --- a/tensorflow/compiler/xla/legacy_flags/llvm_util_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 XLA's llvm_util module. - -#include // NOLINT(build/c++11): only using std::call_once, not mutex. -#include - -#include "tensorflow/compiler/xla/legacy_flags/llvm_util_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 xla { -namespace legacy_flags { - -// Pointers to the parsed value of the flags and flag descriptors, initialized -// via flags_init. -static LlvmUtilFlags* 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 LlvmUtilFlags; - flags->xla_emit_tbaa = true; - flag_list = new std::vector({ - tensorflow::Flag("xla_emit_tbaa", &flags->xla_emit_tbaa, - "Perform type-based alias analysis optimizations for " - "LLVM-based backends."), - }); - ParseFlagsFromEnv(*flag_list); -} - -// Append to *append_to flag definitions associated with XLA's llvm_util -// module. -void AppendLlvmUtilFlags(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 LlvmUtilFlags struct; -// repeated calls return the same pointer. -// This should be called only after Flags::Parse() has returned. -LlvmUtilFlags* GetLlvmUtilFlags() { - std::call_once(flags_init, &AllocateFlags); - return flags; -} - -} // namespace legacy_flags -} // namespace xla diff --git a/tensorflow/compiler/xla/legacy_flags/llvm_util_flags.h b/tensorflow/compiler/xla/legacy_flags/llvm_util_flags.h deleted file mode 100644 index 98da26b4b806dd83c7baf6bdcf60cbf5297457a6..0000000000000000000000000000000000000000 --- a/tensorflow/compiler/xla/legacy_flags/llvm_util_flags.h +++ /dev/null @@ -1,46 +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_LLVM_UTIL_FLAGS_H_ -#define TENSORFLOW_COMPILER_XLA_LEGACY_FLAGS_LLVM_UTIL_FLAGS_H_ - -// Legacy flags for XLA's llvm_util module. - -#include - -#include "tensorflow/core/platform/types.h" -#include "tensorflow/core/util/command_line_flags.h" - -namespace xla { -namespace legacy_flags { - -// Append to *flag_list flag definitions associated with XLA's llvm_util module. -void AppendLlvmUtilFlags(std::vector* flag_list); - -// The values of flags associated with XLA's llvm_util module. -typedef struct { - bool xla_emit_tbaa; // Perform type-based alias analysis optimizations for - // LLVM-based backends. -} LlvmUtilFlags; - -// Return a pointer to the LlvmUtilFlags struct; -// repeated calls return the same pointer. -// This should be called only after Flags::Parse() has returned. -LlvmUtilFlags* GetLlvmUtilFlags(); - -} // namespace legacy_flags -} // namespace xla - -#endif // TENSORFLOW_COMPILER_XLA_LEGACY_FLAGS_LLVM_UTIL_FLAGS_H_ diff --git a/tensorflow/compiler/xla/legacy_flags/parse_flags_from_env_test.cc b/tensorflow/compiler/xla/legacy_flags/parse_flags_from_env_test.cc index 07bbcd802fef1d89f428717a4bd7d669d9f119c2..a3b4286f4c12bf39a44c63dd6e7d303a46a418c3 100644 --- a/tensorflow/compiler/xla/legacy_flags/parse_flags_from_env_test.cc +++ b/tensorflow/compiler/xla/legacy_flags/parse_flags_from_env_test.cc @@ -86,14 +86,14 @@ static const char kTestFlagString[] = "--single_quoted='single quoted \\\\ \n \"' " "--double_quoted=\"double quoted \\\\ \n '\\\"\" "; -// Test that the environent variable is parserd correctly. +// Test that the environent variable is parsed correctly. TEST(ParseFlagsFromEnv, Basic) { // Prepare environment. setenv("TF_XLA_FLAGS", kTestFlagString, true /*overwrite*/); TestParseFlagsFromEnv("(flags in environment variable)"); } -// Test that a file named by the environent variable is parserd correctly. +// Test that a file named by the environent variable is parsed correctly. TEST(ParseFlagsFromEnv, File) { // environment variables where tmp dir may be specified. static const char* kTempVars[] = {"TEST_TMPDIR", "TMP"}; diff --git a/tensorflow/compiler/xla/legacy_flags/service_flags.cc b/tensorflow/compiler/xla/legacy_flags/service_flags.cc deleted file mode 100644 index 41cb8d8bdfc51de1d8fe77906317b4b4a0804802..0000000000000000000000000000000000000000 --- a/tensorflow/compiler/xla/legacy_flags/service_flags.cc +++ /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. -==============================================================================*/ - -// Legacy flags for XLA's service module. - -#include // NOLINT(build/c++11): only using std::call_once, not mutex. -#include - -#include "tensorflow/compiler/xla/legacy_flags/parse_flags_from_env.h" -#include "tensorflow/compiler/xla/legacy_flags/service_flags.h" -#include "tensorflow/core/platform/types.h" -#include "tensorflow/core/util/command_line_flags.h" - -namespace xla { -namespace legacy_flags { - -// Pointers to the parsed value of the flags and flag descriptors, initialized -// via flags_init. -static ServiceFlags* 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 ServiceFlags; - flags->xla_hlo_profile = false; - flags->xla_log_hlo_text = ""; - flags->xla_generate_hlo_graph = ""; - flags->xla_hlo_graph_addresses = false; - flags->xla_hlo_graph_layout = false; - flags->xla_hlo_graph_for_compute_constant = false; - flags->xla_dump_computations_to = ""; - flags->xla_dump_hlo_text_to = ""; - flags->xla_dump_executions_to = ""; - flag_list = new std::vector({ - tensorflow::Flag( - "xla_hlo_profile", &flags->xla_hlo_profile, - "Instrument the computation to collect per-HLO cycle counts"), - tensorflow::Flag( - "xla_log_hlo_text", &flags->xla_log_hlo_text, - "If non-empty, print the text format of " - "HLO modules whose name partially matches this regex. E.g. " - "xla_log_hlo_text=.* will dump the text for every module."), - tensorflow::Flag( - "xla_generate_hlo_graph", &flags->xla_generate_hlo_graph, - "If non-empty, dump graph of HLO modules whose name partially " - "matches this regex. E.g. --xla_generate_hlo_graph=.* will dump " - "the graph of every module."), - tensorflow::Flag("xla_hlo_graph_addresses", - &flags->xla_hlo_graph_addresses, - "Show addresses of HLO ops in graph"), - tensorflow::Flag("xla_hlo_graph_layout", &flags->xla_hlo_graph_layout, - "Show layout of HLO ops in graph"), - tensorflow::Flag( - "xla_hlo_graph_for_compute_constant", - &flags->xla_hlo_graph_for_compute_constant, - "If true, include hlo dumps of graphs from ComputeConstant." - "Such graphs still need to be matched via xla_generate_hlo_graph."), - tensorflow::Flag("xla_dump_computations_to", - &flags->xla_dump_computations_to, - "Dumps computations that XLA executes into the provided " - "directory path"), - tensorflow::Flag("xla_dump_hlo_text_to", &flags->xla_dump_hlo_text_to, - "Dumps HLO modules that XLA executes into the provided " - "directory path"), - tensorflow::Flag("xla_dump_executions_to", &flags->xla_dump_executions_to, - "Dumps parameters and results of computations that XLA " - "executes into the provided directory path"), - }); - ParseFlagsFromEnv(*flag_list); -} - -// Append to *append_to flag definitions associated with XLA's service module. -void AppendServiceFlags(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 ServiceFlags struct; -// repeated calls return the same pointer. -// This should be called only after Flags::Parse() has returned. -ServiceFlags* GetServiceFlags() { - std::call_once(flags_init, &AllocateFlags); - return flags; -} - -} // namespace legacy_flags -} // namespace xla diff --git a/tensorflow/compiler/xla/legacy_flags/service_flags.h b/tensorflow/compiler/xla/legacy_flags/service_flags.h deleted file mode 100644 index d982506944daed41eb6e7c4a238d540b38cf8be3..0000000000000000000000000000000000000000 --- a/tensorflow/compiler/xla/legacy_flags/service_flags.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_COMPILER_XLA_LEGACY_FLAGS_SERVICE_FLAGS_H_ -#define TENSORFLOW_COMPILER_XLA_LEGACY_FLAGS_SERVICE_FLAGS_H_ - -// Legacy flags for XLA's service module. - -#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 { - -// Append to *flag_list flag definitions associated with XLA's service module. -void AppendServiceFlags(std::vector* flag_list); - -// The values of flags associated with XLA's service module. -typedef struct { - bool xla_hlo_profile; // Instrument the computation to collect per-HLO cycle - // counts - string xla_log_hlo_text; // If non-empty, print the text format of the HLO - // modules whose name partially - // matches this regex. E.g. xla_log_hlo_text=.* - // will dump the text for every module. - string xla_generate_hlo_graph; // If non-empty, dump graph of HLO modules - // whose name partially matches this regex. - // E.g. --xla_generate_hlo_graph=.* will dump - // the graph of every module. - bool xla_hlo_graph_addresses; // Show addresses of HLO ops in graph - bool xla_hlo_graph_layout; // Show layout of HLO ops in graph - bool xla_hlo_graph_for_compute_constant; // If true, include hlo dumps of - // graphs from ComputeConstant. - // Such graphs still need to be - // matched via - // xla_generate_hlo_graph. - string xla_dump_hlo_text_to; // Dumps HLO text for each HLO module that is - // executed into the provided directory path - string xla_dump_computations_to; // Dumps computations that XLA executes - // into the provided directory path - // Dumps parameters and results of computations that XLA executes into - // the provided directory path - string xla_dump_executions_to; -} ServiceFlags; - -// Return a pointer to the ServiceFlags struct; -// repeated calls return the same pointer. -// This should be called only after Flags::Parse() has returned. -ServiceFlags* GetServiceFlags(); - -} // namespace legacy_flags -} // namespace xla - -#endif // TENSORFLOW_COMPILER_XLA_LEGACY_FLAGS_SERVICE_FLAGS_H_ diff --git a/tensorflow/compiler/xla/legacy_flags/stream_assignment_flags.cc b/tensorflow/compiler/xla/legacy_flags/stream_assignment_flags.cc deleted file mode 100644 index 6506175777ccd262b6467f8fbe6de8bb24eff945..0000000000000000000000000000000000000000 --- a/tensorflow/compiler/xla/legacy_flags/stream_assignment_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 XLA's stream_assignment module. - -#include // NOLINT(build/c++11): only using std::call_once, not mutex. -#include - -#include "tensorflow/compiler/xla/legacy_flags/parse_flags_from_env.h" -#include "tensorflow/compiler/xla/legacy_flags/stream_assignment_flags.h" -#include "tensorflow/core/platform/types.h" -#include "tensorflow/core/util/command_line_flags.h" - -namespace xla { -namespace legacy_flags { - -// Pointers to the parsed value of the flags and flag descriptors, initialized -// via flags_init. -static StreamAssignmentFlags* 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 StreamAssignmentFlags; - flags->xla_gpu_disable_multi_streaming = false; - flag_list = new std::vector({ - tensorflow::Flag("xla_gpu_disable_multi_streaming", - &flags->xla_gpu_disable_multi_streaming, - "Disable multi-streaming in XLA's GPU backend"), - }); - ParseFlagsFromEnv(*flag_list); -} - -// Append to *append_to flag definitions associated with XLA's stream_assignment -// module. -void AppendStreamAssignmentFlags(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 StreamAssignmentFlags struct; -// repeated calls return the same pointer. -// This should be called only after Flags::Parse() has returned. -StreamAssignmentFlags* GetStreamAssignmentFlags() { - std::call_once(flags_init, &AllocateFlags); - return flags; -} - -} // namespace legacy_flags -} // namespace xla diff --git a/tensorflow/compiler/xla/legacy_flags/stream_assignment_flags.h b/tensorflow/compiler/xla/legacy_flags/stream_assignment_flags.h deleted file mode 100644 index a98f9b34584b43161aa8e3248c28d520403f3f3a..0000000000000000000000000000000000000000 --- a/tensorflow/compiler/xla/legacy_flags/stream_assignment_flags.h +++ /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. -==============================================================================*/ - -#ifndef TENSORFLOW_COMPILER_XLA_LEGACY_FLAGS_STREAM_ASSIGNMENT_FLAGS_H_ -#define TENSORFLOW_COMPILER_XLA_LEGACY_FLAGS_STREAM_ASSIGNMENT_FLAGS_H_ - -// Legacy flags for XLA's stream_assignment module. - -#include - -#include "tensorflow/core/platform/types.h" -#include "tensorflow/core/util/command_line_flags.h" - -namespace xla { -namespace legacy_flags { - -// Append to *flag_list flag definitions associated with XLA's stream_assignment -// module. -void AppendStreamAssignmentFlags(std::vector* flag_list); - -// The values of flags associated with XLA's stream_assignment module. -typedef struct { - bool xla_gpu_disable_multi_streaming; // Disable multi-streaming in XLA's GPU - // backend -} StreamAssignmentFlags; - -// Return a pointer to the StreamAssignmentFlags struct; -// repeated calls return the same pointer. -// This should be called only after Flags::Parse() has returned. -StreamAssignmentFlags* GetStreamAssignmentFlags(); - -} // namespace legacy_flags -} // namespace xla - -#endif // TENSORFLOW_COMPILER_XLA_LEGACY_FLAGS_STREAM_ASSIGNMENT_FLAGS_H_ diff --git a/tensorflow/compiler/xla/legacy_flags/util_flags.cc b/tensorflow/compiler/xla/legacy_flags/util_flags.cc deleted file mode 100644 index e6df19ddd2afbbf14149d77a1e0652df209f58fe..0000000000000000000000000000000000000000 --- a/tensorflow/compiler/xla/legacy_flags/util_flags.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. -==============================================================================*/ - -// Legacy flags for XLA's util module. - -#include // NOLINT(build/c++11): only using std::call_once, not mutex. -#include - -#include "tensorflow/compiler/xla/legacy_flags/parse_flags_from_env.h" -#include "tensorflow/compiler/xla/legacy_flags/util_flags.h" -#include "tensorflow/core/platform/types.h" -#include "tensorflow/core/util/command_line_flags.h" - -namespace xla { -namespace legacy_flags { - -// Pointers to the parsed value of the flags and flag descriptors, initialized -// via flags_init. -static UtilFlags* 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 UtilFlags; - flags->xla_status_add_backtrace = false; - flag_list = new std::vector({ - tensorflow::Flag("xla_status_add_backtrace", - &flags->xla_status_add_backtrace, - "add backtraces to XLA-produced status values"), - }); - ParseFlagsFromEnv(*flag_list); -} - -// Append to *append_to flag definitions associated with XLA's util module. -void AppendUtilFlags(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 UtilFlags struct; -// repeated calls return the same pointer. -// This should be called only after Flags::Parse() has returned. -UtilFlags* GetUtilFlags() { - std::call_once(flags_init, &AllocateFlags); - return flags; -} - -} // namespace legacy_flags -} // namespace xla diff --git a/tensorflow/compiler/xla/legacy_flags/util_flags.h b/tensorflow/compiler/xla/legacy_flags/util_flags.h deleted file mode 100644 index 03bffcd726f0544a185f5e8403ad2c45318bd0ad..0000000000000000000000000000000000000000 --- a/tensorflow/compiler/xla/legacy_flags/util_flags.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_COMPILER_XLA_LEGACY_FLAGS_UTIL_FLAGS_H_ -#define TENSORFLOW_COMPILER_XLA_LEGACY_FLAGS_UTIL_FLAGS_H_ - -// Legacy flags for the XLA's util module. - -#include - -#include "tensorflow/core/platform/types.h" -#include "tensorflow/core/util/command_line_flags.h" - -namespace xla { -namespace legacy_flags { - -// Append to *flag_list flag definitions associated with XLA's util module. -void AppendUtilFlags(std::vector* flag_list); - -// The values of flags associated with XLA's util module. -typedef struct { - bool xla_status_add_backtrace; // add backtraces to XLA-produced statuses -} UtilFlags; - -// Return a pointer to the UtilFlags struct; -// repeated calls return the same pointer. -// This should be called only after Flags::Parse() has returned. -UtilFlags* GetUtilFlags(); - -} // namespace legacy_flags -} // namespace xla - -#endif // TENSORFLOW_COMPILER_XLA_LEGACY_FLAGS_UTIL_FLAGS_H_ diff --git a/tensorflow/compiler/xla/literal_util.cc b/tensorflow/compiler/xla/literal_util.cc index 7091c324d14552d8b7603c3872d0ffc59771d8f7..71995b2307e3a34ac5d1f3307ccea42b4cf230a5 100644 --- a/tensorflow/compiler/xla/literal_util.cc +++ b/tensorflow/compiler/xla/literal_util.cc @@ -16,12 +16,15 @@ limitations under the License. #include "tensorflow/compiler/xla/literal_util.h" #include +#include +#include #include #include #include #include "tensorflow/compiler/xla/index_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/errors.h" @@ -33,31 +36,161 @@ limitations under the License. namespace xla { -/* static */ Literal LiteralUtil::Zero(PrimitiveType primitive_type) { +Literal::StrideConfig::StrideConfig( + const Shape& source_shape, const Shape& dest_shape, + tensorflow::gtl::ArraySlice dimensions) + : dimensions(dimensions), + base(dimensions.size(), 0), + step(dimensions.size(), 1) { + if (!dimensions.empty()) { + // Selects the shape with the largest minor dimension as the one upon + // which to run the tight stride loop. + if (dimensions[source_shape.layout().minor_to_major()[0]] >= + dimensions[dest_shape.layout().minor_to_major()[0]]) { + minor_dimension = source_shape.layout().minor_to_major()[0]; + dest_stride = IndexUtil::GetDimensionStride(dest_shape, minor_dimension); + } else { + minor_dimension = dest_shape.layout().minor_to_major()[0]; + source_stride = + IndexUtil::GetDimensionStride(source_shape, minor_dimension); + } + minor_loop_size = dimensions[minor_dimension]; + step[minor_dimension] = minor_loop_size; + } +} + +std::unique_ptr Literal::CreateFromShape(const Shape& shape) { + auto literal = MakeUnique(); + *literal->mutable_shape() = shape; + if (ShapeUtil::IsTuple(shape)) { + int64 num_elements = ShapeUtil::TupleElementCount(shape); + literal->tuple_literals_.resize(num_elements); + for (int i = 0; i < num_elements; ++i) { + std::unique_ptr elem = + CreateFromShape(ShapeUtil::GetTupleElementShape(shape, i)); + literal->tuple_literals_[i] = std::move(*elem); + } + } else { + literal->Reserve(ShapeUtil::ElementsIn(literal->shape())); + } + return literal; +} + +/* static */ std::unique_ptr Literal::CreateFromDimensions( + PrimitiveType primitive_type, + tensorflow::gtl::ArraySlice dimensions) { + return CreateFromShape(ShapeUtil::MakeShape(primitive_type, dimensions)); +} + +template +Status Literal::CopyRange(const Literal& src_literal, + tensorflow::gtl::ArraySlice src_base, + tensorflow::gtl::ArraySlice dest_base, + tensorflow::gtl::ArraySlice copy_size) { + const Shape& src_shape = src_literal.shape(); + const Shape& dest_shape = shape(); + tensorflow::gtl::ArraySlice src_data = src_literal.GetArraySlice(); + tensorflow::gtl::MutableArraySlice dest_data = GetMutableArraySlice(); + + TF_RET_CHECK(ShapeUtil::Rank(src_shape) == src_base.size()); + TF_RET_CHECK(ShapeUtil::Rank(dest_shape) == dest_base.size()); + if (ShapeUtil::Rank(src_shape) == 0 || ShapeUtil::Rank(dest_shape) == 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()); + StridedCopy(dest_data, LinearIndex(dest_base), 0, src_data, + src_literal.LinearIndex(src_base), 0, 1); + } else if (!ShapeUtil::HasZeroElements(dest_shape)) { + TF_RET_CHECK(!ShapeUtil::HasZeroElements(src_shape)); + TF_RET_CHECK(src_base.size() == dest_base.size()); + TF_RET_CHECK(src_base.size() == copy_size.size()); + + // Scan the source from minor, stepping in copy size blocks, then within + // the index enumaration functor, do a strided copy advancing source index + // by one (walking through the minor dimension), and destination index by + // proper stride size at the matching dimension. + DimensionVector src_indexes(src_base.size(), 0); + DimensionVector dest_indexes(dest_base.size(), 0); + StrideConfig stride_config(src_shape, dest_shape, copy_size); + + auto copy_proc = [&](const std::vector& indexes) { + // Map from multi-dimensional index, to source index. + std::transform(indexes.begin(), indexes.end(), src_base.begin(), + src_indexes.begin(), std::plus()); + // Map from multi-dimensional index, to destination index. + std::transform(indexes.begin(), indexes.end(), dest_base.begin(), + dest_indexes.begin(), std::plus()); + + int64 src_index = src_literal.LinearIndex(src_indexes); + int64 dest_index = LinearIndex(dest_indexes); + + StridedCopy(dest_data, dest_index, stride_config.dest_stride, src_data, + src_index, stride_config.source_stride, + stride_config.minor_loop_size); + return true; + }; + + ShapeUtil::ForEachIndex(src_shape, stride_config.base, + stride_config.dimensions, stride_config.step, + copy_proc); + } + return Status::OK(); +} + +Status Literal::Copy(const Literal& src_literal, + tensorflow::gtl::ArraySlice src_base, + tensorflow::gtl::ArraySlice dest_base, + tensorflow::gtl::ArraySlice copy_size) { + TF_RET_CHECK(ShapeUtil::SameElementType(src_literal.shape(), shape())); + switch (src_literal.shape().element_type()) { + case U32: + return CopyRange(src_literal, src_base, dest_base, copy_size); + case U64: + return CopyRange(src_literal, src_base, dest_base, copy_size); + case S32: + return CopyRange(src_literal, src_base, dest_base, copy_size); + case S64: + return CopyRange(src_literal, src_base, dest_base, copy_size); + case F16: + return CopyRange(src_literal, src_base, dest_base, copy_size); + case F32: + return CopyRange(src_literal, src_base, dest_base, copy_size); + case F64: + return CopyRange(src_literal, src_base, dest_base, copy_size); + case PRED: + return CopyRange(src_literal, src_base, dest_base, copy_size); + default: + break; + } + return Unimplemented("Unhandled primitive type %d", + src_literal.shape().element_type()); +} + +/* static */ Literal Literal::Zero(PrimitiveType primitive_type) { switch (primitive_type) { case U8: - return *LiteralUtil::CreateR0(0); + return *Literal::CreateR0(0); case U32: - return *LiteralUtil::CreateR0(0); + return *Literal::CreateR0(0); case U64: - return *LiteralUtil::CreateR0(0); + return *Literal::CreateR0(0); case S8: - return *LiteralUtil::CreateR0(0); + return *Literal::CreateR0(0); case S32: - return *LiteralUtil::CreateR0(0); + return *Literal::CreateR0(0); case S64: - return *LiteralUtil::CreateR0(0); + return *Literal::CreateR0(0); + case F16: + return *Literal::CreateR0(static_cast(0.0f)); case F32: - return *LiteralUtil::CreateR0(0); + return *Literal::CreateR0(0); case F64: - return *LiteralUtil::CreateR0(0); + return *Literal::CreateR0(0); case PRED: - return *LiteralUtil::CreateR0(false); + return *Literal::CreateR0(false); case S16: case U16: LOG(FATAL) << "u16/s16 literals not yet implemented"; - case F16: - LOG(FATAL) << "f16 literals not yet implemented"; case TUPLE: LOG(FATAL) << "tuple element type cannot take on value of 0"; case OPAQUE: @@ -67,31 +200,31 @@ namespace xla { } } -/* static */ Literal LiteralUtil::One(PrimitiveType primitive_type) { +/* static */ Literal Literal::One(PrimitiveType primitive_type) { switch (primitive_type) { case U8: - return *LiteralUtil::CreateR0(1); + return *Literal::CreateR0(1); case U32: - return *LiteralUtil::CreateR0(1); + return *Literal::CreateR0(1); case U64: - return *LiteralUtil::CreateR0(1); + return *Literal::CreateR0(1); case S8: - return *LiteralUtil::CreateR0(1); + return *Literal::CreateR0(1); case S32: - return *LiteralUtil::CreateR0(1); + return *Literal::CreateR0(1); case S64: - return *LiteralUtil::CreateR0(1); + return *Literal::CreateR0(1); case F32: - return *LiteralUtil::CreateR0(1); + return *Literal::CreateR0(1); case F64: - return *LiteralUtil::CreateR0(1); + return *Literal::CreateR0(1); case PRED: - return *LiteralUtil::CreateR0(true); + return *Literal::CreateR0(true); case S16: case U16: LOG(FATAL) << "u16/s16 literals not yet implemented"; case F16: - LOG(FATAL) << "f16 literals not yet implemented"; + return *Literal::CreateR0(static_cast(1.0f)); case TUPLE: LOG(FATAL) << "tuple element type cannot take on value of 1"; case OPAQUE: @@ -101,33 +234,33 @@ namespace xla { } } -/* static */ Literal LiteralUtil::MinValue(PrimitiveType primitive_type) { +/* static */ Literal Literal::MinValue(PrimitiveType primitive_type) { switch (primitive_type) { case U8: - return *LiteralUtil::CreateR0(std::numeric_limits::min()); + return *Literal::CreateR0(std::numeric_limits::min()); case U32: - return *LiteralUtil::CreateR0(std::numeric_limits::min()); + return *Literal::CreateR0(std::numeric_limits::min()); case U64: - return *LiteralUtil::CreateR0(std::numeric_limits::min()); + return *Literal::CreateR0(std::numeric_limits::min()); case S8: - return *LiteralUtil::CreateR0(std::numeric_limits::min()); + return *Literal::CreateR0(std::numeric_limits::min()); case S32: - return *LiteralUtil::CreateR0(std::numeric_limits::min()); + return *Literal::CreateR0(std::numeric_limits::min()); case S64: - return *LiteralUtil::CreateR0(std::numeric_limits::min()); + return *Literal::CreateR0(std::numeric_limits::min()); case F32: - return *LiteralUtil::CreateR0( - -std::numeric_limits::infinity()); + return *Literal::CreateR0(-std::numeric_limits::infinity()); case F64: - return *LiteralUtil::CreateR0( + return *Literal::CreateR0( -std::numeric_limits::infinity()); case PRED: - return *LiteralUtil::CreateR0(false); + return *Literal::CreateR0(false); case S16: case U16: LOG(FATAL) << "u16/s16 literals not yet implemented"; case F16: - LOG(FATAL) << "f16 literals not yet implemented"; + return *Literal::CreateR0( + static_cast(-std::numeric_limits::infinity())); case TUPLE: LOG(FATAL) << "tuple element type has no minimum value"; case OPAQUE: @@ -137,33 +270,33 @@ namespace xla { } } -/* static */ Literal LiteralUtil::MaxValue(PrimitiveType primitive_type) { +/* static */ Literal Literal::MaxValue(PrimitiveType primitive_type) { switch (primitive_type) { case U8: - return *LiteralUtil::CreateR0(std::numeric_limits::max()); + return *Literal::CreateR0(std::numeric_limits::max()); case U32: - return *LiteralUtil::CreateR0(std::numeric_limits::max()); + return *Literal::CreateR0(std::numeric_limits::max()); case U64: - return *LiteralUtil::CreateR0(std::numeric_limits::max()); + return *Literal::CreateR0(std::numeric_limits::max()); case S8: - return *LiteralUtil::CreateR0(std::numeric_limits::max()); + return *Literal::CreateR0(std::numeric_limits::max()); case S32: - return *LiteralUtil::CreateR0(std::numeric_limits::max()); + return *Literal::CreateR0(std::numeric_limits::max()); case S64: - return *LiteralUtil::CreateR0(std::numeric_limits::max()); + return *Literal::CreateR0(std::numeric_limits::max()); case F32: - return *LiteralUtil::CreateR0( - std::numeric_limits::infinity()); + return *Literal::CreateR0(std::numeric_limits::infinity()); case F64: - return *LiteralUtil::CreateR0( + return *Literal::CreateR0( std::numeric_limits::infinity()); case PRED: - return *LiteralUtil::CreateR0(true); + return *Literal::CreateR0(true); case S16: case U16: LOG(FATAL) << "u16/s16 literals not yet implemented"; case F16: - LOG(FATAL) << "f16 literals not yet implemented"; + return *Literal::CreateR0( + static_cast(std::numeric_limits::infinity())); case TUPLE: LOG(FATAL) << "tuple element type has no maximum value"; case OPAQUE: @@ -173,224 +306,162 @@ namespace xla { } } -/* static */ std::unique_ptr LiteralUtil::CreateR1( +/* static */ std::unique_ptr Literal::CreateR1( const tensorflow::core::Bitmap& values) { auto literal = MakeUnique(); - PopulateR1(values, literal.get()); + literal->PopulateR1(values); return literal; } -/* static */ std::unique_ptr LiteralUtil::CreateR1U8( +/* static */ std::unique_ptr Literal::CreateR1U8( tensorflow::StringPiece value) { auto literal = MakeUnique(); *literal->mutable_shape() = ShapeUtil::MakeShape(U8, {static_cast(value.size())}); - literal->set_u8s(value.ToString()); + literal->set_u8s(tensorflow::StringPiece(value.ToString())); return literal; } -/* static */ std::unique_ptr LiteralUtil::CreateR2F32Linspace( - float from, float to, int64 rows, int64 cols) { +/* static */ std::unique_ptr Literal::CreateR2F32Linspace(float from, + float to, + int64 rows, + int64 cols) { auto value = MakeLinspaceArray2D(from, to, rows, cols); return CreateR2FromArray2D(*value); } -/* static */ std::unique_ptr LiteralUtil::Relayout( - const Literal& original, const Layout& layout) { - // Note: if this were a performance bottleneck, we avoid cloning and just make - // an uninitialized array instead, since all values are clobbered below. - std::unique_ptr result = CloneToUnique(original); +std::unique_ptr Literal::Relayout(const Layout& layout) const { + CHECK(ShapeUtil::IsArray(shape())); + std::unique_ptr result = CloneToUnique(); *result->mutable_shape()->mutable_layout() = layout; - const PrimitiveType primitive_type = original.shape().element_type(); - switch (primitive_type) { - case F32: - LiteralUtil::EachCell( - original, - [&](tensorflow::gtl::ArraySlice indices, float value) { - LiteralUtil::Set(result.get(), indices, value); - }); - return result; - case S32: - LiteralUtil::EachCell( - original, - [&](tensorflow::gtl::ArraySlice indices, int32 value) { - LiteralUtil::Set(result.get(), indices, value); - }); - return result; - case U32: - LiteralUtil::EachCell( - original, - [&](tensorflow::gtl::ArraySlice indices, uint32 value) { - LiteralUtil::Set(result.get(), indices, value); - }); - return result; - default: - LOG(FATAL) << "not yet implemented: " - << PrimitiveType_Name(primitive_type); - } + + DimensionVector base(ShapeUtil::Rank(shape()), 0); + DimensionVector copy_size(shape().dimensions().begin(), + shape().dimensions().end()); + + TF_CHECK_OK(result->Copy(*this, base, base, copy_size)); + return result; } -/* static */ StatusOr> LiteralUtil::Reshape( - const xla::Literal& input, tensorflow::gtl::ArraySlice dimensions) { - if (ShapeUtil::IsTuple(input.shape())) { +StatusOr> Literal::Reshape( + tensorflow::gtl::ArraySlice dimensions) const { + if (ShapeUtil::IsTuple(shape())) { return InvalidArgument("Reshape does not support tuples."); } - - if (!LayoutUtil::IsMonotonicWithDim0Major(input.shape().layout())) { - return Unimplemented( - "Input shape must have a monotonic layout where dimension 0 is major, " - "was: %s", - LayoutUtil::HumanString(input.shape().layout()).c_str()); + std::unique_ptr output; + if (!LayoutUtil::IsMonotonicWithDim0Major(shape().layout())) { + std::vector minor_to_major(ShapeUtil::Rank(shape())); + std::iota(minor_to_major.rbegin(), minor_to_major.rend(), + static_cast(0)); + output = Relayout(LayoutUtil::MakeLayout(minor_to_major)); + } else { + output = CloneToUnique(); } - std::vector layout(dimensions.size()); - std::iota(layout.rbegin(), layout.rend(), 0); - // Because the layout is monotonic, we can simply reuse the same sequence of // values without changing their order. - std::unique_ptr output = CloneToUnique(input); - output->clear_shape(); - output->mutable_shape()->set_element_type(input.shape().element_type()); - for (int64 dimension : dimensions) { - output->mutable_shape()->add_dimensions(dimension); - } - *output->mutable_shape()->mutable_layout() = LayoutUtil::MakeLayout(layout); + *output->mutable_shape() = + ShapeUtil::MakeShape(shape().element_type(), dimensions); - int64 elements_before = ShapeUtil::ElementsIn(input.shape()); + int64 elements_before = ShapeUtil::ElementsIn(shape()); int64 elements_after = ShapeUtil::ElementsIn(output->shape()); if (elements_before != elements_after) { return InvalidArgument( - "Shapes before and after LiteralUtil::Reshape have different numbers " + "Shapes before and after Literal::Reshape have different numbers " "of elements: %s vs %s.", - ShapeUtil::HumanString(input.shape()).c_str(), + ShapeUtil::HumanString(shape()).c_str(), ShapeUtil::HumanString(output->shape()).c_str()); } return std::move(output); } -namespace { - -template -void TransposeLiteralInternal(const Literal& original, - tensorflow::gtl::ArraySlice permutation, - Literal* result) { - std::vector new_indices(ShapeUtil::Rank(original.shape())); - LiteralUtil::EachCell( - original, [&](tensorflow::gtl::ArraySlice indices, T value) { - for (int64 i = 0; i < indices.size(); ++i) { - new_indices[i] = indices[permutation[i]]; - } - LiteralUtil::Set(result, new_indices, value); - }); -} -} // namespace - -/* static */ std::unique_ptr LiteralUtil::Transpose( - const Literal& original, tensorflow::gtl::ArraySlice permutation) { - CHECK(!ShapeUtil::IsTuple(original.shape())) - << "tuple is not supported for transpose"; - std::vector dimension_numbers(ShapeUtil::Rank(original.shape())); - std::iota(dimension_numbers.begin(), dimension_numbers.end(), 0); - CHECK(std::is_permutation(permutation.begin(), permutation.end(), - dimension_numbers.begin())) - << "given permutation is not a permutation of dimension numbers"; - std::vector new_dimension_sizes; - for (const int64 dim : permutation) { - new_dimension_sizes.push_back(original.shape().dimensions(dim)); - } - const auto result_shape = ShapeUtil::MakeShape( - original.shape().element_type(), new_dimension_sizes); - std::unique_ptr result = CloneToUnique(original); - *result->mutable_shape() = result_shape; - const PrimitiveType primitive_type = original.shape().element_type(); - switch (primitive_type) { - case F32: - TransposeLiteralInternal(original, permutation, result.get()); - return result; - case F64: - TransposeLiteralInternal(original, permutation, result.get()); - return result; - case PRED: - TransposeLiteralInternal(original, permutation, result.get()); - return result; - case S8: - TransposeLiteralInternal(original, permutation, result.get()); - return result; - case U8: - TransposeLiteralInternal(original, permutation, result.get()); - return result; - case S32: - TransposeLiteralInternal(original, permutation, result.get()); - return result; - case U32: - TransposeLiteralInternal(original, permutation, result.get()); - return result; - case S64: - TransposeLiteralInternal(original, permutation, result.get()); - return result; - case U64: - TransposeLiteralInternal(original, permutation, result.get()); - return result; - default: - LOG(FATAL) << "not yet implemented: " - << PrimitiveType_Name(primitive_type); +std::unique_ptr Literal::Transpose( + tensorflow::gtl::ArraySlice permutation) const { + CHECK(!ShapeUtil::IsTuple(shape())) << "Tuple is not supported for transpose"; + CHECK(IsPermutation(permutation, ShapeUtil::Rank(shape()))) + << "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. + // This is considerably faster than iterating over every array element using + // the EachCell<>() and Set<>() APIs. + std::vector inverse_permutation = InversePermutation(permutation); + Shape permuted_shape = + ShapeUtil::PermuteDimensions(inverse_permutation, shape()); + // Replace the layout with one affine to this shape, such that a + // transpose operation can be performed by leaving the flat values + // representation intact. + // For example, consider the shape F32[11,8]{1,0} under a {1,0} permutation. + // The shape with affine layout resulting from that operation will be + // F32[8,11]{0,1}, since it leaves the original most minor (the 8 sized), the + // most minor. + // Essentially, given MinMaj(Di) the position of the Di dimension within the + // minor to major vector, and given T(Di) the index that the original Di + // dimension has within the transposed array, a layout is affine if + // MinMaj(Di) == TMinMaj(T(Di)), with TMinMaj() being the minor to major + // vector of the affine layout. + Layout* layout = permuted_shape.mutable_layout(); + layout->clear_minor_to_major(); + for (auto index : shape().layout().minor_to_major()) { + layout->add_minor_to_major(inverse_permutation[index]); } + std::unique_ptr new_literal = CreateFromShape(permuted_shape); + DCHECK_GE(ShapeUtil::ByteSizeOf(new_literal->shape()), + ShapeUtil::ByteSizeOf(shape())); + std::memcpy(new_literal->MutableInternalData(), InternalData(), + ShapeUtil::ByteSizeOf(shape())); + return new_literal; } -/* static */ std::unique_ptr LiteralUtil::Slice( - const Literal& literal, tensorflow::gtl::ArraySlice start_indices, - tensorflow::gtl::ArraySlice limit_indices) { - CHECK(!ShapeUtil::IsTuple(literal.shape())) - << "tuple is not supported for reshape"; +std::unique_ptr Literal::Slice( + tensorflow::gtl::ArraySlice start_indices, + tensorflow::gtl::ArraySlice limit_indices) const { + CHECK(!ShapeUtil::IsTuple(shape())) << "tuple is not supported for reshape"; - std::vector result_dimensions; - for (int64 dnum = 0; dnum < ShapeUtil::Rank(literal.shape()); ++dnum) { + DimensionVector result_dimensions; + for (int64 dnum = 0; dnum < ShapeUtil::Rank(shape()); ++dnum) { CHECK_GE(start_indices[dnum], 0); - CHECK_LE(limit_indices[dnum], literal.shape().dimensions(dnum)); + CHECK_LE(limit_indices[dnum], shape().dimensions(dnum)); int64 dimension = limit_indices[dnum] - start_indices[dnum]; CHECK_GT(dimension, 0); result_dimensions.push_back(dimension); } const auto result_shape = ShapeUtil::MakeShapeWithLayout( - literal.shape().element_type(), result_dimensions, - AsInt64Slice(literal.shape().layout().minor_to_major())); + shape().element_type(), result_dimensions, + AsInt64Slice(shape().layout().minor_to_major())); auto result_literal = MakeUnique(); *result_literal->mutable_shape() = result_shape; - Reserve(ShapeUtil::ElementsIn(result_shape), result_literal.get()); + result_literal->Reserve(ShapeUtil::ElementsIn(result_shape)); - std::vector new_indices(ShapeUtil::Rank(result_shape)); + DimensionVector new_indices(ShapeUtil::Rank(result_shape)); switch (result_shape.element_type()) { case F32: - LiteralUtil::EachCell( - *result_literal, + result_literal->EachCell( [&](tensorflow::gtl::ArraySlice indices, float /*value*/) { for (int64 i = 0; i < ShapeUtil::Rank(result_shape); ++i) { new_indices[i] = indices[i] + start_indices[i]; } - float value = LiteralUtil::Get(literal, new_indices); - LiteralUtil::Set(result_literal.get(), indices, value); + float value = Get(new_indices); + result_literal->Set(indices, value); }); return result_literal; case S32: - LiteralUtil::EachCell( - *result_literal, + result_literal->EachCell( [&](tensorflow::gtl::ArraySlice indices, int32 /*value*/) { for (int64 i = 0; i < ShapeUtil::Rank(result_shape); ++i) { new_indices[i] = indices[i] + start_indices[i]; } - int32 value = LiteralUtil::Get(literal, new_indices); - LiteralUtil::Set(result_literal.get(), indices, value); + int32 value = Get(new_indices); + result_literal->Set(indices, value); }); return result_literal; case U32: - LiteralUtil::EachCell( - *result_literal, + result_literal->EachCell( [&](tensorflow::gtl::ArraySlice indices, uint32 /*value*/) { for (int64 i = 0; i < ShapeUtil::Rank(result_shape); ++i) { new_indices[i] = indices[i] + start_indices[i]; } - uint32 value = LiteralUtil::Get(literal, new_indices); - LiteralUtil::Set(result_literal.get(), indices, value); + uint32 value = Get(new_indices); + result_literal->Set(indices, value); }); return result_literal; default: @@ -399,96 +470,117 @@ void TransposeLiteralInternal(const Literal& original, } } -/* static */ std::unique_ptr LiteralUtil::CloneToUnique( - const Literal& literal) { +std::unique_ptr Literal::CloneToUnique() const { auto unique = MakeUnique(); - *unique = literal; + *unique = *this; return unique; } -/* static */ string LiteralUtil::GetAsString( - const Literal& literal, tensorflow::gtl::ArraySlice multi_index) { - switch (literal.shape().element_type()) { +string Literal::GetAsString( + tensorflow::gtl::ArraySlice multi_index) const { + switch (shape().element_type()) { case PRED: - return Get(literal, multi_index) ? "true" : "false"; + return Get(multi_index) ? "true" : "false"; case U8: - return tensorflow::strings::StrCat(Get(literal, multi_index)); + return tensorflow::strings::StrCat(Get(multi_index)); case S32: - return tensorflow::strings::StrCat(Get(literal, multi_index)); + return tensorflow::strings::StrCat(Get(multi_index)); case S64: - return tensorflow::strings::StrCat(Get(literal, multi_index)); + return tensorflow::strings::StrCat(Get(multi_index)); case U32: - return tensorflow::strings::StrCat(Get(literal, multi_index)); + return tensorflow::strings::StrCat(Get(multi_index)); case U64: - return tensorflow::strings::StrCat(Get(literal, multi_index)); + return tensorflow::strings::StrCat(Get(multi_index)); case F32: - return tensorflow::strings::StrCat(Get(literal, multi_index)); + return tensorflow::strings::StrCat(Get(multi_index)); case F64: - return tensorflow::strings::StrCat(Get(literal, multi_index)); + return tensorflow::strings::StrCat(Get(multi_index)); + case F16: + return tensorflow::strings::StrCat(Get(multi_index)); default: return tensorflow::strings::StrCat( - "[", PrimitiveType_Name(literal.shape().element_type()), "]"); + "[", PrimitiveType_Name(shape().element_type()), "]"); + } +} + +StatusOr Literal::GetIntegralAsS64( + tensorflow::gtl::ArraySlice multi_index) const { + switch (shape().element_type()) { + case PRED: + return Get(multi_index); + case U8: + return Get(multi_index); + case S32: + return Get(multi_index); + case S64: + return Get(multi_index); + case U32: + return Get(multi_index); + case U64: + return Get(multi_index); + default: + return FailedPrecondition( + "Array element type is not integral: %s", + PrimitiveType_Name(shape().element_type()).c_str()); } } -/* static */ int64 LiteralUtil::LinearIndex( - const Literal& literal, tensorflow::gtl::ArraySlice multi_index) { - return IndexUtil::MultidimensionalIndexToLinearIndex(literal.shape(), - multi_index); +int64 Literal::LinearIndex( + tensorflow::gtl::ArraySlice multi_index) const { + return IndexUtil::MultidimensionalIndexToLinearIndex(shape(), multi_index); } -/* static */ string LiteralUtil::ToString(const Literal& literal) { - const Shape& shape = literal.shape(); +string Literal::ToString() const { std::vector pieces; auto element_to_string = - [&literal](tensorflow::gtl::ArraySlice indices) -> string { - PrimitiveType element_type = literal.shape().element_type(); + [this](tensorflow::gtl::ArraySlice indices) -> string { + PrimitiveType element_type = shape().element_type(); if (element_type == PRED) { // We display predicates in a densely packed form. - return Get(literal, indices) ? "1" : "0"; + return Get(indices) ? "1" : "0"; } return ((!indices.empty() && indices.back() > 0) ? ", " : "") + - GetAsString(literal, indices); + GetAsString(indices); }; // TODO(b/32894291): refactor this code to reduce code duplication. - if (ShapeUtil::IsTuple(shape)) { - pieces.push_back(ShapeUtil::HumanString(shape)); + if (ShapeUtil::IsTuple(shape())) { + pieces.push_back(ShapeUtil::HumanString(shape())); pieces.push_back(" (\n"); - for (const auto& element_literal : literal.tuple_literals()) { - pieces.push_back(ToString(element_literal)); + for (const auto& element_literal : tuple_literals()) { + pieces.push_back(element_literal.ToString()); pieces.push_back(",\n"); } pieces.push_back(")"); - } else if (ShapeUtil::Rank(shape) == 0) { - pieces.push_back(GetAsString(literal, {})); - } else if (ShapeUtil::Rank(shape) == 1) { + } else if (ShapeUtil::Rank(shape()) == 0) { + pieces.push_back(GetAsString({})); + } else if (ShapeUtil::Rank(shape()) == 1) { pieces.push_back("{"); - for (int64 i0 = 0; i0 < shape.dimensions(0); ++i0) { + for (int64 i0 = 0; i0 < shape().dimensions(0); ++i0) { pieces.push_back(element_to_string({i0})); } pieces.push_back("}"); - } else if (ShapeUtil::Rank(shape) == 2) { - pieces.push_back(ShapeUtil::HumanString(shape)); + } else if (ShapeUtil::Rank(shape()) == 2) { + pieces.push_back(ShapeUtil::HumanString(shape())); pieces.push_back(" {\n"); - for (int64 i0 = 0; i0 < shape.dimensions(0); ++i0) { + for (int64 i0 = 0; i0 < shape().dimensions(0); ++i0) { pieces.push_back(" { "); - for (int64 i1 = 0; i1 < shape.dimensions(1); ++i1) { + for (int64 i1 = 0; i1 < shape().dimensions(1); ++i1) { pieces.push_back(element_to_string({i0, i1})); } pieces.push_back(" "); pieces.push_back("},\n"); } pieces.push_back("}"); - } else if (ShapeUtil::Rank(shape) == 3) { - pieces.push_back(ShapeUtil::HumanString(shape)); + } else if (ShapeUtil::Rank(shape()) == 3) { + pieces.push_back(ShapeUtil::HumanString(shape())); pieces.push_back(" {\n"); - for (int64 i0 = 0; i0 < shape.dimensions(0); ++i0) { + for (int64 i0 = 0; i0 < shape().dimensions(0); ++i0) { pieces.push_back(i0 > 0 ? ",\n{" : "{"); - for (int64 i1 = 0; i1 < shape.dimensions(1); ++i1) { + for (int64 i1 = 0; i1 < shape().dimensions(1); ++i1) { pieces.push_back(i1 > 0 ? ",\n { " : " { "); - for (int64 i2 = 0; i2 < shape.dimensions(2); ++i2) { + for (int64 i2 = 0; i2 < shape().dimensions(2); ++i2) { pieces.push_back(element_to_string({i0, i1, i2})); } pieces.push_back(" }"); @@ -496,17 +588,17 @@ void TransposeLiteralInternal(const Literal& original, pieces.push_back(" }"); } pieces.push_back("\n}"); - } else if (ShapeUtil::Rank(shape) == 4) { - pieces.push_back(ShapeUtil::HumanString(shape)); + } else if (ShapeUtil::Rank(shape()) == 4) { + pieces.push_back(ShapeUtil::HumanString(shape())); pieces.push_back(" {\n"); - for (int64 i0 = 0; i0 < shape.dimensions(0); ++i0) { + for (int64 i0 = 0; i0 < shape().dimensions(0); ++i0) { pieces.push_back(tensorflow::strings::Printf(" { // i0=%lld\n", i0)); - for (int64 i1 = 0; i1 < shape.dimensions(1); ++i1) { + for (int64 i1 = 0; i1 < shape().dimensions(1); ++i1) { pieces.push_back( tensorflow::strings::Printf(" { // i1=%lld\n", i1)); - for (int64 i2 = 0; i2 < shape.dimensions(2); ++i2) { + for (int64 i2 = 0; i2 < shape().dimensions(2); ++i2) { pieces.push_back(" {"); - for (int64 i3 = 0; i3 < shape.dimensions(3); ++i3) { + for (int64 i3 = 0; i3 < shape().dimensions(3); ++i3) { pieces.push_back(element_to_string({i0, i1, i2, i3})); } pieces.push_back("},\n"); @@ -516,20 +608,20 @@ void TransposeLiteralInternal(const Literal& original, pieces.push_back(" },\n"); } pieces.push_back("}"); - } else if (ShapeUtil::Rank(shape) == 5) { - pieces.push_back(ShapeUtil::HumanString(shape)); + } else if (ShapeUtil::Rank(shape()) == 5) { + pieces.push_back(ShapeUtil::HumanString(shape())); pieces.push_back(" {\n"); - for (int64 i0 = 0; i0 < shape.dimensions(0); ++i0) { + for (int64 i0 = 0; i0 < shape().dimensions(0); ++i0) { pieces.push_back(tensorflow::strings::Printf(" { // i0=%lld\n", i0)); - for (int64 i1 = 0; i1 < shape.dimensions(1); ++i1) { + for (int64 i1 = 0; i1 < shape().dimensions(1); ++i1) { pieces.push_back( tensorflow::strings::Printf(" { // i1=%lld\n", i1)); - for (int64 i2 = 0; i2 < shape.dimensions(2); ++i2) { + for (int64 i2 = 0; i2 < shape().dimensions(2); ++i2) { pieces.push_back( tensorflow::strings::Printf(" { // i2=%lld\n", i2)); - for (int64 i3 = 0; i3 < shape.dimensions(3); ++i3) { + for (int64 i3 = 0; i3 < shape().dimensions(3); ++i3) { pieces.push_back(" {"); - for (int64 i4 = 0; i4 < shape.dimensions(4); ++i4) { + for (int64 i4 = 0; i4 < shape().dimensions(4); ++i4) { pieces.push_back(element_to_string({i0, i1, i2, i3, i4})); } pieces.push_back("},\n"); @@ -542,14 +634,14 @@ void TransposeLiteralInternal(const Literal& original, } pieces.push_back("}"); } else { - pieces.push_back(ShapeUtil::HumanString(shape)); + pieces.push_back(ShapeUtil::HumanString(shape())); pieces.push_back(" {...}"); } return tensorflow::str_util::Join(pieces, ""); } -/* static */ std::unique_ptr LiteralUtil::MakeTuple( +/* static */ std::unique_ptr Literal::MakeTuple( tensorflow::gtl::ArraySlice elements) { auto literal = MakeUnique(); std::vector shape; @@ -561,169 +653,229 @@ void TransposeLiteralInternal(const Literal& original, return literal; } -/* static */ const void* LiteralUtil::InternalData(const Literal& literal) { - switch (literal.shape().element_type()) { +/* static */ std::unique_ptr Literal::MakeTupleOwned( + std::vector> elements) { + auto literal = MakeUnique(); + std::vector shape; + for (auto& tuple_element : elements) { + shape.push_back(tuple_element->shape()); + literal->add_tuple_literals()->Swap(tuple_element.get()); + } + *literal->mutable_shape() = ShapeUtil::MakeTupleShape(shape); + return literal; +} + +const void* Literal::InternalData() const { + return const_cast( + const_cast(this)->MutableInternalData()); +} + +void* Literal::MutableInternalData() { + // NOTE: We access the vectors directly to avoid the const reference + // created by the accessor functions. + switch (shape().element_type()) { case PRED: - return reinterpret_cast(literal.preds().data()); case U8: - return reinterpret_cast(literal.u8s().data()); + return reinterpret_cast(u8s_.data()); case S32: - return reinterpret_cast(literal.s32s().data()); + return reinterpret_cast(s32s_.data()); case S64: - return reinterpret_cast(literal.s64s().data()); + return reinterpret_cast(s64s_.data()); case U32: - return reinterpret_cast(literal.u32s().data()); + return reinterpret_cast(u32s_.data()); case U64: - return reinterpret_cast(literal.u64s().data()); + return reinterpret_cast(u64s_.data()); case F32: - return reinterpret_cast(literal.f32s().data()); + return reinterpret_cast(f32s_.data()); case F64: - return reinterpret_cast(literal.f64s().data()); + return reinterpret_cast(f64s_.data()); + case F16: + return reinterpret_cast(f16s_.data()); default: LOG(FATAL) << "primitive type not supported in literals: " - << PrimitiveType_Name(literal.shape().element_type()); + << PrimitiveType_Name(shape().element_type()); } } -/* static */ void* LiteralUtil::MutableInternalData(Literal* literal) { - return const_cast(LiteralUtil::InternalData(*literal)); -} - -/* static */ void LiteralUtil::Reserve(int64 num_elements, Literal* literal) { - CHECK_EQ(ShapeUtil::ElementsIn(literal->shape()), num_elements); - switch (literal->shape().element_type()) { +void Literal::Reserve(int64 num_elements) { + CHECK_EQ(ShapeUtil::ElementsIn(shape()), num_elements); + switch (shape().element_type()) { case PRED: - GetMutableRepeatedField(literal)->Resize(num_elements, false); + Resize(num_elements, false); + break; + case S8: + Resize(num_elements, 0); break; case U8: - // u8s is an optional "bytes", rather than a repeated field. Therefore its - // access methods are somewhat different from the others. - literal->mutable_u8s()->resize(num_elements, 0); + Resize(num_elements, 0); break; case S32: - GetMutableRepeatedField(literal)->Resize(num_elements, - /*value=*/0); + Resize(num_elements, 0); break; case S64: - GetMutableRepeatedField(literal)->Resize( - num_elements, - /*value=*/0); + Resize(num_elements, 0); break; case U32: - GetMutableRepeatedField(literal)->Resize(num_elements, - /*value=*/0); + Resize(num_elements, 0); break; case U64: - GetMutableRepeatedField(literal)->Resize( - num_elements, - /*value=*/0); + Resize(num_elements, 0); break; case F32: - GetMutableRepeatedField(literal)->Resize(num_elements, - /*value=*/0.0f); + Resize(num_elements, 0); break; case F64: - GetMutableRepeatedField(literal)->Resize(num_elements, - /*value=*/0.0); + Resize(num_elements, 0); + break; + case F16: + Resize(num_elements, static_cast(0.0f)); break; default: LOG(FATAL) << "primitive type not supported in literals: " - << PrimitiveType_Name(literal->shape().element_type()); + << PrimitiveType_Name(shape().element_type()); } } -/* static */ tensorflow::Status LiteralUtil::ValidateLiteral( - const Literal& literal) { - TF_CHECK_OK(ShapeUtil::ValidateShape(literal.shape())); - int64 expected = ShapeUtil::ElementsIn(literal.shape()); +tensorflow::Status Literal::ValidateLiteral() const { + TF_CHECK_OK(ShapeUtil::ValidateShape(shape())); + int64 expected = ShapeUtil::ElementsIn(shape()); int64 actual = -1; - switch (literal.shape().element_type()) { + switch (shape().element_type()) { case PRED: - actual = literal.preds().size(); - break; case U8: - actual = literal.u8s().size(); + actual = u8s_size(); break; case S32: - actual = literal.s32s_size(); + actual = s32s_size(); break; case U32: - actual = literal.u32s_size(); + actual = u32s_size(); break; case S64: - actual = literal.s64s_size(); + actual = s64s_size(); break; case U64: - actual = literal.u64s_size(); + actual = u64s_size(); break; case F32: - actual = literal.f32s_size(); + actual = f32s_size(); break; case F64: - actual = literal.f64s_size(); + actual = f64s_size(); + break; + case F16: + actual = f16s().size() / sizeof(half); break; default: return tensorflow::errors::Unimplemented( "unhandled element type for literal validation: " + - PrimitiveType_Name(literal.shape().element_type())); + PrimitiveType_Name(shape().element_type())); } if (expected != actual) { return tensorflow::errors::InvalidArgument(tensorflow::strings::Printf( "literal has bad number of elements for its shape %s: want %lld " "got %lld", - ShapeUtil::HumanString(literal.shape()).c_str(), expected, actual)); + ShapeUtil::HumanString(shape()).c_str(), expected, actual)); } return tensorflow::Status::OK(); } -/* static */ void LiteralUtil::EachCellAsString( - const Literal& literal, - std::function indices, - const string& value)> - per_cell) { - if (ShapeUtil::Rank(literal.shape()) == 1) { - for (int64 i0 = 0; i0 < literal.shape().dimensions(0); ++i0) { - per_cell({i0}, GetAsString(literal, {i0})); - } +void Literal::EachCellAsString( + const std::function indices, + const string& value)>& per_cell) const { + if (ShapeUtil::HasZeroElements(shape())) { return; } + std::vector indices = IndexUtil::LinearIndexToMultidimensionalIndex( + shape(), /*linear_index=*/0); + do { + per_cell(indices, GetAsString(indices)); + } while (IndexUtil::BumpIndices(shape(), &indices)); +} - if (ShapeUtil::Rank(literal.shape()) == 2) { - for (int64 i0 = 0; i0 < literal.shape().dimensions(0); ++i0) { - for (int64 i1 = 0; i1 < literal.shape().dimensions(1); ++i1) { - per_cell({i0, i1}, GetAsString(literal, {i0, i1})); - } - } - return; +namespace { +template +std::unique_ptr ConvertBetweenNativeTypes(const Literal& src_literal) { + auto result_literal = MakeUnique(); + Shape* result_shape = result_literal->mutable_shape(); + *result_shape = src_literal.shape(); + result_shape->set_element_type( + primitive_util::NativeToPrimitiveType()); + result_literal->Reserve(ShapeUtil::ElementsIn(*result_shape)); + tensorflow::gtl::ArraySlice src_data = + src_literal.GetArraySlice(); + tensorflow::gtl::MutableArraySlice dest_data = + result_literal->GetMutableArraySlice(); + int64 num_elements = ShapeUtil::ElementsIn(src_literal.shape()); + + for (int64 i = 0; i < num_elements; ++i) { + dest_data[i] = static_cast(src_data[i]); } + return result_literal; +} - if (ShapeUtil::Rank(literal.shape()) == 3) { - for (int64 i0 = 0; i0 < literal.shape().dimensions(0); ++i0) { - for (int64 i1 = 0; i1 < literal.shape().dimensions(1); ++i1) { - for (int64 i2 = 0; i2 < literal.shape().dimensions(2); ++i2) { - per_cell({i0, i1, i2}, GetAsString(literal, {i0, i1, i2})); - } - } - } - return; - } +template +std::unique_ptr ConvertIfTypesMatch(const Literal& src_literal) { + CHECK_EQ(primitive_src_type, src_literal.shape().element_type()); + return ConvertBetweenNativeTypes< + typename primitive_util::PrimitiveTypeToNative::type, + typename primitive_util::PrimitiveTypeToNative< + primitive_dest_type>::type>(src_literal); +} - if (ShapeUtil::Rank(literal.shape()) == 4) { - for (int64 i0 = 0; i0 < literal.shape().dimensions(0); ++i0) { - for (int64 i1 = 0; i1 < literal.shape().dimensions(1); ++i1) { - for (int64 i2 = 0; i2 < literal.shape().dimensions(2); ++i2) { - for (int64 i3 = 0; i3 < literal.shape().dimensions(3); ++i3) { - per_cell({i0, i1, i2, i3}, GetAsString(literal, {i0, i1, i2, i3})); - } - } - } - } - return; +template +StatusOr> ConvertIfDestTypeMatches( + const Literal& src_literal, PrimitiveType primitive_dest_type) { + switch (primitive_dest_type) { +#define CONVERT_IF_TYPES_MATCH(type) \ + case (type): \ + return ConvertIfTypesMatch(src_literal); + CONVERT_IF_TYPES_MATCH(PRED) + CONVERT_IF_TYPES_MATCH(S8) + CONVERT_IF_TYPES_MATCH(S32) + CONVERT_IF_TYPES_MATCH(S64) + CONVERT_IF_TYPES_MATCH(U8) + CONVERT_IF_TYPES_MATCH(U32) + CONVERT_IF_TYPES_MATCH(U64) + CONVERT_IF_TYPES_MATCH(F16) + CONVERT_IF_TYPES_MATCH(F32) + CONVERT_IF_TYPES_MATCH(F64) +#undef CONVERT_IF_TYPES_MATCH + // Other types are not yet supported. + default: + return InvalidArgument( + "Unimplemented: Convert from type %s to type %s", + PrimitiveType_Name(src_literal.shape().element_type()).c_str(), + PrimitiveType_Name(primitive_dest_type).c_str()); } +} +} // namespace - LOG(FATAL) << "unhandled rank: " << ShapeUtil::Rank(literal.shape()); +StatusOr> Literal::Convert( + PrimitiveType primitive_dest_type) const { + switch (shape().element_type()) { +#define CONVERT_IF_DEST_TYPE_MATCHES(type) \ + case (type): \ + return ConvertIfDestTypeMatches<(type)>(*this, primitive_dest_type); + CONVERT_IF_DEST_TYPE_MATCHES(PRED) + CONVERT_IF_DEST_TYPE_MATCHES(S8) + CONVERT_IF_DEST_TYPE_MATCHES(S32) + CONVERT_IF_DEST_TYPE_MATCHES(S64) + CONVERT_IF_DEST_TYPE_MATCHES(U8) + CONVERT_IF_DEST_TYPE_MATCHES(U32) + CONVERT_IF_DEST_TYPE_MATCHES(U64) + CONVERT_IF_DEST_TYPE_MATCHES(F16) + CONVERT_IF_DEST_TYPE_MATCHES(F32) + CONVERT_IF_DEST_TYPE_MATCHES(F64) +#undef CONVERT_IF_DEST_TYPE_MATCHES + // Other types are not yet supported. + default: + return InvalidArgument("Unimplemented: Convert from type %s to type %s", + PrimitiveType_Name(shape().element_type()).c_str(), + PrimitiveType_Name(primitive_dest_type).c_str()); + } } namespace { @@ -737,8 +889,8 @@ template bool EqualElements(const Literal& literal1, const Literal& literal2, int dimension, std::vector* multi_index) { if (dimension == ShapeUtil::Rank(literal1.shape())) { - return (LiteralUtil::Get(literal1, *multi_index) == - LiteralUtil::Get(literal2, *multi_index)); + return (literal1.Get(*multi_index) == + literal2.Get(*multi_index)); } for (int64 i = 0; i < literal1.shape().dimensions(dimension); ++i) { (*multi_index)[dimension] = i; @@ -752,138 +904,191 @@ bool EqualElements(const Literal& literal1, const Literal& literal2, } // namespace -/* static */ bool LiteralUtil::Equal(const Literal& literal1, - const Literal& literal2) { - if (!ShapeUtil::Compatible(literal1.shape(), literal2.shape())) { +bool Literal::Equal(const Literal& literal2) const { + if (!ShapeUtil::Compatible(shape(), literal2.shape())) { return false; } - if (ShapeUtil::IsTuple(literal1.shape())) { + if (ShapeUtil::IsTuple(shape())) { // Because the shapes are compatible, they must have the same number of // tuple elements. - CHECK_EQ(literal1.tuple_literals_size(), literal2.tuple_literals_size()); - for (int i = 0; i < literal1.tuple_literals_size(); ++i) { - if (!Equal(literal1.tuple_literals(i), literal2.tuple_literals(i))) { + CHECK_EQ(tuple_literals_size(), literal2.tuple_literals_size()); + for (int i = 0; i < tuple_literals_size(); ++i) { + if (!tuple_literals(i).Equal(literal2.tuple_literals(i))) { return false; } } return true; } else { - std::vector multi_index(ShapeUtil::Rank(literal1.shape()), 0); - switch (literal1.shape().element_type()) { + std::vector multi_index(ShapeUtil::Rank(shape()), 0); + switch (shape().element_type()) { case PRED: - return EqualElements(literal1, literal2, 0, &multi_index); + return EqualElements(*this, literal2, 0, &multi_index); case U8: - return EqualElements(literal1, literal2, 0, &multi_index); + return EqualElements(*this, literal2, 0, &multi_index); case S32: - return EqualElements(literal1, literal2, 0, &multi_index); + return EqualElements(*this, literal2, 0, &multi_index); case S64: - return EqualElements(literal1, literal2, 0, &multi_index); + return EqualElements(*this, literal2, 0, &multi_index); case U32: - return EqualElements(literal1, literal2, 0, &multi_index); + return EqualElements(*this, literal2, 0, &multi_index); case U64: - return EqualElements(literal1, literal2, 0, &multi_index); + return EqualElements(*this, literal2, 0, &multi_index); case F32: - return EqualElements(literal1, literal2, 0, &multi_index); + return EqualElements(*this, literal2, 0, &multi_index); case F64: - return EqualElements(literal1, literal2, 0, &multi_index); + return EqualElements(*this, literal2, 0, &multi_index); + case F16: + return EqualElements(*this, literal2, 0, &multi_index); default: - LOG(FATAL) << "Unimplemented: LiteralUtil::Equal for type " - << PrimitiveType_Name(literal1.shape().element_type()); + LOG(FATAL) << "Unimplemented: Literal::Equal for type " + << PrimitiveType_Name(shape().element_type()); } } } template <> -/* static */ tensorflow::gtl::ArraySlice LiteralUtil::GetArraySlice( - const Literal& literal) { - CHECK(literal.shape().element_type() == PRED); - return literal.preds(); +tensorflow::gtl::MutableArraySlice Literal::GetMutableArraySlice() { + auto values = mutable_preds(); + return tensorflow::gtl::MutableArraySlice( + reinterpret_cast(values->data()), values->size()); +} + +template <> +tensorflow::gtl::MutableArraySlice Literal::GetMutableArraySlice() { + auto values = mutable_u8s(); + return tensorflow::gtl::MutableArraySlice( + reinterpret_cast(values->data()), values->size()); +} + +template <> +tensorflow::gtl::MutableArraySlice Literal::GetMutableArraySlice() { + auto values = mutable_u8s(); + return tensorflow::gtl::MutableArraySlice(values->data(), + values->size()); +} + +template <> +tensorflow::gtl::MutableArraySlice Literal::GetMutableArraySlice() { + auto values = mutable_s32s(); + return tensorflow::gtl::MutableArraySlice(values->data(), + values->size()); +} + +template <> +tensorflow::gtl::MutableArraySlice Literal::GetMutableArraySlice() { + auto values = mutable_u32s(); + return tensorflow::gtl::MutableArraySlice(values->data(), + values->size()); +} + +template <> +tensorflow::gtl::MutableArraySlice Literal::GetMutableArraySlice() { + static_assert(sizeof(int64) == sizeof(tensorflow::protobuf_int64) && + alignof(int64) == alignof(tensorflow::protobuf_int64), + "The int64 and tensorflow::protobuf_int64 types are not " + "compatible"); + auto values = mutable_s64s(); + // Because of the fact that tensorflow::protobuf_int64 is defined as int64_t + // while tensorflow::int64 is defined as long long, a reinterpret_cast<> is + // necessary from the raw data pointer returned by the mutable_data() API. + return tensorflow::gtl::MutableArraySlice( + reinterpret_cast(values->data()), values->size()); +} + +template <> +tensorflow::gtl::MutableArraySlice Literal::GetMutableArraySlice() { + static_assert(sizeof(uint64) == sizeof(tensorflow::protobuf_uint64) && + alignof(uint64) == alignof(tensorflow::protobuf_uint64), + "The uint64 and tensorflow::protobuf_uint64 types are not " + "compatible"); + auto values = mutable_u64s(); + // Because of the fact that tensorflow::protobuf_uint64 is defined as uint64_t + // while tensorflow::uint64 is defined as unsigned long long, a + // reinterpret_cast<> is necessary from the raw data pointer returned by the + // mutable_data() API. + return tensorflow::gtl::MutableArraySlice( + reinterpret_cast(values->data()), values->size()); } template <> -/* static */ tensorflow::protobuf::RepeatedField* -LiteralUtil::GetMutableRepeatedField(Literal* literal) { - CHECK(literal->shape().element_type() == PRED); - return literal->mutable_preds(); +tensorflow::gtl::MutableArraySlice Literal::GetMutableArraySlice() { + auto values = mutable_f32s(); + return tensorflow::gtl::MutableArraySlice(values->data(), + values->size()); } template <> -/* static */ tensorflow::gtl::ArraySlice -LiteralUtil::GetArraySlice(const Literal& literal) { - CHECK(literal.shape().element_type() == U32); - return literal.u32s(); +tensorflow::gtl::MutableArraySlice Literal::GetMutableArraySlice() { + auto values = mutable_f64s(); + return tensorflow::gtl::MutableArraySlice(values->data(), + values->size()); } template <> -/* static */ tensorflow::protobuf::RepeatedField* -LiteralUtil::GetMutableRepeatedField(Literal* literal) { - CHECK(literal->shape().element_type() == U32); - return literal->mutable_u32s(); +tensorflow::gtl::MutableArraySlice Literal::GetMutableArraySlice() { + // TODO - there is an endianess problem here. fix it, or wait for uint16 + // support in protobuf + auto values = mutable_f16s(); + return tensorflow::gtl::MutableArraySlice(values->data(), + values->size()); } template <> -/* static */ tensorflow::gtl::ArraySlice -LiteralUtil::GetArraySlice(const Literal& literal) { - CHECK(literal.shape().element_type() == U64); - return AsUInt64Slice(literal.u64s()); +tensorflow::gtl::ArraySlice Literal::GetArraySlice() const { + CHECK_EQ(shape().element_type(), PRED); + return tensorflow::gtl::ArraySlice( + reinterpret_cast(preds().data()), preds().size()); } template <> -/* static */ tensorflow::protobuf::RepeatedField* -LiteralUtil::GetMutableRepeatedField( - Literal* literal) { - CHECK(literal->shape().element_type() == U64); - return literal->mutable_u64s(); +tensorflow::gtl::ArraySlice Literal::GetArraySlice() const { + CHECK_EQ(shape().element_type(), U8); + return tensorflow::gtl::ArraySlice( + reinterpret_cast(u8s().data()), u8s().size()); } template <> -/* static */ tensorflow::gtl::ArraySlice -LiteralUtil::GetArraySlice(const Literal& literal) { - CHECK(literal.shape().element_type() == S32); - return literal.s32s(); +tensorflow::gtl::ArraySlice Literal::GetArraySlice() const { + CHECK_EQ(shape().element_type(), S8); + return tensorflow::gtl::ArraySlice( + reinterpret_cast(u8s().data()), u8s().size()); } template <> -/* static */ tensorflow::protobuf::RepeatedField* -LiteralUtil::GetMutableRepeatedField(Literal* literal) { - CHECK(literal->shape().element_type() == S32); - return literal->mutable_s32s(); +tensorflow::gtl::ArraySlice Literal::GetArraySlice() const { + CHECK_EQ(shape().element_type(), U32); + return u32s(); } template <> -/* static */ tensorflow::gtl::ArraySlice -LiteralUtil::GetArraySlice(const Literal& literal) { - CHECK(literal.shape().element_type() == S64); - return AsInt64Slice(literal.s64s()); +tensorflow::gtl::ArraySlice Literal::GetArraySlice() const { + CHECK_EQ(shape().element_type(), U64); + return u64s(); } template <> -/* static */ tensorflow::protobuf::RepeatedField* -LiteralUtil::GetMutableRepeatedField( - Literal* literal) { - CHECK(literal->shape().element_type() == S64); - return literal->mutable_s64s(); +tensorflow::gtl::ArraySlice Literal::GetArraySlice() const { + CHECK_EQ(shape().element_type(), S32); + return s32s(); } template <> -/* static */ tensorflow::protobuf::RepeatedField* -LiteralUtil::GetMutableRepeatedField(Literal* literal) { - CHECK(literal->shape().element_type() == F32); - return literal->mutable_f32s(); +tensorflow::gtl::ArraySlice Literal::GetArraySlice() const { + CHECK_EQ(shape().element_type(), S64); + return s64s(); } template <> -/* static */ tensorflow::gtl::ArraySlice -LiteralUtil::GetArraySlice(const Literal& literal) { - CHECK(literal.shape().element_type() == F64); - return literal.f64s(); +tensorflow::gtl::ArraySlice Literal::GetArraySlice() const { + CHECK_EQ(shape().element_type(), F64); + return f64s(); } template <> -/* static */ tensorflow::protobuf::RepeatedField* -LiteralUtil::GetMutableRepeatedField(Literal* literal) { - CHECK(literal->shape().element_type() == F64); - return literal->mutable_f64s(); +tensorflow::gtl::ArraySlice Literal::GetArraySlice() const { + CHECK_EQ(shape().element_type(), F16); + return tensorflow::gtl::ArraySlice(f16s().data(), + f16s().size() / sizeof(half)); } template @@ -891,46 +1096,48 @@ static bool AllElementsEqualValue(const Literal& literal, NativeT value) { for (int64 i = 0; i < ShapeUtil::ElementsIn(literal.shape()); ++i) { auto multi_index = IndexUtil::LinearIndexToMultidimensionalIndex(literal.shape(), i); - if (LiteralUtil::Get(literal, multi_index) != value) { + if (literal.Get(multi_index) != value) { return false; } } return true; } -/* static */ bool LiteralUtil::IsAll(const Literal& literal, int8 value) { - switch (literal.shape().element_type()) { +bool Literal::IsAll(int8 value) const { + switch (shape().element_type()) { case U8: if (value >= 0) { - return AllElementsEqualValue(literal, value); + return AllElementsEqualValue(*this, value); } return false; case U32: if (value >= 0) { - return AllElementsEqualValue(literal, value); + return AllElementsEqualValue(*this, value); } return false; case U64: if (value >= 0) { - return AllElementsEqualValue(literal, value); + return AllElementsEqualValue(*this, value); } return false; case S8: - return AllElementsEqualValue(literal, value); + return AllElementsEqualValue(*this, value); case S32: - return AllElementsEqualValue(literal, value); + return AllElementsEqualValue(*this, value); case S64: - return AllElementsEqualValue(literal, value); + return AllElementsEqualValue(*this, value); case F32: - return AllElementsEqualValue(literal, value); + return AllElementsEqualValue(*this, value); case F64: - return AllElementsEqualValue(literal, value); + return AllElementsEqualValue(*this, value); + case F16: + return AllElementsEqualValue(*this, static_cast(value)); case PRED: if (value == 0) { - return AllElementsEqualValue(literal, false); + return AllElementsEqualValue(*this, false); } if (value == 1) { - return AllElementsEqualValue(literal, true); + return AllElementsEqualValue(*this, true); } return false; default: @@ -938,89 +1145,210 @@ static bool AllElementsEqualValue(const Literal& literal, NativeT value) { } } -/* static */ bool LiteralUtil::IsAllFloat(const Literal& literal, float value) { - switch (literal.shape().element_type()) { +bool Literal::IsAllFloat(float value) const { + switch (shape().element_type()) { case F32: - return AllElementsEqualValue(literal, value); + return AllElementsEqualValue(*this, value); case F64: - return AllElementsEqualValue(literal, value); + return AllElementsEqualValue(*this, value); + case F16: + return AllElementsEqualValue(*this, static_cast(value)); default: return false; } } -/* static */ bool LiteralUtil::IsZero( - const Literal& literal, tensorflow::gtl::ArraySlice indices) { - switch (literal.shape().element_type()) { +bool Literal::IsZero(tensorflow::gtl::ArraySlice indices) const { + switch (shape().element_type()) { case U8: - return Get(literal, indices) == 0; + return Get(indices) == 0; case U32: - return Get(literal, indices) == 0; + return Get(indices) == 0; case U64: - return Get(literal, indices) == 0; + return Get(indices) == 0; case S8: - return Get(literal, indices) == 0; + return Get(indices) == 0; case S32: - return Get(literal, indices) == 0; + return Get(indices) == 0; case S64: - return Get(literal, indices) == 0; + return Get(indices) == 0; case F32: - return Get(literal, indices) == 0.0f; + return Get(indices) == 0.0f; case F64: - return Get(literal, indices) == 0.0; + return Get(indices) == 0.0; + case F16: + return Get(indices) == static_cast(0.0f); case PRED: - return Get(literal, indices) == false; + return Get(indices) == false; default: LOG(FATAL) << "Input literal must be an array."; } } template <> -/* static */ void LiteralUtil::PopulateWithValue( - int64 value, tensorflow::gtl::ArraySlice dimensions, - Literal* literal) { - *literal->mutable_shape() = ShapeUtil::MakeShape( - primitive_util::NativeToPrimitiveType(), dimensions); - tensorflow::protobuf::RepeatedField* - repeated_field = - GetMutableRepeatedField(literal); - for (int64 i = 0; i < ShapeUtil::ElementsIn(literal->shape()); ++i) { - repeated_field->Add(value); - } +/* static */ void Literal::Resize(int64 num_elements, bool value) { + CHECK_EQ(ShapeUtil::ElementsIn(shape()), num_elements); + mutable_preds()->resize(num_elements, value); } template <> -/* static */ void LiteralUtil::PopulateWithValue( - uint64 value, tensorflow::gtl::ArraySlice dimensions, - Literal* literal) { - *literal->mutable_shape() = ShapeUtil::MakeShape( - primitive_util::NativeToPrimitiveType(), dimensions); - tensorflow::protobuf::RepeatedField* - repeated_field = - GetMutableRepeatedField(literal); - for (int64 i = 0; i < ShapeUtil::ElementsIn(literal->shape()); ++i) { - repeated_field->Add(value); - } +void Literal::Resize(int64 num_elements, int8 value) { + CHECK_EQ(ShapeUtil::ElementsIn(shape()), num_elements); + mutable_u8s()->resize(num_elements, value); +} + +template <> +void Literal::Resize(int64 num_elements, uint8 value) { + CHECK_EQ(ShapeUtil::ElementsIn(shape()), num_elements); + mutable_u8s()->resize(num_elements, value); +} + +template <> +void Literal::Resize(int64 num_elements, int32 value) { + CHECK_EQ(ShapeUtil::ElementsIn(shape()), num_elements); + mutable_s32s()->resize(num_elements, value); +} + +template <> +void Literal::Resize(int64 num_elements, uint32 value) { + CHECK_EQ(ShapeUtil::ElementsIn(shape()), num_elements); + mutable_u32s()->resize(num_elements, value); +} + +template <> +void Literal::Resize(int64 num_elements, int64 value) { + CHECK_EQ(ShapeUtil::ElementsIn(shape()), num_elements); + mutable_s64s()->resize(num_elements, value); +} + +template <> +void Literal::Resize(int64 num_elements, uint64 value) { + CHECK_EQ(ShapeUtil::ElementsIn(shape()), num_elements); + mutable_u64s()->resize(num_elements, value); +} + +template <> +void Literal::Resize(int64 num_elements, float value) { + CHECK_EQ(ShapeUtil::ElementsIn(shape()), num_elements); + mutable_f32s()->resize(num_elements, value); } template <> -/* static */ void LiteralUtil::Resize(int64 num_elements, int64 value, - Literal* literal) { - CHECK_EQ(ShapeUtil::ElementsIn(literal->shape()), num_elements); - tensorflow::protobuf::RepeatedField* - repeated_field = - GetMutableRepeatedField(literal); - repeated_field->Resize(num_elements, value); +void Literal::Resize(int64 num_elements, double value) { + CHECK_EQ(ShapeUtil::ElementsIn(shape()), num_elements); + mutable_f64s()->resize(num_elements, value); } template <> -/* static */ void LiteralUtil::Resize(int64 num_elements, uint64 value, - Literal* literal) { - CHECK_EQ(ShapeUtil::ElementsIn(literal->shape()), num_elements); - tensorflow::protobuf::RepeatedField* - repeated_field = - GetMutableRepeatedField(literal); - repeated_field->Resize(num_elements, value); +void Literal::Resize(int64 num_elements, half value) { + CHECK_EQ(ShapeUtil::ElementsIn(shape()), num_elements); + mutable_f16s()->resize(num_elements, value); +} + +template +static void CopyToRepeatedField(RepeatedFieldT* dest, + const std::vector& src) { + *dest = RepeatedFieldT(src.begin(), src.end()); +} + +LiteralProto Literal::ToProto() const { + LiteralProto proto; + proto.Clear(); + *proto.mutable_shape() = shape(); + switch (shape().element_type()) { + case PRED: + CopyToRepeatedField(proto.mutable_preds(), preds()); + break; + case U8: + *proto.mutable_u8s() = u8s_string(); + break; + case S32: + CopyToRepeatedField(proto.mutable_s32s(), s32s()); + break; + case S64: + CopyToRepeatedField(proto.mutable_s64s(), s64s()); + break; + case U32: + CopyToRepeatedField(proto.mutable_u32s(), u32s()); + break; + case U64: + CopyToRepeatedField(proto.mutable_u64s(), u64s()); + break; + case F16: + *proto.mutable_f16s() = + string(reinterpret_cast(f16s_.data()), + f16s_.size() * sizeof(half)); + break; + case F32: + CopyToRepeatedField(proto.mutable_f32s(), f32s()); + break; + case F64: + CopyToRepeatedField(proto.mutable_f64s(), f64s()); + break; + case TUPLE: + for (const auto& tuple : tuple_literals()) { + *proto.add_tuple_literals() = tuple.ToProto(); + } + break; + default: + LOG(FATAL) << "Unhandled primitive type " << shape().element_type(); + } + + return proto; +} + +template +static void CopyFromRepeatedField(std::vector* dest, + const RepeatedFieldT& src) { + *dest = std::vector(src.begin(), src.end()); +} + +void Literal::CopyFromProto(const LiteralProto& literal_proto) { + if (!literal_proto.has_shape()) { + return; + } + + *mutable_shape() = literal_proto.shape(); + switch (shape().element_type()) { + case PRED: + CopyFromRepeatedField(mutable_preds(), literal_proto.preds()); + break; + case U8: + set_u8s(literal_proto.u8s()); + break; + case S32: + CopyFromRepeatedField(mutable_s32s(), literal_proto.s32s()); + break; + case S64: + CopyFromRepeatedField(mutable_s64s(), literal_proto.s64s()); + break; + case U32: + CopyFromRepeatedField(mutable_u32s(), literal_proto.u32s()); + break; + case U64: + CopyFromRepeatedField(mutable_u64s(), literal_proto.u64s()); + break; + case F16: { + const string& s(literal_proto.f16s()); + CHECK_EQ(0, s.size() % sizeof(half)); + f16s_ = std::vector(s.size() / sizeof(half)); + memcpy(f16s_.data(), s.data(), s.size()); + break; + } + case F32: + CopyFromRepeatedField(mutable_f32s(), literal_proto.f32s()); + break; + case F64: + CopyFromRepeatedField(mutable_f64s(), literal_proto.f64s()); + break; + case TUPLE: + for (const auto& proto : literal_proto.tuple_literals()) { + mutable_tuple_literals()->push_back(Literal(proto)); + } + break; + default: + LOG(FATAL) << "Unhandled primitive type " << shape().element_type(); + } } } // namespace xla diff --git a/tensorflow/compiler/xla/literal_util.h b/tensorflow/compiler/xla/literal_util.h index 21bb2e46cf2ebcd72bcce393a1e5526f41757544..447c494bfca01cd268bd3f4a335984a33bafb045 100644 --- a/tensorflow/compiler/xla/literal_util.h +++ b/tensorflow/compiler/xla/literal_util.h @@ -20,6 +20,7 @@ limitations under the License. #include #include +#include #include #include #include @@ -33,6 +34,7 @@ limitations under the License. #include "tensorflow/compiler/xla/primitive_util.h" #include "tensorflow/compiler/xla/ptr_util.h" #include "tensorflow/compiler/xla/shape_util.h" +#include "tensorflow/compiler/xla/status_macros.h" #include "tensorflow/compiler/xla/types.h" #include "tensorflow/compiler/xla/util.h" #include "tensorflow/compiler/xla/xla_data.pb.h" @@ -52,10 +54,126 @@ namespace xla { // PrimitiveType. See ComputationBuilder for details. Not all primitive types // defined in xla_data.proto have a corresponding native type or even have a // storage location in the Literal proto yet (for example, primitive type F16). -class LiteralUtil { +class Literal { public: - // Create new literal of a given rank. To minimize ambiguity (for users and - // the compiler) these CreateR[0-2] methods should explicitly specify the + Literal() {} + + Literal(const Literal& other) = default; + Literal(Literal&&) = default; + + explicit Literal(const LiteralProto& other) { CopyFromProto(other); } + + Literal& operator=(const Literal& other) = default; + Literal& operator=(Literal&&) = default; + + LiteralProto ToProto() const; + + bool has_shape() const { + return shape_.element_type() != PRIMITIVE_TYPE_INVALID; + } + + // Basic accessor functions. Names mirror the original protobuf + // functions for convenience. + string DebugString() const { return ToProto().DebugString(); } + string ShortDebugString() const { return ToProto().ShortDebugString(); } + + void Clear() { + shape_.Clear(); + u8s_.clear(); + s32s_.clear(); + s64s_.clear(); + u32s_.clear(); + u64s_.clear(); + f16s_.clear(); + f32s_.clear(); + f64s_.clear(); + tuple_literals_.clear(); + } + + int preds_size() const { return u8s().size(); } + const std::vector& preds() const { + static_assert(sizeof(uint8) == sizeof(bool), + "The uint8 and bool types should be the same size"); + return u8s_; + } + std::vector* mutable_preds() { + static_assert(sizeof(uint8) == sizeof(bool), + "The uint8 and bool types should be the same size"); + return &u8s_; + } + + int s32s_size() const { return s32s().size(); } + int32 s32s(int i) const { return s32s_[i]; } + const std::vector& s32s() const { return s32s_; } + std::vector* mutable_s32s() { return &s32s_; } + + int s64s_size() const { return s64s().size(); } + void add_s64s(int64 value) { s64s_.push_back(value); } + const std::vector& s64s() const { return s64s_; } + std::vector* mutable_s64s() { return &s64s_; } + + int u32s_size() const { return u32s().size(); } + uint32 u32s(int i) const { return u32s_[i]; } + const std::vector& u32s() const { return u32s_; } + std::vector* mutable_u32s() { return &u32s_; } + + int u64s_size() const { return u64s().size(); } + const std::vector& u64s() const { return u64s_; } + std::vector* mutable_u64s() { return &u64s_; } + + int f16s_size() const { return f16s().size(); } + half f16s(int i) const { return f16s_[i]; } + const std::vector& f16s() const { return f16s_; } + std::vector* mutable_f16s() { return &f16s_; } + + int f32s_size() const { return f32s().size(); } + float f32s(int i) const { return f32s_[i]; } + void add_f32s(float value) { f32s_.push_back(value); } + const std::vector& f32s() const { return f32s_; } + std::vector& f32s() { return f32s_; } + std::vector* mutable_f32s() { return &f32s_; } + + int f64s_size() const { return f64s().size(); } + const std::vector& f64s() const { return f64s_; } + std::vector* mutable_f64s() { return &f64s_; } + + int tuple_literals_size() const { return tuple_literals().size(); } + const Literal& tuple_literals(int i) const { return tuple_literals_[i]; } + Literal* add_tuple_literals() { + tuple_literals_.push_back(Literal()); + return &tuple_literals_.back(); + } + std::vector* mutable_tuple_literals() { return &tuple_literals_; } + const std::vector& tuple_literals() const { return tuple_literals_; } + + int u8s_size() const { return u8s().size(); } + const std::vector& u8s() const { return u8s_; } + void set_u8s(const std::vector& value) { u8s_ = value; } + void set_u8s(tensorflow::StringPiece value) { + u8s_ = std::vector(value.size()); + u8s_.clear(); + append_u8s(value); + } + + void append_u8s(tensorflow::StringPiece value) { + u8s_.insert(u8s_.end(), value.begin(), value.end()); + } + + string u8s_string() const { return string(u8s().begin(), u8s().end()); } + + std::vector* mutable_u8s() { return &u8s_; } + + const Shape& shape() const { return shape_; } + Shape* mutable_shape() { return &shape_; } + + void Swap(Literal* other) { + Literal temp = *this; + *this = *other; + *other = temp; + } + + // Creates a new literal of a given rank. To minimize ambiguity (for users + // and the compiler) these CreateR[0-2] methods should explicitly specify the // native type. For example: // // CreateR1({1.0, 42.0}); @@ -100,75 +218,98 @@ class LiteralUtil { values, const Layout& layout); - // Creates a new value that has the equivalent value as literal, but conforms - // to new_layout; e.g. a literal matrix that was in {0, 1} minor-to-major - // dimension layout can be re-layed-out as {1, 0} minor-to-major dimension - // layout and the value in the cell at any given logical index (i0, i1) will - // be the same. + // Creates a new Literal object with the shape specified as parameter. + // The content of the literal values is the default value of the primitive + // type of literal itself (0 for numeric types, and false for predicates). + static std::unique_ptr CreateFromShape(const Shape& shape); + + // Creates a new Literal object with its values havings the primitive_type + // type, and with dimensions defined by the dimensions parameter. + // The content of the literal values is the default value of the primitive + // type of literal itself (0 for numeric types, and false for predicates). + static std::unique_ptr CreateFromDimensions( + PrimitiveType primitive_type, + tensorflow::gtl::ArraySlice dimensions); + + // Copies the values from src_literal, starting at src_base shape indexes, + // to this literal, starting at dest_base, where the copy size in each + // dimension is specified by copy_size. + // The src_literal and this literal must have the same primitive type, + // src_base+copy_size must fit the source literal dimensions, as well as + // dest_base+copy_size must fit the destination literal dimensions. + Status Copy(const Literal& src_literal, + tensorflow::gtl::ArraySlice src_base, + tensorflow::gtl::ArraySlice dest_base, + tensorflow::gtl::ArraySlice copy_size); + + // Creates a new value that has the equivalent value as this literal, but + // conforms to new_layout; e.g. a literal matrix that was in {0, 1} + // minor-to-major dimension layout can be re-layed-out as {1, 0} + // minor-to-major dimension layout and the value in the cell at any given + // logical index (i0, i1) will be the same. // // Note: this is useful when the client wants to ensure that a value placed in // the XLA allocation tracker has a particular layout; for efficiency // purposes or avoiding unimplemented operation/layout combinations. - static std::unique_ptr Relayout(const Literal& literal, - const Layout& new_layout); + std::unique_ptr Relayout(const Layout& new_layout) const; - // Reshapes literal 'input' to have 'shape'. Both the original shape and - // 'shape' must contain the same number of elements. The implementation - // currently only supports monotonic dim0-major layouts. - static StatusOr> Reshape( - const xla::Literal& input, tensorflow::gtl::ArraySlice shape); + // Creates a new literal by reshaping this literal to have 'shape'. Both the + // original shape and 'shape' must contain the same number of elements. The + // implementation currently only supports monotonic dim0-major layouts. + StatusOr> Reshape( + tensorflow::gtl::ArraySlice shape) const; - // Creates a new literal by reordering the dimensions of the original literal. + // Creates a new literal by reordering the dimensions of this literal. // The given `permutation` must be a permutation of the dimension numbers // in the original literal, and it specifies the order of the new dimensions // in the result literal (i.e., new_order[i] = old_order[permutation[i]]). // For example, a transpose call on a literal of shape [3 x 8 x 4] and // `permutation` = {2, 0, 1} returns a new literal of shape [4 x 3 x 8]. - static std::unique_ptr Transpose( - const Literal& literal, tensorflow::gtl::ArraySlice permutation); + std::unique_ptr Transpose( + tensorflow::gtl::ArraySlice permutation) const; - // Creates a sub-array from the given literal by extracting the indices + // Creates a sub-array from this literal by extracting the indices // [start_index, limit_index) of each dimension. The result literal has the // same rank and layout as for the given literal. The number of indices in // start_indices and limit_indices must be the rank of the literal, and the // indices follow the order of the dimensions. - static std::unique_ptr Slice( - const Literal& literal, tensorflow::gtl::ArraySlice start_indices, - tensorflow::gtl::ArraySlice limit_indices); + std::unique_ptr Slice( + tensorflow::gtl::ArraySlice start_indices, + tensorflow::gtl::ArraySlice limit_indices) const; // Creates a literal with a prepended dimension with bound "times"; e.g. a - // f32[3x2] with times=4 will produce a f32[4x3x2] with the 3x2 from the input + // f32[3x2] with times=4 will produce a f32[4x3x2] with the 3x2 from this // literal replicated four times. template - static std::unique_ptr Replicate(const Literal& input, int64 times); + std::unique_ptr Replicate(int64 times) const; - // Create a literal by converting each element in an original literal to a new - // type. - template - static std::unique_ptr Convert(const Literal& literal); + // Converts this literal to another primitive type. Returns an error if the + // conversion is not possible. + StatusOr> Convert( + PrimitiveType primitive_dest_type) const; - // Create a literal value zero of the given primitive type. + // Creates a literal value zero of the given primitive type. static Literal Zero(PrimitiveType primitive_type); - // Create a literal value one of the given primitive type. + // Creates a literal value one of the given primitive type. static Literal One(PrimitiveType primitive_type); // Creates a literal value containing the minimum value of the given // primitive type. For floating-point types, returns -inf. static Literal MinValue(PrimitiveType primitive_type); - // Create a literal value containing the maximum value of the given + // Creates a literal value containing the maximum value of the given // primitive type. For floating-point types, returns inf. static Literal MaxValue(PrimitiveType primitive_type); - // Create a literal of the given shape where each element is `value`. + // Creates a literal of the given shape where each element is `value`. template static std::unique_ptr CreateFullWithMonotonicDim0MajorLayout( tensorflow::gtl::ArraySlice dimensions, NativeT value); - // Create a new literal from an array. The variants not ending with WithLayout - // use the default XLA layout for the literal's linear representation in - // memory. + // Creates a new literal from an array. The variants not ending with + // WithLayout use the default XLA layout for the literal's linear + // representation in memory. template static std::unique_ptr CreateR2FromArray2D( const Array2D& values); @@ -210,28 +351,49 @@ class LiteralUtil { std::initializer_list> values, int64 projection_p, int64 projection_z); - // Clones literal into an owned unique_ptr version. - static std::unique_ptr CloneToUnique(const Literal& literal); + // Clones this literal into an owned unique_ptr version. + std::unique_ptr CloneToUnique() const; + + // Returns the linear index of the given index within this literal's + // element_type repeated field. + int64 LinearIndex(tensorflow::gtl::ArraySlice multi_index) const; // Gets or sets an element in the literal at the given index. The index is // CHECKed against the dimension sizes. template - static NativeT Get(const Literal& literal, - tensorflow::gtl::ArraySlice multi_index); + NativeT Get(tensorflow::gtl::ArraySlice multi_index) const; + template + void Set(tensorflow::gtl::ArraySlice multi_index, NativeT value); + + // Returns a (Mutable)ArraySlice view of the array for this literal for the + // given NativeT (e.g., float). These functions map native type to XLA + // PrimitiveType via template specialization. The unspecialized forms below + // aborts to handle the error case where the given native type does not map to + // an XLA primitive type. + template + tensorflow::gtl::ArraySlice GetArraySlice() const { + static_assert(!std::is_same::value, + "Cannot map native type to primitive type."); + } template - static void Set(Literal* literal, - tensorflow::gtl::ArraySlice multi_index, - NativeT value); + tensorflow::gtl::MutableArraySlice GetMutableArraySlice() { + static_assert(!std::is_same::value, + "Cannot map native type to primitive type."); + } // Returns the element value at index (0, ..., 0), however many zeroes are // required for that index. template - static NativeT GetFirstElement(const Literal& literal); + NativeT GetFirstElement() const; // As Get(), but determines the correct type and converts the value // into text. - static string GetAsString(const Literal& literal, - tensorflow::gtl::ArraySlice multi_index); + string GetAsString(tensorflow::gtl::ArraySlice multi_index) const; + + // As Get(), but determines the correct type and converts the value into + // int64. + StatusOr GetIntegralAsS64( + tensorflow::gtl::ArraySlice multi_index) const; // Returns an identity matrix (rank 2) with the given row and column count. template @@ -241,12 +403,22 @@ class LiteralUtil { static std::unique_ptr MakeTuple( tensorflow::gtl::ArraySlice elements); + // As above, but intended to be invoked with move semantics; i.e. + // + // std::vector> elements = ...; + // auto result = Literal::MakeTupleOwned(std::move(elements)); + // + // This would have been declared as an overload, but there is ambiguity + // in invocation between the above signature and this one. + static std::unique_ptr MakeTupleOwned( + std::vector> elements); + // Validates that the data payload of the literal matches the literal shape; // if it does not, an appropriate status is returned. - static tensorflow::Status ValidateLiteral(const Literal& literal); + tensorflow::Status ValidateLiteral() const; // Returns a string representation of the literal value. - static string ToString(const Literal& literal); + string ToString() const; // Invokes the "per cell" callback for each element in the provided // literal with the element's indices and a string representation of @@ -255,105 +427,98 @@ class LiteralUtil { // This function is useful if you want a polymorphic representation // of the tensor's elements (turning it to a string for something // like representation in a protobuf). - static void EachCellAsString( - const Literal& literal, - std::function indices, - const string& value)> - per_cell); + void EachCellAsString( + const std::function indices, + const string& value)>& per_cell) const; template - static void EachCell( - const Literal& literal, - std::function indices, - NativeT value)> - per_cell); - - // Templated methods which populate the given repeated field in the Literal - // proto with the given value(s). The Shape field of the Literal proto is set + void EachCell(std::function indices, + NativeT value)> + per_cell) const; + + // Templated methods which populate the given repeated field in this literal + // with the given value(s). The Shape field of this literal is set // to match the array dimensions and type. Examples: // // // Populate with floats. // Array2D float_values = ... - // PopulateR2FromArray2D(values, literal); + // literal.PopulateR2FromArray2D(values); // // // Populate with int32s. - // PopulateR2({{1, 2}, {3, 4}}, literal); + // literal.PopulateR2({{1, 2}, {3, 4}}); // template - static void PopulateR0(NativeT values, Literal* literal); + void PopulateR0(NativeT values); template - static void PopulateR1(tensorflow::gtl::ArraySlice values, - Literal* literal); - static void PopulateR1(const tensorflow::core::Bitmap& values, - Literal* literal); + void PopulateR1(tensorflow::gtl::ArraySlice values); + void PopulateR1(const tensorflow::core::Bitmap& values); template - static void PopulateR2( - std::initializer_list> values, - Literal* literal); + void PopulateR2(std::initializer_list> values); template - static void PopulateR2WithLayout( + void PopulateR2WithLayout( std::initializer_list> values, - const Layout& layout, Literal* literal); + const Layout& layout); template - static void PopulateR2FromArray2D(const Array2D& values, - Literal* literal); + void PopulateR2FromArray2D(const Array2D& values); template - static void PopulateR2FromArray2DWithLayout(const Array2D& values, - const Layout& layout, - Literal* literal); + void PopulateR2FromArray2DWithLayout(const Array2D& values, + const Layout& layout); template - static void PopulateR3FromArray3D(const Array3D& values, - Literal* literal); + void PopulateR3FromArray3D(const Array3D& values); template - static void PopulateR3FromArray3DWithLayout(const Array3D& values, - const Layout& layout, - Literal* literal); + void PopulateR3FromArray3DWithLayout(const Array3D& values, + const Layout& layout); template - static void PopulateR4FromArray4D(const Array4D& values, - Literal* literal); + void PopulateR4FromArray4D(const Array4D& values); template - static void PopulateR4FromArray4DWithLayout(const Array4D& values, - const Layout& layout, - Literal* literal); + void PopulateR4FromArray4DWithLayout(const Array4D& values, + const Layout& layout); + + // Populates literal values by calling the generator function for every cell + // in this literal object. + // + // generator must be a callable of the type + // NativeT(tensorflow::gtl::ArraySlice indexes) or compatible. + template + Status Populate(const FnType& generator); // Creates a Literal of the given dimensions with all elements set to the // given value. template - static void PopulateWithValue(NativeT value, - tensorflow::gtl::ArraySlice dimensions, - Literal* literal); - - // Returns a pointer to the underlying buffer in the protobuf containing the - // array data. Use with care. - static const void* InternalData(const Literal& literal); - static void* MutableInternalData(Literal* literal); - - // Allocates space in the repeated_field of the literal sufficient to hold - // num_elements of the literal's primitive type. Values in the buffer are set - // to zero. num_elements must equal the number of elements in the literals + void PopulateWithValue(NativeT value, + tensorflow::gtl::ArraySlice dimensions); + + // Returns a pointer to the underlying vector corresponding to the Literal's + // shape. + const void* InternalData() const; + void* MutableInternalData(); + + // Allocates space in the underlying vector of this literal sufficient to hold + // num_elements of this literal's primitive type. Values in the vector are set + // to zero. num_elements must equal the number of elements in the literal's // shape. - static void Reserve(int64 num_elements, Literal* literal); + void Reserve(int64 num_elements); - // Allocates space in the repeated_field of the literal sufficient to hold - // num_elements of the literal's primitive type and sets each element in the + // Allocates space in the underlying vector of this literal sufficient to hold + // num_elements of this literal's primitive type and sets each element in this // literal to the given value. num_elements must equal the number of elements - // in the literals shape. + // in this literal's shape. template - static void Resize(int64 num_elements, NativeT value, Literal* literal); + void Resize(int64 num_elements, NativeT value); - // Returns true if the two given literals have the same shape and - // values. Layout is not considered in the comparison. - static bool Equal(const Literal& literal1, const Literal& literal2); + // Returns true if this literal has the same shape and value as the given + // literal. Layout is not considered in the comparison. + bool Equal(const Literal& literal2) const; - // Returns whether every element in the given literal is equal to value. + // Returns whether every element in this literal is equal to value. // // value is an int8 because we expect this to be called with small // compile-time constants (0, -1, etc.) and so that whatever value you pass // can be represented exactly by floating-point types as small as 16 bits. // - // If value doesn't fit in literal's type, returns false. Values of 1/0 are - // considered equal to true/false; other values are not considered equal to - // true. - static bool IsAll(const Literal& literal, int8 value); + // If value doesn't fit in this literal's type, returns false. Values of 1/0 + // are considered equal to true/false; other values are not considered equal + // to true. + bool IsAll(int8 value) const; // Like IsAll(const Literal&, int8), except we check whether the literal is // equal to a particular floating-point number. @@ -364,140 +529,188 @@ class LiteralUtil { // admonishments about floating-point equality checks apply. We expect you to // use this to check for values that can be expressed precisely as a float, // e.g. -0.5. - static bool IsAllFloat(const Literal& literal, float value); + bool IsAllFloat(float value) const; - // Returns whether the literal is zero at the specified index. The literal + // Returns whether this literal is zero at the specified index. This literal // must be an array. - static bool IsZero(const Literal& literal, - tensorflow::gtl::ArraySlice indices); + bool IsZero(tensorflow::gtl::ArraySlice indices) const; private: - // Returns an ArraySlice view of the array for the given literal for the - // given NativeT (e.g., float). These - // functions map native type to XLA PrimitiveType via template - // specialization. The unspecialized forms below aborts to handle the error - // case where the given native type does not map to an XLA primitive type. - template - static tensorflow::gtl::ArraySlice GetArraySlice( - const Literal& literal) { - static_assert(!std::is_same::value, - "Cannot map native type to primitive type."); - } - template - static tensorflow::protobuf::RepeatedField* GetMutableRepeatedField( - Literal* literal) { - // Make the expression depend on the template parameter NativeT so - // that this compile-time error only apperas if this function is - // instantiated with some concrete type that is not specialized - // below. - static_assert(!std::is_same::value, - "Cannot map native type to primitive type."); - } - - // Returns the linear index of the given index within the literal's - // element_type repeated field. - static int64 LinearIndex(const Literal& literal, - tensorflow::gtl::ArraySlice multi_index); - - TF_DISALLOW_COPY_AND_ASSIGN(LiteralUtil); + // Copy from a LiteralProto instance. + void CopyFromProto(const LiteralProto& literal_proto); + + // Internal template helper for the Copy() API, matching its arguments one by + // one. + template + Status CopyRange(const Literal& src_literal, + tensorflow::gtl::ArraySlice src_base, + tensorflow::gtl::ArraySlice dest_base, + tensorflow::gtl::ArraySlice copy_size); + + // Utility structure which is used to create the optimal configuration for + // a ShapeUtil::ForEachIndex() scan across two literals. + struct StrideConfig { + StrideConfig(const Shape& source_shape, const Shape& dest_shape, + tensorflow::gtl::ArraySlice dimensions); + + // The dimensions of the stride operation. Essentially every dimension + // will be iterated from base[i] to base[i]+dimensions[i], in step[i] + // steps. + tensorflow::gtl::ArraySlice dimensions; + DimensionVector base; + DimensionVector step; + int64 minor_dimension = 0; + // The size of the strides for source and destination. One of the two + // (the one looping through its most minor dimension) will be 1, while + // the other will be the stride size at the dimension matching the other + // shape most minor dimension being scanned. + int64 dest_stride = 1; + int64 source_stride = 1; + // The size of the inner loop on the most minor dimension. + int64 minor_loop_size = 1; + }; + + Shape shape_; + std::vector u8s_; + std::vector s32s_; + std::vector s64s_; + std::vector u32s_; + std::vector u64s_; + std::vector f16s_; + std::vector f32s_; + std::vector f64s_; + std::vector tuple_literals_; }; // Declarations of template specializations for GetArraySlice and -// GetMutableRepeatedField. The specializations map native type to XLA primitive +// GetMutableArraySlice. The specializations map native type to XLA primitive // type. template <> -/* static */ tensorflow::gtl::ArraySlice LiteralUtil::GetArraySlice( - const Literal& literal); +tensorflow::gtl::ArraySlice Literal::GetArraySlice() const; template <> -/* static */ tensorflow::protobuf::RepeatedField* -LiteralUtil::GetMutableRepeatedField(Literal* literal); +tensorflow::gtl::ArraySlice Literal::GetArraySlice() const; template <> -/* static */ tensorflow::gtl::ArraySlice -LiteralUtil::GetArraySlice(const Literal& literal); +tensorflow::gtl::ArraySlice Literal::GetArraySlice() const; template <> -/* static */ tensorflow::protobuf::RepeatedField* -LiteralUtil::GetMutableRepeatedField(Literal* literal); +tensorflow::gtl::ArraySlice Literal::GetArraySlice() const; template <> -/* static */ tensorflow::gtl::ArraySlice -LiteralUtil::GetArraySlice(const Literal& literal); +tensorflow::gtl::ArraySlice Literal::GetArraySlice() const; template <> -/* static */ tensorflow::protobuf::RepeatedField* -LiteralUtil::GetMutableRepeatedField( - Literal* literal); +tensorflow::gtl::ArraySlice Literal::GetArraySlice() const; template <> -/* static */ tensorflow::gtl::ArraySlice -LiteralUtil::GetArraySlice(const Literal& literal); +tensorflow::gtl::ArraySlice Literal::GetArraySlice() const; template <> -/* static */ tensorflow::protobuf::RepeatedField* -LiteralUtil::GetMutableRepeatedField(Literal* literal); +inline tensorflow::gtl::ArraySlice Literal::GetArraySlice() + const { + DCHECK(shape().element_type() == F32); + return f32s(); +} template <> -/* static */ tensorflow::gtl::ArraySlice -LiteralUtil::GetArraySlice(const Literal& literal); +tensorflow::gtl::ArraySlice Literal::GetArraySlice() const; template <> -/* static */ tensorflow::protobuf::RepeatedField* -LiteralUtil::GetMutableRepeatedField( - Literal* literal); +tensorflow::gtl::ArraySlice Literal::GetArraySlice() const; template <> -/* static */ inline tensorflow::gtl::ArraySlice -LiteralUtil::GetArraySlice(const Literal& literal) { - DCHECK(literal.shape().element_type() == F32); - return literal.f32s(); -} +tensorflow::gtl::MutableArraySlice Literal::GetMutableArraySlice(); + +template <> +tensorflow::gtl::MutableArraySlice Literal::GetMutableArraySlice(); + +template <> +tensorflow::gtl::MutableArraySlice Literal::GetMutableArraySlice(); + +template <> +tensorflow::gtl::MutableArraySlice Literal::GetMutableArraySlice(); + +template <> +tensorflow::gtl::MutableArraySlice Literal::GetMutableArraySlice(); template <> -/* static */ tensorflow::protobuf::RepeatedField* -LiteralUtil::GetMutableRepeatedField(Literal* literal); +tensorflow::gtl::MutableArraySlice Literal::GetMutableArraySlice(); template <> -/* static */ tensorflow::gtl::ArraySlice -LiteralUtil::GetArraySlice(const Literal& literal); +tensorflow::gtl::MutableArraySlice Literal::GetMutableArraySlice(); template <> -/* static */ tensorflow::protobuf::RepeatedField* -LiteralUtil::GetMutableRepeatedField(Literal* literal); +tensorflow::gtl::MutableArraySlice Literal::GetMutableArraySlice(); + +template <> +tensorflow::gtl::MutableArraySlice Literal::GetMutableArraySlice(); + +template <> +tensorflow::gtl::MutableArraySlice Literal::GetMutableArraySlice(); + +template <> +void Literal::Resize(int64 num_elements, bool value); + +template <> +void Literal::Resize(int64 num_elements, int8 value); + +template <> +void Literal::Resize(int64 num_elements, uint8 value); + +template <> +void Literal::Resize(int64 num_elements, int32 value); + +template <> +void Literal::Resize(int64 num_elements, uint32 value); + +template <> +void Literal::Resize(int64 num_elements, int64 value); + +template <> +void Literal::Resize(int64 num_elements, uint64 value); + +template <> +void Literal::Resize(int64 num_elements, float value); + +template <> +void Literal::Resize(int64 num_elements, double value); + +template <> +void Literal::Resize(int64 num_elements, half value); template -/* static */ std::unique_ptr LiteralUtil::CreateR0(NativeT value) { +/* static */ std::unique_ptr Literal::CreateR0(NativeT value) { auto literal = MakeUnique(); - PopulateR0(value, literal.get()); + literal->PopulateR0(value); return literal; } template -/* static */ std::unique_ptr LiteralUtil::CreateR1( +/* static */ std::unique_ptr Literal::CreateR1( tensorflow::gtl::ArraySlice values) { auto literal = MakeUnique(); - PopulateR1(values, literal.get()); + literal->PopulateR1(values); return literal; } template -/* static */ std::unique_ptr LiteralUtil::CreateR2WithLayout( +/* static */ std::unique_ptr Literal::CreateR2WithLayout( std::initializer_list> values, const Layout& layout) { auto literal = MakeUnique(); - PopulateR2WithLayout(values, layout, literal.get()); + literal->PopulateR2WithLayout(values, layout); return literal; } template -/* static */ std::unique_ptr LiteralUtil::CreateR2( +/* static */ std::unique_ptr Literal::CreateR2( std::initializer_list> values) { return CreateR2WithLayout(values, LayoutUtil::GetDefaultLayoutForR2()); } template -/* static */ std::unique_ptr LiteralUtil::CreateR3WithLayout( +/* static */ std::unique_ptr Literal::CreateR3WithLayout( std::initializer_list>> values, const Layout& layout) { @@ -522,14 +735,14 @@ template } template -/* static */ std::unique_ptr LiteralUtil::CreateR3( +/* static */ std::unique_ptr Literal::CreateR3( std::initializer_list>> values) { return CreateR3WithLayout(values, LayoutUtil::GetDefaultLayoutForR3()); } template -/* static */ std::unique_ptr LiteralUtil::CreateR4WithLayout( +/* static */ std::unique_ptr Literal::CreateR4WithLayout( std::initializer_list>>> values, @@ -560,7 +773,7 @@ template } template -/* static */ std::unique_ptr LiteralUtil::CreateR4( +/* static */ std::unique_ptr Literal::CreateR4( std::initializer_list>>> values) { @@ -568,38 +781,37 @@ template } template -/* static */ std::unique_ptr -LiteralUtil::CreateR2FromArray2DWithLayout(const Array2D& values, - const Layout& layout) { +/* static */ std::unique_ptr Literal::CreateR2FromArray2DWithLayout( + const Array2D& values, const Layout& layout) { auto literal = MakeUnique(); - PopulateR2FromArray2DWithLayout(values, layout, literal.get()); + literal->PopulateR2FromArray2DWithLayout(values, layout); return literal; } template -/* static */ std::unique_ptr LiteralUtil::CreateR2FromArray2D( +/* static */ std::unique_ptr Literal::CreateR2FromArray2D( const Array2D& values) { return CreateR2FromArray2DWithLayout(values, LayoutUtil::GetDefaultLayoutForR2()); } + template -/* static */ std::unique_ptr -LiteralUtil::CreateR3FromArray3DWithLayout(const Array3D& values, - const Layout& layout) { +/* static */ std::unique_ptr Literal::CreateR3FromArray3DWithLayout( + const Array3D& values, const Layout& layout) { auto literal = MakeUnique(); - PopulateR3FromArray3DWithLayout(values, layout, literal.get()); + literal->PopulateR3FromArray3DWithLayout(values, layout); return literal; } template -/* static */ std::unique_ptr LiteralUtil::CreateR3FromArray3D( +/* static */ std::unique_ptr Literal::CreateR3FromArray3D( const Array3D& values) { return CreateR3FromArray3DWithLayout(values, LayoutUtil::GetDefaultLayoutForR3()); } template -/* static */ std::unique_ptr LiteralUtil::CreateR3Projected( +/* static */ std::unique_ptr Literal::CreateR3Projected( std::initializer_list> values, int64 projection) { int64 dim0_size = projection; @@ -624,7 +836,7 @@ template } template -/* static */ std::unique_ptr LiteralUtil::CreateR4Projected( +/* static */ std::unique_ptr Literal::CreateR4Projected( std::initializer_list> values, int64 projection_p, int64 projection_z) { int64 dim0_size = projection_p; @@ -652,91 +864,92 @@ template } template -/* static */ std::unique_ptr LiteralUtil::CreateR4FromArray4D( +/* static */ std::unique_ptr Literal::CreateR4FromArray4D( const Array4D& values) { return CreateR4FromArray4DWithLayout(values, LayoutUtil::GetDefaultLayoutForR4()); } template -/* static */ std::unique_ptr -LiteralUtil::CreateR4FromArray4DWithLayout(const Array4D& values, - const Layout& layout) { +/* static */ std::unique_ptr Literal::CreateR4FromArray4DWithLayout( + const Array4D& values, const Layout& layout) { auto literal = MakeUnique(); - PopulateR4FromArray4DWithLayout(values, layout, literal.get()); + literal->PopulateR4FromArray4DWithLayout(values, layout); return literal; } template -/* static */ NativeT LiteralUtil::Get( - const Literal& literal, tensorflow::gtl::ArraySlice multi_index) { - int64 linear_index = LinearIndex(literal, multi_index); - return GetArraySlice(literal).at(linear_index); +NativeT Literal::Get(tensorflow::gtl::ArraySlice multi_index) const { + int64 linear_index = LinearIndex(multi_index); + return GetArraySlice().at(linear_index); } template -/* static */ NativeT LiteralUtil::GetFirstElement(const Literal& literal) { - return GetArraySlice(literal).at(0); +NativeT Literal::GetFirstElement() const { + return GetArraySlice().at(0); +} + +template <> +inline uint8 Literal::Get( + tensorflow::gtl::ArraySlice multi_index) const { + CHECK(shape().element_type() == U8); + int64 linear_index = LinearIndex(multi_index); + return u8s()[linear_index]; } template <> -/* static */ inline uint8 LiteralUtil::Get( - const Literal& literal, tensorflow::gtl::ArraySlice multi_index) { - CHECK(literal.shape().element_type() == U8); - int64 linear_index = LinearIndex(literal, multi_index); - return literal.u8s()[linear_index]; +inline int8 Literal::Get( + tensorflow::gtl::ArraySlice multi_index) const { + CHECK(shape().element_type() == S8); + int64 linear_index = LinearIndex(multi_index); + return u8s()[linear_index]; } template <> -/* static */ inline int8 LiteralUtil::Get( - const Literal& literal, tensorflow::gtl::ArraySlice multi_index) { - CHECK(literal.shape().element_type() == S8); - int64 linear_index = LinearIndex(literal, multi_index); - return literal.u8s()[linear_index]; +inline half Literal::Get( + tensorflow::gtl::ArraySlice multi_index) const { + CHECK(shape().element_type() == F16); + int64 linear_index = LinearIndex(multi_index); + return GetArraySlice()[linear_index]; } template -/* static */ void LiteralUtil::Set( - Literal* literal, tensorflow::gtl::ArraySlice multi_index, - NativeT value) { - int64 linear_index = LinearIndex(*literal, multi_index); - GetMutableRepeatedField(literal)->Set(linear_index, value); +void Literal::Set(tensorflow::gtl::ArraySlice multi_index, + NativeT value) { + int64 linear_index = LinearIndex(multi_index); + GetMutableArraySlice().at(linear_index) = value; } template <> -/* static */ inline void LiteralUtil::Set( - Literal* literal, tensorflow::gtl::ArraySlice multi_index, - uint8 value) { - int64 linear_index = LinearIndex(*literal, multi_index); - (*literal->mutable_u8s())[linear_index] = value; +inline void Literal::Set(tensorflow::gtl::ArraySlice multi_index, + uint8 value) { + int64 linear_index = LinearIndex(multi_index); + (*mutable_u8s())[linear_index] = value; } template <> -/* static */ inline void LiteralUtil::Set( - Literal* literal, tensorflow::gtl::ArraySlice multi_index, - int8 value) { - return Set(literal, multi_index, value); +inline void Literal::Set(tensorflow::gtl::ArraySlice multi_index, + int8 value) { + return Set(multi_index, value); } template <> -/* static */ inline void LiteralUtil::Set( - Literal* literal, tensorflow::gtl::ArraySlice multi_index, - int64 value) { - int64 linear_index = LinearIndex(*literal, multi_index); - (*literal->mutable_s64s())[linear_index] = value; +inline void Literal::Set(tensorflow::gtl::ArraySlice multi_index, + int64 value) { + int64 linear_index = LinearIndex(multi_index); + (*mutable_s64s())[linear_index] = value; } template <> -/* static */ inline void LiteralUtil::Set( - Literal* literal, tensorflow::gtl::ArraySlice multi_index, - uint64 value) { - int64 linear_index = LinearIndex(*literal, multi_index); - (*literal->mutable_u64s())[linear_index] = value; +/* static */ inline void Literal::Set( + tensorflow::gtl::ArraySlice multi_index, uint64 value) { + int64 linear_index = LinearIndex(multi_index); + (*mutable_u64s())[linear_index] = value; } // Returns an identity matrix (rank 2) with the given row and column count. template -/* static */ std::unique_ptr LiteralUtil::MakeIdentityR2(int64 size) { +/* static */ std::unique_ptr Literal::MakeIdentityR2(int64 size) { Array2D array(size, size, 0); for (int64 i = 0; i < size; ++i) { array(i, i) = 1; @@ -745,88 +958,51 @@ template } template -/* static */ void LiteralUtil::EachCell( - const Literal& literal, +void Literal::EachCell( std::function indices, NativeT value)> - per_cell) { - if (ShapeUtil::HasZeroElements(literal.shape())) { + per_cell) const { + if (ShapeUtil::HasZeroElements(shape())) { return; } - std::vector indices(ShapeUtil::Rank(literal.shape()), 0); + std::vector indices(ShapeUtil::Rank(shape()), 0); do { - per_cell(indices, Get(literal, indices)); - } while (IndexUtil::BumpIndices(literal.shape(), &indices)); + per_cell(indices, Get(indices)); + } while (IndexUtil::BumpIndices(shape(), &indices)); } template -/* static */ void LiteralUtil::PopulateR0(NativeT value, Literal* literal) { - *literal->mutable_shape() = ShapeUtil::MakeShape( +inline void Literal::PopulateR0(NativeT value) { + *mutable_shape() = ShapeUtil::MakeShape( primitive_util::NativeToPrimitiveType(), {}); - tensorflow::protobuf::RepeatedField* repeated_field = - GetMutableRepeatedField(literal); - repeated_field->Add(value); -} - -template <> -/* static */ inline void LiteralUtil::PopulateR0(uint8 value, - Literal* literal) { - *literal->mutable_shape() = - ShapeUtil::MakeShape(primitive_util::NativeToPrimitiveType(), {}); - literal->mutable_u8s()->push_back(value); -} - -template <> -/* static */ inline void LiteralUtil::PopulateR0(int8 value, - Literal* literal) { - *literal->mutable_shape() = - ShapeUtil::MakeShape(primitive_util::NativeToPrimitiveType(), {}); - literal->mutable_u8s()->push_back(value); -} - -template <> -/* static */ inline void LiteralUtil::PopulateR0(uint64 value, - Literal* literal) { - *literal->mutable_shape() = - ShapeUtil::MakeShape(primitive_util::NativeToPrimitiveType(), {}); - literal->mutable_u64s()->Add(value); -} - -template <> -/* static */ inline void LiteralUtil::PopulateR0(int64 value, - Literal* literal) { - *literal->mutable_shape() = - ShapeUtil::MakeShape(primitive_util::NativeToPrimitiveType(), {}); - literal->mutable_s64s()->Add(value); + Resize(1, value); } template -/* static */ void LiteralUtil::PopulateR1( - tensorflow::gtl::ArraySlice values, Literal* literal) { - *literal->mutable_shape() = +inline void Literal::PopulateR1(tensorflow::gtl::ArraySlice values) { + *mutable_shape() = ShapeUtil::MakeShape(primitive_util::NativeToPrimitiveType(), {static_cast(values.size())}); - Reserve(values.size(), literal); + Reserve(values.size()); for (int64 i = 0; i < values.size(); ++i) { - Set(literal, {i}, values[i]); + Set({i}, values[i]); } } -/* static */ inline void LiteralUtil::PopulateR1( - const tensorflow::core::Bitmap& values, Literal* literal) { - *literal->mutable_shape() = +inline void Literal::PopulateR1(const tensorflow::core::Bitmap& values) { + *mutable_shape() = ShapeUtil::MakeShape(PRED, {static_cast(values.bits())}); - Reserve(values.bits(), literal); + Reserve(values.bits()); for (int64 i = 0; i < static_cast(values.bits()); ++i) { - Set(literal, {i}, values.get(i)); + Set({i}, values.get(i)); } } template -/* static */ void LiteralUtil::PopulateR2WithLayout( +void Literal::PopulateR2WithLayout( std::initializer_list> values, - const Layout& layout, Literal* literal) { - *literal->mutable_shape() = ShapeUtil::MakeShapeWithLayout( + const Layout& layout) { + *mutable_shape() = ShapeUtil::MakeShapeWithLayout( primitive_util::NativeToPrimitiveType(), {static_cast(values.size()), static_cast(values.begin()->size())}, @@ -834,17 +1010,17 @@ template const int64 dim0_size = values.size(); const int64 dim1_size = values.begin()->size(); - CHECK_EQ(dim0_size, literal->shape().dimensions(0)); - CHECK_EQ(dim1_size, literal->shape().dimensions(1)); + CHECK_EQ(dim0_size, shape().dimensions(0)); + CHECK_EQ(dim1_size, shape().dimensions(1)); const int64 num_elements = dim1_size * dim0_size; - Reserve(num_elements, literal); + Reserve(num_elements); int64 dim0 = 0; for (auto inner_list : values) { int64 dim1 = 0; for (auto value : inner_list) { - Set(literal, {dim0, dim1}, value); + Set({dim0, dim1}, value); ++dim1; } CHECK_EQ(dim1_size, dim1); @@ -853,84 +1029,79 @@ template } template -/* static */ void LiteralUtil::PopulateR2( - std::initializer_list> values, - Literal* literal) { - PopulateR2WithLayout(values, LayoutUtil::GetDefaultLayoutForR2(), literal); +void Literal::PopulateR2( + std::initializer_list> values) { + PopulateR2WithLayout(values, LayoutUtil::GetDefaultLayoutForR2()); } template -/* static */ void LiteralUtil::PopulateR2FromArray2DWithLayout( - const Array2D& values, const Layout& layout, Literal* literal) { - *literal->mutable_shape() = ShapeUtil::MakeShapeWithLayout( +void Literal::PopulateR2FromArray2DWithLayout(const Array2D& values, + const Layout& layout) { + *mutable_shape() = ShapeUtil::MakeShapeWithLayout( primitive_util::NativeToPrimitiveType(), {values.height(), values.width()}, AsInt64Slice(layout.minor_to_major())); const int64 dim1_size = values.width(); const int64 dim0_size = values.height(); - CHECK_EQ(dim0_size, literal->shape().dimensions(0)); - CHECK_EQ(dim1_size, literal->shape().dimensions(1)); - Reserve(dim1_size * dim0_size, literal); + CHECK_EQ(dim0_size, shape().dimensions(0)); + CHECK_EQ(dim1_size, shape().dimensions(1)); + Reserve(dim1_size * dim0_size); for (int64 dim0 = 0; dim0 < dim0_size; ++dim0) { for (int64 dim1 = 0; dim1 < dim1_size; ++dim1) { - Set(literal, {dim0, dim1}, values(dim0, dim1)); + Set({dim0, dim1}, values(dim0, dim1)); } } } template -/* static */ void LiteralUtil::PopulateR2FromArray2D( - const Array2D& values, Literal* literal) { - PopulateR2FromArray2DWithLayout(values, LayoutUtil::GetDefaultLayoutForR2(), - literal); +void Literal::PopulateR2FromArray2D(const Array2D& values) { + PopulateR2FromArray2DWithLayout(values, LayoutUtil::GetDefaultLayoutForR2()); } + template -/* static */ void LiteralUtil::PopulateR3FromArray3DWithLayout( - const Array3D& values, const Layout& layout, Literal* literal) { - *literal->mutable_shape() = ShapeUtil::MakeShapeWithLayout( +void Literal::PopulateR3FromArray3DWithLayout(const Array3D& values, + const Layout& layout) { + *mutable_shape() = ShapeUtil::MakeShapeWithLayout( primitive_util::NativeToPrimitiveType(), {values.n1(), values.n2(), values.n3()}, AsInt64Slice(layout.minor_to_major())); - CHECK_EQ(values.n1(), literal->shape().dimensions(0)); - CHECK_EQ(values.n2(), literal->shape().dimensions(1)); - CHECK_EQ(values.n3(), literal->shape().dimensions(2)); - Reserve(values.n1() * values.n2() * values.n3(), literal); + CHECK_EQ(values.n1(), shape().dimensions(0)); + CHECK_EQ(values.n2(), shape().dimensions(1)); + CHECK_EQ(values.n3(), shape().dimensions(2)); + Reserve(values.n1() * values.n2() * values.n3()); for (int64 dim0 = 0; dim0 < values.n1(); ++dim0) { for (int64 dim1 = 0; dim1 < values.n2(); ++dim1) { for (int64 dim2 = 0; dim2 < values.n3(); ++dim2) { - Set(literal, {dim0, dim1, dim2}, values(dim0, dim1, dim2)); + Set({dim0, dim1, dim2}, values(dim0, dim1, dim2)); } } } } template -/* static */ void LiteralUtil::PopulateR3FromArray3D( - const Array3D& values, Literal* literal) { - PopulateR3FromArray3DWithLayout(values, LayoutUtil::GetDefaultLayoutForR3(), - literal); +void Literal::PopulateR3FromArray3D(const Array3D& values) { + PopulateR3FromArray3DWithLayout(values, LayoutUtil::GetDefaultLayoutForR3()); } template -/* static */ void LiteralUtil::PopulateR4FromArray4DWithLayout( - const Array4D& values, const Layout& layout, Literal* literal) { - *literal->mutable_shape() = ShapeUtil::MakeShapeWithLayout( +void Literal::PopulateR4FromArray4DWithLayout(const Array4D& values, + const Layout& layout) { + *mutable_shape() = ShapeUtil::MakeShapeWithLayout( primitive_util::NativeToPrimitiveType(), {values.planes(), values.depth(), values.height(), values.width()}, AsInt64Slice(layout.minor_to_major())); - CHECK_EQ(values.n1(), literal->shape().dimensions(0)); - CHECK_EQ(values.n2(), literal->shape().dimensions(1)); - CHECK_EQ(values.n3(), literal->shape().dimensions(2)); - CHECK_EQ(values.n4(), literal->shape().dimensions(3)); - Reserve(values.n1() * values.n2() * values.n3() * values.n4(), literal); + CHECK_EQ(values.n1(), shape().dimensions(0)); + CHECK_EQ(values.n2(), shape().dimensions(1)); + CHECK_EQ(values.n3(), shape().dimensions(2)); + CHECK_EQ(values.n4(), shape().dimensions(3)); + Reserve(values.n1() * values.n2() * values.n3() * values.n4()); for (int64 dim0 = 0; dim0 < values.n1(); ++dim0) { for (int64 dim1 = 0; dim1 < values.n2(); ++dim1) { for (int64 dim2 = 0; dim2 < values.n3(); ++dim2) { for (int64 dim3 = 0; dim3 < values.n4(); ++dim3) { - Set(literal, {dim0, dim1, dim2, dim3}, - values(dim0, dim1, dim2, dim3)); + Set({dim0, dim1, dim2, dim3}, values(dim0, dim1, dim2, dim3)); } } } @@ -938,116 +1109,92 @@ template } template -/* static */ void LiteralUtil::PopulateR4FromArray4D( - const Array4D& values, Literal* literal) { - PopulateR4FromArray4DWithLayout(values, LayoutUtil::GetDefaultLayoutForR4(), - literal); +void Literal::PopulateR4FromArray4D(const Array4D& values) { + PopulateR4FromArray4DWithLayout(values, LayoutUtil::GetDefaultLayoutForR4()); } -template -/* static */ void LiteralUtil::PopulateWithValue( - NativeT value, tensorflow::gtl::ArraySlice dimensions, - Literal* literal) { - *literal->mutable_shape() = ShapeUtil::MakeShape( - primitive_util::NativeToPrimitiveType(), dimensions); - tensorflow::protobuf::RepeatedField* repeated_field = - GetMutableRepeatedField(literal); - for (int64 i = 0; i < ShapeUtil::ElementsIn(literal->shape()); ++i) { - repeated_field->Add(value); +template +Status Literal::Populate(const FnType& generator) { + const Shape& this_shape = shape(); + const int64 rank = ShapeUtil::Rank(this_shape); + TF_RET_CHECK(this_shape.element_type() == + primitive_util::NativeToPrimitiveType()); + tensorflow::gtl::MutableArraySlice data = + GetMutableArraySlice(); + if (rank > 0) { + StrideConfig stride_config(this_shape, this_shape, + AsInt64Slice(this_shape.dimensions())); + DimensionVector minor_scan_indexes(rank, 0); + int64 minor_dimension_size = + ShapeUtil::GetDimension(this_shape, stride_config.minor_dimension); + + auto init_function = [&](const std::vector& indexes) { + const int64 index = LinearIndex(indexes); + std::copy(indexes.begin(), indexes.end(), minor_scan_indexes.begin()); + for (int64 i = 0; i < minor_dimension_size; ++i) { + minor_scan_indexes[stride_config.minor_dimension] = i; + data.at(index + i) = generator(minor_scan_indexes); + } + return true; + }; + ShapeUtil::ForEachIndex(this_shape, stride_config.base, + stride_config.dimensions, stride_config.step, + init_function); + } else { + // For scalars. + data.at(0) = generator({}); } -} - -template <> -/* static */ void LiteralUtil::PopulateWithValue( - int64 value, tensorflow::gtl::ArraySlice dimensions, - Literal* literal); - -template <> -/* static */ void LiteralUtil::PopulateWithValue( - uint64 value, tensorflow::gtl::ArraySlice dimensions, - Literal* literal); - -template -/* static */ std::unique_ptr LiteralUtil::Convert( - const Literal& literal) { - auto result_literal = MakeUnique(); - Shape result_shape = literal.shape(); - result_shape.set_element_type( - primitive_util::NativeToPrimitiveType()); - *result_literal->mutable_shape() = result_shape; - LiteralUtil::Reserve(ShapeUtil::ElementsIn(result_shape), - result_literal.get()); - LiteralUtil::EachCell( - literal, - [&](tensorflow::gtl::ArraySlice indices, NativeSrcT value) { - LiteralUtil::Set(result_literal.get(), indices, - static_cast(value)); - }); - return result_literal; + return Status::OK(); } template -/* static */ void LiteralUtil::Resize(int64 num_elements, NativeT value, - Literal* literal) { - CHECK_EQ(ShapeUtil::ElementsIn(literal->shape()), num_elements); - tensorflow::protobuf::RepeatedField* repeated_field = - GetMutableRepeatedField(literal); - repeated_field->Resize(num_elements, value); +void Literal::PopulateWithValue(NativeT value, + tensorflow::gtl::ArraySlice dimensions) { + *mutable_shape() = ShapeUtil::MakeShape( + primitive_util::NativeToPrimitiveType(), dimensions); + Resize(ShapeUtil::ElementsIn(shape()), value); } -template <> -/* static */ void LiteralUtil::Resize(int64 num_elements, int64 value, - Literal* literal); - -template <> -/* static */ void LiteralUtil::Resize(int64 num_elements, uint64 value, - Literal* literal); - template /* static */ std::unique_ptr -LiteralUtil::CreateFullWithMonotonicDim0MajorLayout( +Literal::CreateFullWithMonotonicDim0MajorLayout( tensorflow::gtl::ArraySlice dimensions, NativeT value) { - Shape shape = ShapeUtil::MakeShapeWithMonotonicDim0MajorLayout( + Shape this_shape = ShapeUtil::MakeShapeWithMonotonicDim0MajorLayout( primitive_util::NativeToPrimitiveType(), dimensions); auto literal = MakeUnique(); - *literal->mutable_shape() = shape; - Reserve(ShapeUtil::ElementsIn(shape), literal.get()); + *literal->mutable_shape() = this_shape; + literal->Reserve(ShapeUtil::ElementsIn(this_shape)); std::vector index(dimensions.size(), 0); do { - Set(literal.get(), index, value); - } while (IndexUtil::BumpIndices(shape, &index)); + literal->Set(index, value); + } while (IndexUtil::BumpIndices(this_shape, &index)); return literal; } template -/* static */ std::unique_ptr LiteralUtil::Replicate( - const Literal& input, int64 times) { - // Ranks greater than 8 are very rare, so use InlinedVector to store - // the bounds and indices. - static constexpr int kInlineRank = 8; - tensorflow::gtl::InlinedVector bounds = {times}; - bounds.reserve(input.shape().dimensions_size() + 1); - for (int64 bound : input.shape().dimensions()) { +std::unique_ptr Literal::Replicate(int64 times) const { + DimensionVector bounds = {times}; + bounds.reserve(shape().dimensions_size() + 1); + for (int64 bound : shape().dimensions()) { bounds.push_back(bound); } auto literal = MakeUnique(); *literal->mutable_shape() = - ShapeUtil::MakeShape(input.shape().element_type(), bounds); + ShapeUtil::MakeShape(shape().element_type(), bounds); int64 elements = ShapeUtil::ElementsIn(literal->shape()); if (elements == 0) { return literal; } - Reserve(elements, literal.get()); + literal->Reserve(elements); - tensorflow::gtl::InlinedVector output_indices( - bounds.size(), 0); + DimensionVector output_indices(bounds.size(), 0); tensorflow::gtl::ArraySlice input_indices = output_indices; input_indices.remove_prefix(1); bool done = false; while (!done) { - const auto element = Get(input, input_indices); - Set(literal.get(), output_indices, element); + const auto element = Get(input_indices); + literal->Set(output_indices, element); done = true; for (int n = 0; n < output_indices.size(); ++n) { diff --git a/tensorflow/compiler/xla/literal_util_test.cc b/tensorflow/compiler/xla/literal_util_test.cc index e53763376bfe58b7c5a811987161cac966d14222..a33c0fe09dd95734c7ff98ff4f7955c0124415b9 100644 --- a/tensorflow/compiler/xla/literal_util_test.cc +++ b/tensorflow/compiler/xla/literal_util_test.cc @@ -21,14 +21,17 @@ limitations under the License. #include "tensorflow/compiler/xla/array4d.h" #include "tensorflow/compiler/xla/layout_util.h" #include "tensorflow/compiler/xla/shape_util.h" -#include "tensorflow/compiler/xla/test_helpers.h" +#include "tensorflow/compiler/xla/test.h" #include "tensorflow/compiler/xla/types.h" -#include "tensorflow/core/platform/test.h" +#include "tensorflow/core/lib/core/status_test_util.h" +#include "tensorflow/core/platform/macros.h" #include "tensorflow/core/platform/types.h" namespace xla { namespace { +using ::testing::ElementsAre; + class LiteralUtilTest : public ::testing::Test { protected: LiteralUtilTest() { @@ -69,11 +72,11 @@ class LiteralUtilTest : public ::testing::Test { layout_r4_dim0minor_ = LayoutUtil::MakeLayout({0, 1, 2, 3}); literal_r4_2x2x3x3_dim0major_ = - LiteralUtil::CreateR4FromArray4DWithLayout(arr4d, - layout_r4_dim0major_); + Literal::CreateR4FromArray4DWithLayout(arr4d, + layout_r4_dim0major_); literal_r4_2x2x3x3_dim0minor_ = - LiteralUtil::CreateR4FromArray4DWithLayout(arr4d, - layout_r4_dim0minor_); + Literal::CreateR4FromArray4DWithLayout(arr4d, + layout_r4_dim0minor_); } Layout layout_r2_dim0major_; @@ -87,40 +90,42 @@ class LiteralUtilTest : public ::testing::Test { }; TEST_F(LiteralUtilTest, LiteralScalarToString) { - auto true_lit = LiteralUtil::CreateR0(true); - ASSERT_EQ("true", LiteralUtil::ToString(*true_lit)); + auto true_lit = Literal::CreateR0(true); + ASSERT_EQ("true", true_lit->ToString()); + + auto false_lit = Literal::CreateR0(false); + ASSERT_EQ("false", false_lit->ToString()); - auto false_lit = LiteralUtil::CreateR0(false); - ASSERT_EQ("false", LiteralUtil::ToString(*false_lit)); + auto u32_lit = Literal::CreateR0(42); + ASSERT_EQ("42", u32_lit->ToString()); - auto u32_lit = LiteralUtil::CreateR0(42); - ASSERT_EQ("42", LiteralUtil::ToString(*u32_lit)); + auto s32_lit = Literal::CreateR0(-999); + ASSERT_EQ("-999", s32_lit->ToString()); - auto s32_lit = LiteralUtil::CreateR0(-999); - ASSERT_EQ("-999", LiteralUtil::ToString(*s32_lit)); + auto f32_lit = Literal::CreateR0(3.14f); + ASSERT_EQ("3.14", f32_lit->ToString()); - auto f32_lit = LiteralUtil::CreateR0(3.14f); - ASSERT_EQ("3.14", LiteralUtil::ToString(*f32_lit)); + auto f16_lit = Literal::CreateR0(static_cast(0.5f)); + ASSERT_EQ("0.5", f16_lit->ToString()); } TEST_F(LiteralUtilTest, LiteralVectorToString) { - auto pred_vec = LiteralUtil::CreateR1({true, false, true}); - ASSERT_EQ("{101}", LiteralUtil::ToString(*pred_vec)); + auto pred_vec = Literal::CreateR1({true, false, true}); + ASSERT_EQ("{101}", pred_vec->ToString()); } TEST_F(LiteralUtilTest, R2ToString) { - const auto literal = LiteralUtil::CreateR2({{1, 2}, {3, 4}, {5, 6}}); + const auto literal = Literal::CreateR2({{1, 2}, {3, 4}, {5, 6}}); const string expected = R"(s32[3,2] { { 1, 2 }, { 3, 4 }, { 5, 6 }, })"; - ASSERT_EQ(expected, LiteralUtil::ToString(*literal)); + ASSERT_EQ(expected, literal->ToString()); } TEST_F(LiteralUtilTest, R3ToString) { - const auto literal = - LiteralUtil::CreateR3({{{1}, {2}}, {{3}, {4}}, {{5}, {6}}}); + const auto literal = Literal::CreateR3({{{1}, {2}}, {{3}, {4}}, {{5}, {6}}}); const string expected = R"(s32[3,2,1] { { { 1 }, { 2 } }, @@ -129,13 +134,13 @@ TEST_F(LiteralUtilTest, R3ToString) { { { 5 }, { 6 } } })"; - ASSERT_EQ(expected, LiteralUtil::ToString(*literal)); + ASSERT_EQ(expected, literal->ToString()); } 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.get(), matrix.get()}); + auto scalar = Literal::CreateR0(1.0); + auto matrix = Literal::CreateR2({{1.0, 2.0}, {3.0, 4.0}}); + auto tuple = Literal::MakeTuple({scalar.get(), matrix.get()}); const string expected = R"((f32[], f32[2,2]) ( 1, f32[2,2] { @@ -143,7 +148,7 @@ f32[2,2] { { 3, 4 }, }, ))"; - ASSERT_EQ(expected, LiteralUtil::ToString(*tuple)); + ASSERT_EQ(expected, tuple->ToString()); } TEST_F(LiteralUtilTest, CreateR3FromArray3d) { @@ -158,11 +163,9 @@ TEST_F(LiteralUtilTest, CreateR3FromArray3d) { }); // clang-format on - auto literal = LiteralUtil::CreateR3FromArray3D(array_3d); - EXPECT_MATCH(testing::PBToVec( - literal->shape().dimensions()), - testing::VectorMatcher({2, 3, 2})); - string result = LiteralUtil::ToString(*literal); + auto literal = Literal::CreateR3FromArray3D(array_3d); + EXPECT_THAT(literal->shape().dimensions(), ElementsAre(2, 3, 2)); + string result = literal->ToString(); const string expected = R"(f32[2,3,2] { { { 1, 2 }, { 3, 4 }, @@ -176,16 +179,14 @@ TEST_F(LiteralUtilTest, CreateR3FromArray3d) { TEST_F(LiteralUtilTest, LiteralR4F32ProjectedStringifies) { // clang-format off - auto literal = LiteralUtil::CreateR4Projected({ + auto literal = Literal::CreateR4Projected({ {1, 2}, {1001, 1002}, {2001, 2002}, }, /*projection_p=*/1, /*projection_z=*/2); // clang-format on - EXPECT_MATCH( - testing::PBToVec(literal->shape().dimensions()), - testing::VectorMatcher({1, 2, 3, 2})); - string result = LiteralUtil::ToString(*literal); + 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 @@ -204,11 +205,9 @@ TEST_F(LiteralUtilTest, LiteralR4F32ProjectedStringifies) { } TEST_F(LiteralUtilTest, LiteralR4F32Stringifies) { - EXPECT_MATCH( - testing::PBToVec( - literal_r4_2x2x3x3_dim0major_->shape().dimensions()), - testing::VectorMatcher({2, 2, 3, 3})); - string result = LiteralUtil::ToString(*literal_r4_2x2x3x3_dim0major_); + EXPECT_THAT(literal_r4_2x2x3x3_dim0major_->shape().dimensions(), + 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 @@ -240,14 +239,13 @@ TEST_F(LiteralUtilTest, LiteralR4F32Stringifies) { TEST_F(LiteralUtilTest, EachCellR2F32) { // clang-format off - auto literal = LiteralUtil::CreateR2({ + auto literal = Literal::CreateR2({ {3.1f, 4.2f}, {9.3f, 12.4f}, }); // clang-format on std::vector> seen; - LiteralUtil::EachCellAsString( - *literal, + literal->EachCellAsString( [&seen](tensorflow::gtl::ArraySlice indices, const string& value) { seen.emplace_back(indices[0], indices[1], value); }); @@ -259,167 +257,171 @@ TEST_F(LiteralUtilTest, EachCellR2F32) { } TEST_F(LiteralUtilTest, ScalarEquality) { - // Test LiteralUtil::Equal with scalars. - auto f32_42 = LiteralUtil::CreateR0(42.0); - auto f32_42_clone = LiteralUtil::CreateR0(42.0); + // Test Literal::Equal with scalars. + auto f32_42 = Literal::CreateR0(42.0); + auto f32_42_clone = Literal::CreateR0(42.0); - EXPECT_TRUE(LiteralUtil::Equal(*f32_42, *f32_42)); - EXPECT_TRUE(LiteralUtil::Equal(*f32_42, *f32_42_clone)); + EXPECT_TRUE(f32_42->Equal(*f32_42)); + EXPECT_TRUE(f32_42->Equal(*f32_42_clone)); - auto f32_123 = LiteralUtil::CreateR0(123.0); - EXPECT_FALSE(LiteralUtil::Equal(*f32_42, *f32_123)); + auto f32_123 = Literal::CreateR0(123.0); + EXPECT_FALSE(f32_42->Equal(*f32_123)); - auto f64_42 = LiteralUtil::CreateR0(42.0); - EXPECT_FALSE(LiteralUtil::Equal(*f32_42, *f64_42)); + auto f64_42 = Literal::CreateR0(42.0); + EXPECT_FALSE(f32_42->Equal(*f64_42)); } TEST_F(LiteralUtilTest, NonScalarEquality) { - // Test LiteralUtil::Equal with nonscalars. - auto matrix = LiteralUtil::CreateR2({{1.0, 2.0}, {3.0, 4.0}}); - auto matrix_clone = LiteralUtil::CreateR2({{1.0, 2.0}, {3.0, 4.0}}); - auto matrix_different = - LiteralUtil::CreateR2({{4.0, 3.0}, {1.0, 2.0}}); - auto vector_literal = LiteralUtil::CreateR1({1.0, 2.0, 3.0, 4.0}); - auto scalar = LiteralUtil::CreateR0(1.0); - - EXPECT_TRUE(LiteralUtil::Equal(*matrix, *matrix)); - EXPECT_TRUE(LiteralUtil::Equal(*matrix, *matrix_clone)); - EXPECT_FALSE(LiteralUtil::Equal(*matrix, *matrix_different)); - EXPECT_FALSE(LiteralUtil::Equal(*matrix, *vector_literal)); - EXPECT_FALSE(LiteralUtil::Equal(*matrix, *scalar)); + // Test Literal::Equal with nonscalars. + auto matrix = Literal::CreateR2({{1.0, 2.0}, {3.0, 4.0}}); + auto matrix_clone = Literal::CreateR2({{1.0, 2.0}, {3.0, 4.0}}); + auto matrix_different = Literal::CreateR2({{4.0, 3.0}, {1.0, 2.0}}); + auto vector_literal = Literal::CreateR1({1.0, 2.0, 3.0, 4.0}); + auto scalar = Literal::CreateR0(1.0); + + EXPECT_TRUE(matrix->Equal(*matrix)); + EXPECT_TRUE(matrix->Equal(*matrix_clone)); + EXPECT_FALSE(matrix->Equal(*matrix_different)); + EXPECT_FALSE(matrix->Equal(*vector_literal)); + EXPECT_FALSE(matrix->Equal(*scalar)); } TEST_F(LiteralUtilTest, DifferentLayoutEquality) { - // Test LiteralUtil::Equal with literals which have different layouts. + // Test Literal::Equal with literals which have different layouts. auto colmajor = MakeUnique(); *colmajor->mutable_shape() = ShapeUtil::MakeShape(F32, {2, 2}); *colmajor->mutable_shape()->mutable_layout() = LayoutUtil::MakeLayout({0, 1}); - LiteralUtil::Reserve(4, colmajor.get()); - LiteralUtil::Set(colmajor.get(), {0, 0}, 1.0); - LiteralUtil::Set(colmajor.get(), {0, 1}, 2.0); - LiteralUtil::Set(colmajor.get(), {1, 0}, 3.0); - LiteralUtil::Set(colmajor.get(), {1, 1}, 4.0); + colmajor->Reserve(4); + colmajor->Set({0, 0}, 1.0); + colmajor->Set({0, 1}, 2.0); + colmajor->Set({1, 0}, 3.0); + colmajor->Set({1, 1}, 4.0); auto rowmajor = MakeUnique(); *rowmajor->mutable_shape() = ShapeUtil::MakeShape(F32, {2, 2}); *rowmajor->mutable_shape()->mutable_layout() = LayoutUtil::MakeLayout({1, 0}); - LiteralUtil::Reserve(4, rowmajor.get()); - LiteralUtil::Set(rowmajor.get(), {0, 0}, 1.0); - LiteralUtil::Set(rowmajor.get(), {0, 1}, 2.0); - LiteralUtil::Set(rowmajor.get(), {1, 0}, 3.0); - LiteralUtil::Set(rowmajor.get(), {1, 1}, 4.0); + rowmajor->Reserve(4); + rowmajor->Set({0, 0}, 1.0); + rowmajor->Set({0, 1}, 2.0); + rowmajor->Set({1, 0}, 3.0); + rowmajor->Set({1, 1}, 4.0); - EXPECT_TRUE(LiteralUtil::Equal(*rowmajor, *colmajor)); + EXPECT_TRUE(rowmajor->Equal(*colmajor)); } TEST_F(LiteralUtilTest, TupleEquality) { - // Test LiteralUtil::Equal with tuples. - auto scalar = LiteralUtil::CreateR0(1.0); - auto matrix = LiteralUtil::CreateR2({{1.0, 2.0}, {3.0, 4.0}}); - auto tuple1 = LiteralUtil::MakeTuple({scalar.get(), matrix.get()}); + // Test Literal::Equal with tuples. + auto scalar = Literal::CreateR0(1.0); + auto matrix = Literal::CreateR2({{1.0, 2.0}, {3.0, 4.0}}); + auto tuple1 = Literal::MakeTuple({scalar.get(), matrix.get()}); // 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 scalar_clone = LiteralUtil::CreateR0(1.0); - auto tuple2 = LiteralUtil::MakeTuple({scalar_clone.get(), matrix.get()}); - EXPECT_TRUE(LiteralUtil::Equal(*tuple1, *tuple2)); + auto scalar_clone = Literal::CreateR0(1.0); + auto tuple2 = Literal::MakeTuple({scalar_clone.get(), matrix.get()}); + EXPECT_TRUE(tuple1->Equal(*tuple2)); // Tuple with elements reversed. - auto reversed_tuple = LiteralUtil::MakeTuple({matrix.get(), scalar.get()}); - EXPECT_FALSE(LiteralUtil::Equal(*tuple1, *reversed_tuple)); + auto reversed_tuple = Literal::MakeTuple({matrix.get(), scalar.get()}); + EXPECT_FALSE(tuple1->Equal(*reversed_tuple)); // Tuple with different value. - auto scalar_42 = LiteralUtil::CreateR0(42.0); - auto different_tuple = - LiteralUtil::MakeTuple({scalar_42.get(), matrix.get()}); - EXPECT_FALSE(LiteralUtil::Equal(*tuple1, *different_tuple)); + auto scalar_42 = Literal::CreateR0(42.0); + auto different_tuple = Literal::MakeTuple({scalar_42.get(), matrix.get()}); + EXPECT_FALSE(tuple1->Equal(*different_tuple)); } TEST_F(LiteralUtilTest, IsAllTuple) { - auto element1 = LiteralUtil::CreateR0(0.0); - auto element2 = LiteralUtil::CreateR2({{0.0, 0.0}, {0.0, 0.0}}); - auto tuple = LiteralUtil::MakeTuple({element1.get(), element1.get()}); + auto element1 = Literal::CreateR0(0.0); + auto element2 = Literal::CreateR2({{0.0, 0.0}, {0.0, 0.0}}); + auto tuple = Literal::MakeTuple({element1.get(), element1.get()}); // Tuples should always return false for IsAll. - EXPECT_FALSE(LiteralUtil::IsAll(*tuple, 0)); - EXPECT_FALSE(LiteralUtil::IsAll(*tuple, 1)); + EXPECT_FALSE(tuple->IsAll(0)); + EXPECT_FALSE(tuple->IsAll(1)); +} + +// Verifies that CreateFromShape works for tuples. +TEST_F(LiteralUtilTest, CreateFromShapeTuple) { + auto scalar = Literal::CreateR0(0.0); + auto matrix = Literal::CreateR2({{0, 0}, {0, 0}}); + auto tuple = Literal::MakeTuple({scalar.get(), matrix.get()}); + + auto x = Literal::CreateFromShape(tuple->shape()); + EXPECT_TRUE(tuple->Equal(*x)); } TEST_F(LiteralUtilTest, IsAll) { - EXPECT_TRUE(LiteralUtil::IsAll(*LiteralUtil::CreateR0(false), 0)); - EXPECT_TRUE(LiteralUtil::IsAll(*LiteralUtil::CreateR0(true), 1)); - EXPECT_FALSE(LiteralUtil::IsAll(*LiteralUtil::CreateR0(false), 1)); - EXPECT_FALSE(LiteralUtil::IsAll(*LiteralUtil::CreateR0(false), 2)); - EXPECT_FALSE(LiteralUtil::IsAll(*LiteralUtil::CreateR0(true), 0)); - EXPECT_FALSE(LiteralUtil::IsAll(*LiteralUtil::CreateR0(true), 2)); - EXPECT_FALSE(LiteralUtil::IsAll(*LiteralUtil::CreateR0(true), -1)); + EXPECT_TRUE(Literal::CreateR0(false)->IsAll(0)); + EXPECT_TRUE(Literal::CreateR0(true)->IsAll(1)); + EXPECT_FALSE(Literal::CreateR0(false)->IsAll(1)); + EXPECT_FALSE(Literal::CreateR0(false)->IsAll(2)); + EXPECT_FALSE(Literal::CreateR0(true)->IsAll(0)); + EXPECT_FALSE(Literal::CreateR0(true)->IsAll(2)); + EXPECT_FALSE(Literal::CreateR0(true)->IsAll(-1)); // We shouldn't reinterpret int8_min as an unsigned type and then decide that // it is equal to 255. auto int8_min = std::numeric_limits::min(); - EXPECT_FALSE( - LiteralUtil::IsAll(*LiteralUtil::CreateR0(255), int8_min)); + EXPECT_FALSE(Literal::CreateR0(255)->IsAll(int8_min)); - EXPECT_TRUE(LiteralUtil::IsAll(*LiteralUtil::CreateR0(42.0), 42)); - EXPECT_FALSE(LiteralUtil::IsAll(*LiteralUtil::CreateR0(42.0001), 42)); + EXPECT_TRUE(Literal::CreateR0(42.0)->IsAll(42)); + EXPECT_FALSE(Literal::CreateR0(42.0001)->IsAll(42)); - EXPECT_TRUE( - LiteralUtil::IsAll(*LiteralUtil::CreateR1({100, 100, 100}), 100)); - EXPECT_FALSE(LiteralUtil::IsAll( - *LiteralUtil::CreateR1({100, 100, 100.001}), 100)); + EXPECT_TRUE(Literal::CreateR1({100, 100, 100})->IsAll(100)); + EXPECT_FALSE(Literal::CreateR1({100, 100, 100.001})->IsAll(100)); - EXPECT_TRUE( - LiteralUtil::IsAll(*LiteralUtil::CreateR2({{8, 8}, {8, 8}}), 8)); - EXPECT_FALSE( - LiteralUtil::IsAll(*LiteralUtil::CreateR2({{8, 8}, {8, 9}}), 8)); - EXPECT_FALSE( - LiteralUtil::IsAll(*LiteralUtil::CreateR2({{9, 8}, {8, 8}}), 8)); + EXPECT_TRUE(Literal::CreateR2({{8, 8}, {8, 8}})->IsAll(8)); + EXPECT_FALSE(Literal::CreateR2({{8, 8}, {8, 9}})->IsAll(8)); + EXPECT_FALSE(Literal::CreateR2({{9, 8}, {8, 8}})->IsAll(8)); + + half h8(8.0f); + half h9(9.0f); + EXPECT_TRUE(Literal::CreateR2({{h8}, {h8}})->IsAll(8)); + EXPECT_FALSE(Literal::CreateR2({{h8}, {h9}})->IsAll(8)); + EXPECT_FALSE(Literal::CreateR2({{h9}, {h8}})->IsAll(8)); auto uint64_max = std::numeric_limits::max(); - EXPECT_FALSE(LiteralUtil::IsAll( - *LiteralUtil::CreateR2( - {{uint64_max, uint64_max}, {uint64_max, uint64_max}}), - -1)); + EXPECT_FALSE(Literal::CreateR2( + {{uint64_max, uint64_max}, {uint64_max, uint64_max}}) + ->IsAll(-1)); } TEST_F(LiteralUtilTest, IsAllFloat) { // IsAllFloat always returns false when the literal is not floating-point. - EXPECT_FALSE(LiteralUtil::IsAllFloat(*LiteralUtil::CreateR0(false), 0)); - EXPECT_FALSE(LiteralUtil::IsAllFloat(*LiteralUtil::CreateR0(0), 0)); - EXPECT_FALSE(LiteralUtil::IsAllFloat(*LiteralUtil::CreateR0(0), 0)); - EXPECT_FALSE(LiteralUtil::IsAllFloat(*LiteralUtil::CreateR0(0), 0)); - - EXPECT_TRUE(LiteralUtil::IsAllFloat(*LiteralUtil::CreateR0(0), 0)); - EXPECT_TRUE(LiteralUtil::IsAllFloat(*LiteralUtil::CreateR0(.5), .5)); - EXPECT_TRUE(LiteralUtil::IsAllFloat(*LiteralUtil::CreateR0(-.5), -.5)); + EXPECT_FALSE(Literal::CreateR0(false)->IsAllFloat(0)); + EXPECT_FALSE(Literal::CreateR0(0)->IsAllFloat(0)); + EXPECT_FALSE(Literal::CreateR0(0)->IsAllFloat(0)); + EXPECT_FALSE(Literal::CreateR0(0)->IsAllFloat(0)); + + EXPECT_TRUE(Literal::CreateR0(0)->IsAllFloat(0)); + EXPECT_TRUE(Literal::CreateR0(.5)->IsAllFloat(.5)); + EXPECT_TRUE(Literal::CreateR0(-.5)->IsAllFloat(-.5)); + EXPECT_FALSE(Literal::CreateR0(-.5)->IsAllFloat(-.49)); EXPECT_FALSE( - LiteralUtil::IsAllFloat(*LiteralUtil::CreateR0(-.5), -.49)); - EXPECT_FALSE(LiteralUtil::IsAllFloat( - *LiteralUtil::CreateR2({{0, 0, 0}, {0, .1, 0}}), 0)); - EXPECT_TRUE(LiteralUtil::IsAllFloat( - *LiteralUtil::CreateR2({{.5, .5, .5}, {.5, .5, .5}}), .5)); - - EXPECT_TRUE(LiteralUtil::IsAllFloat(*LiteralUtil::CreateR0(0), 0)); - EXPECT_TRUE(LiteralUtil::IsAllFloat(*LiteralUtil::CreateR0(.5), .5)); + Literal::CreateR2({{0, 0, 0}, {0, .1, 0}})->IsAllFloat(0)); EXPECT_TRUE( - LiteralUtil::IsAllFloat(*LiteralUtil::CreateR0(-.5), -.5)); + Literal::CreateR2({{.5, .5, .5}, {.5, .5, .5}})->IsAllFloat(.5)); + + EXPECT_TRUE(Literal::CreateR0(0)->IsAllFloat(0)); + EXPECT_TRUE(Literal::CreateR0(.5)->IsAllFloat(.5)); + EXPECT_TRUE(Literal::CreateR0(-.5)->IsAllFloat(-.5)); + EXPECT_FALSE(Literal::CreateR0(-.5)->IsAllFloat(-.49)); EXPECT_FALSE( - LiteralUtil::IsAllFloat(*LiteralUtil::CreateR0(-.5), -.49)); - EXPECT_FALSE(LiteralUtil::IsAllFloat( - *LiteralUtil::CreateR2({{0, 0, 0}, {0, .1, 0}}), 0)); + Literal::CreateR2({{0, 0, 0}, {0, .1, 0}})->IsAllFloat(0)); } TEST_F(LiteralUtilTest, IsZero) { - auto scalar_zero = LiteralUtil::CreateR0(0.0f); - auto scalar_one = LiteralUtil::CreateR0(1.0f); - EXPECT_TRUE(LiteralUtil::IsZero(*scalar_zero, {})); - EXPECT_FALSE(LiteralUtil::IsZero(*scalar_one, {})); - - auto array = LiteralUtil::CreateR2({{1, 2, 0, 3}, {1, 0, 1, 2}}); - EXPECT_FALSE(LiteralUtil::IsZero(*array, {0, 1})); - EXPECT_TRUE(LiteralUtil::IsZero(*array, {0, 2})); - EXPECT_TRUE(LiteralUtil::IsZero(*array, {1, 1})); - EXPECT_FALSE(LiteralUtil::IsZero(*array, {1, 2})); + auto scalar_zero = Literal::CreateR0(0.0f); + auto scalar_one = Literal::CreateR0(1.0f); + EXPECT_TRUE(scalar_zero->IsZero({})); + EXPECT_FALSE(scalar_one->IsZero({})); + + auto array = Literal::CreateR2({{1, 2, 0, 3}, {1, 0, 1, 2}}); + EXPECT_FALSE(array->IsZero({0, 1})); + EXPECT_TRUE(array->IsZero({0, 2})); + EXPECT_TRUE(array->IsZero({1, 1})); + EXPECT_FALSE(array->IsZero({1, 2})); } template @@ -431,112 +433,123 @@ TYPED_TEST_CASE(LiteralUtilTestTemplated, TestedTypes); TYPED_TEST(LiteralUtilTestTemplated, Relayout2x2) { // Make a non-integer for floating point types. TypeParam half = TypeParam(1) / TypeParam(2); - auto data = LiteralUtil::CreateR2({{half, 2}, {3, 4}}); + auto data = Literal::CreateR2({{half, 2}, {3, 4}}); const Layout layout01 = LayoutUtil::MakeLayout({0, 1}); const Layout layout10 = LayoutUtil::MakeLayout({1, 0}); - auto data01 = LiteralUtil::Relayout(*data, layout01); + auto data01 = data->Relayout(layout01); EXPECT_TRUE(LayoutUtil::Equal(data01->shape().layout(), layout01)); - EXPECT_TRUE(LiteralUtil::Equal(*data, *data01)); + EXPECT_TRUE(data->Equal(*data01)); - auto data10 = LiteralUtil::Relayout(*data, layout10); + auto data10 = data->Relayout(layout10); EXPECT_TRUE(LayoutUtil::Equal(data10->shape().layout(), layout10)); - EXPECT_TRUE(LiteralUtil::Equal(*data, *data10)); + EXPECT_TRUE(data->Equal(*data10)); } TEST_F(LiteralUtilTest, ReshapeR0) { - auto original = LiteralUtil::CreateR0(1.7f); - auto reshape = - LiteralUtil::Reshape(*original, /*shape=*/{}).ConsumeValueOrDie(); - EXPECT_TRUE(LiteralUtil::Equal(*original, *reshape)); + auto original = Literal::CreateR0(1.7f); + auto reshape = original->Reshape(/*shape=*/{}).ConsumeValueOrDie(); + EXPECT_TRUE(original->Equal(*reshape)); } TEST_F(LiteralUtilTest, ReshapeR4) { // clang-format off // F32[1x3x2x4] - auto original = LiteralUtil::CreateR4WithLayout({{ + auto original = Literal::CreateR4WithLayout({{ {{10, 11, 12, 13}, {14, 15, 16, 17}}, {{18, 19, 20, 21}, {22, 23, 24, 25}}, {{26, 27, 28, 29}, {30, 31, 32, 33}}, }}, layout_r4_dim0major_); // F32[1x3x4x2] - auto expected = LiteralUtil::CreateR3WithLayout({ + auto expected = Literal::CreateR3WithLayout({ {{10, 11}, {12, 13}, {14, 15}, {16, 17}}, {{18, 19}, {20, 21}, {22, 23}, {24, 25}}, {{26, 27}, {28, 29}, {30, 31}, {32, 33}}, }, layout_r3_dim0major_); // clang-format on - auto reshape = LiteralUtil::Reshape(*original, {3, 4, 2}).ConsumeValueOrDie(); + auto reshape = original->Reshape({3, 4, 2}).ConsumeValueOrDie(); - EXPECT_TRUE(LiteralUtil::Equal(*expected, *reshape)); + EXPECT_TRUE(expected->Equal(*reshape)); +} + +TEST_F(LiteralUtilTest, ReshapeR4Dim0Minor) { + // clang-format off + // F32[1x3x2x4] + auto original = Literal::CreateR4WithLayout({{ + {{10, 11, 12, 13}, {14, 15, 16, 17}}, + {{18, 19, 20, 21}, {22, 23, 24, 25}}, + {{26, 27, 28, 29}, {30, 31, 32, 33}}, + }}, layout_r4_dim0minor_); + // F32[1x3x4x2] + auto expected = Literal::CreateR3WithLayout({ + {{10, 11}, {12, 13}, {14, 15}, {16, 17}}, + {{18, 19}, {20, 21}, {22, 23}, {24, 25}}, + {{26, 27}, {28, 29}, {30, 31}, {32, 33}}, + }, layout_r3_dim0major_); + // clang-format on + auto reshape = original->Reshape({3, 4, 2}).ConsumeValueOrDie(); + + EXPECT_TRUE(expected->Equal(*reshape)); } TEST_F(LiteralUtilTest, TransposeR0) { - auto original = LiteralUtil::CreateR0(1.7f); - auto reshape = LiteralUtil::Transpose(*original, /*permutation=*/{}); - EXPECT_TRUE(LiteralUtil::Equal(*original, *reshape)); + auto original = Literal::CreateR0(1.7f); + auto reshape = original->Transpose(/*permutation=*/{}); + EXPECT_TRUE(original->Equal(*reshape)); } TEST_F(LiteralUtilTest, TransposeR4) { // clang-format off // F32[1x3x2x4] - auto original = LiteralUtil::CreateR4({{ + auto original = Literal::CreateR4({{ {{10, 11, 12, 13}, {14, 15, 16, 17}}, {{18, 19, 20, 21}, {22, 23, 24, 25}}, {{26, 27, 28, 29}, {30, 31, 32, 33}}, }}); // clang-format on - auto reshape = - LiteralUtil::Transpose(*original, /*permutation=*/{2, 3, 0, 1}); - - LiteralUtil::EachCell( - *reshape, [&](tensorflow::gtl::ArraySlice indices, float value) { - EXPECT_EQ(value, - LiteralUtil::Get(*original, {indices[2], indices[3], - indices[0], indices[1]})); + auto reshape = original->Transpose(/*permutation=*/{2, 3, 0, 1}); + + reshape->EachCell( + [&](tensorflow::gtl::ArraySlice indices, float value) { + EXPECT_EQ(value, original->Get( + {indices[2], indices[3], indices[0], indices[1]})); }); } TEST_F(LiteralUtilTest, TestR4RelayoutEquivalence) { // Tests that using Relayout on an array is equivalent to creating it in the // target layout in the first place. - auto dim0minor_relaid_to_dim0major = LiteralUtil::Relayout( - *literal_r4_2x2x3x3_dim0minor_, layout_r4_dim0major_); - EXPECT_TRUE(LiteralUtil::Equal(*literal_r4_2x2x3x3_dim0major_, - *dim0minor_relaid_to_dim0major)); + auto dim0minor_relaid_to_dim0major = + literal_r4_2x2x3x3_dim0minor_->Relayout(layout_r4_dim0major_); + EXPECT_TRUE( + literal_r4_2x2x3x3_dim0major_->Equal(*dim0minor_relaid_to_dim0major)); - auto dim0major_relaid_to_dim0minor = LiteralUtil::Relayout( - *literal_r4_2x2x3x3_dim0major_, layout_r4_dim0minor_); - EXPECT_TRUE(LiteralUtil::Equal(*literal_r4_2x2x3x3_dim0minor_, - *dim0major_relaid_to_dim0minor)); + auto dim0major_relaid_to_dim0minor = + literal_r4_2x2x3x3_dim0major_->Relayout(layout_r4_dim0minor_); + EXPECT_TRUE( + literal_r4_2x2x3x3_dim0minor_->Equal(*dim0major_relaid_to_dim0minor)); } TEST_F(LiteralUtilTest, TestR2LinearLayout) { // Test expected memory layout of R2 dim0-minor (column-major) literal. - auto mat_dim0minor = LiteralUtil::CreateR2WithLayout( - {{1, 2, 3}, {4, 5, 6}}, layout_r2_dim0minor_); + auto mat_dim0minor = Literal::CreateR2WithLayout({{1, 2, 3}, {4, 5, 6}}, + layout_r2_dim0minor_); EXPECT_EQ(mat_dim0minor->s32s_size(), 6); - EXPECT_MATCH(testing::PBToVec(mat_dim0minor->s32s()), - testing::VectorMatcher({1, 4, 2, 5, 3, 6})); + EXPECT_THAT(mat_dim0minor->s32s(), ElementsAre(1, 4, 2, 5, 3, 6)); // Test expected memory layout when using Relayout to row major. - auto relaid_mat_to_dim0major = - LiteralUtil::Relayout(*mat_dim0minor, layout_r2_dim0major_); - EXPECT_MATCH(testing::PBToVec(relaid_mat_to_dim0major->s32s()), - testing::VectorMatcher({1, 2, 3, 4, 5, 6})); + auto relaid_mat_to_dim0major = mat_dim0minor->Relayout(layout_r2_dim0major_); + EXPECT_THAT(relaid_mat_to_dim0major->s32s(), ElementsAre(1, 2, 3, 4, 5, 6)); // Test expected memory layout of R2 created with dim0-major (row-major). - auto mat_dim0major = LiteralUtil::CreateR2WithLayout( - {{1, 2, 3}, {4, 5, 6}}, layout_r2_dim0major_); + auto mat_dim0major = Literal::CreateR2WithLayout({{1, 2, 3}, {4, 5, 6}}, + layout_r2_dim0major_); EXPECT_EQ(mat_dim0major->s32s_size(), 6); - EXPECT_MATCH(testing::PBToVec(mat_dim0major->s32s()), - testing::VectorMatcher({1, 2, 3, 4, 5, 6})); + EXPECT_THAT(mat_dim0major->s32s(), ElementsAre(1, 2, 3, 4, 5, 6)); // Test expected memory layout when using Relayout to column major. - auto relaid_mat_to_dim0minor = - LiteralUtil::Relayout(*mat_dim0major, layout_r2_dim0minor_); - EXPECT_MATCH(testing::PBToVec(relaid_mat_to_dim0minor->s32s()), - testing::VectorMatcher({1, 4, 2, 5, 3, 6})); + auto relaid_mat_to_dim0minor = mat_dim0major->Relayout(layout_r2_dim0minor_); + EXPECT_THAT(relaid_mat_to_dim0minor->s32s(), ElementsAre(1, 4, 2, 5, 3, 6)); } TEST_F(LiteralUtilTest, TestR3LinearLayout) { @@ -553,107 +566,458 @@ TEST_F(LiteralUtilTest, TestR3LinearLayout) { {10, 11, 12}, }, }); // clang-format on - auto lit_dim0minor = LiteralUtil::CreateR3FromArray3DWithLayout( - arr3d, layout_r3_dim0minor_); + auto lit_dim0minor = + Literal::CreateR3FromArray3DWithLayout(arr3d, layout_r3_dim0minor_); EXPECT_EQ(lit_dim0minor->s32s_size(), 12); std::vector expected_dim0minor{1, 7, 4, 10, 2, 8, 5, 11, 3, 9, 6, 12}; - EXPECT_MATCH(testing::PBToVec(lit_dim0minor->s32s()), - testing::VectorMatcher(expected_dim0minor)); + EXPECT_THAT(lit_dim0minor->s32s(), + testing::ElementsAreArray(expected_dim0minor)); // Test expected memory layout when using Relayout to row major. - auto relaid_lit_to_dim0major = - LiteralUtil::Relayout(*lit_dim0minor, layout_r3_dim0major_); + auto relaid_lit_to_dim0major = lit_dim0minor->Relayout(layout_r3_dim0major_); std::vector expected_dim0major{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12}; - EXPECT_MATCH(testing::PBToVec(relaid_lit_to_dim0major->s32s()), - testing::VectorMatcher(expected_dim0major)); + EXPECT_THAT(relaid_lit_to_dim0major->s32s(), + testing::ElementsAreArray(expected_dim0major)); // Test expected memory layout of R3 created with dim0-major (row-major). - auto lit_dim0major = LiteralUtil::CreateR3FromArray3DWithLayout( - arr3d, layout_r3_dim0major_); + auto lit_dim0major = + Literal::CreateR3FromArray3DWithLayout(arr3d, layout_r3_dim0major_); EXPECT_EQ(lit_dim0major->s32s_size(), 12); - EXPECT_MATCH(testing::PBToVec(lit_dim0major->s32s()), - testing::VectorMatcher(expected_dim0major)); + EXPECT_THAT(lit_dim0major->s32s(), + testing::ElementsAreArray(expected_dim0major)); // Test expected memory layout when using Relayout to column major. - auto relaid_lit_to_dim0minor = - LiteralUtil::Relayout(*lit_dim0major, layout_r3_dim0minor_); - EXPECT_MATCH(testing::PBToVec(relaid_lit_to_dim0minor->s32s()), - testing::VectorMatcher(expected_dim0minor)); + auto relaid_lit_to_dim0minor = lit_dim0major->Relayout(layout_r3_dim0minor_); + EXPECT_THAT(relaid_lit_to_dim0minor->s32s(), + testing::ElementsAreArray(expected_dim0minor)); } TEST_F(LiteralUtilTest, SliceR0S32) { - auto input = LiteralUtil::CreateR0(1); - auto result = LiteralUtil::Slice(*input, {}, {}); - EXPECT_TRUE(LiteralUtil::Equal(*input, *result)); + auto input = Literal::CreateR0(1); + auto result = input->Slice({}, {}); + EXPECT_TRUE(input->Equal(*result)); } TEST_F(LiteralUtilTest, SliceR1F32) { - auto input = LiteralUtil::CreateR1({1.0, 2.0, 3.0, 4.0, 5.0}); - auto result = LiteralUtil::Slice(*input, {3}, {4}); - auto expected = LiteralUtil::CreateR1({4.0}); - EXPECT_TRUE(LiteralUtil::Equal(*expected, *result)); + auto input = Literal::CreateR1({1.0, 2.0, 3.0, 4.0, 5.0}); + auto result = input->Slice({3}, {4}); + auto expected = Literal::CreateR1({4.0}); + EXPECT_TRUE(expected->Equal(*result)); } TEST_F(LiteralUtilTest, SliceR2U32) { - auto input_3x4 = LiteralUtil::CreateR2( - {{1, 2, 3, 4}, {5, 6, 7, 8}, {9, 10, 11, 12}}); - auto result = LiteralUtil::Slice(*input_3x4, {0, 2}, {2, 4}); - auto expected = LiteralUtil::CreateR2({{3, 4}, {7, 8}}); - EXPECT_TRUE(LiteralUtil::Equal(*expected, *result)); + auto input_3x4 = + Literal::CreateR2({{1, 2, 3, 4}, {5, 6, 7, 8}, {9, 10, 11, 12}}); + auto result = input_3x4->Slice({0, 2}, {2, 4}); + auto expected = Literal::CreateR2({{3, 4}, {7, 8}}); + EXPECT_TRUE(expected->Equal(*result)); } TEST_F(LiteralUtilTest, SliceR3U32Full) { - auto input_2x3x2 = LiteralUtil::CreateR3( + auto input_2x3x2 = Literal::CreateR3( {{{1, 2}, {3, 4}, {5, 6}}, {{7, 8}, {9, 10}, {11, 12}}}); - auto result = LiteralUtil::Slice(*input_2x3x2, {0, 0, 0}, {2, 3, 2}); - EXPECT_TRUE(LiteralUtil::Equal(*input_2x3x2, *result)); + auto result = input_2x3x2->Slice({0, 0, 0}, {2, 3, 2}); + EXPECT_TRUE(input_2x3x2->Equal(*result)); } TEST_F(LiteralUtilTest, PopulateR1S64) { Literal output; - LiteralUtil::PopulateR1({77}, &output); - auto expected = LiteralUtil::CreateR1({77}); - EXPECT_TRUE(LiteralUtil::Equal(output, *expected)); + output.PopulateR1({77}); + auto expected = Literal::CreateR1({77}); + EXPECT_TRUE(output.Equal(*expected)); } TEST_F(LiteralUtilTest, PopulateR2U64) { Literal output; - LiteralUtil::PopulateR1({{77, 88}}, &output); - auto expected = LiteralUtil::CreateR1({{77, 88}}); - EXPECT_TRUE(LiteralUtil::Equal(output, *expected)); + output.PopulateR1({{77, 88}}); + auto expected = Literal::CreateR1({{77, 88}}); + EXPECT_TRUE(output.Equal(*expected)); } TEST_F(LiteralUtilTest, PopulateWithValueR0F32) { Literal output; - LiteralUtil::PopulateWithValue(2.5f, {}, &output); - auto expected = LiteralUtil::CreateR0(2.5f); - EXPECT_TRUE(LiteralUtil::Equal(output, *expected)); + output.PopulateWithValue(2.5f, {}); + auto expected = Literal::CreateR0(2.5f); + EXPECT_TRUE(output.Equal(*expected)); } TEST_F(LiteralUtilTest, PopulateWithValueR1S64) { Literal output; - LiteralUtil::PopulateWithValue(-7, {3}, &output); - auto expected = LiteralUtil::CreateR1({-7, -7, -7}); - EXPECT_TRUE(LiteralUtil::Equal(output, *expected)); + output.PopulateWithValue(-7, {3}); + auto expected = Literal::CreateR1({-7, -7, -7}); + EXPECT_TRUE(output.Equal(*expected)); } TEST_F(LiteralUtilTest, PopulateWithValueR2U64) { Literal output; - LiteralUtil::PopulateWithValue(42, {2, 2}, &output); - auto expected = LiteralUtil::CreateR2({{42, 42}, {42, 42}}); - EXPECT_TRUE(LiteralUtil::Equal(output, *expected)); + output.PopulateWithValue(42, {2, 2}); + auto expected = Literal::CreateR2({{42, 42}, {42, 42}}); + EXPECT_TRUE(output.Equal(*expected)); +} + +TEST_F(LiteralUtilTest, PopulateWithValueR0F16) { + Literal output; + half h(0.25f); + output.PopulateWithValue(h, {}); + auto expected = Literal::CreateR0(h); + EXPECT_TRUE(output.Equal(*expected)); +} + +TEST_F(LiteralUtilTest, PopulateWithValueR1F16) { + Literal output; + half h(0.5f); + output.PopulateWithValue(h, {3}); + auto expected = Literal::CreateR1({h, h, h}); + EXPECT_TRUE(output.Equal(*expected)); +} + +TEST_F(LiteralUtilTest, PopulateWithValueR2F16) { + Literal output; + half h(2.0f); + output.PopulateWithValue(h, {2, 2}); + auto expected = Literal::CreateR2({{h, h}, {h, h}}); + EXPECT_TRUE(output.Equal(*expected)); } TEST_F(LiteralUtilTest, ReplicateR2U32) { - auto input = LiteralUtil::CreateR2( - {{1, 2, 3, 4}, {5, 6, 7, 8}, {9, 10, 11, 12}}); - auto output = LiteralUtil::Replicate(*input, 3); - auto expected = LiteralUtil::CreateR3( + auto input = + Literal::CreateR2({{1, 2, 3, 4}, {5, 6, 7, 8}, {9, 10, 11, 12}}); + auto output = input->Replicate(3); + auto expected = Literal::CreateR3( {{{1, 2, 3, 4}, {5, 6, 7, 8}, {9, 10, 11, 12}}, {{1, 2, 3, 4}, {5, 6, 7, 8}, {9, 10, 11, 12}}, {{1, 2, 3, 4}, {5, 6, 7, 8}, {9, 10, 11, 12}}}); - EXPECT_TRUE(LiteralUtil::Equal(*output, *expected)); + EXPECT_TRUE(output->Equal(*expected)); +} + +TEST_F(LiteralUtilTest, Copy) { + const int64 dimensions[] = {17, 15, 34, 21}; + const int64 layouts[][4] = { + {3, 2, 1, 0}, {0, 2, 1, 3}, {0, 1, 2, 3}, {2, 0, 3, 1}, {1, 3, 0, 2}}; + for (const auto& layout : layouts) { + Shape shape = ShapeUtil::MakeShapeWithLayout( + primitive_util::NativeToPrimitiveType(), dimensions, layout); + auto blank = Literal::CreateFromShape(shape); + auto source = Literal::CreateFromShape(shape); + const int64 zero_base[] = {0, 0, 0, 0}; + const int64 step[] = {1, 1, 1, 1}; + uint32 seqnr = 0; + auto init_proc = [&](const std::vector& indexes) { + source->Set(indexes, ++seqnr); + return true; + }; + + ShapeUtil::ForEachIndex(source->shape(), zero_base, dimensions, step, + init_proc); + + const int64 src_base[] = {3, 1, 5, 7}; + const int64 dest_base[] = {6, 4, 12, 2}; + const int64 copy_size[] = {7, 8, 11, 9}; + + TF_EXPECT_OK(blank->Copy(*source, src_base, dest_base, copy_size)); + std::vector source_indexes(TF_ARRAYSIZE(dimensions), 0); + std::vector blank_indexes(TF_ARRAYSIZE(dimensions), 0); + bool matched = true; + auto check_proc = [&](const std::vector& indexes) { + std::copy(indexes.begin(), indexes.end(), source_indexes.begin()); + std::transform(source_indexes.begin(), source_indexes.end(), src_base, + source_indexes.begin(), std::plus()); + std::copy(indexes.begin(), indexes.end(), blank_indexes.begin()); + std::transform(blank_indexes.begin(), blank_indexes.end(), dest_base, + blank_indexes.begin(), std::plus()); + auto bval = blank->Get(blank_indexes); + matched = (bval != 0 && bval == source->Get(source_indexes)); + return matched; + }; + ShapeUtil::ForEachIndex(source->shape(), zero_base, copy_size, step, + check_proc); + EXPECT_TRUE(matched); + } +} + +TEST_F(LiteralUtilTest, CopyScalars) { + auto zero = Literal::CreateR0(0); + auto nine = Literal::CreateR0(9); + TF_EXPECT_OK(zero->Copy(*nine, {}, {}, {})); + EXPECT_TRUE(zero->Equal(*nine)); + + auto vect = Literal::CreateR1({3, 4, 9, 12, 5, 17, 21}); + TF_EXPECT_OK(zero->Copy(*vect, {5}, {}, {})); + EXPECT_EQ(zero->Get({}), 17); + TF_EXPECT_OK(vect->Copy(*zero, {}, {4}, {})); + EXPECT_EQ(vect->Get({4}), 17); +} + +TEST_F(LiteralUtilTest, F16) { + // Verify that the internal data views are consistent and that they + // are in little endian format + // TODO - modify if we make the data format machine endianess dependent + auto m1 = Literal::CreateFromShape(ShapeUtil::MakeShape(F16, {2, 2})); + Literal* l1 = m1.get(); + const char* d1 = static_cast(l1->InternalData()); + EXPECT_EQ(d1[0], 0); + EXPECT_EQ(d1[1], 0); + EXPECT_EQ(d1[2], 0); + EXPECT_EQ(d1[3], 0); + EXPECT_EQ(d1[4], 0); + EXPECT_EQ(d1[5], 0); + EXPECT_EQ(d1[6], 0); + EXPECT_EQ(d1[7], 0); + EXPECT_EQ(l1->InternalData(), l1->MutableInternalData()); + + half h1(1.0f); + half h2(2.0f); + auto m2 = Literal::CreateR2({{h1, h2}, {h2, h1}}); + Literal* l2 = m2.get(); + const char* d2 = static_cast(l2->InternalData()); + EXPECT_EQ(d2[0], 0); + EXPECT_EQ(d2[1], 0x3C); + EXPECT_EQ(d2[2], 0); + EXPECT_EQ(d2[3], 0x40); + EXPECT_EQ(d2[4], 0); + EXPECT_EQ(d2[5], 0x40); + EXPECT_EQ(d2[6], 0); + EXPECT_EQ(d2[7], 0x3C); + EXPECT_EQ(l2->InternalData(), l2->MutableInternalData()); +} + +TEST_F(LiteralUtilTest, Populate) { + struct PopulateData { + std::vector dimensions; + std::vector layout; + } populate_data[] = { + {{}, {}}, + {{0}, {0}}, + {{16}, {0}}, + {{2, 0}, {1, 0}}, + {{4, 16}, {1, 0}}, + {{21, 12}, {0, 1}}, + {{6, 11, 17}, {2, 0, 1}}, + {{6, 11, 5, 17}, {3, 2, 0, 1}}, + }; + for (const auto& data : populate_data) { + Shape shape = ShapeUtil::MakeShapeWithLayout( + primitive_util::NativeToPrimitiveType(), data.dimensions, + data.layout); + auto literal = Literal::CreateFromShape(shape); + auto generator = [&](tensorflow::gtl::ArraySlice indexes) -> uint32 { + // Offsets from linear index just to avoid R0 literals to be initialized + // with zero. + return literal->LinearIndex(indexes) + 17; + }; + TF_EXPECT_OK(literal->Populate(generator)); + + std::vector zero_base(data.dimensions.size(), 0); + std::vector step(data.dimensions.size(), 1); + bool matched = true; + auto check_function = [&](const std::vector& indexes) { + auto value = literal->Get(indexes); + matched = matched && (value == generator(indexes)); + return matched; + }; + ShapeUtil::ForEachIndex(literal->shape(), zero_base, data.dimensions, step, + check_function); + EXPECT_TRUE(matched); + } +} + +TEST_F(LiteralUtilTest, ConvertR4) { + // clang-format off + auto original = Literal::CreateR4WithLayout({{ + {{10, 11, 12, 13}, {14, 15, 16, 17}}, + {{18, 19, 20, 21}, {22, 23, 24, 25}}, + {{26, 27, 28, 29}, {30, 31, 32, 33}}, + }}, layout_r4_dim0major_); + auto expected = Literal::CreateR4WithLayout({{ + {{10, 11, 12, 13}, {14, 15, 16, 17}}, + {{18, 19, 20, 21}, {22, 23, 24, 25}}, + {{26, 27, 28, 29}, {30, 31, 32, 33}}, + }}, layout_r4_dim0major_); + // clang-format on + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr converted, + original->Convert(U32)); + + EXPECT_TRUE(expected->Equal(*converted)); +} + +TEST_F(LiteralUtilTest, ConvertIfTypesMatch) { + // clang-format off + auto s8 = Literal::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 = Literal::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 = Literal::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 s64 = Literal::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 u64 = Literal::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 pred = Literal::CreateR4WithLayout({{ + {{true, false, true, false}, {false, true, false, true}}, + {{false, true, false, true}, {true, false, true, false}}, + {{true, false, true, false}, {false, true, false, true}}, + }}, layout_r4_dim0major_); + auto int32_pred = Literal::CreateR4WithLayout({{ + {{1, 0, 1, 0}, {0, 1, 0, 1}}, + {{0, 1, 0, 1}, {1, 0, 1, 0}}, + {{1, 0, 1, 0}, {0, 1, 0, 1}}, + }}, layout_r4_dim0major_); + auto f16 = Literal::CreateR4WithLayout({{ + {{half(10.0), half(0.0), half(12.0), half(0.0)}, + {half(0.0), half(15.0), half(0.0), half(17.0)}}, + {{half(0.0), half(19.0), half(0.0), half(21.0)}, + {half(22.0), half(0.0), half(24.0), half(0.0)}}, + {{half(26.0), half(0.0), half(28.0), half(0.0)}, + {half(0.0), half(31.0), half(0.0), half(33.0)}}, + }}, layout_r4_dim0major_); + auto f32 = Literal::CreateR4WithLayout({{ + {{10.0f, 0.0f, 12.0f, 0.0f}, {0.0f, 15.0f, 0.0f, 17.0f}}, + {{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_); + auto f64 = Literal::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 + std::unique_ptr conv; + + conv = s8->Convert(U32).ConsumeValueOrDie(); + EXPECT_TRUE(conv->Equal(*u32)); + + conv = s8->Convert(S32).ConsumeValueOrDie(); + EXPECT_TRUE(conv->Equal(*s32)); + + conv = s8->Convert(U64).ConsumeValueOrDie(); + EXPECT_TRUE(conv->Equal(*u64)); + + conv = s8->Convert(S64).ConsumeValueOrDie(); + EXPECT_TRUE(conv->Equal(*s64)); + + conv = s8->Convert(PRED).ConsumeValueOrDie(); + EXPECT_TRUE(conv->Equal(*pred)); + + conv = pred->Convert(S32).ConsumeValueOrDie(); + EXPECT_TRUE(conv->Equal(*int32_pred)); + + conv = f32->Convert(S32).ConsumeValueOrDie(); + EXPECT_TRUE(conv->Equal(*s32)); + + conv = f64->Convert(S32).ConsumeValueOrDie(); + EXPECT_TRUE(conv->Equal(*s32)); + + conv = s32->Convert(F32).ConsumeValueOrDie(); + EXPECT_TRUE(conv->Equal(*f32)); + + conv = f32->Convert(F16).ConsumeValueOrDie(); + EXPECT_TRUE(conv->Equal(*f16)); + + conv = f64->Convert(F16).ConsumeValueOrDie(); + EXPECT_TRUE(conv->Equal(*f16)); + + conv = s32->Convert(F16).ConsumeValueOrDie(); + EXPECT_TRUE(conv->Equal(*f16)); + + conv = u32->Convert(F16).ConsumeValueOrDie(); + EXPECT_TRUE(conv->Equal(*f16)); + + EXPECT_EQ(s32->Convert(TUPLE).status().code(), + tensorflow::error::INVALID_ARGUMENT); + EXPECT_EQ(s32->Convert(S16).status().code(), + tensorflow::error::INVALID_ARGUMENT); + EXPECT_EQ(s32->Convert(U16).status().code(), + tensorflow::error::INVALID_ARGUMENT); +} + +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); + p.clear_preds(); + for (int i = 0; i < len; ++i) { + p.add_preds((i % 2) == (len % 2)); + } + + Literal literal(p); + ASSERT_EQ(len, literal.preds_size()); + int i = 0; + for (auto it = literal.preds().begin(); it < literal.preds().end(); ++it) { + EXPECT_EQ((i % 2) == (len % 2), *it); + ++i; + } + } +} + +// Note that f16 is currently stored in a byte array in little endian byte order +TEST_F(LiteralUtilTest, ToProto_f16) { + half h1(1.0f); + half h2(2.0f); + + auto m = Literal::CreateR2({{h1, h2}, {h2, h1}}); + Literal* l = m.get(); + EXPECT_EQ(4, ShapeUtil::ElementsIn(l->shape())); + EXPECT_EQ(4, l->f16s().size()); + EXPECT_EQ(4, l->f16s_size()); + + LiteralProto p = l->ToProto(); + EXPECT_EQ(4, ShapeUtil::ElementsIn(p.shape())); + EXPECT_EQ(8, p.f16s().size()); + const char* d = p.f16s().data(); + EXPECT_EQ(d[0], 0); + EXPECT_EQ(d[1], 0x3C); + EXPECT_EQ(d[2], 0); + EXPECT_EQ(d[3], 0x40); + EXPECT_EQ(d[4], 0); + EXPECT_EQ(d[5], 0x40); + EXPECT_EQ(d[6], 0); + EXPECT_EQ(d[7], 0x3C); +} + +// Note that f16 is currently stored in a byte array in little endian byte order +TEST_F(LiteralUtilTest, CopyFromProto_f16) { + half h1(1.0f); + half h2(2.0f); + + const char half_vals[8] = {0x00, 0x3C, 0x00, 0x40, 0x00, 0x40, 0x00, 0x3C}; + LiteralProto p; + p.mutable_shape()->set_element_type(F16); + p.mutable_shape()->clear_dimensions(); + p.mutable_shape()->add_dimensions(4); + p.clear_f16s(); + p.set_f16s(half_vals, 8); + + Literal literal(p); + ASSERT_EQ(4, literal.f16s_size()); + ASSERT_EQ(h1, literal.f16s(0)); + ASSERT_EQ(h2, literal.f16s(1)); + ASSERT_EQ(h2, literal.f16s(2)); + ASSERT_EQ(h1, literal.f16s(3)); + + const std::vector& r = literal.f16s(); + ASSERT_EQ(4, r.size()); + ASSERT_EQ(h1, r[0]); + ASSERT_EQ(h2, r[1]); + ASSERT_EQ(h2, r[2]); + ASSERT_EQ(h1, r[3]); } } // namespace diff --git a/tensorflow/compiler/xla/metric_table_report.cc b/tensorflow/compiler/xla/metric_table_report.cc index cd7c42f6e17e15b5e1c6ebfa1f24a40a9003a63e..fed0e58e66a04df2ff9554cb0dd0053b7c669803 100644 --- a/tensorflow/compiler/xla/metric_table_report.cc +++ b/tensorflow/compiler/xla/metric_table_report.cc @@ -38,7 +38,8 @@ void MetricTableReport::SetEntryName(string entry_name) { void MetricTableReport::SetShowAllEntries() { max_entries_to_show_ = std::numeric_limits::max(); - max_metric_proportion_to_show = 1.1; // more than 100% + max_entries_per_category_to_show_ = std::numeric_limits::max(); + max_metric_proportion_to_show_ = 1.1; // more than 100% } void MetricTableReport::SetShowCategoryTable() { show_category_table_ = true; } @@ -141,7 +142,7 @@ void MetricTableReport::AppendCategoryTable() { int64 categories_shown = 0; for (const auto& category : categories) { if (categories_shown >= max_entries_to_show_ || - metric_sum / expected_metric_sum_ > max_metric_proportion_to_show) { + metric_sum / expected_metric_sum_ > max_metric_proportion_to_show_) { break; } ++categories_shown; @@ -149,22 +150,21 @@ void MetricTableReport::AppendCategoryTable() { // Show the category. string text = category.category_text; - if (text == "") { + if (text.empty()) { text = "[no category]"; } tensorflow::strings::StrAppend(&text, " (", category.entries.size(), " ", entry_name_, ")"); AppendTableRow(text, category.metric_sum, metric_sum); - // Show the top few entries in the category. - const int64 kMaxToShow = 5; + // Show the top entries in the category. const char* const kIndentPrefix = " * "; - int64 entries_to_show = - std::min(kMaxToShow, category.entries.size()); - if (category.entries.size() == kMaxToShow + 1) { + int64 entries_to_show = std::min(max_entries_per_category_to_show_, + category.entries.size()); + if (category.entries.size() == entries_to_show + 1) { // May as well show the last entry on the line that would otherwise say // that there is a single entry not shown. - entries_to_show = category.entries.size(); + ++entries_to_show; } for (int64 i = 0; i < entries_to_show; ++i) { AppendLine(kIndentPrefix, MetricPercent(category.entries[i]->metric), " ", @@ -193,14 +193,14 @@ void MetricTableReport::AppendEntryTable() { int64 entries_shown = 0; for (const auto& entry : entries_) { if (entries_shown >= max_entries_to_show_ || - metric_sum / expected_metric_sum_ > max_metric_proportion_to_show) { + metric_sum / expected_metric_sum_ > max_metric_proportion_to_show_) { break; } ++entries_shown; metric_sum += entry.metric; string text = entry.text; - if (text == "") { + if (text.empty()) { text = "[no entry text]"; } AppendTableRow(text, entry.metric, metric_sum); @@ -220,7 +220,14 @@ void MetricTableReport::AppendTableRow(const string& text, const double metric, const int64 max_metric_string_size = MetricString(expected_metric_sum_).size(); string metric_string = MetricString(metric); - string padding(max_metric_string_size - metric_string.size() + 1, ' '); + + // Don't try to make a gigantic string and crash if expected_metric_sum_ is + // wrong somehow. + int64 padding_len = 1; + if (max_metric_string_size >= metric_string.size()) { + padding_len += max_metric_string_size - metric_string.size(); + } + string padding(padding_len, ' '); AppendLine(padding, metric_string, " (", MetricPercent(metric), " Σ", MetricPercent(running_metric_sum), ") ", text); } diff --git a/tensorflow/compiler/xla/metric_table_report.h b/tensorflow/compiler/xla/metric_table_report.h index e967627bff4446a695bfae514faac4b1acca4968..818fb1d3fe0b8bbe1a8eba363ff6445e2f3df9d2 100644 --- a/tensorflow/compiler/xla/metric_table_report.h +++ b/tensorflow/compiler/xla/metric_table_report.h @@ -103,6 +103,7 @@ class MetricTableReport { private: static constexpr double kDefaultMaxMetricProportionToShow = 0.99; static constexpr int64 kDefaultMaxEntriesToShow = 100; + static constexpr int64 kDefaultMaxEntriesPerCategoryToShow = 5; // Append all parameters to the report. template @@ -162,7 +163,8 @@ class MetricTableReport { // These members control how many categories and entries to show in tables. int64 max_entries_to_show_ = kDefaultMaxEntriesToShow; - double max_metric_proportion_to_show = kDefaultMaxMetricProportionToShow; + int64 max_entries_per_category_to_show_ = kDefaultMaxEntriesPerCategoryToShow; + double max_metric_proportion_to_show_ = kDefaultMaxMetricProportionToShow; // The report that is being created. string report_; diff --git a/tensorflow/compiler/xla/packed_literal_reader.cc b/tensorflow/compiler/xla/packed_literal_reader.cc index 21766a2a0c890c95e02b48a015470fe2489b9bd4..70e0f5a74711c8ceef1b6d4225141aa1cc9c6219 100644 --- a/tensorflow/compiler/xla/packed_literal_reader.cc +++ b/tensorflow/compiler/xla/packed_literal_reader.cc @@ -58,10 +58,9 @@ StatusOr> PackedLiteralReader::Read( } int64 elements = ShapeUtil::ElementsIn(shape); - LiteralUtil::Resize(elements, std::numeric_limits::quiet_NaN(), - result.get()); - tensorflow::protobuf::RepeatedField* field = result->mutable_f32s(); - char* data = tensorflow::bit_cast(field->mutable_data()); + result->Resize(elements, std::numeric_limits::quiet_NaN()); + std::vector* field = result->mutable_f32s(); + char* data = tensorflow::bit_cast(field->data()); uint64 bytes = elements * sizeof(float); tensorflow::StringPiece sp; auto s = file_->Read(offset_, bytes, &sp, data); diff --git a/tensorflow/compiler/xla/packed_literal_reader.h b/tensorflow/compiler/xla/packed_literal_reader.h index 563d978cf5ddaa96905f26eb833f474f11daad12..45a9fe012784d3e4168e7549240dec962aa1a17a 100644 --- a/tensorflow/compiler/xla/packed_literal_reader.h +++ b/tensorflow/compiler/xla/packed_literal_reader.h @@ -18,6 +18,7 @@ limitations under the License. #include +#include "tensorflow/compiler/xla/literal_util.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/port/BUILD b/tensorflow/compiler/xla/port/BUILD deleted file mode 100644 index 6fc5f1185c9d56075f18928e4b2c8e3819cf9ddd..0000000000000000000000000000000000000000 --- a/tensorflow/compiler/xla/port/BUILD +++ /dev/null @@ -1,33 +0,0 @@ -licenses(["notice"]) # Apache 2.0 - -# Filegroup used to collect source files for dependency checking. -filegroup( - name = "c_srcs", - data = glob([ - "**/*.cc", - "**/*.h", - ]), - visibility = ["//tensorflow/compiler/xla:internal"], -) - -cc_library( - name = "initialize", - hdrs = ["initialize.h"], - visibility = [ - "//tensorflow/compiler/xla:__subpackages__", - ], -) - -# ----------------------------------------------------------------------------- - -filegroup( - name = "all_files", - srcs = glob( - ["**/*"], - exclude = [ - "**/METADATA", - "**/OWNERS", - ], - ), - visibility = ["//tensorflow:__subpackages__"], -) diff --git a/tensorflow/compiler/xla/port/initialize.h b/tensorflow/compiler/xla/port/initialize.h deleted file mode 100644 index 13d9632f97c72296e9a335c2a10edefa9abc0e17..0000000000000000000000000000000000000000 --- a/tensorflow/compiler/xla/port/initialize.h +++ /dev/null @@ -1,39 +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_PORT_INITIALIZE_H_ -#define TENSORFLOW_COMPILER_XLA_PORT_INITIALIZE_H_ - -#undef REGISTER_MODULE_INITIALIZER - -namespace xla { - -class Initializer { - public: - typedef void (*InitializerFunc)(); - explicit Initializer(InitializerFunc func) { func(); } -}; - -} // namespace xla - -#define REGISTER_INITIALIZER(type, name, body) \ - static void google_init_##type##_##name() { body; } \ - xla::Initializer google_initializer_##type##_##name( \ - google_init_##type##_##name) - -#define REGISTER_MODULE_INITIALIZER(name, body) \ - REGISTER_INITIALIZER(module, name, body) - -#endif // TENSORFLOW_COMPILER_XLA_PORT_INITIALIZE_H_ diff --git a/tensorflow/compiler/xla/primitive_util.cc b/tensorflow/compiler/xla/primitive_util.cc index e3909ae8e9736351d3ee91332572b5db62727289..e4e37177a2d74e6da20300f1439942a146ad8d49 100644 --- a/tensorflow/compiler/xla/primitive_util.cc +++ b/tensorflow/compiler/xla/primitive_util.cc @@ -78,6 +78,11 @@ PrimitiveType NativeToPrimitiveType() { return F64; } +template <> +PrimitiveType NativeToPrimitiveType() { + return F16; +} + bool IsFloatingPointType(PrimitiveType type) { return type == F16 || type == F32 || type == F64; } diff --git a/tensorflow/compiler/xla/primitive_util.h b/tensorflow/compiler/xla/primitive_util.h index 78f0ee6f592d9b9ec2ed85f23297634c5e2e4d41..162a11c7d2966346979b98c804917203f82c806c 100644 --- a/tensorflow/compiler/xla/primitive_util.h +++ b/tensorflow/compiler/xla/primitive_util.h @@ -75,6 +75,8 @@ template <> PrimitiveType NativeToPrimitiveType(); template <> PrimitiveType NativeToPrimitiveType(); +template <> +PrimitiveType NativeToPrimitiveType(); bool IsFloatingPointType(PrimitiveType type); @@ -150,6 +152,10 @@ template <> struct PrimitiveTypeToNative { using type = double; }; +template <> +struct PrimitiveTypeToNative { + using type = half; +}; } // namespace primitive_util } // namespace xla diff --git a/tensorflow/compiler/xla/protobuf_util.cc b/tensorflow/compiler/xla/protobuf_util.cc index adb2e99ad25ae87b6411707c2259a44ab699886d..cdc4139cd69c3d6eb4afc2e5d25f9446ffad0a11 100644 --- a/tensorflow/compiler/xla/protobuf_util.cc +++ b/tensorflow/compiler/xla/protobuf_util.cc @@ -14,7 +14,13 @@ limitations under the License. ==============================================================================*/ #include "tensorflow/compiler/xla/protobuf_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/io/path.h" +#include "tensorflow/core/platform/env.h" +#include "tensorflow/core/platform/protobuf.h" namespace xla { namespace protobuf_util { @@ -31,5 +37,35 @@ bool ProtobufEquals(const tensorflow::protobuf::Message& m1, return (serialized1 == serialized2); } +StatusOr ToJson(const tensorflow::protobuf::Message& message) { + string json_output; + tensorflow::protobuf::util::JsonPrintOptions json_options; + json_options.add_whitespace = true; + json_options.always_print_primitive_fields = true; + auto status = tensorflow::protobuf::util::MessageToJsonString( + message, &json_output, json_options); + if (!status.ok()) { + return InternalError("MessageToJsonString failed: %s", + status.error_message().data()); + } + return json_output; +} + +Status DumpJsonToDirectory(const tensorflow::protobuf::Message& message, + const string& directory, const string& file_name) { + TF_ASSIGN_OR_RETURN(const string json_output, ToJson(message)); + + tensorflow::Env* env = tensorflow::Env::Default(); + TF_RETURN_IF_ERROR(env->RecursivelyCreateDir(directory)); + string safe_file_name = file_name + ".json"; + for (char& c : safe_file_name) { + if (c == '/' || c == '\\') { + c = '_'; + } + } + const string path = tensorflow::io::JoinPath(directory, safe_file_name); + return tensorflow::WriteStringToFile(env, path, json_output); +} + } // namespace protobuf_util } // namespace xla diff --git a/tensorflow/compiler/xla/protobuf_util.h b/tensorflow/compiler/xla/protobuf_util.h index 36247f1bdec2589b98b039d5aee83591dac34258..1a895c3585902e8fbc0d20475c2817ef4caa4c71 100644 --- a/tensorflow/compiler/xla/protobuf_util.h +++ b/tensorflow/compiler/xla/protobuf_util.h @@ -16,6 +16,8 @@ limitations under the License. #ifndef TENSORFLOW_COMPILER_XLA_PROTOBUF_UTIL_H_ #define TENSORFLOW_COMPILER_XLA_PROTOBUF_UTIL_H_ +#include "tensorflow/compiler/xla/statusor.h" +#include "tensorflow/compiler/xla/types.h" #include "tensorflow/core/platform/protobuf.h" namespace xla { @@ -29,6 +31,17 @@ namespace protobuf_util { // base, this form of equality checking is sufficient. extern bool ProtobufEquals(const tensorflow::protobuf::Message& m1, const tensorflow::protobuf::Message& m2); + +// Returns 'message' as a JSON string. +StatusOr ToJson(const tensorflow::protobuf::Message& message); + +// Converts 'message' to JSON, and dumps it to the path formed by joining +// 'directory/file_name.json'. The 'directory' is recursively created if it +// doesn't already exist, and the 'file_name' is sanitized by replacing illegal +// characters with underscore '_'. +Status DumpJsonToDirectory(const tensorflow::protobuf::Message& message, + const string& directory, const string& file_name); + } // namespace protobuf_util } // namespace xla diff --git a/tensorflow/compiler/xla/reference_util.cc b/tensorflow/compiler/xla/reference_util.cc index 86c9c3b1ac38d755effad733590f78aafa9571db..90aa9720a1e18bad06842adeead46fc3120d01dd 100644 --- a/tensorflow/compiler/xla/reference_util.cc +++ b/tensorflow/compiler/xla/reference_util.cc @@ -16,9 +16,13 @@ limitations under the License. #include "tensorflow/compiler/xla/reference_util.h" #include +#include #include "tensorflow/compiler/xla/client/computation_builder.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" #include "tensorflow/compiler/xla/window_util.h" #include "tensorflow/compiler/xla/xla_data.pb.h" #include "tensorflow/core/lib/math/math_util.h" @@ -85,6 +89,53 @@ namespace xla { return result; } +/* static */ std::unique_ptr> ReferenceUtil::ConvArray3D( + const Array3D& lhs, const Array3D& rhs, int64 kernel_stride, + Padding padding) { + return ConvArray3DGeneralDimensionsDilated( + lhs, rhs, kernel_stride, padding, 1, 1, + ComputationBuilder::CreateDefaultConvDimensionNumbers(1)); +} + +/*static*/ std::unique_ptr> +ReferenceUtil::ConvArray3DGeneralDimensionsDilated( + const Array3D& lhs, const Array3D& rhs, int64 kernel_stride, + Padding padding, int64 lhs_dilation, int64 rhs_dilation, + const ConvolutionDimensionNumbers& dnums) { + CHECK_EQ(dnums.spatial_dimensions_size(), 1); + // Reuse the code for Array4D-convolution by extending the 3D input into a 4D + // array by adding a fourth dummy dimension of size 1 without stride, padding + // and dilation. + Array4D a4dlhs(lhs.n1(), lhs.n2(), lhs.n3(), 1); + a4dlhs.Each( + [&](tensorflow::gtl::ArraySlice indices, float* value_ptr) { + CHECK_EQ(indices[3], 0); + *value_ptr = lhs.operator()(indices[0], indices[1], indices[2]); + }); + Array4D a4drhs(rhs.n1(), rhs.n2(), rhs.n3(), 1); + a4drhs.Each( + [&](tensorflow::gtl::ArraySlice indices, float* value_ptr) { + CHECK_EQ(indices[3], 0); + *value_ptr = rhs.operator()(indices[0], indices[1], indices[2]); + }); + // Add a second dummy spatial dimensions. + ConvolutionDimensionNumbers dnums2d = dnums; + dnums2d.add_spatial_dimensions(3); + dnums2d.add_kernel_spatial_dimensions(3); + std::unique_ptr> convr4 = ConvArray4DGeneralDimensionsDilated( + a4dlhs, a4drhs, {kernel_stride, 1}, padding, {lhs_dilation, 1}, + {rhs_dilation, 1}, dnums2d); + + auto convr3 = MakeUnique>(convr4->planes(), convr4->depth(), + convr4->height()); + convr4->Each( + [&](tensorflow::gtl::ArraySlice indices, float* value_ptr) { + CHECK_EQ(indices[3], 0); + convr3->operator()(indices[0], indices[1], indices[2]) = *value_ptr; + }); + return convr3; +} + /* static */ std::unique_ptr> ReferenceUtil::ConvArray4D( const Array4D& lhs, const Array4D& rhs, std::pair kernel_stride, Padding padding) { @@ -134,6 +185,49 @@ ReferenceUtil::SeparableConvArray4D(const Array4D& input, return tensorflow::MathUtil::CeilOfRatio(unpadded_width, stride); } +/* static */ std::unique_ptr> +ReferenceUtil::ReduceWindow1DGeneric( + const tensorflow::gtl::ArraySlice& operand, float init, + const std::function& reduce_func, + const tensorflow::gtl::ArraySlice& window, + const tensorflow::gtl::ArraySlice& stride, Padding padding) { + std::vector dim_lengths{static_cast(operand.size())}; + auto padding_both = xla::MakePadding(dim_lengths, window, stride, padding); + + std::vector window_counts(window.size(), 0); + std::vector pad_low(window.size(), 0); + for (int64 i = 0; i < window.size(); ++i) { + window_counts[i] = + WindowCount(dim_lengths[i], window[i], stride[i], padding); + pad_low[i] = padding_both[i].first; + } + auto result = MakeUnique>(window_counts[0]); + + // Do a full 1D reduce window. + for (int64 i0 = 0; i0 < window_counts[0]; ++i0) { + int64 i0_base = i0 * stride[0] - pad_low[0]; + + float val = init; + for (int64 i0_win = 0; i0_win < window[0]; ++i0_win) { + if (i0_base + i0_win >= 0 && i0_base + i0_win < dim_lengths[0]) { + val = reduce_func(val, operand[i0_base + i0_win]); + } + } + (*result)[i0] = val; + } + return result; +} + +/* static */ std::unique_ptr> +ReferenceUtil::ReduceWindow1DAdd( + const tensorflow::gtl::ArraySlice& operand, float init, + const tensorflow::gtl::ArraySlice& window, + const tensorflow::gtl::ArraySlice& stride, Padding padding) { + const auto add_reduce = [](float arg1, float arg2) { return arg1 + arg2; }; + return ReduceWindow1DGeneric(operand, init, add_reduce, window, stride, + padding); +} + /* static */ std::unique_ptr> ReferenceUtil::ReduceWindow2DAdd( const Array2D& operand, float init, const tensorflow::gtl::ArraySlice& window, @@ -180,14 +274,28 @@ ReferenceUtil::ReduceWindow4DGeneric( const tensorflow::gtl::ArraySlice& stride, Padding padding) { std::vector dim_lengths{operand.n1(), operand.n2(), operand.n3(), operand.n4()}; - auto padding_both = xla::MakePadding(dim_lengths, window, stride, padding); + return ReduceWindow4DGeneric( + operand, init, reduce_func, window, stride, + xla::MakePadding(dim_lengths, window, stride, padding)); +} + +/* static */ std::unique_ptr> +ReferenceUtil::ReduceWindow4DGeneric( + const Array4D& operand, float init, + const std::function& reduce_func, + const tensorflow::gtl::ArraySlice& window, + const tensorflow::gtl::ArraySlice& stride, + const tensorflow::gtl::ArraySlice>& padding) { + std::vector dim_lengths{operand.n1(), operand.n2(), operand.n3(), + operand.n4()}; std::vector window_counts(window.size(), 0); std::vector pad_low(window.size(), 0); for (int64 i = 0; i < window.size(); ++i) { + int64 padded_width = padding[i].first + dim_lengths[i] + padding[i].second; window_counts[i] = - WindowCount(dim_lengths[i], window[i], stride[i], padding); - pad_low[i] = padding_both[i].first; + window_util::StridedBound(padded_width, window[i], stride[i]); + pad_low[i] = padding[i].first; } auto result = MakeUnique>(window_counts[0], window_counts[1], window_counts[2], window_counts[3]); @@ -237,6 +345,20 @@ ReferenceUtil::ReduceWindow4DGeneric( padding); } +/* static */ std::unique_ptr> ReferenceUtil::BatchNorm4D( + const Array4D& input, const Array4D& mean, + const Array4D& var, const Array4D& scale, + const Array4D& offset, float epsilon) { + auto normalized = + *MapArray4D(input, mean, [](float a, float b) { return a - b; }); + normalized = *MapArray4D(normalized, var, [&](float a, float b) { + return a / std::sqrt(b + epsilon); + }); + normalized = + *MapArray4D(normalized, scale, [](float a, float b) { return a * b; }); + return MapArray4D(normalized, offset, [](float a, float b) { return a + b; }); +} + /* static */ std::unique_ptr> ReferenceUtil::SelectAndScatter4DGePlus( const Array4D& operand, const Array4D& source, float init, @@ -317,7 +439,8 @@ ReferenceUtil::ConvArray4DGeneralDimensions( std::pair kernel_stride, Padding padding, ConvolutionDimensionNumbers dimension_numbers) { return ConvArray4DGeneralDimensionsDilated(lhs, rhs, kernel_stride, padding, - {1, 1}, {1, 1}, dimension_numbers); + {1, 1}, {1, 1}, + std::move(dimension_numbers)); } /* static */ std::unique_ptr> @@ -326,162 +449,94 @@ ReferenceUtil::ConvArray4DGeneralDimensionsDilated( std::pair kernel_stride, Padding padding, std::pair lhs_dilation, std::pair rhs_dilation, ConvolutionDimensionNumbers dnums) { - std::array lhs_dimensions{{lhs.n1(), lhs.n2(), lhs.n3(), lhs.n4()}}; - std::array rhs_dimensions{{rhs.n1(), rhs.n2(), rhs.n3(), rhs.n4()}}; - - const int64 ksy = kernel_stride.first; - const int64 ksx = kernel_stride.second; - const int64 dy = lhs_dilation.first; - const int64 dx = lhs_dilation.second; - const int64 dky = rhs_dilation.first; - const int64 dkx = rhs_dilation.second; - CHECK_GE(dky, 1); - CHECK_GE(dkx, 1); - CHECK_GE(dy, 1); - CHECK_GE(dx, 1); - - // Get all dimension sizes in lhs and rhs based on the given convolution - // dimension configuration. - const int64 ix = window_util::DilatedBound( - lhs_dimensions[dnums.spatial_dimensions(1)], dx); - const int64 iy = window_util::DilatedBound( - lhs_dimensions[dnums.spatial_dimensions(0)], dy); - const int64 iz = lhs_dimensions[dnums.feature_dimension()]; - const int64 samples = lhs_dimensions[dnums.batch_dimension()]; - const int64 kx = window_util::DilatedBound( - rhs_dimensions[dnums.kernel_spatial_dimensions(1)], dkx); - const int64 ky = window_util::DilatedBound( - rhs_dimensions[dnums.kernel_spatial_dimensions(0)], dky); - const int64 oz = rhs_dimensions[dnums.kernel_output_feature_dimension()]; - { - const int64 kiz = rhs_dimensions[dnums.kernel_input_feature_dimension()]; - CHECK_EQ(kiz, iz); + HloComputation::Builder b("ConvArray4DGeneralDimensionDilated"); + auto lhs_literal = Literal::CreateR4FromArray4D(lhs); + auto rhs_literal = Literal::CreateR4FromArray4D(rhs); + + std::array ordered_kernel_strides; + std::array ordered_input_dimensions; + std::array ordered_kernel_dimensions; + if (dnums.kernel_spatial_dimensions(0) > dnums.kernel_spatial_dimensions(1)) { + ordered_kernel_strides[0] = kernel_stride.second; + ordered_kernel_strides[1] = kernel_stride.first; + } else { + ordered_kernel_strides[0] = kernel_stride.first; + ordered_kernel_strides[1] = kernel_stride.second; } - if (padding == Padding::kSame) { - // We reject same padding with kernel striding, since it's somewhat - // nonsensical. We can always follow up to implement this with the desired - // semantics if anybody actually uses it. - CHECK_EQ(1, ksy); - CHECK_EQ(1, ksx); - } - - const int64 ox = - padding == Padding::kSame ? ix : window_util::StridedBound(ix, kx, ksx); - const int64 oy = - padding == Padding::kSame ? iy : window_util::StridedBound(iy, ky, ksy); - const int64 istartx = - padding == Padding::kValid ? 0 : kx % 2 == 0 ? -(kx / 2 - 1) : -kx / 2; - const int64 istarty = - padding == Padding::kValid ? 0 : ky % 2 == 0 ? -(ky / 2 - 1) : -ky / 2; - // Create the output result array and reset the values to 0. - std::array result_dimensions; - result_dimensions[dnums.batch_dimension()] = samples; - result_dimensions[dnums.feature_dimension()] = oz; - result_dimensions[dnums.spatial_dimensions(0)] = oy; - result_dimensions[dnums.spatial_dimensions(1)] = ox; + ordered_input_dimensions[0] = + lhs_literal->shape().dimensions(dnums.spatial_dimensions(0)); + ordered_input_dimensions[1] = + lhs_literal->shape().dimensions(dnums.spatial_dimensions(1)); + ordered_kernel_dimensions[0] = + rhs_literal->shape().dimensions(dnums.kernel_spatial_dimensions(0)); + ordered_kernel_dimensions[1] = + rhs_literal->shape().dimensions(dnums.kernel_spatial_dimensions(1)); + + std::vector> paddings = + MakePadding(ordered_input_dimensions, ordered_kernel_dimensions, + ordered_kernel_strides, padding); + CHECK_EQ(paddings.size(), 2); + + Window window; + + WindowDimension dim; + dim.set_size( + rhs_literal->shape().dimensions(dnums.kernel_spatial_dimensions(0))); + dim.set_stride(kernel_stride.first); + dim.set_padding_low(paddings[0].first); + dim.set_padding_high(paddings[0].second); + dim.set_window_dilation(rhs_dilation.first); + dim.set_base_dilation(lhs_dilation.first); + *window.add_dimensions() = dim; + + WindowDimension dim2; + dim2.set_size( + rhs_literal->shape().dimensions(dnums.kernel_spatial_dimensions(1))); + dim2.set_stride(kernel_stride.second); + dim2.set_padding_low(paddings[1].first); + dim2.set_padding_high(paddings[1].second); + dim2.set_window_dilation(rhs_dilation.second); + dim2.set_base_dilation(lhs_dilation.second); + *window.add_dimensions() = dim2; + + const Shape& shape = + ShapeInference::InferConvolveShape(lhs_literal->shape(), + rhs_literal->shape(), window, dnums) + .ConsumeValueOrDie(); + + HloInstruction* lhs_instruction = + b.AddInstruction(HloInstruction::CreateConstant(std::move(lhs_literal))); + HloInstruction* rhs_instruction = + b.AddInstruction(HloInstruction::CreateConstant(std::move(rhs_literal))); + + b.AddInstruction(HloInstruction::CreateConvolve( + shape, lhs_instruction, rhs_instruction, window, dnums)); + HloModule module("ReferenceUtil"); + auto computation = module.AddEntryComputation(b.Build()); + + HloEvaluator evaluator; + std::unique_ptr result_literal = + evaluator.Evaluate(*computation, {}).ConsumeValueOrDie(); + + CHECK_EQ(ShapeUtil::Rank(result_literal->shape()), 4); auto result = - MakeUnique>(result_dimensions[0], result_dimensions[1], - result_dimensions[2], result_dimensions[3]); - result->Fill(0.0); - - const auto is_int32 = [](int64 x) { - return x >= std::numeric_limits::min() && - x <= std::numeric_limits::max(); - }; - - // 64-bit idiv/mod are much more expensive x86-64 than 32-bit idiv/imod (at - // least on x86-64), so we avoid them where possible. - const auto fast_idiv64 = [&](int64 a, int64 b) { - if (is_int32(a) && is_int32(b)) { - return static_cast(static_cast(a) / static_cast(b)); - } - return a / b; - }; - const auto fast_imod64 = [&](int64 a, int64 b) { - if (is_int32(a) && is_int32(b)) { - return static_cast(static_cast(a) % static_cast(b)); - } - return a % b; - }; - - // Lambda to access the lhs operand at the given 4D index. - const auto lhs_element = [&](int64 batch, int64 feature, int64 height, - int64 width) { - if (fast_imod64(height, dy) != 0 || fast_imod64(width, dx) != 0) { - return 0.0f; - } + MakeUnique>(result_literal->shape().dimensions(0), + result_literal->shape().dimensions(1), + result_literal->shape().dimensions(2), + result_literal->shape().dimensions(3)); + + result->Each([&](tensorflow::gtl::ArraySlice indices, float* value) { + *value = result_literal->Get(indices); + }); - std::array index; - index[dnums.batch_dimension()] = batch; - index[dnums.feature_dimension()] = feature; - index[dnums.spatial_dimensions(0)] = fast_idiv64(height, dy); - index[dnums.spatial_dimensions(1)] = fast_idiv64(width, dx); - return lhs(index[0], index[1], index[2], index[3]); - }; - - // Lambda to access the rhs operand at the given 4D index. height_over_dky - // should be equal to height / dky, and width_over_dkx should be equal to - // width / dkx. (This is an optimization to avoid doing divisions.) - const auto rhs_element = [&]( - int64 kernel_output_feature, int64 kernel_input_feature, int64 height, - int64 width, int64 height_over_dky, int64 width_over_dkx) { - DCHECK_EQ(height % dky, 0); - DCHECK_EQ(width % dkx, 0); - DCHECK_EQ(height / dky, height_over_dky); - DCHECK_EQ(width / dkx, width_over_dkx); - - std::array index; - index[dnums.kernel_output_feature_dimension()] = kernel_output_feature; - index[dnums.kernel_input_feature_dimension()] = kernel_input_feature; - index[dnums.kernel_spatial_dimensions(0)] = height_over_dky; - index[dnums.kernel_spatial_dimensions(1)] = width_over_dkx; - return rhs(index[0], index[1], index[2], index[3]); - }; - - // Lambda to access the result data at the given 4D index. - const auto result_element = [&](int64 batch, int64 kernel_output_feature, - int64 height, int64 width) -> float& { - std::array index; - index[dnums.batch_dimension()] = batch; - index[dnums.feature_dimension()] = kernel_output_feature; - index[dnums.spatial_dimensions(0)] = height; - index[dnums.spatial_dimensions(1)] = width; - return (*result)(index[0], index[1], index[2], index[3]); - }; - - for (int64 oyi = 0; oyi < oy; ++oyi) { - for (int64 oxi = 0; oxi < ox; ++oxi) { - for (int64 sample = 0; sample < samples; ++sample) { - for (int64 izi = 0; izi < iz; ++izi) { - for (int64 ozi = 0; ozi < oz; ++ozi) { - for (int64 kyi = 0, kyi_over_dky = 0; kyi < ky; - kyi += dky, kyi_over_dky++) { - for (int64 kxi = 0, kxi_over_dkx = 0; kxi < kx; - kxi += dkx, kxi_over_dkx++) { - int64 iyi = istarty + ksy * oyi + kyi; - int64 ixi = istartx + ksx * oxi + kxi; - float input = (iyi >= iy || ixi >= ix || iyi < 0 || ixi < 0) - ? 0.0 - : lhs_element(sample, izi, iyi, ixi); - float gain = - rhs_element(ozi, izi, kyi, kxi, kyi_over_dky, kxi_over_dkx); - float addend = input * gain; - result_element(sample, ozi, oyi, oxi) += addend; - } - } - } - } - } - } - } return result; } /* static */ std::unique_ptr> ReferenceUtil::ReduceToColArray2D( const Array2D& matrix, float init, - std::function reduce_function) { + const std::function& reduce_function) { int64 rows = matrix.height(); int64 cols = matrix.width(); auto result = MakeUnique>(); @@ -498,7 +553,7 @@ ReferenceUtil::ReduceToColArray2D( /* static */ std::unique_ptr> ReferenceUtil::ReduceToRowArray2D( const Array2D& matrix, float init, - std::function reduce_function) { + const std::function& reduce_function) { int64 rows = matrix.height(); int64 cols = matrix.width(); auto result = MakeUnique>(); @@ -515,7 +570,7 @@ ReferenceUtil::ReduceToRowArray2D( /*static*/ std::vector ReferenceUtil::Reduce4DTo1D( const Array4D& array, float init, tensorflow::gtl::ArraySlice dims, - std::function reduce_function) { + const std::function& reduce_function) { std::vector result; CHECK_EQ(dims.size(), 3); const std::set dim_set(dims.begin(), dims.end()); @@ -550,10 +605,42 @@ ReferenceUtil::ReduceToRowArray2D( return result; } +/* static */ std::unique_ptr> ReferenceUtil::Broadcast1DTo4D( + const std::vector& array, const std::vector& bounds, + int64 broadcast_from_dim) { + auto result = + MakeUnique>(bounds[0], bounds[1], bounds[2], bounds[3]); + for (int64 i = 0; i < result->n1(); ++i) { + for (int64 j = 0; j < result->n2(); ++j) { + for (int64 k = 0; k < result->n3(); ++k) { + for (int64 l = 0; l < result->n4(); ++l) { + switch (broadcast_from_dim) { + case 0: + (*result)(i, j, k, l) = array[i]; + break; + case 1: + (*result)(i, j, k, l) = array[j]; + break; + case 2: + (*result)(i, j, k, l) = array[k]; + break; + case 3: + (*result)(i, j, k, l) = array[l]; + break; + default: + break; + } + } + } + } + } + return result; +} + /* static */ std::unique_ptr> ReferenceUtil::Reduce3DTo2D( const Array3D& array, float init, tensorflow::gtl::ArraySlice dims, - std::function reduce_function) { + const std::function& reduce_function) { CHECK_EQ(dims.size(), 1); int64 rows = dims[0] == 0 ? array.n2() : array.n1(); int64 cols = dims[0] == 2 ? array.n2() : array.n3(); @@ -649,4 +736,104 @@ ReferenceUtil::ReduceToRowArray2D( return result; } +/* static */ Array3D ReferenceUtil::PadArray3D( + const Array3D& operand, const PaddingConfig& padding, + const float pad) { + CHECK_EQ(padding.dimensions_size(), 3); + + const std::vector input_bounds = {operand.n1(), operand.n2(), + operand.n3()}; + std::vector pad_low(3); + std::vector pad_high(3); + std::vector pad_interior(3); + std::vector output_bounds(3); + for (int64 i = 0; i < 3; ++i) { + pad_low[i] = padding.dimensions(i).edge_padding_low(); + pad_high[i] = padding.dimensions(i).edge_padding_high(); + CHECK_LE(0, pad_low[i]); + CHECK_LE(0, pad_high[i]); + CHECK_LE(0, padding.dimensions(i).interior_padding()) << "not implemented"; + pad_interior[i] = padding.dimensions(i).interior_padding(); + + output_bounds[i] = pad_low[i] + input_bounds[i] + pad_high[i] + + (input_bounds[i] - 1) * pad_interior[i]; + } + + Array3D result(output_bounds[0], output_bounds[1], output_bounds[2]); + std::vector indices = {0, 0, 0}; + for (indices[0] = 0; indices[0] < output_bounds[0]; ++indices[0]) { + for (indices[1] = 0; indices[1] < output_bounds[1]; ++indices[1]) { + for (indices[2] = 0; indices[2] < output_bounds[2]; ++indices[2]) { + float* value = &result(indices[0], indices[1], indices[2]); + bool value_padded = false; + for (int i = 0; i < 3; ++i) { + bool in_low_padding = indices[i] < pad_low[i]; + bool in_high_padding = indices[i] >= output_bounds[i] - pad_high[i]; + if (in_low_padding || in_high_padding) { + *value = pad; + value_padded = true; + } + if (pad_interior[i] && + (indices[i] - pad_low[i]) % (pad_interior[i] + 1)) { + *value = pad; + value_padded = true; + } + } + if (value_padded) { + continue; + } + *value = operand((indices[0] - pad_low[0]) / (pad_interior[0] + 1), + (indices[1] - pad_low[1]) / (pad_interior[1] + 1), + (indices[2] - pad_low[2]) / (pad_interior[2] + 1)); + } + } + } + return result; +} + +/* static */ Array4D ReferenceUtil::PadArray4D( + const Array4D& operand, const PaddingConfig& padding, + const float pad) { + CHECK_EQ(padding.dimensions_size(), 4); + + const std::vector input_bounds = {operand.n1(), operand.n2(), + operand.n3(), operand.n4()}; + std::vector pad_low(4); + std::vector pad_high(4); + std::vector pad_interior(4); + std::vector output_bounds(4); + for (int64 i = 0; i < 4; ++i) { + pad_low[i] = padding.dimensions(i).edge_padding_low(); + pad_high[i] = padding.dimensions(i).edge_padding_high(); + CHECK_LE(0, padding.dimensions(i).interior_padding()) << "not implemented"; + pad_interior[i] = padding.dimensions(i).interior_padding(); + + output_bounds[i] = pad_low[i] + input_bounds[i] + pad_high[i] + + (input_bounds[i] - 1) * pad_interior[i]; + } + + Array4D result(output_bounds[0], output_bounds[1], output_bounds[2], + output_bounds[3]); + result.Each([&](tensorflow::gtl::ArraySlice indices, float* value) { + for (int i = 0; i < 4; ++i) { + bool in_low_padding = indices[i] < pad_low[i]; + bool in_high_padding = indices[i] >= output_bounds[i] - pad_high[i]; + if (in_low_padding || in_high_padding) { + *value = pad; + return; + } + if (pad_interior[i] && + (indices[i] - pad_low[i]) % (pad_interior[i] + 1)) { + *value = pad; + return; + } + } + *value = operand((indices[0] - pad_low[0]) / (pad_interior[0] + 1), + (indices[1] - pad_low[1]) / (pad_interior[1] + 1), + (indices[2] - pad_low[2]) / (pad_interior[2] + 1), + (indices[3] - pad_low[3]) / (pad_interior[3] + 1)); + }); + return result; +} + } // namespace xla diff --git a/tensorflow/compiler/xla/reference_util.h b/tensorflow/compiler/xla/reference_util.h index 9e0f247203866d544595a877fabd33af148cc307..2da17307817858eea60e868f4be1ab8138784385 100644 --- a/tensorflow/compiler/xla/reference_util.h +++ b/tensorflow/compiler/xla/reference_util.h @@ -27,6 +27,7 @@ limitations under the License. #include "tensorflow/compiler/xla/array4d.h" #include "tensorflow/compiler/xla/client/padding.h" #include "tensorflow/compiler/xla/ptr_util.h" +#include "tensorflow/compiler/xla/util.h" #include "tensorflow/compiler/xla/xla_data.pb.h" #include "tensorflow/core/lib/gtl/array_slice.h" #include "tensorflow/core/platform/macros.h" @@ -73,6 +74,20 @@ class ReferenceUtil { std::pair lhs_dilation, std::pair rhs_dilation, ConvolutionDimensionNumbers dnums); + // Returns the result of a convolution `lhs rhs`, with the default + // convolution dimension numbers returned from + // ComputationBuilder::CreateDefaultConvDimensionNumbers(). + static std::unique_ptr> ConvArray3D(const Array3D& lhs, + const Array3D& rhs, + int64 kernel_stride, + Padding padding); + + // Returns the result of a convolution `lhs rhs`. + static std::unique_ptr> ConvArray3DGeneralDimensionsDilated( + const Array3D& lhs, const Array3D& rhs, int64 kernel_stride, + Padding padding, int64 lhs_dilation, int64 rhs_dilation, + const ConvolutionDimensionNumbers& dnums); + // Returns the result of a separable convolution with the given parameters. // kernel_stride and padding applies to the depthwise convolution during // the separable convolution. pointwise_weights.depth() must be equal to @@ -87,21 +102,21 @@ class ReferenceUtil { // to apply for each reduction step. static std::unique_ptr> ReduceToColArray2D( const Array2D& matrix, float init, - std::function reduce_function); + const std::function& reduce_function); // Returns the result of reducing a matrix to a row vector. init is the // initial value for the reduce operation, and reduce_function is the function // to apply for each reduction step. static std::unique_ptr> ReduceToRowArray2D( const Array2D& matrix, float init, - std::function reduce_function); + const std::function& reduce_function); // Performs a R2=>R1 reduction by reducing away the dimension specified in // 'dimension_to_reduce'. template static std::vector ReduceR2ToR1(const Array2D& input, int dimension_to_reduce, T init, - std::function freduce) { + const std::function& freduce) { std::vector result(dimension_to_reduce == 0 ? input.n2() : input.n1(), init); for (int i0 = 0; i0 < input.n1(); ++i0) { @@ -118,14 +133,19 @@ class ReferenceUtil { static std::vector Reduce4DTo1D( const Array4D& array, float init, tensorflow::gtl::ArraySlice dims, - std::function reduce_function); + const std::function& reduce_function); + + // Broadcast 1D dimension to 4D, from the dimension `broadcast_from_dim`. + static std::unique_ptr> Broadcast1DTo4D( + const std::vector& array, const std::vector& bounds, + int64 broadcast_from_dim); // Returns the result of reducing the 3D array to a 2D array, reducing away // the dimensions specified in dims. static std::unique_ptr> Reduce3DTo2D( const Array3D& array, float init, tensorflow::gtl::ArraySlice dims, - std::function reduce_function); + const std::function& reduce_function); // Applies map_function to each element in the input (2D array) and returns // the result. @@ -144,24 +164,43 @@ class ReferenceUtil { static int64 WindowCount(int64 unpadded_width, int64 window_len, int64 stride, Padding padding); - // Performs a 2D window reduction with Add as the function to apply. + // Windowed reductions with Add as the function to apply. + static std::unique_ptr> ReduceWindow1DAdd( + const tensorflow::gtl::ArraySlice& operand, float init, + const tensorflow::gtl::ArraySlice& window, + const tensorflow::gtl::ArraySlice& stride, Padding padding); static std::unique_ptr> ReduceWindow2DAdd( const Array2D& operand, float init, const tensorflow::gtl::ArraySlice& window, const tensorflow::gtl::ArraySlice& stride, Padding padding); - - // Performs a 4D window reduction with Add as the function to apply. static std::unique_ptr> ReduceWindow4DAdd( const Array4D& operand, float init, const tensorflow::gtl::ArraySlice& window, const tensorflow::gtl::ArraySlice& stride, Padding padding); - // Performs a 4D window reduction with a generic reduce function. + // Windowed reductions with a generic reduce function. + static std::unique_ptr> ReduceWindow1DGeneric( + const tensorflow::gtl::ArraySlice& operand, float init, + const std::function& reduce_func, + const tensorflow::gtl::ArraySlice& window, + const tensorflow::gtl::ArraySlice& stride, Padding padding); static std::unique_ptr> ReduceWindow4DGeneric( const Array4D& operand, float init, const std::function& reduce_func, const tensorflow::gtl::ArraySlice& window, const tensorflow::gtl::ArraySlice& stride, Padding padding); + static std::unique_ptr> ReduceWindow4DGeneric( + const Array4D& operand, float init, + const std::function& reduce_func, + const tensorflow::gtl::ArraySlice& window, + const tensorflow::gtl::ArraySlice& stride, + const tensorflow::gtl::ArraySlice>& padding); + + // Batch normalize data. + static std::unique_ptr> BatchNorm4D( + const Array4D& input, const Array4D& mean, + const Array4D& var, const Array4D& scale, + const Array4D& offset, float epsilon); // Performs select and scatter with Greater Than or equal as the select, plus // as the scatter, and Same Padding. @@ -277,48 +316,56 @@ class ReferenceUtil { return result; } - // Slices the input array given starting indices in each dimension and limit - // indices in each dimension. + // Slices the input array given starting indices, limit indices, and strides + // in each dimension. template static std::unique_ptr> Slice2D(const Array2D& input, std::array starts, - std::array limits) { + std::array limits, + std::array strides) { CHECK_LE(starts[0], input.n1()); CHECK_LE(starts[1], input.n2()); CHECK_LE(limits[0], input.n1()); CHECK_LE(limits[1], input.n2()); + CHECK_GE(strides[0], 1); + CHECK_GE(strides[1], 1); auto result = - MakeUnique>(limits[0] - starts[0], limits[1] - starts[1]); + MakeUnique>(CeilOfRatio(limits[0] - starts[0], strides[0]), + CeilOfRatio(limits[1] - starts[1], strides[1])); for (int64 i0 = 0; i0 < result->n1(); ++i0) { for (int64 i1 = 0; i1 < result->n2(); ++i1) { - (*result)(i0, i1) = input(starts[0] + i0, starts[1] + i1); + (*result)(i0, i1) = + input(starts[0] + i0 * strides[0], starts[1] + i1 * strides[1]); } } return result; } template - static std::unique_ptr> Slice4D(const Array4D& input, - std::array starts, - std::array limits) { + static std::unique_ptr> Slice3D(const Array3D& input, + std::array starts, + std::array limits, + std::array strides) { CHECK_LE(starts[0], input.n1()); CHECK_LE(starts[1], input.n2()); CHECK_LE(starts[2], input.n3()); - CHECK_LE(starts[3], input.n4()); CHECK_LE(limits[0], input.n1()); CHECK_LE(limits[1], input.n2()); CHECK_LE(limits[2], input.n3()); - CHECK_LE(limits[3], input.n4()); + CHECK_GE(strides[0], 1); + CHECK_GE(strides[1], 1); + CHECK_GE(strides[2], 1); auto result = - MakeUnique>(limits[0] - starts[0], limits[1] - starts[1], - limits[2] - starts[2], limits[3] - starts[3]); + MakeUnique>(CeilOfRatio(limits[0] - starts[0], strides[0]), + CeilOfRatio(limits[1] - starts[1], strides[1]), + CeilOfRatio(limits[2] - starts[2], strides[2])); + for (int64 i0 = 0; i0 < result->n1(); ++i0) { for (int64 i1 = 0; i1 < result->n2(); ++i1) { for (int64 i2 = 0; i2 < result->n3(); ++i2) { - for (int64 i3 = 0; i3 < result->n4(); ++i3) { - (*result)(i0, i1, i2, i3) = input(starts[0] + i0, starts[1] + i1, - starts[2] + i2, starts[3] + i3); - } + (*result)(i0, i1, i2) = + input(starts[0] + i0 * strides[0], starts[1] + i1 * strides[1], + starts[2] + i2 * strides[2]); } } } @@ -326,22 +373,35 @@ class ReferenceUtil { } template - static std::unique_ptr> Slice3D(const Array3D& input, - std::array starts, - std::array limits) { + static std::unique_ptr> Slice4D(const Array4D& input, + std::array starts, + std::array limits, + std::array strides) { CHECK_LE(starts[0], input.n1()); CHECK_LE(starts[1], input.n2()); CHECK_LE(starts[2], input.n3()); + CHECK_LE(starts[3], input.n4()); CHECK_LE(limits[0], input.n1()); CHECK_LE(limits[1], input.n2()); CHECK_LE(limits[2], input.n3()); - auto result = MakeUnique>( - limits[0] - starts[0], limits[1] - starts[1], limits[2] - starts[2]); + CHECK_LE(limits[3], input.n4()); + CHECK_GE(strides[0], 1); + CHECK_GE(strides[1], 1); + CHECK_GE(strides[2], 1); + CHECK_GE(strides[3], 1); + auto result = + MakeUnique>(CeilOfRatio(limits[0] - starts[0], strides[0]), + CeilOfRatio(limits[1] - starts[1], strides[1]), + CeilOfRatio(limits[2] - starts[2], strides[2]), + CeilOfRatio(limits[3] - starts[3], strides[3])); for (int64 i0 = 0; i0 < result->n1(); ++i0) { for (int64 i1 = 0; i1 < result->n2(); ++i1) { for (int64 i2 = 0; i2 < result->n3(); ++i2) { - (*result)(i0, i1, i2) = - input(starts[0] + i0, starts[1] + i1, starts[2] + i2); + for (int64 i3 = 0; i3 < result->n4(); ++i3) { + (*result)(i0, i1, i2, i3) = + input(starts[0] + i0 * strides[0], starts[1] + i1 * strides[1], + starts[2] + i2 * strides[2], starts[3] + i3 * strides[3]); + } } } } @@ -390,12 +450,96 @@ class ReferenceUtil { return result; } + // Applies map_function to each pair of elements in the input lhs and rhs + // (4D array) and returns the result. + template + static std::unique_ptr> MapArray4D(const Array4D& lhs, + const Array4D& rhs, + F&& map_function) { + return MapWithIndexArray4D( + lhs, rhs, [&](float lhs, float rhs, int64, int64, int64, int64) { + return map_function(lhs, rhs); + }); + } + + // Applies map_function to each pair of element in lhs and rhs (4D array) and + // returns the result. + // (plane, depth, height, width) index of each element is also provided as + // arguments to map_function. + template + static std::unique_ptr> MapWithIndexArray4D( + const Array4D& lhs, const Array4D& rhs, F&& map_function) { + auto result = MakeUnique>(lhs.planes(), lhs.depth(), + lhs.height(), lhs.width()); + for (int64 plane = 0; plane < lhs.planes(); ++plane) { + for (int64 depth = 0; depth < lhs.depth(); ++depth) { + for (int64 height = 0; height < lhs.height(); ++height) { + for (int64 width = 0; width < lhs.width(); ++width) { + (*result)(plane, depth, height, width) = map_function( + lhs(plane, depth, height, width), + rhs(plane, depth, height, width), plane, depth, height, width); + } + } + } + } + return result; + } + // Returns the result of a 2D pad on an input matrix. static std::unique_ptr> PadArray2D( const Array2D& operand, const PaddingConfig& padding, const float pad); + // Returns the result of a 3D pad on an input matrix. + static Array3D PadArray3D(const Array3D& operand, + const PaddingConfig& padding, + const float pad); + + // Returns the result of a 4D pad on an input array. + static Array4D PadArray4D(const Array4D& operand, + const PaddingConfig& padding, + const float pad); + + // ApplyElementwise2D(f, x, y, ...) returns the Array2D formed by running + // f(x[i], y[i], ...) for each array element in the Array2Ds x, y, .... + // + // The given arrays must have the same size and element type, and the return + // type of f must be implicitly convertible to the arrays' element type. + // + // Example usage: + // + // Array2D x, y, z = ...; + // std::unique_ptr result = ReferenceUtil::ApplyElementwise2D( + // [](float a, float b, float c) { return a * b + c; }, x, y, z); + // + template + static std::unique_ptr> ApplyElementwise2D( + F&& f, const Array2D& array1, const Array2D&... arrays) { + AssertSameSize2D(array1, arrays...); + auto result = MakeUnique>(array1.n1(), array1.n2()); + for (int64 i = 0; i < array1.n1(); ++i) { + for (int64 j = 0; j < array1.n2(); ++j) { + (*result)(i, j) = f(array1(i, j), arrays(i, j)...); + } + } + return result; + } + private: + template + static void AssertSameSize2D(const Array2D& array1, + const Array2D& array2, + const Array2D&... arrays) { + static_assert(std::is_same::value, "Args must be same type."); + CHECK_EQ(array1.n1(), array2.n1()); + CHECK_EQ(array1.n2(), array2.n2()); + AssertSameSize2D(array2, arrays...); + } + + // Recursive base case for AssertSameSize2D. + template + static void AssertSameSize2D(const Array1& array1) {} + TF_DISALLOW_COPY_AND_ASSIGN(ReferenceUtil); }; diff --git a/tensorflow/compiler/xla/reference_util_test.cc b/tensorflow/compiler/xla/reference_util_test.cc index c53351ca93e81f70920291019798f16f0f1c6a57..35b5e8cd52ab0ec21a4bd2df3e9fa0538ae60816 100644 --- a/tensorflow/compiler/xla/reference_util_test.cc +++ b/tensorflow/compiler/xla/reference_util_test.cc @@ -19,13 +19,14 @@ limitations under the License. #include #include "tensorflow/compiler/xla/array2d.h" +#include "tensorflow/compiler/xla/array3d.h" #include "tensorflow/compiler/xla/array4d.h" #include "tensorflow/compiler/xla/client/padding.h" #include "tensorflow/compiler/xla/literal_util.h" #include "tensorflow/compiler/xla/ptr_util.h" +#include "tensorflow/compiler/xla/test.h" #include "tensorflow/compiler/xla/tests/literal_test_util.h" #include "tensorflow/compiler/xla/xla_data.pb.h" -#include "tensorflow/core/platform/test.h" namespace xla { namespace { @@ -52,9 +53,9 @@ class ReferenceUtilTest : public ::testing::Test { TEST_F(ReferenceUtilTest, TransposeArray2D) { auto result = ReferenceUtil::TransposeArray2D(*matrix_); - auto result_literal = LiteralUtil::CreateR2FromArray2D(*result); + auto actual_literal = Literal::CreateR2FromArray2D(*result); LiteralTestUtil::ExpectR2Near({{1.f, 4.f}, {2.f, 5.f}, {3.f, 6.f}}, - *result_literal, ErrorSpec(0.0001)); + *actual_literal, ErrorSpec(0.0001)); } TEST_F(ReferenceUtilTest, MatmulArray2D) { @@ -62,32 +63,32 @@ TEST_F(ReferenceUtilTest, MatmulArray2D) { {7.f, 8.f}, {9.f, 10.f}, {11.f, 12.f}, }); auto result = ReferenceUtil::MatmulArray2D(*matrix_, rhs); - auto result_literal = LiteralUtil::CreateR2FromArray2D(*result); + auto actual_literal = Literal::CreateR2FromArray2D(*result); LiteralTestUtil::ExpectR2Near({{58.f, 64.f}, {139.f, 154.f}}, - *result_literal, ErrorSpec(0.0001)); + *actual_literal, ErrorSpec(0.0001)); } TEST_F(ReferenceUtilTest, ReduceToColArray2D) { auto add = [](float lhs, float rhs) { return lhs + rhs; }; auto result = ReferenceUtil::ReduceToColArray2D(*matrix_, 0.0f, add); - auto result_literal = LiteralUtil::CreateR1(*result); - LiteralTestUtil::ExpectR1Near({6.f, 15.f}, *result_literal, + auto actual_literal = Literal::CreateR1(*result); + LiteralTestUtil::ExpectR1Near({6.f, 15.f}, *actual_literal, ErrorSpec(0.0001)); } TEST_F(ReferenceUtilTest, ReduceToRowArray2D) { auto add = [](float lhs, float rhs) { return lhs + rhs; }; auto result = ReferenceUtil::ReduceToRowArray2D(*matrix_, 0.0f, add); - auto result_literal = LiteralUtil::CreateR1(*result); - LiteralTestUtil::ExpectR1Near({5.f, 7.f, 9.f}, *result_literal, + auto actual_literal = Literal::CreateR1(*result); + LiteralTestUtil::ExpectR1Near({5.f, 7.f, 9.f}, *actual_literal, ErrorSpec(0.0001)); } TEST_F(ReferenceUtilTest, MapArray2D) { auto identity = [](float value) { return log(exp(value)); }; auto result = ReferenceUtil::MapArray2D(*matrix_, identity); - auto result_literal = LiteralUtil::CreateR2FromArray2D(*result); - LiteralTestUtil::ExpectR2NearArray2D(*matrix_, *result_literal, + auto actual_literal = Literal::CreateR2FromArray2D(*result); + LiteralTestUtil::ExpectR2NearArray2D(*matrix_, *actual_literal, ErrorSpec(0.0001)); } @@ -96,9 +97,9 @@ TEST_F(ReferenceUtilTest, MapWithIndexArray2D) { return value + row + col; }; auto result = ReferenceUtil::MapWithIndexArray2D(*matrix_, add_index); - auto result_literal = LiteralUtil::CreateR2FromArray2D(*result); + auto actual_literal = Literal::CreateR2FromArray2D(*result); LiteralTestUtil::ExpectR2Near({{1.f, 3.f, 5.f}, {5.f, 7.f, 9.f}}, - *result_literal, ErrorSpec(0.0001)); + *actual_literal, ErrorSpec(0.0001)); } TEST_F(ReferenceUtilTest, MapArray4D) { @@ -107,11 +108,11 @@ TEST_F(ReferenceUtilTest, MapArray4D) { input->FillWithMultiples(1.0f); auto multiply_by_two = [](float value) { return 2 * value; }; auto result = ReferenceUtil::MapArray4D(*input, multiply_by_two); - auto result_literal = LiteralUtil::CreateR4FromArray4D(*result); + auto actual_literal = Literal::CreateR4FromArray4D(*result); Array4D expected(/*planes=*/2, /*depth=*/3, /*height=*/4, /*width=*/5); expected.FillWithMultiples(2.0f); - LiteralTestUtil::ExpectR4NearArray4D(expected, *result_literal, + LiteralTestUtil::ExpectR4NearArray4D(expected, *actual_literal, ErrorSpec(0.0001)); } @@ -124,14 +125,109 @@ TEST_F(ReferenceUtilTest, MapWithIndexArray4D) { return value - (3 * 4 * 5 * plane + 4 * 5 * depth + 5 * height + width); }; auto result = ReferenceUtil::MapWithIndexArray4D(*input, subtract_index); - auto result_literal = LiteralUtil::CreateR4FromArray4D(*result); + auto actual_literal = Literal::CreateR4FromArray4D(*result); Array4D expected(/*planes=*/2, /*depth=*/3, /*height=*/4, /*width=*/5); expected.Fill(0.0f); - LiteralTestUtil::ExpectR4NearArray4D(expected, *result_literal, + LiteralTestUtil::ExpectR4NearArray4D(expected, *actual_literal, ErrorSpec(0.0001)); } +TEST_F(ReferenceUtilTest, SliceArray2D) { + auto result = ReferenceUtil::Slice2D(*matrix_, {{0, 0}}, {{2, 2}}, {{1, 1}}); + auto actual_literal = Literal::CreateR2FromArray2D(*result); + + LiteralTestUtil::ExpectR2Near({{1.f, 2.f}, {4.f, 5.f}}, + *actual_literal, ErrorSpec(0.0001)); +} + +TEST_F(ReferenceUtilTest, SliceStridedArray2D) { + auto result = ReferenceUtil::Slice2D(*matrix_, {{0, 0}}, {{2, 3}}, {{1, 2}}); + auto actual_literal = Literal::CreateR2FromArray2D(*result); + + LiteralTestUtil::ExpectR2Near({{1.f, 3.f}, {4.f, 6.f}}, + *actual_literal, ErrorSpec(0.0001)); +} + +TEST_F(ReferenceUtilTest, SliceArray3D) { + Array3D input(2, 3, 4); + input.FillIota(0); + + auto result = + ReferenceUtil::Slice3D(input, {{0, 0, 0}}, {{2, 2, 2}}, {{1, 1, 1}}); + auto actual_literal = Literal::CreateR3FromArray3D(*result); + + LiteralTestUtil::ExpectR3Near( + {{{0.f, 1.f}, {4.f, 5.f}}, {{12.f, 13.f}, {16.f, 17.f}}}, *actual_literal, + ErrorSpec(0.0001)); +} + +TEST_F(ReferenceUtilTest, SliceStridedArray3D) { + Array3D input(2, 3, 4); + input.FillIota(0); + + auto result = + ReferenceUtil::Slice3D(input, {{0, 0, 0}}, {{2, 3, 4}}, {{1, 2, 2}}); + auto actual_literal = Literal::CreateR3FromArray3D(*result); + + LiteralTestUtil::ExpectR3Near( + {{{0.f, 2.f}, {8.f, 10.f}}, {{12.f, 14.f}, {20.f, 22.f}}}, + *actual_literal, ErrorSpec(0.0001)); +} + +TEST_F(ReferenceUtilTest, SliceArray4D) { + Array4D input(2, 3, 4, 5); + input.FillIota(0); + + auto result = ReferenceUtil::Slice4D(input, {{1, 0, 0, 0}}, {{2, 2, 2, 2}}, + {{1, 1, 1, 1}}); + auto actual_literal = Literal::CreateR4FromArray4D(*result); + + LiteralTestUtil::ExpectR4Near( + {{{{60.f, 61.f}, {65.f, 66.f}}, {{80.f, 81.f}, {85.f, 86.f}}}}, + *actual_literal, ErrorSpec(0.0001)); +} + +TEST_F(ReferenceUtilTest, SliceStridedArray4D) { + Array4D input(2, 3, 4, 5); + input.FillIota(0); + + auto result = ReferenceUtil::Slice4D(input, {{1, 0, 0, 0}}, {{2, 3, 4, 5}}, + {{1, 2, 2, 2}}); + auto actual_literal = Literal::CreateR4FromArray4D(*result); + + LiteralTestUtil::ExpectR4Near( + {{{{60.f, 62.f, 64.f}, {70.f, 72.f, 74.f}}, + {{100.f, 102.f, 104.f}, {110.f, 112.f, 114.f}}}}, + *actual_literal, ErrorSpec(0.0001)); +} + +TEST_F(ReferenceUtilTest, ConvArray3DWithSamePadding) { + Array3D input = {{{1, 2, 3, 4}}}; + Array3D weights = {{{5, 6}}}; + std::unique_ptr> actual = + ReferenceUtil::ConvArray3D(input, weights, 1, Padding::kSame); + Array3D expected = {{{17, 28, 39, 20}}}; + + auto actual_literal = Literal::CreateR3FromArray3D(*actual); + + LiteralTestUtil::ExpectR3NearArray3D(expected, *actual_literal, + ErrorSpec(0.0001)); +} + +TEST_F(ReferenceUtilTest, ConvArray3DWithValidPadding) { + Array3D input = {{{1, 2, 3, 4}}}; + Array3D weights = {{{5, 6}}}; + std::unique_ptr> actual = + ReferenceUtil::ConvArray3D(input, weights, 1, Padding::kValid); + Array3D expected = {{{17, 28, 39}}}; + + auto actual_literal = Literal::CreateR3FromArray3D(*actual); + + LiteralTestUtil::ExpectR3NearArray3D(expected, *actual_literal, + ErrorSpec(0.0001)); +} + TEST_F(ReferenceUtilTest, ConvWithSamePadding) { Array4D input(1, 1, 4, 4); // clang-format off @@ -161,7 +257,7 @@ TEST_F(ReferenceUtilTest, ConvWithSamePadding) { })); // clang-format on - auto actual_literal = LiteralUtil::CreateR4FromArray4D(*actual); + auto actual_literal = Literal::CreateR4FromArray4D(*actual); LiteralTestUtil::ExpectR4NearArray4D(expected, *actual_literal, ErrorSpec(0.0001)); @@ -195,7 +291,7 @@ TEST_F(ReferenceUtilTest, ConvWithValidPadding) { })); // clang-format on - auto actual_literal = LiteralUtil::CreateR4FromArray4D(*actual); + auto actual_literal = Literal::CreateR4FromArray4D(*actual); LiteralTestUtil::ExpectR4NearArray4D(expected, *actual_literal, ErrorSpec(0.0001)); @@ -247,7 +343,7 @@ TEST_F(ReferenceUtilTest, ConvGeneralDimensionsWithSamePadding) { }}); // clang-format on - auto actual_literal = LiteralUtil::CreateR4FromArray4D(*actual); + auto actual_literal = Literal::CreateR4FromArray4D(*actual); LiteralTestUtil::ExpectR4NearArray4D(expected, *actual_literal, ErrorSpec(0.0001)); @@ -296,11 +392,23 @@ TEST_F(ReferenceUtilTest, ConvGeneralDimensionsWithValidPadding) { Array4D expected({{{{2514, 2685}}}}); // clang-format on - auto actual_literal = LiteralUtil::CreateR4FromArray4D(*actual); + auto actual_literal = Literal::CreateR4FromArray4D(*actual); LiteralTestUtil::ExpectR4NearArray4D(expected, *actual_literal, ErrorSpec(0.0001)); } +TEST_F(ReferenceUtilTest, ApplyElementwise2D) { + Array2D a({{1, 2}, {3, 4}}); + Array2D b({{10, 20}, {30, 40}}); + Array2D c({{100, 200}, {300, 400}}); + + auto actual = ReferenceUtil::ApplyElementwise2D( + [](float x, float y, float z) { return 100 * x + 10 * y + z; }, a, b, c); + auto actual_literal = Literal::CreateR2FromArray2D(*actual); + LiteralTestUtil::ExpectR2Near({{300.f, 600.f}, {900.f, 1200.f}}, + *actual_literal, ErrorSpec(0.0001)); +} + } // namespace } // namespace xla diff --git a/tensorflow/compiler/xla/service/BUILD b/tensorflow/compiler/xla/service/BUILD index b9118fab2549689d045a1caf826b9d3937019e1c..98cc3401c14f93cc2f209806baa0d97f41582f4b 100644 --- a/tensorflow/compiler/xla/service/BUILD +++ b/tensorflow/compiler/xla/service/BUILD @@ -21,6 +21,12 @@ xla_proto_library( deps = ["//tensorflow/compiler/xla:xla_data_proto"], ) +xla_proto_library( + name = "hlo_proto", + srcs = ["hlo.proto"], + deps = ["//tensorflow/compiler/xla:xla_data_proto"], +) + # Filegroup used to collect source files for dependency checking. filegroup( name = "c_srcs", @@ -35,6 +41,7 @@ cc_library( srcs = ["shape_inference.cc"], hdrs = ["shape_inference.h"], deps = [ + ":hlo", "//tensorflow/compiler/xla:shape_util", "//tensorflow/compiler/xla:status_macros", "//tensorflow/compiler/xla:statusor", @@ -48,10 +55,12 @@ cc_library( cc_test( name = "shape_inference_test", + size = "small", srcs = ["shape_inference_test.cc"], deps = [ ":shape_inference", "//tensorflow/compiler/xla:shape_util", + "//tensorflow/compiler/xla:test", "//tensorflow/compiler/xla:test_helpers", "//tensorflow/compiler/xla:types", "//tensorflow/compiler/xla:xla_data_proto", @@ -61,11 +70,55 @@ cc_test( cc_test( name = "hlo_opcode_test", + size = "small", srcs = ["hlo_opcode_test.cc"], deps = [ ":hlo", + "//tensorflow/compiler/xla:test", "//tensorflow/compiler/xla:types", - "//tensorflow/core:test", + "//tensorflow/core:test_main", + ], +) + +cc_library( + name = "hlo_evaluator", + srcs = ["hlo_evaluator.cc"], + hdrs = ["hlo_evaluator.h"], + deps = [ + ":hlo", + ":hlo_query", + ":shape_inference", + "//tensorflow/compiler/xla:literal_util", + "//tensorflow/compiler/xla:shape_util", + "//tensorflow/compiler/xla:statusor", + "//tensorflow/compiler/xla:types", + "//tensorflow/compiler/xla:util", + "//tensorflow/compiler/xla:window_util", + "//tensorflow/compiler/xla:xla_data_proto", + "//tensorflow/core:lib", + ], +) + +cc_test( + name = "hlo_evaluator_test", + size = "small", + srcs = ["hlo_evaluator_test.cc"], + deps = [ + ":hlo", + ":hlo_evaluator", + "//tensorflow/compiler/xla:literal_util", + "//tensorflow/compiler/xla:reference_util", + "//tensorflow/compiler/xla:shape_util", + "//tensorflow/compiler/xla:status_macros", + "//tensorflow/compiler/xla:statusor", + "//tensorflow/compiler/xla:test", + "//tensorflow/compiler/xla:types", + "//tensorflow/compiler/xla:util", + "//tensorflow/compiler/xla:xla_data_proto", + "//tensorflow/compiler/xla/client:computation_builder", + "//tensorflow/compiler/xla/tests:hlo_test_base", + "//tensorflow/compiler/xla/tests:literal_test_util", + "//tensorflow/core:lib", "//tensorflow/core:test_main", ], ) @@ -88,10 +141,14 @@ cc_library( "hlo_opcode.h", ], deps = [ + ":hlo_module_config", + ":hlo_proto", + ":hlo_reachability", ":name_uniquer", ":versioned_computation_handle", "//tensorflow/compiler/xla:literal_util", "//tensorflow/compiler/xla:protobuf_util", + "//tensorflow/compiler/xla:shape_tree", "//tensorflow/compiler/xla:shape_util", "//tensorflow/compiler/xla:status", "//tensorflow/compiler/xla:status_macros", @@ -105,6 +162,53 @@ cc_library( ], ) +cc_library( + name = "hlo_reachability", + srcs = ["hlo_reachability.cc"], + hdrs = ["hlo_reachability.h"], + deps = [ + "//tensorflow/compiler/xla:types", + "//tensorflow/compiler/xla:util", + "//tensorflow/core:lib", + "//tensorflow/core:lib_internal", + ], +) + +cc_test( + name = "hlo_reachability_test", + srcs = ["hlo_reachability_test.cc"], + deps = [ + ":hlo", + ":hlo_reachability", + "//tensorflow/compiler/xla:test", + "//tensorflow/compiler/xla:test_helpers", + "//tensorflow/compiler/xla/tests:hlo_test_base", + "//tensorflow/core:test_main", + ], +) + +cc_library( + name = "hlo_matchers", + testonly = 1, + srcs = ["hlo_matchers.cc"], + hdrs = ["hlo_matchers.h"], + deps = [ + ":hlo", + "//tensorflow/compiler/xla:test", + "//tensorflow/core:test_main", + ], +) + +cc_test( + name = "hlo_matchers_test", + size = "small", + srcs = ["hlo_matchers_test.cc"], + deps = [ + ":hlo_matchers", + "//tensorflow/compiler/xla:shape_util", + ], +) + cc_library( name = "versioned_computation_handle", srcs = ["versioned_computation_handle.cc"], @@ -118,15 +222,17 @@ cc_library( cc_test( name = "hlo_instruction_test", + size = "small", srcs = ["hlo_instruction_test.cc"], deps = [ ":hlo", "//tensorflow/compiler/xla:literal_util", + "//tensorflow/compiler/xla:protobuf_util", "//tensorflow/compiler/xla:shape_util", + "//tensorflow/compiler/xla:test", "//tensorflow/compiler/xla:test_helpers", "//tensorflow/compiler/xla:util", "//tensorflow/compiler/xla/tests:hlo_test_base", - "//tensorflow/core:test_main", ], ) @@ -137,7 +243,6 @@ cc_library( deps = [ ":hlo", "//tensorflow/compiler/xla:status_macros", - "//tensorflow/compiler/xla:statusor", "//tensorflow/compiler/xla:util", "//tensorflow/core:lib", ], @@ -145,19 +250,86 @@ cc_library( cc_test( name = "call_graph_test", + size = "small", srcs = ["call_graph_test.cc"], deps = [ ":call_graph", "//tensorflow/compiler/xla:literal_util", "//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/service:hlo", + "//tensorflow/compiler/xla/tests:hlo_test_base", + "//tensorflow/core:test", + ], +) + +cc_library( + name = "flatten_call_graph", + srcs = ["flatten_call_graph.cc"], + hdrs = ["flatten_call_graph.h"], + deps = [ + ":call_graph", + ":hlo", + ":hlo_pass", + "//tensorflow/compiler/xla:statusor", + "//tensorflow/compiler/xla:util", + "//tensorflow/core:lib", + ], +) + +cc_library( + name = "call_inliner", + srcs = ["call_inliner.cc"], + hdrs = ["call_inliner.h"], + deps = [ + ":hlo_pass", + "//tensorflow/compiler/xla:statusor", + "//tensorflow/core:lib", + ], +) + +cc_test( + name = "call_inliner_test", + size = "small", + srcs = ["call_inliner_test.cc"], + deps = [ + ":call_inliner", + ":hlo", + ":hlo_matchers", + ":hlo_pass", + "//tensorflow/compiler/xla:literal_util", + "//tensorflow/compiler/xla:shape_util", + "//tensorflow/compiler/xla:test", + "//tensorflow/compiler/xla:types", + "//tensorflow/compiler/xla:util", + "//tensorflow/compiler/xla:xla_data_proto", + "//tensorflow/compiler/xla/tests:hlo_test_base", + "//tensorflow/core:lib", + "//tensorflow/core:test", + ], +) + +cc_test( + name = "flatten_call_graph_test", + size = "small", + srcs = ["flatten_call_graph_test.cc"], + deps = [ + ":call_graph", + ":flatten_call_graph", + "//tensorflow/compiler/xla:literal_util", + "//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/service:hlo", "//tensorflow/compiler/xla/tests:hlo_test_base", "//tensorflow/core:test", - "//tensorflow/core:test_main", ], ) @@ -177,23 +349,26 @@ cc_library( "//tensorflow/compiler/xla:types", "//tensorflow/compiler/xla:util", "//tensorflow/compiler/xla:xla_data_proto", + "//tensorflow/compiler/xla:xla_proto", "//tensorflow/core:lib", ], ) cc_test( name = "user_computation_test", + size = "small", srcs = ["user_computation_test.cc"], deps = [ + ":hlo_matchers", ":user_computation", "//tensorflow/compiler/xla:literal_util", "//tensorflow/compiler/xla:shape_util", "//tensorflow/compiler/xla:status_macros", + "//tensorflow/compiler/xla:test", "//tensorflow/compiler/xla:test_helpers", "//tensorflow/compiler/xla:xla_data_proto", "//tensorflow/compiler/xla/service:hlo", "//tensorflow/core:test", - "//tensorflow/core:test_main", ], ) @@ -218,6 +393,7 @@ cc_library( hdrs = ["backend.h"], deps = [ ":compiler", + ":computation_placer", ":device_memory_allocator", ":platform_util", ":pool", @@ -226,7 +402,6 @@ cc_library( "//tensorflow/compiler/xla:statusor", "//tensorflow/compiler/xla:types", "//tensorflow/compiler/xla:util", - "//tensorflow/compiler/xla/legacy_flags:backend_flags", "//tensorflow/core:core_cpu_internal", "//tensorflow/core:lib", "//tensorflow/core:stream_executor_no_cuda", @@ -250,10 +425,11 @@ cc_library( ":device_memory_allocator", ":executable", ":execution_tracker", + ":gpu_transfer_manager", ":hlo", ":hlo_cost_analysis", + ":hlo_evaluator", ":hlo_execution_profile", - ":hlo_graph_dumper", ":hlo_module_config", ":platform_util", ":session_proto", @@ -261,6 +437,7 @@ cc_library( ":user_computation", ":versioned_computation_handle", "//tensorflow/compiler/xla:executable_run_options", + "//tensorflow/compiler/xla:execution_options_util", "//tensorflow/compiler/xla:service_interface", "//tensorflow/compiler/xla:shape_layout", "//tensorflow/compiler/xla:shape_util", @@ -270,10 +447,9 @@ cc_library( "//tensorflow/compiler/xla:util", "//tensorflow/compiler/xla:xla_data_proto", "//tensorflow/compiler/xla:xla_proto", - "//tensorflow/compiler/xla/legacy_flags:service_flags", + "//tensorflow/compiler/xla/legacy_flags:debug_options_flags", "//tensorflow/compiler/xla/service/cpu:cpu_compiler", "//tensorflow/core:lib", - "//tensorflow/core:regexp_internal", "//tensorflow/core:stream_executor_no_cuda", ], alwayslink = 1, @@ -298,6 +474,7 @@ cc_library( ":shaped_buffer", ":user_computation", ":versioned_computation_handle", + "//tensorflow/compiler/xla:execution_options_util", "//tensorflow/compiler/xla:shape_layout", "//tensorflow/compiler/xla:shape_util", "//tensorflow/compiler/xla:status_macros", @@ -305,7 +482,28 @@ cc_library( "//tensorflow/compiler/xla:types", "//tensorflow/compiler/xla:util", "//tensorflow/compiler/xla:xla_data_proto", - "//tensorflow/compiler/xla/legacy_flags:service_flags", + "//tensorflow/core:lib", + "//tensorflow/core:stream_executor_no_cuda", + ], +) + +cc_library( + name = "compile_only_service", + srcs = ["compile_only_service.cc"], + hdrs = ["compile_only_service.h"], + deps = [ + ":backend", + ":compiler", + ":computation_layout", + ":computation_tracker", + ":platform_util", + ":service", + "//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:stream_executor_no_cuda", ], @@ -324,7 +522,7 @@ cc_library( cc_library( name = "gpu_plugin", deps = [ - ":generic_transfer_manager", + ":gpu_transfer_manager", ":service", "//tensorflow/compiler/xla/service/gpu:gpu_compiler", "//tensorflow/core:stream_executor_no_cuda", @@ -361,8 +559,9 @@ cc_library( ":computation_layout", ":device_memory_allocator", ":hlo", + ":hlo_cost_analysis", ":hlo_execution_profile", - ":hlo_module_config", + ":hlo_graph_dumper", ":pool", ":session_proto", ":shaped_buffer", @@ -370,8 +569,9 @@ cc_library( "//tensorflow/compiler/xla:executable_run_options", "//tensorflow/compiler/xla:status_macros", "//tensorflow/compiler/xla:statusor", + "//tensorflow/compiler/xla:util", "//tensorflow/compiler/xla:xla_data_proto", - "//tensorflow/compiler/xla/legacy_flags:service_flags", + "//tensorflow/compiler/xla/legacy_flags:debug_options_flags", "//tensorflow/core:lib", "//tensorflow/core:stream_executor_no_cuda", "//tensorflow/stream_executor", @@ -386,6 +586,7 @@ cc_library( ":executable", ":hlo", ":hlo_module_config", + ":logical_buffer", "//tensorflow/compiler/xla:statusor", "//tensorflow/compiler/xla:types", "//tensorflow/compiler/xla:util", @@ -394,11 +595,21 @@ cc_library( ], ) +cc_library( + name = "llvm_compiler", + hdrs = ["llvm_compiler.h"], + deps = [ + ":compiler", + "@llvm//:core", + ], +) + cc_library( name = "transfer_manager", srcs = ["transfer_manager.cc"], hdrs = ["transfer_manager.h"], deps = [ + "//tensorflow/compiler/xla:literal_util", "//tensorflow/compiler/xla:shape_util", "//tensorflow/compiler/xla:status_macros", "//tensorflow/compiler/xla:statusor", @@ -451,6 +662,7 @@ cc_library( hdrs = ["computation_tracker.h"], deps = [ ":hlo", + ":hlo_module_config", ":session_proto", ":user_computation", ":versioned_computation_handle", @@ -504,22 +716,18 @@ cc_library( "//tensorflow/compiler/xla:shape_util", "//tensorflow/compiler/xla:types", "//tensorflow/compiler/xla:util", - "//tensorflow/core:lib", ], ) cc_test( name = "liveness_util_test", + size = "small", srcs = ["liveness_util_test.cc"], deps = [ ":hlo", ":liveness_util", ":tuple_points_to_analysis", - "//tensorflow/compiler/xla:shape_util", - "//tensorflow/compiler/xla:types", - "//tensorflow/compiler/xla:util", "//tensorflow/compiler/xla/tests:hlo_test_base", - "//tensorflow/core:test_main", ], ) @@ -548,10 +756,10 @@ cc_library( cc_test( name = "buffer_liveness_test", + size = "small", srcs = ["buffer_liveness_test.cc"], deps = [ ":buffer_liveness", - ":cpu_plugin", ":hlo", "//tensorflow/compiler/xla:shape_util", "//tensorflow/compiler/xla:types", @@ -574,6 +782,8 @@ cc_library( ":buffer_liveness", ":heap_simulator", ":hlo", + ":hlo_proto", + ":hlo_scheduling", ":logical_buffer", ":tuple_points_to_analysis", "//tensorflow/compiler/xla:shape_util", @@ -582,7 +792,6 @@ cc_library( "//tensorflow/compiler/xla:types", "//tensorflow/compiler/xla:util", "//tensorflow/compiler/xla:xla_data_proto", - "//tensorflow/compiler/xla/legacy_flags:buffer_assignment_flags", "//tensorflow/core:lib", "//tensorflow/core:lib_internal", ], @@ -590,33 +799,72 @@ cc_library( cc_test( name = "buffer_assignment_test", + size = "small", srcs = ["buffer_assignment_test.cc"], deps = [ ":buffer_assignment", + ":call_graph", ":computation_tracker", - ":cpu_plugin", + ":copy_insertion", + ":flatten_call_graph", ":hlo", + ":hlo_ordering", + ":hlo_scheduling", "//tensorflow/compiler/xla:literal_util", "//tensorflow/compiler/xla:shape_util", + "//tensorflow/compiler/xla:test", + "//tensorflow/compiler/xla:test_helpers", "//tensorflow/compiler/xla:types", "//tensorflow/compiler/xla:util", "//tensorflow/compiler/xla:xla_data_proto", "//tensorflow/compiler/xla/tests:hlo_test_base", "//tensorflow/core:lib", - "//tensorflow/core:test_main", ], ) cc_library( - name = "heap_simulator", - srcs = [ - "heap_simulator.cc", + name = "hlo_ordering", + srcs = ["hlo_ordering.cc"], + hdrs = ["hlo_ordering.h"], + deps = [ + ":call_graph", + ":hlo", + ":hlo_proto", + ":hlo_value", + ":liveness_util", + "//tensorflow/compiler/xla:shape_util", + "//tensorflow/compiler/xla:status_macros", + "//tensorflow/compiler/xla:statusor", + "//tensorflow/compiler/xla:types", + "//tensorflow/compiler/xla:util", + "//tensorflow/core:lib", ], - hdrs = [ - "heap_simulator.h", +) + +cc_test( + name = "hlo_ordering_test", + size = "small", + srcs = ["hlo_ordering_test.cc"], + deps = [ + ":hlo", + ":hlo_dataflow_analysis", + ":hlo_ordering", + ":hlo_scheduling", + "//tensorflow/compiler/xla:shape_util", + "//tensorflow/compiler/xla:types", + "//tensorflow/compiler/xla:xla_data_proto", + "//tensorflow/compiler/xla/tests:hlo_test_base", ], +) + +cc_library( + name = "heap_simulator", + srcs = ["heap_simulator.cc"], + hdrs = ["heap_simulator.h"], deps = [ ":hlo", + ":hlo_ordering", + ":hlo_proto", ":liveness_util", ":logical_buffer", ":tuple_points_to_analysis", @@ -628,18 +876,16 @@ cc_library( cc_test( name = "heap_simulator_test", + size = "small", srcs = ["heap_simulator_test.cc"], deps = [ ":heap_simulator", ":hlo", + ":hlo_ordering", ":logical_buffer", ":tuple_points_to_analysis", "//tensorflow/compiler/xla:literal_util", - "//tensorflow/compiler/xla:shape_util", "//tensorflow/compiler/xla:status_macros", - "//tensorflow/compiler/xla:types", - "//tensorflow/compiler/xla:util", - "//tensorflow/compiler/xla:xla_data_proto", "//tensorflow/compiler/xla/tests:hlo_test_base", "//tensorflow/core:lib", "//tensorflow/core:test_main", @@ -647,16 +893,13 @@ cc_test( ) cc_library( - name = "hlo_ordering", - srcs = [ - "hlo_ordering.cc", - ], - hdrs = [ - "hlo_ordering.h", - ], + name = "hlo_scheduling", + srcs = ["hlo_scheduling.cc"], + hdrs = ["hlo_scheduling.h"], deps = [ ":heap_simulator", ":hlo", + ":hlo_ordering", ":logical_buffer", ":tuple_points_to_analysis", "//tensorflow/compiler/xla:shape_util", @@ -669,17 +912,17 @@ cc_library( ) cc_test( - name = "hlo_ordering_test", - srcs = ["hlo_ordering_test.cc"], + name = "hlo_scheduling_test", + size = "small", + srcs = ["hlo_scheduling_test.cc"], deps = [ - ":cpu_plugin", ":hlo", ":hlo_ordering", + ":hlo_scheduling", "//tensorflow/compiler/xla:shape_util", "//tensorflow/compiler/xla:types", "//tensorflow/compiler/xla:xla_data_proto", "//tensorflow/compiler/xla/tests:hlo_test_base", - "//tensorflow/core:test_main", ], ) @@ -708,11 +951,52 @@ cc_library( cc_test( name = "instruction_fusion_test", + size = "small", srcs = ["instruction_fusion_test.cc"], deps = [ + ":hlo_matchers", ":instruction_fusion", "//tensorflow/compiler/xla/tests:hlo_test_base", - "//tensorflow/core:test_main", + ], +) + +cc_library( + name = "batchnorm_rewriter", + srcs = ["batchnorm_rewriter.cc"], + hdrs = ["batchnorm_rewriter.h"], + deps = [ + ":hlo", + ":hlo_pass", + ":hlo_query", + ":shape_inference", + "//tensorflow/compiler/xla:literal_util", + "//tensorflow/compiler/xla:shape_util", + "//tensorflow/compiler/xla:status_macros", + "//tensorflow/compiler/xla:types", + "//tensorflow/compiler/xla:util", + "//tensorflow/compiler/xla:window_util", + "//tensorflow/compiler/xla:xla_data_proto", + "//tensorflow/core:lib", + ], +) + +cc_test( + name = "batchnorm_rewriter_test", + size = "small", + srcs = ["batchnorm_rewriter_test.cc"], + deps = [ + ":batchnorm_rewriter", + ":hlo", + ":hlo_matchers", + ":hlo_pass", + "//tensorflow/compiler/xla:literal_util", + "//tensorflow/compiler/xla:shape_util", + "//tensorflow/compiler/xla:test", + "//tensorflow/compiler/xla:types", + "//tensorflow/compiler/xla:util", + "//tensorflow/compiler/xla:xla_data_proto", + "//tensorflow/compiler/xla/tests:hlo_test_base", + "//tensorflow/core:lib", ], ) @@ -738,21 +1022,21 @@ cc_library( cc_test( name = "algebraic_simplifier_test", + size = "small", srcs = ["algebraic_simplifier_test.cc"], deps = [ ":algebraic_simplifier", - ":cpu_plugin", ":hlo", + ":hlo_matchers", ":hlo_pass", "//tensorflow/compiler/xla:literal_util", "//tensorflow/compiler/xla:shape_util", - "//tensorflow/compiler/xla:test_helpers", + "//tensorflow/compiler/xla:test", "//tensorflow/compiler/xla:types", "//tensorflow/compiler/xla:util", "//tensorflow/compiler/xla:xla_data_proto", "//tensorflow/compiler/xla/tests:hlo_test_base", "//tensorflow/core:lib", - "//tensorflow/core:test_main", ], ) @@ -764,25 +1048,29 @@ cc_library( ":hlo_pass", "//tensorflow/compiler/xla:literal_util", "//tensorflow/compiler/xla:shape_util", + "//tensorflow/compiler/xla:status_macros", "//tensorflow/compiler/xla:util", + "//tensorflow/core:lib", ], ) cc_test( name = "reshape_mover_test", + size = "small", srcs = ["reshape_mover_test.cc"], deps = [ ":hlo", + ":hlo_matchers", ":reshape_mover", "//tensorflow/compiler/xla:literal_util", "//tensorflow/compiler/xla:shape_util", + "//tensorflow/compiler/xla:test", "//tensorflow/compiler/xla:test_helpers", "//tensorflow/compiler/xla:types", "//tensorflow/compiler/xla:util", "//tensorflow/compiler/xla:xla_data_proto", "//tensorflow/compiler/xla/tests:hlo_test_base", "//tensorflow/core:lib", - "//tensorflow/core:test_main", ], ) @@ -802,14 +1090,15 @@ cc_library( cc_test( name = "inliner_test", + size = "small", srcs = ["inliner_test.cc"], deps = [ - ":cpu_plugin", ":hlo", + ":hlo_matchers", ":inliner", "//tensorflow/compiler/xla:literal_util", "//tensorflow/compiler/xla:shape_util", - "//tensorflow/compiler/xla:test_helpers", + "//tensorflow/compiler/xla:test", "//tensorflow/compiler/xla:util", "//tensorflow/compiler/xla:xla_data_proto", "//tensorflow/compiler/xla/tests:hlo_test_base", @@ -818,6 +1107,38 @@ cc_test( ], ) +cc_library( + name = "computation_placer", + srcs = ["computation_placer.cc"], + hdrs = ["computation_placer.h"], + deps = [ + "//tensorflow/compiler/xla:array2d", + "//tensorflow/compiler/xla:literal_util", + "//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/core:lib", + "//tensorflow/core:stream_executor_no_cuda", + ], + alwayslink = True, # Contains per-platform computation placer registration +) + +cc_library( + name = "human_readable_profile_builder", + srcs = ["human_readable_profile_builder.cc"], + hdrs = ["human_readable_profile_builder.h"], + deps = [ + "//tensorflow/compiler/xla:metric_table_report", + "//tensorflow/compiler/xla:types", + "//tensorflow/compiler/xla:util", + "//tensorflow/core:lib", + ], +) + cc_library( name = "generic_transfer_manager", srcs = ["generic_transfer_manager.cc"], @@ -858,8 +1179,30 @@ cc_library( alwayslink = True, # Contains per-platform transfer manager registration ) +cc_library( + name = "gpu_transfer_manager", + srcs = ["gpu_transfer_manager.cc"], + hdrs = ["gpu_transfer_manager.h"], + deps = [ + ":generic_transfer_manager", + ":transfer_manager", + "//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:util", + "//tensorflow/compiler/xla:xla_data_proto", + "//tensorflow/compiler/xla/service/gpu:infeed_manager", + "//tensorflow/core:lib", + "//tensorflow/core:stream_executor_no_cuda", + ], + alwayslink = True, # Contains per-platform transfer manager registration +) + cc_test( name = "transfer_manager_test", + size = "small", srcs = ["transfer_manager_test.cc"], deps = [ ":cpu_transfer_manager", @@ -878,12 +1221,8 @@ cc_test( cc_library( name = "hlo_cost_analysis", - srcs = [ - "hlo_cost_analysis.cc", - ], - hdrs = [ - "hlo_cost_analysis.h", - ], + srcs = ["hlo_cost_analysis.cc"], + hdrs = ["hlo_cost_analysis.h"], deps = [ ":hlo", "//tensorflow/compiler/xla:shape_util", @@ -897,6 +1236,7 @@ cc_library( cc_test( name = "hlo_cost_analysis_test", + size = "small", srcs = ["hlo_cost_analysis_test.cc"], deps = [ ":computation_tracker", @@ -915,6 +1255,7 @@ cc_test( "//tensorflow/compiler/xla/client:computation_builder", "//tensorflow/compiler/xla/client:local_client", "//tensorflow/compiler/xla/client:padding", + "//tensorflow/compiler/xla/tests:hlo_test_base", "//tensorflow/core:lib", "//tensorflow/core:test_main", ], @@ -927,7 +1268,7 @@ cc_library( deps = [ ":hlo", ":hlo_cost_analysis", - "//tensorflow/compiler/xla:metric_table_report", + ":human_readable_profile_builder", "//tensorflow/compiler/xla:types", "//tensorflow/compiler/xla:util", "//tensorflow/core:lib", @@ -937,12 +1278,14 @@ cc_library( cc_test( name = "hlo_computation_test", + size = "small", srcs = ["hlo_computation_test.cc"], deps = [ - ":cpu_plugin", ":hlo", + ":hlo_matchers", "//tensorflow/compiler/xla:literal_util", "//tensorflow/compiler/xla:shape_util", + "//tensorflow/compiler/xla:test", "//tensorflow/compiler/xla:test_helpers", "//tensorflow/compiler/xla/tests:hlo_test_base", "//tensorflow/core:test_main", @@ -966,49 +1309,171 @@ cc_binary( cc_test( name = "hlo_module_test", + size = "small", srcs = ["hlo_module_test.cc"], deps = [ - ":cpu_plugin", ":hlo", "//tensorflow/compiler/xla:literal_util", "//tensorflow/compiler/xla:shape_util", - "//tensorflow/compiler/xla:test_helpers", + "//tensorflow/compiler/xla:test", "//tensorflow/compiler/xla:util", "//tensorflow/compiler/xla:xla_data_proto", "//tensorflow/compiler/xla/tests:hlo_test_base", "//tensorflow/core:lib", - "//tensorflow/core:test_main", ], ) cc_library( name = "logical_buffer", - srcs = [ - "logical_buffer.cc", + srcs = ["logical_buffer.cc"], + hdrs = ["logical_buffer.h"], + deps = [ + ":hlo", + ":hlo_proto", + "//tensorflow/compiler/xla:shape_util", + "//tensorflow/compiler/xla:types", + "//tensorflow/compiler/xla:xla_data_proto", + "//tensorflow/core:lib", + "//tensorflow/core:lib_internal", ], - hdrs = [ - "logical_buffer.h", +) + +cc_library( + name = "hlo_value", + srcs = ["hlo_value.cc"], + hdrs = ["hlo_value.h"], + deps = [ + ":hlo", + "//tensorflow/compiler/xla:shape_tree", + "//tensorflow/compiler/xla:shape_util", + "//tensorflow/compiler/xla:status", + "//tensorflow/compiler/xla:types", + "//tensorflow/compiler/xla:util", + "//tensorflow/compiler/xla:xla_data_proto", + "//tensorflow/core:lib", ], +) + +cc_library( + name = "hlo_dataflow_analysis", + srcs = ["hlo_dataflow_analysis.cc"], + hdrs = ["hlo_dataflow_analysis.h"], deps = [ + ":call_graph", ":hlo", + ":hlo_ordering", + ":hlo_value", + ":liveness_util", "//tensorflow/compiler/xla:shape_util", + "//tensorflow/compiler/xla:status", + "//tensorflow/compiler/xla:statusor", "//tensorflow/compiler/xla:types", + "//tensorflow/compiler/xla:util", "//tensorflow/compiler/xla:xla_data_proto", "//tensorflow/core:lib", ], ) +cc_test( + name = "hlo_dataflow_analysis_test", + size = "small", + srcs = ["hlo_dataflow_analysis_test.cc"], + deps = [ + ":hlo", + ":hlo_dataflow_analysis", + ":hlo_matchers", + ":instruction_fusion", + "//tensorflow/compiler/xla:literal_util", + "//tensorflow/compiler/xla:shape_util", + "//tensorflow/compiler/xla:status_macros", + "//tensorflow/compiler/xla:test", + "//tensorflow/compiler/xla:test_helpers", + "//tensorflow/compiler/xla:xla_data_proto", + "//tensorflow/compiler/xla/tests:hlo_test_base", + "//tensorflow/core:lib", + "//tensorflow/core:test", + "//tensorflow/core:test_main", + ], +) + cc_library( - name = "tuple_points_to_analysis", - srcs = [ - "tuple_points_to_analysis.cc", + name = "hlo_buffer", + srcs = ["hlo_buffer.cc"], + hdrs = ["hlo_buffer.h"], + deps = [ + ":hlo", + ":hlo_value", + "//tensorflow/compiler/xla:shape_tree", + "//tensorflow/compiler/xla:shape_util", + "//tensorflow/compiler/xla:types", + "//tensorflow/compiler/xla:util", + "//tensorflow/compiler/xla:xla_data_proto", + "//tensorflow/core:lib", ], - hdrs = [ - "tuple_points_to_analysis.h", +) + +cc_library( + name = "hlo_alias_analysis", + srcs = ["hlo_alias_analysis.cc"], + hdrs = ["hlo_alias_analysis.h"], + deps = [ + ":hlo", + ":hlo_buffer", + ":hlo_dataflow_analysis", + ":hlo_value", + "//tensorflow/compiler/xla:shape_util", + "//tensorflow/compiler/xla:status", + "//tensorflow/compiler/xla:statusor", + "//tensorflow/compiler/xla:types", + "//tensorflow/compiler/xla:util", + "//tensorflow/compiler/xla:xla_data_proto", + "//tensorflow/core:lib", ], +) + +cc_test( + name = "hlo_alias_analysis_test", + srcs = ["hlo_alias_analysis_test.cc"], + deps = [ + ":flatten_call_graph", + ":hlo", + ":hlo_alias_analysis", + ":hlo_matchers", + ":instruction_fusion", + "//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/core:lib", + "//tensorflow/core:test", + "//tensorflow/core:test_main", + ], +) + +cc_library( + name = "logical_buffer_analysis", + srcs = ["logical_buffer_analysis.cc"], + hdrs = ["logical_buffer_analysis.h"], + deps = [ + ":hlo", + ":logical_buffer", + "//tensorflow/compiler/xla:shape_util", + "//tensorflow/compiler/xla:statusor", + "//tensorflow/core:lib", + "//tensorflow/core:lib_internal", + ], +) + +cc_library( + name = "tuple_points_to_analysis", + srcs = ["tuple_points_to_analysis.cc"], + hdrs = ["tuple_points_to_analysis.h"], deps = [ ":hlo", ":logical_buffer", + ":logical_buffer_analysis", "//tensorflow/compiler/xla:shape_tree", "//tensorflow/compiler/xla:shape_util", "//tensorflow/compiler/xla:statusor", @@ -1021,13 +1486,16 @@ cc_library( cc_test( name = "tuple_points_to_analysis_test", + size = "small", srcs = ["tuple_points_to_analysis_test.cc"], deps = [ ":hlo", + ":hlo_matchers", ":instruction_fusion", ":tuple_points_to_analysis", "//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", @@ -1039,12 +1507,8 @@ cc_test( cc_library( name = "compilation_cache", - srcs = [ - "compilation_cache.cc", - ], - hdrs = [ - "compilation_cache.h", - ], + srcs = ["compilation_cache.cc"], + hdrs = ["compilation_cache.h"], deps = [ ":executable", ":hlo_module_config", @@ -1089,6 +1553,7 @@ cc_library( ":buffer_liveness", ":hlo", ":hlo_pass", + ":liveness_util", ":logical_buffer", ":tuple_points_to_analysis", "//tensorflow/compiler/xla:status_macros", @@ -1101,19 +1566,19 @@ cc_library( cc_test( name = "copy_insertion_test", + size = "small", srcs = ["copy_insertion_test.cc"], deps = [ - ":buffer_liveness", ":copy_insertion", - ":cpu_plugin", ":hlo", + ":hlo_matchers", ":tuple_points_to_analysis", "//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/core:test_main", ], ) @@ -1138,13 +1603,9 @@ cc_library( srcs = ["hlo_verifier.cc"], hdrs = ["hlo_verifier.h"], deps = [ - ":hlo", ":hlo_pass", - "//tensorflow/compiler/xla:status", + ":shape_inference", "//tensorflow/compiler/xla:status_macros", - "//tensorflow/compiler/xla:statusor", - "//tensorflow/compiler/xla:types", - "//tensorflow/compiler/xla:util", "//tensorflow/core:lib", ], ) @@ -1156,10 +1617,12 @@ cc_library( deps = [ ":buffer_liveness", ":call_graph", + ":flatten_call_graph", ":hlo", - ":hlo_cost_analysis", ":hlo_dce", ":hlo_ordering", + ":hlo_scheduling", + ":liveness_util", ":logical_buffer", ":tuple_points_to_analysis", "//tensorflow/compiler/xla:shape_util", @@ -1173,10 +1636,11 @@ cc_library( cc_test( name = "hlo_rematerialization_test", + size = "small", srcs = ["hlo_rematerialization_test.cc"], deps = [ - ":cpu_plugin", ":hlo", + ":hlo_matchers", ":hlo_ordering", ":hlo_rematerialization", "//tensorflow/compiler/xla:shape_util", @@ -1189,9 +1653,9 @@ cc_test( cc_test( name = "hlo_dce_test", + size = "small", srcs = ["hlo_dce_test.cc"], deps = [ - ":cpu_plugin", ":hlo", ":hlo_dce", "//tensorflow/compiler/xla:literal_util", @@ -1203,29 +1667,30 @@ cc_test( "//tensorflow/compiler/xla/tests:literal_test_util", "//tensorflow/compiler/xla/tests:test_utils", "//tensorflow/core:lib", - "//tensorflow/core:test_main", + "//tensorflow/core:test", ], ) cc_test( name = "layout_assignment_test", + size = "small", srcs = ["layout_assignment_test.cc"], deps = [ ":algebraic_simplifier", ":computation_layout", - ":cpu_plugin", ":hlo", + ":hlo_matchers", ":layout_assignment", "//tensorflow/compiler/xla:literal_util", "//tensorflow/compiler/xla:shape_layout", "//tensorflow/compiler/xla:shape_util", + "//tensorflow/compiler/xla:test", "//tensorflow/compiler/xla:test_helpers", "//tensorflow/compiler/xla:util", "//tensorflow/compiler/xla:xla_data_proto", "//tensorflow/compiler/xla/tests:hlo_test_base", "//tensorflow/compiler/xla/tests:test_utils", "//tensorflow/core:lib", - "//tensorflow/core:test_main", ], ) @@ -1253,14 +1718,13 @@ cc_library( "hlo_pass_pipeline.h", ], deps = [ - ":compiler", ":hlo", + ":hlo_graph_dumper", ":hlo_pass", "//tensorflow/compiler/xla:status_macros", "//tensorflow/compiler/xla:statusor", "//tensorflow/compiler/xla:types", "//tensorflow/compiler/xla:util", - "//tensorflow/compiler/xla/legacy_flags:hlo_pass_pipeline_flags", "//tensorflow/core:lib", ], ) @@ -1274,7 +1738,6 @@ cc_library( ":hlo_pass", "//tensorflow/compiler/xla:literal_util", "//tensorflow/compiler/xla:shape_util", - "//tensorflow/compiler/xla:status_macros", "//tensorflow/compiler/xla:types", "//tensorflow/compiler/xla:xla_data_proto", "//tensorflow/core:lib", @@ -1283,11 +1746,12 @@ cc_library( cc_test( name = "hlo_cse_test", + size = "small", srcs = ["hlo_cse_test.cc"], deps = [ - ":cpu_plugin", ":hlo", ":hlo_cse", + ":hlo_matchers", "//tensorflow/compiler/xla:literal_util", "//tensorflow/compiler/xla:shape_util", "//tensorflow/compiler/xla:types", @@ -1297,7 +1761,6 @@ cc_test( "//tensorflow/compiler/xla/tests:literal_test_util", "//tensorflow/compiler/xla/tests:test_utils", "//tensorflow/core:lib", - "//tensorflow/core:test_main", ], ) @@ -1307,16 +1770,34 @@ cc_library( hdrs = ["hlo_constant_folding.h"], deps = [ ":hlo", + ":hlo_evaluator", ":hlo_pass", + ":hlo_query", "//tensorflow/compiler/xla:literal_util", "//tensorflow/compiler/xla:shape_util", - "//tensorflow/compiler/xla:status_macros", "//tensorflow/compiler/xla:types", - "//tensorflow/compiler/xla:xla_data_proto", "//tensorflow/core:lib", ], ) +cc_test( + name = "hlo_constant_folding_test", + size = "small", + srcs = ["hlo_constant_folding_test.cc"], + deps = [ + ":hlo", + ":hlo_constant_folding", + ":hlo_matchers", + ":hlo_pass", + "//tensorflow/compiler/xla:literal_util", + "//tensorflow/compiler/xla:shape_util", + "//tensorflow/compiler/xla:test", + "//tensorflow/compiler/xla:types", + "//tensorflow/compiler/xla/tests:hlo_test_base", + "//tensorflow/compiler/xla/tests:literal_test_util", + ], +) + cc_library( name = "device_memory_allocator", srcs = ["device_memory_allocator.cc"], @@ -1345,6 +1826,7 @@ cc_library( "//tensorflow/compiler/xla:util", "//tensorflow/compiler/xla:xla_data_proto", "//tensorflow/compiler/xla/service/llvm_ir:ir_array", + "//tensorflow/compiler/xla/service/llvm_ir:llvm_loop", "//tensorflow/compiler/xla/service/llvm_ir:llvm_util", "//tensorflow/compiler/xla/service/llvm_ir:loop_emitter", "//tensorflow/core:lib", @@ -1362,6 +1844,7 @@ cc_library( ":computation_layout", "//tensorflow/compiler/xla:shape_layout", "//tensorflow/compiler/xla:types", + "//tensorflow/compiler/xla:util", "//tensorflow/compiler/xla:xla_data_proto", "//tensorflow/compiler/xla:xla_proto", "//tensorflow/core:lib", @@ -1391,6 +1874,7 @@ cc_library( cc_test( name = "hlo_subcomputation_unification_test", + size = "small", srcs = ["hlo_subcomputation_unification_test.cc"], deps = [ ":hlo", @@ -1399,7 +1883,33 @@ cc_test( "//tensorflow/compiler/xla:shape_util", "//tensorflow/compiler/xla/tests:hlo_test_base", "//tensorflow/compiler/xla/tests:test_utils", - "//tensorflow/core:test_main", + ], +) + +cc_library( + name = "hlo_tfgraph_builder", + srcs = ["hlo_tfgraph_builder.cc"], + hdrs = ["hlo_tfgraph_builder.h"], + visibility = ["//tensorflow/compiler/xla/tools:__pkg__"], + deps = [ + ":hlo", + "//tensorflow/compiler/xla:literal_util", + "//tensorflow/compiler/xla:shape_util", + "//tensorflow/core:framework", + "//tensorflow/core:lib", + "//tensorflow/core:protos_all_cc", + ], +) + +cc_test( + name = "hlo_tfgraph_builder_test", + size = "small", + srcs = ["hlo_tfgraph_builder_test.cc"], + deps = [ + ":hlo_tfgraph_builder", + "//tensorflow/compiler/xla/client:computation_builder", + "//tensorflow/compiler/xla/tests:hlo_test_base", + "//tensorflow/core:protos_all_cc", ], ) @@ -1412,12 +1922,14 @@ cc_library( deps = [ ":hlo", ":hlo_execution_profile", + ":hlo_tfgraph_builder", "//tensorflow/compiler/xla:literal_util", "//tensorflow/compiler/xla:shape_util", "//tensorflow/compiler/xla:types", "//tensorflow/compiler/xla:window_util", - "//tensorflow/compiler/xla/legacy_flags:hlo_graph_dumper_flags", + "//tensorflow/compiler/xla:xla_proto", "//tensorflow/core:lib", + "//tensorflow/core:regexp_internal", ], alwayslink = 1, ) @@ -1429,7 +1941,9 @@ cc_library( deps = [ ":hlo", ":hlo_pass", + "//tensorflow/compiler/xla:shape_util", "//tensorflow/compiler/xla:status_macros", + "//tensorflow/compiler/xla:util", "//tensorflow/compiler/xla/service/gpu:ir_emission_utils", "//tensorflow/core:lib", ], @@ -1437,14 +1951,18 @@ cc_library( cc_test( name = "transpose_folding_test", + size = "small", srcs = ["transpose_folding_test.cc"], deps = [ ":hlo", + ":shape_inference", ":transpose_folding", "//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/client:computation_builder", "//tensorflow/compiler/xla/service/gpu:ir_emission_utils", "//tensorflow/core:lib", "//tensorflow/core:test_main", @@ -1462,6 +1980,7 @@ cc_library( cc_test( name = "pool_test", + size = "small", srcs = ["pool_test.cc"], deps = [ ":pool", @@ -1470,6 +1989,48 @@ cc_test( ], ) +cc_library( + name = "hlo_proto_util", + srcs = ["hlo_proto_util.cc"], + hdrs = ["hlo_proto_util.h"], + deps = [ + ":buffer_assignment", + ":hlo", + ":hlo_proto", + "//tensorflow/compiler/xla:status", + ], +) + +cc_library( + name = "reduce_precision_insertion", + srcs = ["reduce_precision_insertion.cc"], + hdrs = ["reduce_precision_insertion.h"], + deps = [ + ":buffer_liveness", + ":hlo", + ":hlo_pass", + ":hlo_pass_pipeline", + "//tensorflow/compiler/xla:shape_util", + "//tensorflow/core:lib", + ], +) + +cc_test( + name = "reduce_precision_insertion_test", + size = "small", + srcs = ["reduce_precision_insertion_test.cc"], + deps = [ + ":hlo", + ":hlo_matchers", + ":reduce_precision_insertion", + "//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", + ], +) + # ----------------------------------------------------------------------------- filegroup( diff --git a/tensorflow/compiler/xla/service/algebraic_simplifier.cc b/tensorflow/compiler/xla/service/algebraic_simplifier.cc index 415aafe69ac3ffca428524eb6e998febf4f9d75a..74f8e3143d718e46e09e76bf9439b4f96b012226 100644 --- a/tensorflow/compiler/xla/service/algebraic_simplifier.cc +++ b/tensorflow/compiler/xla/service/algebraic_simplifier.cc @@ -48,7 +48,17 @@ namespace { // Returns whether operand is a literal with the given value. bool IsLiteralWithValue(const HloInstruction* operand, int8 value) { return operand->opcode() == HloOpcode::kConstant && - LiteralUtil::IsAll(operand->literal(), value); + operand->literal().IsAll(value); +} + +bool IsAll(const HloInstruction* op, int8 value) { + if (IsLiteralWithValue(op, value)) { + return true; + } + if (op->opcode() == HloOpcode::kBroadcast && IsAll(op->operand(0), value)) { + return true; + } + return false; } // Returns whether the given transpose produces a result which is bit-wise @@ -110,12 +120,20 @@ class AlgebraicSimplifierVisitor : public DfsHloVisitorWithDefault { Status HandleAdd(HloInstruction* add, HloInstruction* lhs, HloInstruction* rhs) override; + Status HandleBitcast(HloInstruction* bitcast) override; + Status HandleBroadcast(HloInstruction* broadcast) override; - Status HandleCopy(HloInstruction* copy, HloInstruction* operand) override; + Status HandleConcatenate( + HloInstruction* concatenate, + tensorflow::gtl::ArraySlice operands) override; - Status HandleConvert(HloInstruction* convert, - HloInstruction* operand) override; + Status HandleConstant(HloInstruction* constant, + const Literal& literal) override; + + Status HandleCopy(HloInstruction* copy) override; + + Status HandleConvert(HloInstruction* convert) override; Status HandleConvolution(HloInstruction* convolution, HloInstruction* lhs, HloInstruction* rhs, const Window& window) override; @@ -146,20 +164,27 @@ class AlgebraicSimplifierVisitor : public DfsHloVisitorWithDefault { tensorflow::gtl::ArraySlice dimensions, HloComputation* function) override; + Status HandleReduceWindow(HloInstruction* reduce_window, + HloInstruction* operand, const Window& window, + HloComputation* function) override; + Status HandleReverse(HloInstruction* reverse, HloInstruction* operand) override; Status HandleSlice(HloInstruction* slice, HloInstruction* operand) override; + Status HandleDynamicSlice(HloInstruction* slice, HloInstruction* operand, + HloInstruction* start_indices) override; + Status HandleDynamicUpdateSlice(HloInstruction* dynamic_update_slice, + HloInstruction* operand, + HloInstruction* update, + HloInstruction* start_indices) override; Status HandleTranspose(HloInstruction* transpose) override; Status HandleSubtract(HloInstruction* sub, HloInstruction* lhs, HloInstruction* rhs) override; - Status HandleMaximum(HloInstruction* maximum, HloInstruction* lhs, - HloInstruction* rhs) override; - - Status HandleMinimum(HloInstruction* minimum, HloInstruction* lhs, - HloInstruction* rhs) override; + Status HandleMaximum(HloInstruction* maximum) override; + Status HandleMinimum(HloInstruction* minimum) override; // Returns whether algebraic simplification has occurred. const bool changed() const { return changed_; } @@ -210,6 +235,35 @@ class AlgebraicSimplifierVisitor : public DfsHloVisitorWithDefault { StatusOr TryToSinkReshapeOrBroadcastAfterOpWithUniqueNonScalarOperand( HloInstruction* reshape_or_broadcast); + // Replaces the existing HLO instruction old_instruction, with + // new_instruction, and marks the optimizer status as changed. + // Returns the Status representing the result of the replace operation. + Status ReplaceWithNewInstruction( + HloInstruction* old_instruction, + std::unique_ptr new_instruction) { + VLOG(3) << "Replacing instruction:"; + VLOG(3) << " old: " << old_instruction->ToString(); + VLOG(3) << " new: " << new_instruction->ToString(); + TF_RETURN_IF_ERROR(computation_->ReplaceWithNewInstruction( + old_instruction, std::move(new_instruction))); + changed_ = true; + return Status::OK(); + } + + // Replaces the existing HLO instruction old_instruction, with + // new_instruction, and marks the optimizer status as changed. + // Returns the Status representing the result of the replace operation. + Status ReplaceInstruction(HloInstruction* old_instruction, + HloInstruction* new_instruction) { + VLOG(3) << "Replacing instruction:"; + VLOG(3) << " old: " << old_instruction->ToString(); + VLOG(3) << " new: " << new_instruction->ToString(); + TF_RETURN_IF_ERROR( + computation_->ReplaceInstruction(old_instruction, new_instruction)); + changed_ = true; + return Status::OK(); + } + // Current HloComputation instance the AlgebraicSimplifierVisitor is // traversing. HloComputation* computation_; @@ -258,8 +312,7 @@ void AlgebraicSimplifierVisitor::ReplaceWithBitcast( auto bitcast = computation_->AddInstruction( HloInstruction::CreateUnary(instruction->shape(), HloOpcode::kBitcast, instruction->mutable_operand(0))); - TF_CHECK_OK(computation_->ReplaceInstruction(instruction, bitcast)); - changed_ = true; + TF_CHECK_OK(ReplaceInstruction(instruction, bitcast)); } bool AlgebraicSimplifierVisitor::ReplaceInstructionIfSameShape( @@ -267,9 +320,7 @@ bool AlgebraicSimplifierVisitor::ReplaceInstructionIfSameShape( if (!SameShape(old_instruction, new_instruction)) { return false; } - TF_CHECK_OK( - computation_->ReplaceInstruction(old_instruction, new_instruction)); - changed_ = true; + TF_CHECK_OK(ReplaceInstruction(old_instruction, new_instruction)); return true; } @@ -278,22 +329,136 @@ Status AlgebraicSimplifierVisitor::HandleAdd(HloInstruction* add, HloInstruction* rhs) { // A + 0 => A VLOG(10) << "trying transform [A + 0 => A]: " << add->ToString(); - if (IsLiteralWithValue(rhs, 0) && ReplaceInstructionIfSameShape(add, lhs)) { + if (IsAll(rhs, 0) && ReplaceInstructionIfSameShape(add, lhs)) { return Status::OK(); } // 0 + A => A VLOG(10) << "trying transform [0 + A => A]: " << add->ToString(); - if (IsLiteralWithValue(lhs, 0) && ReplaceInstructionIfSameShape(add, rhs)) { + if (IsAll(lhs, 0) && ReplaceInstructionIfSameShape(add, rhs)) { return Status::OK(); } return Status::OK(); } -Status AlgebraicSimplifierVisitor::HandleCopy(HloInstruction* copy, - HloInstruction* operand) { +Status AlgebraicSimplifierVisitor::HandleBitcast(HloInstruction* bitcast) { + // If a bitcast feeds a bitcast, make it a single bitcast. + if (bitcast->operand(0)->opcode() == HloOpcode::kBitcast) { + return ReplaceWithNewInstruction( + bitcast, HloInstruction::CreateUnary( + bitcast->shape(), HloOpcode::kBitcast, + bitcast->mutable_operand(0)->mutable_operand(0))); + } + // All bitcasts can be eliminated (assuming layout constraints are + // satisified). + ReplaceInstructionIfSameShape(bitcast, bitcast->mutable_operand(0)); + return Status::OK(); +} + +Status AlgebraicSimplifierVisitor::HandleCopy(HloInstruction* copy) { + // If a copy feeds a copy, make it a single copy. + if (copy->operand(0)->opcode() == HloOpcode::kCopy) { + return ReplaceWithNewInstruction( + copy, HloInstruction::CreateUnary( + copy->shape(), HloOpcode::kCopy, + copy->mutable_operand(0)->mutable_operand(0))); + } // All copies can be eliminated (assuming layout constraints are satisified). - ReplaceInstructionIfSameShape(copy, operand); + ReplaceInstructionIfSameShape(copy, copy->mutable_operand(0)); + return Status::OK(); +} + +Status AlgebraicSimplifierVisitor::HandleConcatenate( + HloInstruction* concatenate, + tensorflow::gtl::ArraySlice operands) { + if (operands.size() == 1) { + // Unary concatenates are useless. + ReplaceInstructionIfSameShape(concatenate, operands[0]); + return Status::OK(); + } + // Filter out and remove empty operands. + std::vector nonempty_operands; + for (HloInstruction* operand : operands) { + if (!ShapeUtil::HasZeroElements(operand->shape())) { + nonempty_operands.push_back(operand); + } + } + if (nonempty_operands.size() < operands.size()) { + HloInstruction* replacement; + if (nonempty_operands.empty()) { + replacement = operands[0]; + } else if (nonempty_operands.size() == 1) { + replacement = nonempty_operands[0]; + } else { + replacement = + computation_->AddInstruction(concatenate->CloneWithNewOperands( + concatenate->shape(), nonempty_operands)); + } + VLOG(10) << "trying to replace " << concatenate->ToString() << " with " + << replacement->ToString(); + ReplaceInstructionIfSameShape(concatenate, replacement); + } else 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 = + operands[0]->opcode() == HloOpcode::kBroadcast && + ShapeUtil::IsScalar(operands[0]->operand(0)->shape()); + bool is_effective_high_pad = + operands[1]->opcode() == HloOpcode::kBroadcast && + ShapeUtil::IsScalar(operands[1]->operand(0)->shape()); + if (!is_effective_low_pad && !is_effective_high_pad) { + return Status::OK(); + } + PaddingConfig padding_config; + for (int64 dim = 0; dim < ShapeUtil::Rank(operands[0]->shape()); ++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 (is_effective_low_pad) { + padding_config_dim->set_edge_padding_low( + operands[0]->shape().dimensions(dim)); + } else { + padding_config_dim->set_edge_padding_high( + operands[1]->shape().dimensions(dim)); + } + } + } + int64 operand_to_pad = is_effective_low_pad ? 1 : 0; + int64 pad_value_operand = is_effective_low_pad ? 0 : 1; + HloInstruction* pad = + computation_->AddInstruction(HloInstruction::CreatePad( + concatenate->shape(), operands[operand_to_pad], + operands[pad_value_operand]->mutable_operand(0), padding_config)); + return ReplaceInstruction(concatenate, pad); + } + return Status::OK(); +} + +static HloInstruction* BuildTupleConstant(HloComputation* computation, + const Literal& literal) { + if (ShapeUtil::IsTuple(literal.shape())) { + std::vector elems; + elems.reserve(ShapeUtil::TupleElementCount(literal.shape())); + for (const Literal& child : literal.tuple_literals()) { + elems.push_back(BuildTupleConstant(computation, child)); + } + return computation->AddInstruction(HloInstruction::CreateTuple(elems)); + } else { + return computation->AddInstruction( + HloInstruction::CreateConstant(MakeUnique(literal))); + } +} + +Status AlgebraicSimplifierVisitor::HandleConstant(HloInstruction* constant, + const Literal& literal) { + // Tuple constants aren't directly supported by any backend. Expand them into + // explicit Tuple instructions. + if (ShapeUtil::IsTuple(constant->shape())) { + return ReplaceInstruction(constant, + BuildTupleConstant(computation_, literal)); + } return Status::OK(); } @@ -302,7 +467,7 @@ Status AlgebraicSimplifierVisitor::HandleSubtract(HloInstruction* sub, HloInstruction* rhs) { // A - 0 => A VLOG(10) << "trying transform [A - 0 => A]: " << sub->ToString(); - if (IsLiteralWithValue(rhs, 0) && ReplaceInstructionIfSameShape(sub, lhs)) { + if (IsAll(rhs, 0) && ReplaceInstructionIfSameShape(sub, lhs)) { return Status::OK(); } @@ -314,8 +479,7 @@ Status AlgebraicSimplifierVisitor::HandleDivide(HloInstruction* divide, HloInstruction* rhs) { // A/1 => A VLOG(10) << "trying transform [A/1 => A]: " << divide->ToString(); - if (IsLiteralWithValue(rhs, 1) && - ReplaceInstructionIfSameShape(divide, lhs)) { + if (IsAll(rhs, 1) && ReplaceInstructionIfSameShape(divide, lhs)) { return Status::OK(); } @@ -326,12 +490,91 @@ Status AlgebraicSimplifierVisitor::HandleDivide(HloInstruction* divide, computation_->AddInstruction(HloInstruction::CreateBinary( divide->shape(), HloOpcode::kSubtract, lhs->mutable_operand(0), rhs->mutable_operand(0))); - changed_ = true; - return computation_->ReplaceWithNewInstruction( + return ReplaceWithNewInstruction( divide, HloInstruction::CreateUnary(divide->shape(), HloOpcode::kExp, subtract)); } + // A/exp(B) => A*exp(-B) + if (rhs->opcode() == HloOpcode::kExp) { + VLOG(10) << "transform [A/exp(B) => A*exp(-B)]: " << divide->ToString(); + HloInstruction* negate = + computation_->AddInstruction(HloInstruction::CreateUnary( + divide->shape(), HloOpcode::kNegate, rhs->mutable_operand(0))); + HloInstruction* new_exp = computation_->AddInstruction( + HloInstruction::CreateUnary(divide->shape(), HloOpcode::kExp, negate)); + return ReplaceWithNewInstruction( + divide, HloInstruction::CreateBinary( + divide->shape(), HloOpcode::kMultiply, lhs, new_exp)); + } + + // A/pow(B,C) => A*pow(B,-C) + if (rhs->opcode() == HloOpcode::kPower) { + VLOG(10) << "transform [A/pow(B,C) => A*pow(B,-C)]: " << divide->ToString(); + HloInstruction* negate = + computation_->AddInstruction(HloInstruction::CreateUnary( + divide->shape(), HloOpcode::kNegate, rhs->mutable_operand(1))); + HloInstruction* new_power = computation_->AddInstruction( + HloInstruction::CreateBinary(divide->shape(), HloOpcode::kPower, + rhs->mutable_operand(0), negate)); + return ReplaceWithNewInstruction( + divide, HloInstruction::CreateBinary( + divide->shape(), HloOpcode::kMultiply, lhs, new_power)); + } + + // Simplifying integral division would produce unexpected results. + if (ShapeUtil::ElementIsIntegral(divide->shape())) { + return Status::OK(); + } + + // (A / B) / (C / D) => (A / B)*(D / C) => (A * D) / (B * C) + if (lhs->opcode() == HloOpcode::kDivide && + rhs->opcode() == HloOpcode::kDivide) { + TF_ASSIGN_OR_RETURN( + const Shape a_times_d_shape, + ShapeInference::InferBinaryOpShape(HloOpcode::kMultiply, + lhs->operand(0), rhs->operand(1))); + auto a_times_d = computation_->AddInstruction(HloInstruction::CreateBinary( + a_times_d_shape, HloOpcode::kMultiply, lhs->mutable_operand(0), + rhs->mutable_operand(1))); + TF_ASSIGN_OR_RETURN( + const Shape b_times_c_shape, + ShapeInference::InferBinaryOpShape(HloOpcode::kMultiply, + lhs->operand(1), rhs->operand(0))); + auto b_times_c = computation_->AddInstruction(HloInstruction::CreateBinary( + b_times_c_shape, HloOpcode::kMultiply, lhs->mutable_operand(1), + rhs->mutable_operand(0))); + return ReplaceWithNewInstruction( + divide, HloInstruction::CreateBinary( + divide->shape(), HloOpcode::kDivide, a_times_d, b_times_c)); + } + + // (A / B) / C => A / (B * C) + if (lhs->opcode() == HloOpcode::kDivide) { + TF_ASSIGN_OR_RETURN(const Shape b_times_c_shape, + ShapeInference::InferBinaryOpShape( + HloOpcode::kMultiply, lhs->operand(1), rhs)); + auto b_times_c = computation_->AddInstruction(HloInstruction::CreateBinary( + b_times_c_shape, HloOpcode::kMultiply, lhs->mutable_operand(1), rhs)); + return ReplaceWithNewInstruction( + divide, + HloInstruction::CreateBinary(divide->shape(), HloOpcode::kDivide, + lhs->mutable_operand(0), b_times_c)); + } + + // A / (B / C) => (A*C) / B + if (rhs->opcode() == HloOpcode::kDivide) { + TF_ASSIGN_OR_RETURN(const Shape a_times_c_shape, + ShapeInference::InferBinaryOpShape( + HloOpcode::kMultiply, lhs, rhs->operand(1))); + auto a_times_c = computation_->AddInstruction(HloInstruction::CreateBinary( + a_times_c_shape, HloOpcode::kMultiply, lhs, rhs->mutable_operand(1))); + return ReplaceWithNewInstruction( + divide, + HloInstruction::CreateBinary(divide->shape(), HloOpcode::kDivide, + a_times_c, rhs->mutable_operand(0))); + } + return Status::OK(); } @@ -353,19 +596,17 @@ Status AlgebraicSimplifierVisitor::HandleDot(HloInstruction* dot, ShapeUtil::HasZeroElements(lhs->shape()) || ShapeUtil::HasZeroElements(rhs->shape())) { auto zero = computation_->AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(0.0f))); - changed_ = true; - return computation_->ReplaceWithNewInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(0.0f))); + return ReplaceWithNewInstruction( dot, HloInstruction::CreateBroadcast(dot->shape(), zero, {})); } - // Simplify dot(transpose(a), transpose(b)) to tranpose(dot(b,a)). + // Simplify dot(transpose(a), transpose(b)) to transpose(dot(b,a)). if (lhs->IsRank2Transpose() && rhs->IsRank2Transpose()) { auto new_dot = computation_->AddInstruction(HloInstruction::CreateBinary( ShapeUtil::PermuteDimensions({1, 0}, dot->shape()), HloOpcode::kDot, rhs->mutable_operand(0), lhs->mutable_operand(0))); - changed_ = true; - return computation_->ReplaceWithNewInstruction( + return ReplaceWithNewInstruction( dot, HloInstruction::CreateTranspose(dot->shape(), new_dot, {1, 0})); } @@ -373,8 +614,7 @@ Status AlgebraicSimplifierVisitor::HandleDot(HloInstruction* dot, // // A dot(a[M, 1], b[1, N]) = multiply(a [M,1], b [1, N]) if (ShapeUtil::Rank(rhs->shape()) == 2 && rhs->shape().dimensions(0) == 1) { - changed_ = true; - return computation_->ReplaceWithNewInstruction( + return ReplaceWithNewInstruction( dot, HloInstruction::CreateBinary(dot->shape(), HloOpcode::kMultiply, lhs, rhs)); } @@ -394,12 +634,11 @@ Status AlgebraicSimplifierVisitor::HandleDot(HloInstruction* dot, HloComputation* add_reduce_computation = CreateScalarBinaryComputation( computation_->parent(), F32, HloOpcode::kAdd); auto zero = computation_->AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(0.0f))); + HloInstruction::CreateConstant(Literal::CreateR0(0.0f))); auto reduce = computation_->AddInstruction(HloInstruction::CreateReduce( ShapeUtil::MakeShape(dot->shape().element_type(), {}), multiply, zero, {0}, add_reduce_computation)); - changed_ = true; - return computation_->ReplaceWithNewInstruction( + return ReplaceWithNewInstruction( dot, HloInstruction::CreateReshape(dot->shape(), reduce)); } @@ -419,7 +658,7 @@ Status AlgebraicSimplifierVisitor::HandleDot(HloInstruction* dot, HloComputation* add_reduce_computation = CreateScalarBinaryComputation( computation_->parent(), F32, HloOpcode::kAdd); auto zero = computation_->AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(0.0f))); + HloInstruction::CreateConstant(Literal::CreateR0(0.0f))); HloInstruction* reduce; if (ShapeUtil::Rank(rhs->shape()) == 1) { auto multiply = computation_->AddInstruction(HloInstruction::CreateBinary( @@ -438,8 +677,7 @@ Status AlgebraicSimplifierVisitor::HandleDot(HloInstruction* dot, {rhs->shape().dimensions(1)}), multiply, zero, {0}, add_reduce_computation)); } - changed_ = true; - return computation_->ReplaceWithNewInstruction( + return ReplaceWithNewInstruction( dot, HloInstruction::CreateReshape(dot->shape(), reduce)); } @@ -460,13 +698,12 @@ Status AlgebraicSimplifierVisitor::HandleDot(HloInstruction* dot, HloComputation* add_reduce_computation = CreateScalarBinaryComputation( computation_->parent(), F32, HloOpcode::kAdd); auto zero = computation_->AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(0.0f))); + HloInstruction::CreateConstant(Literal::CreateR0(0.0f))); auto reduce = computation_->AddInstruction(HloInstruction::CreateReduce( ShapeUtil::MakeShape(dot->shape().element_type(), {lhs->shape().dimensions(0)}), multiply, zero, {1}, add_reduce_computation)); - changed_ = true; - return computation_->ReplaceWithNewInstruction( + return ReplaceWithNewInstruction( dot, HloInstruction::CreateReshape(dot->shape(), reduce)); } return Status::OK(); @@ -477,16 +714,24 @@ Status AlgebraicSimplifierVisitor::HandleMultiply(HloInstruction* multiply, HloInstruction* rhs) { // A*1 => A VLOG(10) << "trying transform [A*1 => A]: " << multiply->ToString(); - if (IsLiteralWithValue(rhs, 1) && - ReplaceInstructionIfSameShape(multiply, lhs)) { + if (IsAll(rhs, 1) && ReplaceInstructionIfSameShape(multiply, lhs)) { return Status::OK(); } // 1*A => A VLOG(10) << "trying transform [1*A => A]: " << multiply->ToString(); - if (IsLiteralWithValue(lhs, 1) && - ReplaceInstructionIfSameShape(multiply, rhs)) { + if (IsAll(lhs, 1) && ReplaceInstructionIfSameShape(multiply, rhs)) { return Status::OK(); } + + // exp(A) * exp(B) => exp(A+B) + if (lhs->opcode() == HloOpcode::kExp && rhs->opcode() == HloOpcode::kExp) { + auto add = computation_->AddInstruction(HloInstruction::CreateBinary( + multiply->shape(), HloOpcode::kAdd, lhs->mutable_operand(0), + rhs->mutable_operand(0))); + return ReplaceWithNewInstruction( + multiply, + HloInstruction::CreateUnary(multiply->shape(), HloOpcode::kExp, add)); + } return Status::OK(); } @@ -498,6 +743,17 @@ Status AlgebraicSimplifierVisitor::HandleLog(HloInstruction* log, ReplaceInstructionIfSameShape(log, operand->mutable_operand(0))) { return Status::OK(); } + + // ln(pow(A,B)) => B*ln(A) + if (operand->opcode() == HloOpcode::kPower) { + auto new_log = computation_->AddInstruction(HloInstruction::CreateUnary( + log->shape(), HloOpcode::kLog, operand->mutable_operand(0))); + return ReplaceWithNewInstruction( + log, + HloInstruction::CreateBinary(log->shape(), HloOpcode::kMultiply, + new_log, operand->mutable_operand(1))); + } + return Status::OK(); } @@ -556,8 +812,9 @@ std::pair> ReshapeLeavesDimensionsUnmodified( return std::make_pair(true, output_dim_indices); } -// Returns true if the output of "instruction" is a permutation of the elements -// of "operand". Precondition: "operand" is an operand of "instruction". +// Returns true if the output of "instruction" is a permutation of the +// elements of "operand". Precondition: "operand" is an operand of +// "instruction". bool OutputIsPermutationOfOperandElements(HloInstruction* instruction, HloInstruction* operand) { DCHECK(!instruction->OperandIndices(operand).empty()); @@ -605,8 +862,7 @@ Status AlgebraicSimplifierVisitor::HandleBroadcast(HloInstruction* broadcast) { ShapeUtil::ElementsIn(operand->shape())) { VLOG(10) << "transform broadcast(X) -> reshape(X) where " "n(broadcast(X)) == n(X)"; - changed_ = true; - return computation_->ReplaceWithNewInstruction( + return ReplaceWithNewInstruction( broadcast, HloInstruction::CreateReshape(broadcast->shape(), operand)); } @@ -618,14 +874,13 @@ Status AlgebraicSimplifierVisitor::HandleBroadcast(HloInstruction* broadcast) { ShapeUtil::ElementsIn(operand->shape())) { VLOG(10) << "transform broadcast(X) -> transpose(X) where " "n(broadcast(X)) == n(X)"; - changed_ = true; - return computation_->ReplaceWithNewInstruction( + return ReplaceWithNewInstruction( broadcast, HloInstruction::CreateTranspose(broadcast->shape(), operand, broadcast->dimensions())); } - // A broadcast of a reshape which merely inserts 1-sized dimensions can elide - // its operand. + // A broadcast of a reshape which merely inserts 1-sized dimensions can + // elide its operand. { bool merely_inserts_or_deletes_1_sized_dimensions; std::vector inserted_indices, deleted_indices; @@ -639,8 +894,7 @@ Status AlgebraicSimplifierVisitor::HandleBroadcast(HloInstruction* broadcast) { for (auto inserted_index : inserted_indices) { dims.erase(dims.begin() + inserted_index); } - changed_ = true; - return computation_->ReplaceWithNewInstruction( + return ReplaceWithNewInstruction( broadcast, HloInstruction::CreateBroadcast(broadcast->shape(), operand->mutable_operand(0), dims)); @@ -683,83 +937,14 @@ Status AlgebraicSimplifierVisitor::HandleBroadcast(HloInstruction* broadcast) { return Status::OK(); } -template -static std::unique_ptr ConvertIfTypesMatch( - const Literal& src_literal) { - CHECK_EQ(primitive_src_type, src_literal.shape().element_type()); - - return HloInstruction::CreateConstant( - LiteralUtil::Convert::type, - typename primitive_util::PrimitiveTypeToNative< - primitive_dest_type>::type>(src_literal)); -} - -template -static std::unique_ptr ConvertIfDestTypeMatches( - const Literal& src_literal, PrimitiveType primitive_dest_type) { - switch (primitive_dest_type) { -#define CONVERT_IF_TYPES_MATCH(type) \ - case (type): \ - return ConvertIfTypesMatch(src_literal); - CONVERT_IF_TYPES_MATCH(PRED) - CONVERT_IF_TYPES_MATCH(S8) - CONVERT_IF_TYPES_MATCH(S32) - CONVERT_IF_TYPES_MATCH(S64) - CONVERT_IF_TYPES_MATCH(U8) - CONVERT_IF_TYPES_MATCH(U32) - CONVERT_IF_TYPES_MATCH(U64) - CONVERT_IF_TYPES_MATCH(F32) - CONVERT_IF_TYPES_MATCH(F64) -#undef CONVERT_IF_TYPES_MATCH - // Other types are not yet supported. - default: - LOG(FATAL) << "Unimplemented: ConvertIfDestTypeMatches for type " - << PrimitiveType_Name(src_literal.shape().element_type()); - } -} - -static std::unique_ptr ConvertIfSrcTypeMatches( - const Literal& src_literal, PrimitiveType primitive_dest_type) { - switch (src_literal.shape().element_type()) { -#define CONVERT_IF_DEST_TYPE_MATCHES(type) \ - case (type): \ - return ConvertIfDestTypeMatches<(type)>(src_literal, primitive_dest_type); - CONVERT_IF_DEST_TYPE_MATCHES(PRED) - CONVERT_IF_DEST_TYPE_MATCHES(S8) - CONVERT_IF_DEST_TYPE_MATCHES(S32) - CONVERT_IF_DEST_TYPE_MATCHES(S64) - CONVERT_IF_DEST_TYPE_MATCHES(U8) - CONVERT_IF_DEST_TYPE_MATCHES(U32) - CONVERT_IF_DEST_TYPE_MATCHES(U64) - CONVERT_IF_DEST_TYPE_MATCHES(F32) - CONVERT_IF_DEST_TYPE_MATCHES(F64) -#undef CONVERT_IF_DEST_TYPE_MATCHES - // Other types are not yet supported. - default: - LOG(FATAL) << "Unimplemented: ConvertIfSrcTypeMatches for type " - << PrimitiveType_Name(src_literal.shape().element_type()); - } -} - // A conversion to the same element type as the operand is a nop and can be // removed. A conversion of a constant can be simplified by making a new // constant. -Status AlgebraicSimplifierVisitor::HandleConvert(HloInstruction* convert, - HloInstruction* operand) { - PrimitiveType src_type = operand->shape().element_type(); +Status AlgebraicSimplifierVisitor::HandleConvert(HloInstruction* convert) { + PrimitiveType src_type = convert->operand(0)->shape().element_type(); PrimitiveType dest_type = convert->shape().element_type(); if (src_type == dest_type) { - changed_ = true; - return computation_->ReplaceInstruction(convert, operand); - } - if (operand->opcode() == HloOpcode::kConstant) { - const Literal& src_literal = operand->literal(); - std::unique_ptr new_constant = - ConvertIfSrcTypeMatches(src_literal, dest_type); - changed_ = true; - return computation_->ReplaceWithNewInstruction(convert, - std::move(new_constant)); + return ReplaceInstruction(convert, convert->mutable_operand(0)); } return Status::OK(); } @@ -822,6 +1007,7 @@ Status AlgebraicSimplifierVisitor::HandlePad(HloInstruction* pad) { // Second, construct the slice instruction to perform the negative padding. std::vector start_indices; std::vector end_indices; + std::vector strides; for (int64 i = 0; i < pad->padding_config().dimensions_size(); ++i) { const PaddingConfig::PaddingConfigDimension& padding_dimension = pad->padding_config().dimensions(i); @@ -835,18 +1021,19 @@ Status AlgebraicSimplifierVisitor::HandlePad(HloInstruction* pad) { } start_indices.push_back(start); end_indices.push_back(end); + strides.push_back(1); } // Verify that the slice shape matches the pad shape. - TF_ASSIGN_OR_RETURN(Shape inferred_slice_shape, - ShapeInference::InferSliceShape( - nonzero_pad_shape, start_indices, end_indices)); + TF_ASSIGN_OR_RETURN( + Shape inferred_slice_shape, + ShapeInference::InferSliceShape(nonzero_pad_shape, start_indices, + end_indices, strides)); TF_RET_CHECK(ShapeUtil::Compatible(inferred_slice_shape, pad->shape())); std::unique_ptr slice = HloInstruction::CreateSlice( - pad->shape(), nonzero_pad, start_indices, end_indices); - changed_ = true; - return computation_->ReplaceWithNewInstruction(pad, std::move(slice)); + pad->shape(), nonzero_pad, start_indices, end_indices, strides); + return ReplaceWithNewInstruction(pad, std::move(slice)); } return Status::OK(); @@ -856,9 +1043,9 @@ Status AlgebraicSimplifierVisitor::HandlePower(HloInstruction* power, HloInstruction* lhs, HloInstruction* rhs) { VLOG(10) << "trying transform [pow(A, 0) => 1]: " << power->ToString(); - if (IsLiteralWithValue(rhs, 0)) { - auto one = HloInstruction::CreateConstant(LiteralUtil::CloneToUnique( - LiteralUtil::One(power->shape().element_type()))); + if (IsAll(rhs, 0)) { + auto one = HloInstruction::CreateConstant( + Literal::One(power->shape().element_type()).CloneToUnique()); std::unique_ptr ones; if (ShapeUtil::IsScalar(power->shape())) { ones = std::move(one); @@ -866,30 +1053,34 @@ Status AlgebraicSimplifierVisitor::HandlePower(HloInstruction* power, ones = HloInstruction::CreateBroadcast( power->shape(), computation_->AddInstruction(std::move(one)), {}); } - changed_ = true; - return computation_->ReplaceWithNewInstruction(power, std::move(ones)); + return ReplaceWithNewInstruction(power, std::move(ones)); } VLOG(10) << "trying transform [pow(A, 1) => A]: " << power->ToString(); - if (IsLiteralWithValue(rhs, 1) && ReplaceInstructionIfSameShape(power, lhs)) { + if (IsAll(rhs, 1) && ReplaceInstructionIfSameShape(power, lhs)) { return Status::OK(); } + // pow(exp(A),B) => exp(A*B) + if (lhs->opcode() == HloOpcode::kExp) { + auto a_times_b = computation_->AddInstruction(HloInstruction::CreateBinary( + power->shape(), HloOpcode::kMultiply, lhs->operands()[0], rhs)); + return ReplaceWithNewInstruction( + power, HloInstruction::CreateUnary(power->shape(), HloOpcode::kExp, + a_times_b)); + } VLOG(10) << "trying transform [pow(A, 2) => A*A]: " << power->ToString(); - if (IsLiteralWithValue(rhs, 2)) { - changed_ = true; - return computation_->ReplaceWithNewInstruction( + if (IsAll(rhs, 2)) { + return ReplaceWithNewInstruction( power, HloInstruction::CreateBinary(power->shape(), HloOpcode::kMultiply, lhs, lhs)); } VLOG(10) << "trying transform [pow(A, -1) => 1/A]: " << power->ToString(); - if (IsLiteralWithValue(rhs, -1)) { - auto* one = computation_->AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CloneToUnique( - LiteralUtil::One(rhs->shape().element_type())))); - changed_ = true; - return computation_->ReplaceWithNewInstruction( + if (IsAll(rhs, -1)) { + auto* one = computation_->AddInstruction(HloInstruction::CreateConstant( + Literal::One(rhs->shape().element_type()).CloneToUnique())); + return ReplaceWithNewInstruction( power, HloInstruction::CreateBinary(power->shape(), HloOpcode::kDivide, one, lhs)); } @@ -900,6 +1091,9 @@ StatusOr AlgebraicSimplifierVisitor:: TryToSinkReshapeOrBroadcastAfterOpWithUniqueNonScalarOperand( HloInstruction* reshape_or_broadcast) { bool changed = false; + if (ShapeUtil::IsScalar(reshape_or_broadcast->shape())) { + return false; + } HloInstruction* operand = reshape_or_broadcast->mutable_operand(0); for (HloInstruction* user : reshape_or_broadcast->users()) { if (user->user_count() == 0 && user != computation_->root_instruction()) { @@ -932,31 +1126,43 @@ StatusOr AlgebraicSimplifierVisitor:: if (scalar_count != user->operand_count() - 1) { continue; } + VLOG(4) << "Sinking reshape or broadcast after user:"; + VLOG(4) << " old reshape/broadcast: " << reshape_or_broadcast->ToString(); + VLOG(4) << " old user: " << user->ToString(); CHECK_EQ(user->operand(reshape_or_broadcast_operand_index), reshape_or_broadcast); - std::vector new_user_operands = user->operands(); + auto new_user_operands = user->operands(); new_user_operands[reshape_or_broadcast_operand_index] = operand; auto new_user = computation_->AddInstruction(user->CloneWithNewOperands( - ShapeUtil::MakeShape(user->shape().element_type(), - AsInt64Slice(operand->shape().dimensions())), + ShapeUtil::MakeShapeWithLayout( + user->shape().element_type(), + AsInt64Slice(operand->shape().dimensions()), + AsInt64Slice(operand->shape().layout().minor_to_major())), new_user_operands)); + VLOG(4) << " new user: " << new_user->ToString(); HloInstruction* new_reshape_or_broadcast = nullptr; if (reshape_or_broadcast->opcode() == HloOpcode::kReshape) { new_reshape_or_broadcast = computation_->AddInstruction(HloInstruction::CreateReshape( - ShapeUtil::MakeShape( + ShapeUtil::MakeShapeWithLayout( user->shape().element_type(), - AsInt64Slice(reshape_or_broadcast->shape().dimensions())), + AsInt64Slice(reshape_or_broadcast->shape().dimensions()), + AsInt64Slice( + reshape_or_broadcast->shape().layout().minor_to_major())), new_user)); } else { TF_RET_CHECK(reshape_or_broadcast->opcode() == HloOpcode::kBroadcast); new_reshape_or_broadcast = computation_->AddInstruction(HloInstruction::CreateBroadcast( - ShapeUtil::MakeShape( + ShapeUtil::MakeShapeWithLayout( user->shape().element_type(), - AsInt64Slice(reshape_or_broadcast->shape().dimensions())), + AsInt64Slice(reshape_or_broadcast->shape().dimensions()), + AsInt64Slice( + reshape_or_broadcast->shape().layout().minor_to_major())), new_user, reshape_or_broadcast->dimensions())); } + VLOG(4) << " new reshape/broadcast: " + << new_reshape_or_broadcast->ToString(); TF_RETURN_IF_ERROR( computation_->ReplaceUsesOfInstruction(user, new_reshape_or_broadcast)); changed = true; @@ -967,17 +1173,24 @@ StatusOr AlgebraicSimplifierVisitor:: Status AlgebraicSimplifierVisitor::HandleReshape(HloInstruction* reshape) { auto operand = reshape->mutable_operand(0); + // Reshape directly to empty constant if the shape contains zero-element + // dimension. + if (ShapeUtil::HasZeroElements(reshape->shape())) { + auto empty_constant = HloInstruction::CreateConstant( + Literal::CreateFromShape(reshape->shape())); + + return ReplaceWithNewInstruction(reshape, std::move(empty_constant)); + } + // Delete no-op reshapes, i.e. where shape = operand shape. if (SameShape(reshape, operand)) { VLOG(10) << "deleting no-op reshape"; - changed_ = true; - return computation_->ReplaceInstruction(reshape, operand); + return ReplaceInstruction(reshape, operand); } // Merge reshapes. if (HloOpcode::kReshape == operand->opcode()) { - changed_ = true; - return computation_->ReplaceWithNewInstruction( + return ReplaceWithNewInstruction( reshape, HloInstruction::CreateReshape(reshape->shape(), operand->mutable_operand(0))); } @@ -986,8 +1199,7 @@ Status AlgebraicSimplifierVisitor::HandleReshape(HloInstruction* reshape) { auto opt_dims = ReshapeLeavesDimensionsUnmodified( reshape, reshape->operand(0)->dimensions()); if (opt_dims.first) { - changed_ = true; - return computation_->ReplaceWithNewInstruction( + return ReplaceWithNewInstruction( reshape, HloInstruction::CreateBroadcast( reshape->shape(), reshape->mutable_operand(0)->mutable_operand(0), @@ -1016,15 +1228,14 @@ Status AlgebraicSimplifierVisitor::HandleReshape(HloInstruction* reshape) { Status AlgebraicSimplifierVisitor::HandleReverse(HloInstruction* reverse, HloInstruction* operand) { - // When all the dimensions to reverse are trivial (i.e. the bound is 1), there - // is nothing to be done. + // When all the dimensions to reverse are trivial (i.e. the bound is 1), + // there is nothing to be done. 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)) { - changed_ = true; - return computation_->ReplaceInstruction(reverse, operand); + return ReplaceInstruction(reverse, operand); } return Status::OK(); } @@ -1038,15 +1249,33 @@ Status AlgebraicSimplifierVisitor::HandleSlice(HloInstruction* slice, return Status::OK(); } +Status AlgebraicSimplifierVisitor::HandleDynamicSlice( + HloInstruction* dynamic_slice, HloInstruction* operand, + HloInstruction* start_indices) { + if (ShapeUtil::IsScalar(dynamic_slice->shape())) { + return ReplaceInstruction(dynamic_slice, operand); + } + return Status::OK(); +} + +Status AlgebraicSimplifierVisitor::HandleDynamicUpdateSlice( + HloInstruction* dynamic_update_slice, HloInstruction* operand, + HloInstruction* update, HloInstruction* start_indices) { + // DynamicUpdateSlice on a scalar just passes through the update argument. + if (ShapeUtil::IsScalar(dynamic_update_slice->shape())) { + return ReplaceInstruction(dynamic_update_slice, update); + } + return Status::OK(); +} + Status AlgebraicSimplifierVisitor::HandleReduce( HloInstruction* reduce, HloInstruction* arg, HloInstruction* init_value, tensorflow::gtl::ArraySlice dimensions, HloComputation* function) { if (ShapeUtil::HasZeroElements(arg->shape()) || ShapeUtil::HasZeroElements(reduce->shape())) { - return computation_->ReplaceWithNewInstruction( + return ReplaceWithNewInstruction( reduce, HloInstruction::CreateBroadcast(reduce->shape(), init_value, {})); - return Status::OK(); } // A Transpose feeding a reduce can simply permute the reduction dimensions // field. @@ -1056,15 +1285,15 @@ Status AlgebraicSimplifierVisitor::HandleReduce( for (auto dim : dimensions) { new_reduce_dimensions.push_back(transpose_dimensions[dim]); } - return computation_->ReplaceWithNewInstruction( + return ReplaceWithNewInstruction( reduce, HloInstruction::CreateReduce( reduce->shape(), arg->mutable_operand(0), init_value, new_reduce_dimensions, function)); } - // A reshape that collapses multiple dimensions into a dimension being reduced - // can just reduce all of those dimensions instead of doing a collapsing - // reshape before a reduction. + // A reshape that collapses multiple dimensions into a dimension being + // reduced can just reduce all of those dimensions instead of doing a + // collapsing reshape before a reduction. if (arg->opcode() == HloOpcode::kReshape) { std::vector> unmodified_dims = ShapeUtil::DimensionsUnmodifiedByReshape(arg->operand(0)->shape(), @@ -1100,7 +1329,7 @@ Status AlgebraicSimplifierVisitor::HandleReduce( new_reduce_dimensions.push_back(i); } } - return computation_->ReplaceWithNewInstruction( + return ReplaceWithNewInstruction( reduce, HloInstruction::CreateReduce( reduce->shape(), arg->mutable_operand(0), init_value, new_reduce_dimensions, function)); @@ -1111,27 +1340,83 @@ Status AlgebraicSimplifierVisitor::HandleReduce( ShapeUtil::HasZeroElements(arg->shape())) { auto reshape = computation_->AddInstruction( HloInstruction::CreateReshape(reduce->shape(), arg)); - changed_ = true; - return computation_->ReplaceWithNewInstruction( + return ReplaceWithNewInstruction( reduce, HloInstruction::CreateMap(reduce->shape(), {reshape, init_value}, function)); } return Status::OK(); } +Status AlgebraicSimplifierVisitor::HandleReduceWindow( + HloInstruction* reduce_window, HloInstruction* operand, + const Window& window, HloComputation* function) { + VLOG(10) << "Considering folding Pad: " << operand->ToString() + << "\ninto reduce-window: " << reduce_window->ToString(); + + // This optimization folds a pad op into reduce_window. + if (operand->opcode() != HloOpcode::kPad) { + VLOG(10) << "Not folding pad into reduce-window as there is no pad."; + return Status::OK(); + } + + // Do not fold interior padding into ReduceWindow since the backends do not + // support it. + const PaddingConfig& pad_config = operand->padding_config(); + if (HasInteriorPadding(pad_config)) { + VLOG(10) << "Not folding pad into reduce-window due to interior padding."; + return Status::OK(); + } + + // If reduce_window already has padding, the pad value of the pad op and the + // init value of reduce_window must match to allow folding the pad. + const HloInstruction* pad_value = operand->operand(1); + const HloInstruction* reduce_init_value = reduce_window->operand(1); + if (pad_value != reduce_init_value) { + // The pad value is usually a constant, so we handle that case and do not + // try to get more fancy about proving equivalence in cases beyond that. + if (pad_value->opcode() != HloOpcode::kConstant || + reduce_init_value->opcode() != HloOpcode::kConstant || + !pad_value->literal().Equal(reduce_init_value->literal())) { + VLOG(10) << "Not folding pad into reduce-window due to different pad " + "values."; + return Status::OK(); + } + } + + // 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()); + TF_RET_CHECK(pad_config.dimensions_size() == rank); + TF_RET_CHECK(window.dimensions_size() == rank); + for (int64 i = 0; i < rank; ++i) { + const auto& pad_dim = pad_config.dimensions(i); + auto& window_dim = *new_window.mutable_dimensions(i); + window_dim.set_padding_low(window_dim.padding_low() + + pad_dim.edge_padding_low()); + window_dim.set_padding_high(window_dim.padding_high() + + pad_dim.edge_padding_high()); + } + return ReplaceWithNewInstruction( + reduce_window, HloInstruction::CreateReduceWindow( + /*shape=*/reduce_window->shape(), + /*operand=*/operand->mutable_operand(0), + /*init_value=*/reduce_window->mutable_operand(1), + /*window=*/new_window, + /*reduce_computation=*/function)); +} + Status AlgebraicSimplifierVisitor::HandleTranspose(HloInstruction* transpose) { auto operand = transpose->mutable_operand(0); if (std::is_sorted(transpose->dimensions().begin(), transpose->dimensions().end())) { VLOG(10) << "deleting no-op transpose"; - changed_ = true; - return computation_->ReplaceInstruction(transpose, operand); + return ReplaceInstruction(transpose, operand); } if (HloOpcode::kTranspose == operand->opcode()) { - changed_ = true; - return computation_->ReplaceWithNewInstruction( + return ReplaceWithNewInstruction( transpose, HloInstruction::CreateTranspose( transpose->shape(), operand->mutable_operand(0), ComposePermutations(operand->dimensions(), @@ -1158,7 +1443,9 @@ Status AlgebraicSimplifierVisitor::HandleConvolution( // bitcasts_ == true. // TODO(cwhipkey): b/31337498, make this layout insensitive. - if (!is_layout_sensitive_) return Status::OK(); + if (!is_layout_sensitive_) { + return Status::OK(); + } const ConvolutionDimensionNumbers& dnums = convolution->convolution_dimension_numbers(); @@ -1248,8 +1535,8 @@ Status AlgebraicSimplifierVisitor::HandleConvolution( // 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_(lhs->shape(), input_shape) || - !valid_bitcast_callback_(rhs->shape(), new_filter_shape) || + 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 Status::OK(); } @@ -1258,9 +1545,7 @@ Status AlgebraicSimplifierVisitor::HandleConvolution( auto new_rhs = add_bitcast(new_filter_shape, rhs); auto dot = computation_->AddInstruction(HloInstruction::CreateBinary( dot_output_shape, HloOpcode::kDot, new_lhs, new_rhs)); - changed_ = true; - return computation_->ReplaceInstruction(convolution, - add_bitcast(convolution_shape, dot)); + return ReplaceInstruction(convolution, add_bitcast(convolution_shape, dot)); } bool AlgebraicSimplifierVisitor::TransformToClampIfSameShape( @@ -1274,14 +1559,11 @@ bool AlgebraicSimplifierVisitor::TransformToClampIfSameShape( auto clamp = HloInstruction::CreateTernary(root->shape(), HloOpcode::kClamp, max_operand, operand, min_operand); - TF_CHECK_OK(computation_->ReplaceWithNewInstruction(root, std::move(clamp))); - changed_ = true; + TF_CHECK_OK(ReplaceWithNewInstruction(root, std::move(clamp))); return true; } -Status AlgebraicSimplifierVisitor::HandleMaximum(HloInstruction* maximum, - HloInstruction* lhs, - HloInstruction* rhs) { +Status AlgebraicSimplifierVisitor::HandleMaximum(HloInstruction* maximum) { // Match the following tree: // min_operand operand // \ / @@ -1312,9 +1594,7 @@ Status AlgebraicSimplifierVisitor::HandleMaximum(HloInstruction* maximum, return Status::OK(); } -Status AlgebraicSimplifierVisitor::HandleMinimum(HloInstruction* minimum, - HloInstruction* lhs, - HloInstruction* rhs) { +Status AlgebraicSimplifierVisitor::HandleMinimum(HloInstruction* minimum) { // Match the following tree: // max_operand operand // \ / @@ -1349,8 +1629,17 @@ StatusOr AlgebraicSimplifier::Run(HloModule* module) { XLA_VLOG_LINES(2, "AlgebraicSimplifier::Run(), before:\n" + module->ToString()); bool changed = false; + // Make a copy of the computations because we may add computations to the + // module, invalidating iteration. + std::vector computations; for (auto& comp : module->computations()) { - if (AlgebraicSimplifierVisitor::Run(comp.get(), is_layout_sensitive_, + if (comp->IsFusionComputation()) { + continue; + } + computations.push_back(comp.get()); + } + for (auto& comp : computations) { + if (AlgebraicSimplifierVisitor::Run(comp, is_layout_sensitive_, valid_bitcast_callback_, enable_dot_simplification_)) { changed = true; diff --git a/tensorflow/compiler/xla/service/algebraic_simplifier.h b/tensorflow/compiler/xla/service/algebraic_simplifier.h index 5d59a27c716d6865e7ea1c7d0fa5392bc403eaec..4295a3227a837ffc8483b3be59994c9e6ac96aec 100644 --- a/tensorflow/compiler/xla/service/algebraic_simplifier.h +++ b/tensorflow/compiler/xla/service/algebraic_simplifier.h @@ -26,16 +26,17 @@ namespace xla { // A pass which performs AlgebraicSimplications. class AlgebraicSimplifier : public HloPassInterface { public: - // Given two shapes, determines if it is valid to bitcast between them 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 = std::function; + // 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 = + std::function; // 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 tranposes with + // returns true, then the pass will replace reshapes and transposes with // bitcasts. AlgebraicSimplifier(bool is_layout_sensitive, ValidBitcastCallback valid_bitcast_callback, diff --git a/tensorflow/compiler/xla/service/algebraic_simplifier_test.cc b/tensorflow/compiler/xla/service/algebraic_simplifier_test.cc index 27a1c0fec8855810cd016b36b1706a17c0204d63..c442e2d0bc962a5e4aae9a563099e9584b41f201 100644 --- a/tensorflow/compiler/xla/service/algebraic_simplifier_test.cc +++ b/tensorflow/compiler/xla/service/algebraic_simplifier_test.cc @@ -23,21 +23,25 @@ limitations under the License. #include "tensorflow/compiler/xla/ptr_util.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_pass_fix.h" #include "tensorflow/compiler/xla/shape_util.h" -#include "tensorflow/compiler/xla/test_helpers.h" +#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_data.pb.h" #include "tensorflow/core/lib/strings/str_util.h" +namespace op = xla::testing::opcode_matchers; + namespace xla { namespace { AlgebraicSimplifier::ValidBitcastCallback bitcasting_callback() { return [](const Shape&, const Shape&) { return true; }; } + AlgebraicSimplifier::ValidBitcastCallback non_bitcasting_callback() { return [](const Shape&, const Shape&) { return false; }; } @@ -51,11 +55,57 @@ TEST_F(AlgebraicSimplifierTest, AddZero) { HloInstruction* param0 = builder.AddInstruction( HloInstruction::CreateParameter(0, r0f32, "param0")); HloInstruction* zero = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(0.0f))); + HloInstruction::CreateConstant(Literal::CreateR0(0.0f))); builder.AddInstruction( HloInstruction::CreateBinary(r0f32, HloOpcode::kAdd, param0, zero)); - auto module = MakeUnique(TestName()); + auto module = CreateNewModule(); + auto computation = module->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.get()).ValueOrDie()); + root = computation->root_instruction(); + EXPECT_EQ(root, param0); +} + +TEST_F(AlgebraicSimplifierTest, AddBroadcastZeroR0Operand) { + Shape r2f32 = ShapeUtil::MakeShape(F32, {3, 2}); + HloComputation::Builder builder(TestName()); + HloInstruction* param0 = builder.AddInstruction( + HloInstruction::CreateParameter(0, r2f32, "param0")); + HloInstruction* zero = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(0.0f))); + HloInstruction* bcast = builder.AddInstruction( + HloInstruction::CreateBroadcast(r2f32, zero, {0, 1})); + builder.AddInstruction( + HloInstruction::CreateBinary(r2f32, HloOpcode::kAdd, bcast, param0)); + + auto module = CreateNewModule(); + auto computation = module->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.get()).ValueOrDie()); + root = computation->root_instruction(); + EXPECT_EQ(root, param0); +} + +TEST_F(AlgebraicSimplifierTest, AddBroadcastZeroR1Operand) { + Shape r2f32 = ShapeUtil::MakeShape(F32, {3, 2}); + HloComputation::Builder builder(TestName()); + HloInstruction* param0 = builder.AddInstruction( + HloInstruction::CreateParameter(0, r2f32, "param0")); + HloInstruction* zero = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR1({0, 0, 0}))); + HloInstruction* bcast = + builder.AddInstruction(HloInstruction::CreateBroadcast(r2f32, zero, {1})); + builder.AddInstruction( + HloInstruction::CreateBinary(r2f32, HloOpcode::kAdd, bcast, param0)); + + auto module = CreateNewModule(); auto computation = module->AddEntryComputation(builder.Build()); HloInstruction* root = computation->root_instruction(); EXPECT_EQ(root->opcode(), HloOpcode::kAdd); @@ -73,11 +123,11 @@ TEST_F(AlgebraicSimplifierTest, SubZero) { HloInstruction* param0 = builder.AddInstruction( HloInstruction::CreateParameter(0, r0f32, "param0")); HloInstruction* zero = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(0.0f))); + HloInstruction::CreateConstant(Literal::CreateR0(0.0f))); builder.AddInstruction( HloInstruction::CreateBinary(r0f32, HloOpcode::kSubtract, param0, zero)); - auto module = MakeUnique(TestName()); + auto module = CreateNewModule(); auto computation = module->AddEntryComputation(builder.Build()); HloInstruction* root = computation->root_instruction(); EXPECT_EQ(root->opcode(), HloOpcode::kSubtract); @@ -88,6 +138,158 @@ TEST_F(AlgebraicSimplifierTest, SubZero) { EXPECT_EQ(root, param0); } +// Test that (A/B)/C is simplified to A/(B*C). +TEST_F(AlgebraicSimplifierTest, LhsDivOfDiv) { + Shape r0f32 = ShapeUtil::MakeShape(F32, {}); + HloComputation::Builder builder(TestName()); + HloInstruction* param0 = builder.AddInstruction( + HloInstruction::CreateParameter(0, r0f32, "param0")); + HloInstruction* param1 = builder.AddInstruction( + HloInstruction::CreateParameter(1, r0f32, "param1")); + HloInstruction* param2 = builder.AddInstruction( + HloInstruction::CreateParameter(2, r0f32, "param2")); + HloInstruction* div = builder.AddInstruction( + HloInstruction::CreateBinary(r0f32, HloOpcode::kDivide, param0, param1)); + builder.AddInstruction( + HloInstruction::CreateBinary(r0f32, HloOpcode::kDivide, div, param2)); + + auto module = CreateNewModule(); + auto computation = module->AddEntryComputation(builder.Build()); + + EXPECT_THAT(computation->root_instruction(), + op::Divide(op::Divide(param0, param1), param2)); + + AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, + non_bitcasting_callback()); + ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); + + EXPECT_THAT(computation->root_instruction(), + op::Divide(param0, op::Multiply(param1, param2))); +} + +// Test that A/(B/C) is simplified to (A*C)/B. +TEST_F(AlgebraicSimplifierTest, RhsDivOfDiv) { + Shape r0f32 = ShapeUtil::MakeShape(F32, {}); + HloComputation::Builder builder(TestName()); + HloInstruction* param0 = builder.AddInstruction( + HloInstruction::CreateParameter(0, r0f32, "param0")); + HloInstruction* param1 = builder.AddInstruction( + HloInstruction::CreateParameter(1, r0f32, "param1")); + HloInstruction* param2 = builder.AddInstruction( + HloInstruction::CreateParameter(2, r0f32, "param2")); + HloInstruction* div = builder.AddInstruction( + HloInstruction::CreateBinary(r0f32, HloOpcode::kDivide, param1, param2)); + builder.AddInstruction( + HloInstruction::CreateBinary(r0f32, HloOpcode::kDivide, param0, div)); + + auto module = CreateNewModule(); + auto computation = module->AddEntryComputation(builder.Build()); + + EXPECT_THAT(computation->root_instruction(), + op::Divide(param0, op::Divide(param1, param2))); + + AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, + non_bitcasting_callback()); + ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); + + EXPECT_THAT(computation->root_instruction(), + op::Divide(op::Multiply(param0, param2), param1)); +} + +// Test that (A/B)/(C/D) is simplified to (A*D)/(B*C). +TEST_F(AlgebraicSimplifierTest, DivOfDivAndDiv) { + Shape r0f32 = ShapeUtil::MakeShape(F32, {}); + Shape r2f32 = ShapeUtil::MakeShape(F32, {42, 123}); + HloComputation::Builder builder(TestName()); + HloInstruction* param0 = builder.AddInstruction( + HloInstruction::CreateParameter(0, r0f32, "param0")); + HloInstruction* param1 = builder.AddInstruction( + HloInstruction::CreateParameter(1, r2f32, "param1")); + HloInstruction* param2 = builder.AddInstruction( + HloInstruction::CreateParameter(2, r2f32, "param2")); + HloInstruction* param3 = builder.AddInstruction( + HloInstruction::CreateParameter(3, r0f32, "param3")); + HloInstruction* div0 = builder.AddInstruction( + HloInstruction::CreateBinary(r2f32, HloOpcode::kDivide, param0, param1)); + HloInstruction* div1 = builder.AddInstruction( + HloInstruction::CreateBinary(r2f32, HloOpcode::kDivide, param2, param3)); + builder.AddInstruction( + HloInstruction::CreateBinary(r2f32, HloOpcode::kDivide, div0, div1)); + + auto module = CreateNewModule(); + auto computation = module->AddEntryComputation(builder.Build()); + + EXPECT_THAT( + computation->root_instruction(), + op::Divide(op::Divide(param0, param1), op::Divide(param2, param3))); + + AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, + non_bitcasting_callback()); + ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); + + EXPECT_THAT( + computation->root_instruction(), + op::Divide(op::Multiply(param0, param3), op::Multiply(param1, param2))); + EXPECT_TRUE( + ShapeUtil::Compatible(computation->root_instruction()->shape(), r2f32)); +} + +// Test that A/exp(B) is simplified to A*exp(-B). +TEST_F(AlgebraicSimplifierTest, DivOfExp) { + Shape r0f32 = ShapeUtil::MakeShape(F32, {}); + HloComputation::Builder builder(TestName()); + HloInstruction* param0 = builder.AddInstruction( + HloInstruction::CreateParameter(0, r0f32, "param0")); + HloInstruction* param1 = builder.AddInstruction( + HloInstruction::CreateParameter(1, r0f32, "param1")); + HloInstruction* exp = builder.AddInstruction( + HloInstruction::CreateUnary(r0f32, HloOpcode::kExp, param1)); + builder.AddInstruction( + HloInstruction::CreateBinary(r0f32, HloOpcode::kDivide, param0, exp)); + + auto module = CreateNewModule(); + auto computation = module->AddEntryComputation(builder.Build()); + + EXPECT_THAT(computation->root_instruction(), + op::Divide(param0, op::Exp(param1))); + + AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, + non_bitcasting_callback()); + ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); + + EXPECT_THAT(computation->root_instruction(), + op::Multiply(param0, op::Exp(op::Negate(param1)))); +} + +// Test that A/pow(B,C) is simplified to A*pow(B,-C). +TEST_F(AlgebraicSimplifierTest, DivOfPower) { + Shape r0f32 = ShapeUtil::MakeShape(F32, {}); + HloComputation::Builder builder(TestName()); + HloInstruction* param0 = builder.AddInstruction( + HloInstruction::CreateParameter(0, r0f32, "param0")); + HloInstruction* param1 = builder.AddInstruction( + HloInstruction::CreateParameter(1, r0f32, "param1")); + HloInstruction* param2 = builder.AddInstruction( + HloInstruction::CreateParameter(2, r0f32, "param2")); + HloInstruction* power = builder.AddInstruction( + HloInstruction::CreateBinary(r0f32, HloOpcode::kPower, param1, param2)); + builder.AddInstruction( + HloInstruction::CreateBinary(r0f32, HloOpcode::kDivide, param0, power)); + + auto module = CreateNewModule(); + auto computation = module->AddEntryComputation(builder.Build()); + + EXPECT_THAT(computation->root_instruction(), + op::Divide(param0, op::Power(param1, param2))); + + AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, + non_bitcasting_callback()); + ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); + + EXPECT_THAT(computation->root_instruction(), + op::Multiply(param0, op::Power(param1, op::Negate(param2)))); +} + // Test that A/1 is simplified to A for a scalar. TEST_F(AlgebraicSimplifierTest, DivOneScalar) { Shape r0f32 = ShapeUtil::MakeShape(F32, {}); @@ -95,11 +297,11 @@ TEST_F(AlgebraicSimplifierTest, DivOneScalar) { HloInstruction* param0 = builder.AddInstruction( HloInstruction::CreateParameter(0, r0f32, "param0")); HloInstruction* one = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(1.0f))); + HloInstruction::CreateConstant(Literal::CreateR0(1.0f))); HloInstruction* div = builder.AddInstruction( HloInstruction::CreateBinary(r0f32, HloOpcode::kDivide, param0, one)); - auto module = MakeUnique(TestName()); + auto module = CreateNewModule(); auto computation = module->AddEntryComputation(builder.Build()); HloInstruction* root = computation->root_instruction(); EXPECT_EQ(root, div); @@ -117,11 +319,11 @@ TEST_F(AlgebraicSimplifierTest, DivOneArray) { HloInstruction* param0 = builder.AddInstruction( HloInstruction::CreateParameter(0, r2f32, "param0")); HloInstruction* one = builder.AddInstruction(HloInstruction::CreateConstant( - LiteralUtil::CreateR2({{1.0, 1.0}, {1.0, 1.0}}))); + Literal::CreateR2({{1.0, 1.0}, {1.0, 1.0}}))); HloInstruction* div = builder.AddInstruction( HloInstruction::CreateBinary(r2f32, HloOpcode::kDivide, param0, one)); - auto module = MakeUnique(TestName()); + auto module = CreateNewModule(); auto computation = module->AddEntryComputation(builder.Build()); HloInstruction* root = computation->root_instruction(); EXPECT_EQ(root, div); @@ -149,7 +351,7 @@ TEST_F(AlgebraicSimplifierTest, SelectMakeTuple) { HloInstruction* add = builder.AddInstruction( HloInstruction::CreateBinary(r0f32, HloOpcode::kAdd, get, param2)); - auto module = MakeUnique(TestName()); + auto module = CreateNewModule(); auto computation = module->AddEntryComputation(builder.Build()); HloInstruction* root = computation->root_instruction(); EXPECT_EQ(root, add); @@ -157,9 +359,7 @@ TEST_F(AlgebraicSimplifierTest, SelectMakeTuple) { non_bitcasting_callback()); ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); root = computation->root_instruction(); - EXPECT_EQ(root, add); - EXPECT_EQ(root->operand(0), param1); - EXPECT_EQ(root->operand(1), param2); + EXPECT_THAT(root, op::Add(param1, param2)); } // Test that exp(A)/exp(B) is simplified to exp(A-B) @@ -177,19 +377,101 @@ TEST_F(AlgebraicSimplifierTest, ExpDiv) { builder.AddInstruction( HloInstruction::CreateBinary(r0f32, HloOpcode::kDivide, exp0, exp1)); - auto module = MakeUnique(TestName()); + auto module = CreateNewModule(); auto computation = module->AddEntryComputation(builder.Build()); - HloInstruction* root = computation->root_instruction(); - EXPECT_EQ(root->opcode(), HloOpcode::kDivide); + + EXPECT_THAT(computation->root_instruction(), + op::Divide(op::Exp(param0), op::Exp(param1))); + AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, non_bitcasting_callback()); ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); - root = computation->root_instruction(); - EXPECT_EQ(root->opcode(), HloOpcode::kExp); - EXPECT_EQ(root->operand_count(), 1); - EXPECT_EQ(root->operand(0)->opcode(), HloOpcode::kSubtract); - EXPECT_EQ(root->operand(0)->operand(0), param0); - EXPECT_EQ(root->operand(0)->operand(1), param1); + + EXPECT_THAT(computation->root_instruction(), + op::Exp(op::Subtract(param0, param1))); +} + +// Test that exp(A)*exp(B) is simplified to exp(A+B) +TEST_F(AlgebraicSimplifierTest, ExpMul) { + Shape r0f32 = ShapeUtil::MakeShape(F32, {}); + HloComputation::Builder builder(TestName()); + HloInstruction* param0 = builder.AddInstruction( + HloInstruction::CreateParameter(0, r0f32, "param0")); + HloInstruction* param1 = builder.AddInstruction( + HloInstruction::CreateParameter(1, r0f32, "param1")); + HloInstruction* exp0 = builder.AddInstruction( + HloInstruction::CreateUnary(r0f32, HloOpcode::kExp, param0)); + HloInstruction* exp1 = builder.AddInstruction( + HloInstruction::CreateUnary(r0f32, HloOpcode::kExp, param1)); + builder.AddInstruction( + HloInstruction::CreateBinary(r0f32, HloOpcode::kMultiply, exp0, exp1)); + + auto module = CreateNewModule(); + auto computation = module->AddEntryComputation(builder.Build()); + + EXPECT_THAT(computation->root_instruction(), + op::Multiply(op::Exp(param0), op::Exp(param1))); + + AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, + non_bitcasting_callback()); + ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); + + EXPECT_THAT(computation->root_instruction(), + op::Exp(op::Add(param0, param1))); +} + +// Test that pow(exp(A), B) is simplified to exp(A*B) +TEST_F(AlgebraicSimplifierTest, PowExp) { + Shape r0f32 = ShapeUtil::MakeShape(F32, {}); + HloComputation::Builder builder(TestName()); + HloInstruction* param0 = builder.AddInstruction( + HloInstruction::CreateParameter(0, r0f32, "param0")); + HloInstruction* param1 = builder.AddInstruction( + HloInstruction::CreateParameter(1, r0f32, "param1")); + HloInstruction* exp0 = builder.AddInstruction( + HloInstruction::CreateUnary(r0f32, HloOpcode::kExp, param0)); + builder.AddInstruction( + HloInstruction::CreateBinary(r0f32, HloOpcode::kPower, exp0, param1)); + + auto module = CreateNewModule(); + auto computation = module->AddEntryComputation(builder.Build()); + + EXPECT_THAT(computation->root_instruction(), + op::Power(op::Exp(param0), param1)); + + AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, + non_bitcasting_callback()); + ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); + + EXPECT_THAT(computation->root_instruction(), + op::Exp(op::Multiply(param0, param1))); +} + +// Test that ln(pow(A, B)) is simplified to ln(A)*B +TEST_F(AlgebraicSimplifierTest, LnPow) { + Shape r0f32 = ShapeUtil::MakeShape(F32, {}); + HloComputation::Builder builder(TestName()); + HloInstruction* param0 = builder.AddInstruction( + HloInstruction::CreateParameter(0, r0f32, "param0")); + HloInstruction* param1 = builder.AddInstruction( + HloInstruction::CreateParameter(1, r0f32, "param1")); + HloInstruction* pow = builder.AddInstruction( + HloInstruction::CreateBinary(r0f32, HloOpcode::kPower, param0, param1)); + builder.AddInstruction( + HloInstruction::CreateUnary(r0f32, HloOpcode::kLog, pow)); + + auto module = CreateNewModule(); + auto computation = module->AddEntryComputation(builder.Build()); + + EXPECT_THAT(computation->root_instruction(), + op::Log(op::Power(param0, param1))); + + AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, + non_bitcasting_callback()); + ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); + + EXPECT_THAT(computation->root_instruction(), + op::Multiply(op::Log(param0), param1)); } // Test that ln(exp(A)) is simplified to A @@ -203,16 +485,16 @@ TEST_F(AlgebraicSimplifierTest, LnExp) { builder.AddInstruction( HloInstruction::CreateUnary(r0f32, HloOpcode::kLog, exp0)); - auto module = MakeUnique(TestName()); + auto module = CreateNewModule(); auto computation = module->AddEntryComputation(builder.Build()); - HloInstruction* root = computation->root_instruction(); - EXPECT_EQ(root->opcode(), HloOpcode::kLog); + + EXPECT_THAT(computation->root_instruction(), op::Log(op::Exp(param0))); + AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, non_bitcasting_callback()); ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); - root = computation->root_instruction(); - EXPECT_EQ(root->opcode(), HloOpcode::kParameter); - EXPECT_EQ(root, param0); + + EXPECT_EQ(computation->root_instruction(), param0); } // Test that ln(exp(A)/exp(B)) is simplified to A-B @@ -232,17 +514,17 @@ TEST_F(AlgebraicSimplifierTest, LnExpDiv) { builder.AddInstruction( HloInstruction::CreateUnary(r0f32, HloOpcode::kLog, div)); - auto module = MakeUnique(TestName()); + auto module = CreateNewModule(); auto computation = module->AddEntryComputation(builder.Build()); - HloInstruction* root = computation->root_instruction(); - EXPECT_EQ(root->opcode(), HloOpcode::kLog); + + EXPECT_THAT(computation->root_instruction(), + op::Log(op::Divide(op::Exp(param0), op::Exp(param1)))); + AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, non_bitcasting_callback()); ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); - root = computation->root_instruction(); - EXPECT_EQ(root->opcode(), HloOpcode::kSubtract); - EXPECT_EQ(root->operand(0), param0); - EXPECT_EQ(root->operand(1), param1); + + EXPECT_THAT(computation->root_instruction(), op::Subtract(param0, param1)); } // Test that pow(A, 0) where A is a scalar is simplified to the scalar @@ -253,18 +535,22 @@ TEST_F(AlgebraicSimplifierTest, Pow0Scalar) { HloInstruction* param0 = builder.AddInstruction( HloInstruction::CreateParameter(0, r0f32, "param0")); HloInstruction* zero = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(0))); + HloInstruction::CreateConstant(Literal::CreateR0(0))); builder.AddInstruction( HloInstruction::CreateBinary(r0f32, HloOpcode::kPower, param0, zero)); - auto module = MakeUnique(TestName()); + auto module = CreateNewModule(); auto computation = module->AddEntryComputation(builder.Build()); + + EXPECT_THAT(computation->root_instruction(), op::Power(param0, zero)); + AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, non_bitcasting_callback()); ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); + HloInstruction* root = computation->root_instruction(); - EXPECT_EQ(root->opcode(), HloOpcode::kConstant); - EXPECT_EQ(LiteralUtil::GetFirstElement(root->literal()), 1); + EXPECT_THAT(root, op::Constant()); + EXPECT_EQ(root->literal().GetFirstElement(), 1); } // Test that pow(A, 0) where A is not a scalar is simplified to broadcast(1). @@ -274,23 +560,26 @@ TEST_F(AlgebraicSimplifierTest, Pow0Vector) { HloInstruction* param0 = builder.AddInstruction( HloInstruction::CreateParameter(0, r1f32, "param0")); HloInstruction* zero = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(0))); + HloInstruction::CreateConstant(Literal::CreateR0(0))); builder.AddInstruction( HloInstruction::CreateBinary(r1f32, HloOpcode::kPower, param0, zero)); - auto module = MakeUnique(TestName()); + auto module = CreateNewModule(); auto computation = module->AddEntryComputation(builder.Build()); + + EXPECT_THAT(computation->root_instruction(), op::Power(param0, zero)); + AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, non_bitcasting_callback()); ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); + HloInstruction* root = computation->root_instruction(); - EXPECT_EQ(root->opcode(), HloOpcode::kBroadcast); + EXPECT_THAT(root, op::Broadcast()); EXPECT_TRUE(ShapeUtil::Equal(root->shape(), r1f32)) << ShapeUtil::HumanString(root->shape()); EXPECT_EQ(root->dimensions().size(), 0); EXPECT_TRUE(ShapeUtil::IsScalar(root->operand(0)->shape())); - EXPECT_EQ(LiteralUtil::GetFirstElement(root->operand(0)->literal()), - 1); + EXPECT_EQ(root->operand(0)->literal().GetFirstElement(), 1); } // Test that pow(A, 1) is simplified to A. @@ -300,18 +589,20 @@ TEST_F(AlgebraicSimplifierTest, Pow1) { HloInstruction* param0 = builder.AddInstruction( HloInstruction::CreateParameter(0, r0f32, "param0")); HloInstruction* one = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(1))); + HloInstruction::CreateConstant(Literal::CreateR0(1))); builder.AddInstruction( HloInstruction::CreateBinary(r0f32, HloOpcode::kPower, param0, one)); - auto module = MakeUnique(TestName()); + auto module = CreateNewModule(); auto computation = module->AddEntryComputation(builder.Build()); + + EXPECT_THAT(computation->root_instruction(), op::Power(param0, one)); + AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, non_bitcasting_callback()); ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); - HloInstruction* root = computation->root_instruction(); - EXPECT_EQ(root->opcode(), HloOpcode::kParameter); - EXPECT_EQ(root, param0); + + EXPECT_EQ(computation->root_instruction(), param0); } // Test that pow(A, 2) is simplified to A*A. @@ -321,19 +612,20 @@ TEST_F(AlgebraicSimplifierTest, Pow2) { HloInstruction* param0 = builder.AddInstruction( HloInstruction::CreateParameter(0, r0f32, "param0")); HloInstruction* two = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(2))); + HloInstruction::CreateConstant(Literal::CreateR0(2))); builder.AddInstruction( HloInstruction::CreateBinary(r0f32, HloOpcode::kPower, param0, two)); - auto module = MakeUnique(TestName()); + auto module = CreateNewModule(); auto computation = module->AddEntryComputation(builder.Build()); + + EXPECT_THAT(computation->root_instruction(), op::Power(param0, two)); + AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, non_bitcasting_callback()); ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); - HloInstruction* root = computation->root_instruction(); - EXPECT_EQ(root->opcode(), HloOpcode::kMultiply); - EXPECT_EQ(root->operand(0), param0); - EXPECT_EQ(root->operand(1), param0); + + EXPECT_THAT(computation->root_instruction(), op::Multiply(param0, param0)); } // Test that pow(A, -1) is simplified to 1/A. @@ -343,21 +635,22 @@ TEST_F(AlgebraicSimplifierTest, PowNegative1) { HloInstruction* param0 = builder.AddInstruction( HloInstruction::CreateParameter(0, r0f32, "param0")); HloInstruction* negative_one = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(-1))); + HloInstruction::CreateConstant(Literal::CreateR0(-1))); builder.AddInstruction(HloInstruction::CreateBinary(r0f32, HloOpcode::kPower, param0, negative_one)); - auto module = MakeUnique(TestName()); + auto module = CreateNewModule(); auto computation = module->AddEntryComputation(builder.Build()); + + EXPECT_THAT(computation->root_instruction(), op::Power(param0, negative_one)); + AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, non_bitcasting_callback()); ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); + HloInstruction* root = computation->root_instruction(); - EXPECT_EQ(root->opcode(), HloOpcode::kDivide); - EXPECT_EQ(root->operand(0)->opcode(), HloOpcode::kConstant); - EXPECT_EQ(LiteralUtil::GetFirstElement(root->operand(0)->literal()), - 1); - EXPECT_EQ(root->operand(1), param0); + EXPECT_THAT(root, op::Divide(op::Constant(), param0)); + EXPECT_EQ(root->operand(0)->literal().GetFirstElement(), 1); } TEST_F(AlgebraicSimplifierTest, ReshapeBroadcast) { @@ -374,124 +667,164 @@ TEST_F(AlgebraicSimplifierTest, ReshapeBroadcast) { ShapeUtil::MakeShape(F32, {3, 2}), broadcast)); auto computation = builder.Build(); - auto module = MakeUnique(TestName()); + auto module = CreateNewModule(); module->AddEntryComputation(std::move(computation)); - HloInstruction* root = module->entry_computation()->root_instruction(); + + EXPECT_THAT(module->entry_computation()->root_instruction(), + op::Reshape(op::Broadcast(op::Reshape(op)))); + HloPassFix simplifier(/*is_layout_sensitive=*/false, non_bitcasting_callback()); ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); - root = module->entry_computation()->root_instruction(); - EXPECT_EQ(root->opcode(), HloOpcode::kParameter); + + EXPECT_THAT(module->entry_computation()->root_instruction(), op); } // Test that convert(A, $TYPE) is simplified to A if A is of type $TYPE. TEST_F(AlgebraicSimplifierTest, ConvertBetweenSameType) { HloComputation::Builder builder(TestName()); HloInstruction* input = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(42.0f))); + HloInstruction::CreateConstant(Literal::CreateR0(42.0f))); builder.AddInstruction( HloInstruction::CreateConvert(ShapeUtil::MakeShape(F32, {}), input)); - auto module = MakeUnique(TestName()); + auto module = CreateNewModule(); auto computation = module->AddEntryComputation(builder.Build()); - EXPECT_EQ(HloOpcode::kConvert, computation->root_instruction()->opcode()); + EXPECT_THAT(computation->root_instruction(), op::Convert(input)); AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, non_bitcasting_callback()); ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); - EXPECT_EQ(HloOpcode::kConstant, computation->root_instruction()->opcode()); + EXPECT_THAT(computation->root_instruction(), input); } -TEST_F(AlgebraicSimplifierTest, ConvertF32ToS64) { +// Test that copies are removed. +TEST_F(AlgebraicSimplifierTest, RemoveCopy) { + Shape r0f32 = ShapeUtil::MakeShape(F32, {}); HloComputation::Builder builder(TestName()); - HloInstruction* input = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(42.0f))); + HloInstruction* param0 = builder.AddInstruction( + HloInstruction::CreateParameter(0, r0f32, "param0")); builder.AddInstruction( - HloInstruction::CreateConvert(ShapeUtil::MakeShape(S64, {}), input)); + HloInstruction::CreateUnary(param0->shape(), HloOpcode::kCopy, param0)); - auto module = MakeUnique(TestName()); + auto module = CreateNewModule(); auto computation = module->AddEntryComputation(builder.Build()); - EXPECT_EQ(HloOpcode::kConvert, computation->root_instruction()->opcode()); + EXPECT_THAT(computation->root_instruction(), op::Copy(param0)); AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, non_bitcasting_callback()); ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); - EXPECT_EQ(HloOpcode::kConstant, computation->root_instruction()->opcode()); - EXPECT_EQ(LiteralUtil::GetFirstElement( - computation->root_instruction()->literal()), - 42); + EXPECT_THAT(computation->root_instruction(), param0); } -TEST_F(AlgebraicSimplifierTest, ConvertS64ToF32) { +// Test that unary concatenates are removed. +TEST_F(AlgebraicSimplifierTest, RemoveUnaryConcatenate) { + Shape r1f32 = ShapeUtil::MakeShape(F32, {100}); HloComputation::Builder builder(TestName()); - HloInstruction* input = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(42))); + HloInstruction* param0 = builder.AddInstruction( + HloInstruction::CreateParameter(0, r1f32, "param0")); builder.AddInstruction( - HloInstruction::CreateConvert(ShapeUtil::MakeShape(F32, {}), input)); + HloInstruction::CreateConcatenate(param0->shape(), {param0}, 0)); - auto module = MakeUnique(TestName()); + auto module = CreateNewModule(); auto computation = module->AddEntryComputation(builder.Build()); - EXPECT_EQ(HloOpcode::kConvert, computation->root_instruction()->opcode()); + EXPECT_THAT(computation->root_instruction(), op::Concatenate(param0)); AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, non_bitcasting_callback()); ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); - EXPECT_EQ(HloOpcode::kConstant, computation->root_instruction()->opcode()); - EXPECT_EQ(LiteralUtil::GetFirstElement( - computation->root_instruction()->literal()), - 42.0f); + EXPECT_THAT(computation->root_instruction(), param0); } -TEST_F(AlgebraicSimplifierTest, ConvertF32ArrayToS64Array) { +// Test that empty operands of concatenates are removed. +TEST_F(AlgebraicSimplifierTest, RemoveEmptyConcatenateOperands) { + const int kParamLength = 100; + Shape r1f32 = ShapeUtil::MakeShape(F32, {kParamLength}); HloComputation::Builder builder(TestName()); - HloInstruction* input = builder.AddInstruction(HloInstruction::CreateConstant( - LiteralUtil::CreateR1({42.0f, 19.0f}))); - builder.AddInstruction( - HloInstruction::CreateConvert(ShapeUtil::MakeShape(S64, {2}), input)); + HloInstruction* param0 = builder.AddInstruction( + HloInstruction::CreateParameter(0, r1f32, "param0")); + HloInstruction* param1 = builder.AddInstruction( + HloInstruction::CreateParameter(1, r1f32, "param1")); + HloInstruction* empty_literal = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR1({}))); + HloInstruction* empty_slice = + builder.AddInstruction(HloInstruction::CreateSlice( + ShapeUtil::MakeShape(F32, {0}), param1, {42}, {42}, {1})); + Shape result_shape = ShapeUtil::MakeShape(F32, {3 * kParamLength}); + builder.AddInstruction(HloInstruction::CreateConcatenate( + result_shape, {empty_literal, param0, param0, empty_slice, param1}, 0)); + + auto module = CreateNewModule(); + auto computation = module->AddEntryComputation(builder.Build()); - auto module = MakeUnique(TestName()); + EXPECT_THAT( + computation->root_instruction(), + op::Concatenate(empty_literal, param0, param0, empty_slice, param1)); + + AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, + non_bitcasting_callback()); + ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); + + EXPECT_THAT(computation->root_instruction(), + op::Concatenate(param0, param0, param1)); +} + +// Test a concatenate with only empty operands is removed. +TEST_F(AlgebraicSimplifierTest, OnlyEmptyConcatenateOperands) { + const int kParamLength = 100; + Shape r1f32 = ShapeUtil::MakeShape(F32, {kParamLength}); + HloComputation::Builder builder(TestName()); + HloInstruction* param0 = builder.AddInstruction( + HloInstruction::CreateParameter(0, r1f32, "param0")); + HloInstruction* empty_literal = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR1({}))); + HloInstruction* empty_slice = + builder.AddInstruction(HloInstruction::CreateSlice( + ShapeUtil::MakeShape(F32, {0}), param0, {42}, {42}, {1})); + Shape result_shape = ShapeUtil::MakeShape(F32, {0}); + builder.AddInstruction(HloInstruction::CreateConcatenate( + result_shape, {empty_literal, empty_slice}, 0)); + + auto module = CreateNewModule(); auto computation = module->AddEntryComputation(builder.Build()); - EXPECT_EQ(HloOpcode::kConvert, computation->root_instruction()->opcode()); + EXPECT_THAT(computation->root_instruction(), + op::Concatenate(empty_literal, empty_slice)); AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, non_bitcasting_callback()); ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); - EXPECT_EQ(HloOpcode::kConstant, computation->root_instruction()->opcode()); - EXPECT_EQ( - LiteralUtil::Get(computation->root_instruction()->literal(), {0}), - 42); - EXPECT_EQ( - LiteralUtil::Get(computation->root_instruction()->literal(), {1}), - 19); + EXPECT_EQ(computation->root_instruction(), empty_literal); } -// Test that copies are removed. -TEST_F(AlgebraicSimplifierTest, RemoveCopy) { +// Test that concat with a scalar broadcast becomes a pad. +TEST_F(AlgebraicSimplifierTest, ConcatenateOfBroadcastBecomesPad) { + Shape r1f32 = ShapeUtil::MakeShape(F32, {100}); Shape r0f32 = ShapeUtil::MakeShape(F32, {}); HloComputation::Builder builder(TestName()); HloInstruction* param0 = builder.AddInstruction( - HloInstruction::CreateParameter(0, r0f32, "param0")); - HloInstruction* copy = builder.AddInstruction( - HloInstruction::CreateUnary(param0->shape(), HloOpcode::kCopy, param0)); + HloInstruction::CreateParameter(0, r1f32, "param0")); + HloInstruction* param1 = builder.AddInstruction( + HloInstruction::CreateParameter(1, r0f32, "param1")); + HloInstruction* broadcast = builder.AddInstruction( + HloInstruction::CreateBroadcast(r1f32, param1, {})); + builder.AddInstruction(HloInstruction::CreateConcatenate( + param0->shape(), {broadcast, param0}, 0)); - auto module = MakeUnique(TestName()); + auto module = CreateNewModule(); auto computation = module->AddEntryComputation(builder.Build()); - EXPECT_EQ(copy, computation->root_instruction()); - AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, non_bitcasting_callback()); ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); - - EXPECT_EQ(param0, computation->root_instruction()); + EXPECT_THAT(computation->root_instruction(), op::Pad(param0, param1)); } // Test that a simplification which changes layouts is not performed if layout @@ -504,21 +837,21 @@ TEST_F(AlgebraicSimplifierTest, CopyWithDifferentLayout) { HloInstruction* copy = builder.AddInstruction( HloInstruction::CreateUnary(param0->shape(), HloOpcode::kCopy, param0)); - auto module = MakeUnique(TestName()); + auto module = CreateNewModule(); auto computation = module->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_EQ(copy, computation->root_instruction()); + EXPECT_THAT(computation->root_instruction(), op::Copy(param0)); AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/true, non_bitcasting_callback()); EXPECT_FALSE(simplifier.Run(module.get()).ValueOrDie()); // Copy has not been removed. - EXPECT_EQ(copy, computation->root_instruction()); + EXPECT_THAT(computation->root_instruction(), op::Copy(param0)); } // Test that a simplification which preserves layouts is performed if layout @@ -531,21 +864,21 @@ TEST_F(AlgebraicSimplifierTest, CopyWithSameLayout) { HloInstruction* copy = builder.AddInstruction( HloInstruction::CreateUnary(param0->shape(), HloOpcode::kCopy, param0)); - auto module = MakeUnique(TestName()); + auto module = CreateNewModule(); auto computation = module->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_EQ(copy, computation->root_instruction()); + EXPECT_THAT(computation->root_instruction(), op::Copy(param0)); AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/true, non_bitcasting_callback()); ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); // Copy has been removed. - EXPECT_EQ(param0, computation->root_instruction()); + EXPECT_THAT(computation->root_instruction(), param0); } // Test that a reshape which could be replaced with a bitcast is not if @@ -563,17 +896,17 @@ TEST_F(AlgebraicSimplifierTest, NoBitcastAdded) { *reshape->mutable_shape()->mutable_layout() = LayoutUtil::MakeLayout({0, 1, 2, 3, 4, 5}); - auto module = MakeUnique(TestName()); + auto module = CreateNewModule(); auto computation = module->AddEntryComputation(builder.Build()); - EXPECT_EQ(reshape, computation->root_instruction()); + EXPECT_THAT(computation->root_instruction(), op::Reshape(param0)); AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/true, non_bitcasting_callback()); EXPECT_FALSE(simplifier.Run(module.get()).ValueOrDie()); // Reshape is not replaced with a bitcast. - EXPECT_EQ(reshape, computation->root_instruction()); + EXPECT_THAT(computation->root_instruction(), op::Reshape(param0)); } // Test transforming reshapes to bitcasts under various conditions. @@ -609,25 +942,21 @@ TEST_F(AlgebraicSimplifierTest, ReshapeReplacedWithBitcast) { builder.AddInstruction(HloInstruction::CreateTuple( {transformable_reshape, dimensions_wrong_reshape, layout_wrong_reshape})); - auto module = MakeUnique(TestName()); + auto module = CreateNewModule(); auto computation = module->AddEntryComputation(builder.Build()); - EXPECT_EQ(transformable_reshape, computation->root_instruction()->operand(0)); - EXPECT_EQ(dimensions_wrong_reshape, - computation->root_instruction()->operand(1)); - EXPECT_EQ(layout_wrong_reshape, computation->root_instruction()->operand(2)); + EXPECT_THAT(computation->root_instruction(), + op::Tuple(transformable_reshape, dimensions_wrong_reshape, + layout_wrong_reshape)); AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/true, bitcasting_callback()); simplifier.Run(module.get()).ValueOrDie(); // Verify that only the first reshape is replaced. - EXPECT_NE(transformable_reshape, computation->root_instruction()->operand(0)); - EXPECT_EQ(HloOpcode::kBitcast, - computation->root_instruction()->operand(0)->opcode()); - EXPECT_EQ(dimensions_wrong_reshape, - computation->root_instruction()->operand(1)); - EXPECT_EQ(layout_wrong_reshape, computation->root_instruction()->operand(2)); + EXPECT_THAT( + computation->root_instruction(), + op::Tuple(op::Bitcast(), dimensions_wrong_reshape, layout_wrong_reshape)); } TEST_F(AlgebraicSimplifierTest, ReshapeAfterEffectiveUnary) { @@ -639,20 +968,50 @@ TEST_F(AlgebraicSimplifierTest, ReshapeAfterEffectiveUnary) { builder.AddInstruction(HloInstruction::CreateReshape( ShapeUtil::MakeShape(F32, {1, 2, 3, 4, 5}), param)); HloInstruction* zero = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(0.0f))); + HloInstruction::CreateConstant(Literal::CreateR0(0.0f))); builder.AddInstruction( HloInstruction::CreateBinary(ShapeUtil::MakeShape(F32, {1, 2, 3, 4, 5}), HloOpcode::kMaximum, movable_reshape, zero)); - auto module = MakeUnique(TestName()); + auto module = CreateNewModule(); auto computation = module->AddEntryComputation(builder.Build()); - HloInstruction* root = computation->root_instruction(); - EXPECT_EQ(root->opcode(), HloOpcode::kMaximum); + + EXPECT_THAT(computation->root_instruction(), + op::Maximum(op::Reshape(param), zero)); + + AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, + bitcasting_callback()); + + simplifier.Run(module.get()).ValueOrDie(); + EXPECT_THAT(computation->root_instruction(), + op::Reshape(op::Maximum(param, zero))); +} + +// Regression test for a bug in the reshape sinking transformation, where +// moving a reshape to a scalar led to a crash. +TEST_F(AlgebraicSimplifierTest, ReshapeToScalarNotHoistedAfterEffectiveUnary) { + HloComputation::Builder builder(TestName()); + HloInstruction* param = + builder.AddInstruction(HloInstruction::CreateParameter( + 0, ShapeUtil::MakeShape(F32, {1, 1}), "param")); + HloInstruction* reshape = builder.AddInstruction( + HloInstruction::CreateReshape(ShapeUtil::MakeShape(F32, {}), param)); + HloInstruction* zero = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR1({1., 2., 3.}))); + builder.AddInstruction(HloInstruction::CreateBinary( + ShapeUtil::MakeShape(F32, {3}), HloOpcode::kMaximum, reshape, zero)); + auto module = CreateNewModule(); + auto computation = module->AddEntryComputation(builder.Build()); + + EXPECT_THAT(computation->root_instruction(), + op::Maximum(op::Reshape(param), zero)); + AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, bitcasting_callback()); + simplifier.Run(module.get()).ValueOrDie(); - EXPECT_EQ(HloOpcode::kReshape, computation->root_instruction()->opcode()); - EXPECT_EQ(HloOpcode::kMaximum, - computation->root_instruction()->operand(0)->opcode()); + + EXPECT_THAT(computation->root_instruction(), + op::Maximum(op::Reshape(param), zero)); } TEST_F(AlgebraicSimplifierTest, TransposeEqualsBitcast1) { @@ -669,16 +1028,17 @@ TEST_F(AlgebraicSimplifierTest, TransposeEqualsBitcast1) { *transpose->mutable_shape()->mutable_layout() = LayoutUtil::MakeLayout({0, 1, 2, 3}); - auto module = MakeUnique(TestName()); + auto module = CreateNewModule(); auto computation = module->AddEntryComputation(builder.Build()); + EXPECT_THAT(computation->root_instruction(), op::Transpose(param)); + AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/true, bitcasting_callback()); ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); // Verify that the reshape is replaced. - EXPECT_EQ(2, computation->instruction_count()); - EXPECT_EQ(HloOpcode::kBitcast, computation->root_instruction()->opcode()); + EXPECT_THAT(computation->root_instruction(), op::Bitcast(param)); } TEST_F(AlgebraicSimplifierTest, TransposeEqualsBitcast2) { @@ -695,16 +1055,17 @@ TEST_F(AlgebraicSimplifierTest, TransposeEqualsBitcast2) { *transpose->mutable_shape()->mutable_layout() = LayoutUtil::MakeLayout({3, 1, 2, 0}); - auto module = MakeUnique(TestName()); + auto module = CreateNewModule(); auto computation = module->AddEntryComputation(builder.Build()); + EXPECT_THAT(computation->root_instruction(), op::Transpose(param)); + AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/true, bitcasting_callback()); ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); // Verify that the reshape is replaced. - EXPECT_EQ(2, computation->instruction_count()); - EXPECT_EQ(HloOpcode::kBitcast, computation->root_instruction()->opcode()); + EXPECT_THAT(computation->root_instruction(), op::Bitcast(param)); } TEST_F(AlgebraicSimplifierTest, ReshapesMerged) { @@ -717,23 +1078,47 @@ TEST_F(AlgebraicSimplifierTest, ReshapesMerged) { builder.AddInstruction(HloInstruction::CreateReshape( ShapeUtil::MakeShape(F32, {2, 1, 2}), param0)); - HloInstruction* reshape2 = - builder.AddInstruction(HloInstruction::CreateReshape( - ShapeUtil::MakeShape(F32, {1, 2, 1, 1, 2, 1}), reshape1)); + builder.AddInstruction(HloInstruction::CreateReshape( + ShapeUtil::MakeShape(F32, {1, 2, 1, 1, 2, 1}), reshape1)); - auto module = MakeUnique(TestName()); + auto module = CreateNewModule(); auto computation = module->AddEntryComputation(builder.Build()); - EXPECT_EQ(reshape2, computation->root_instruction()); - EXPECT_EQ(reshape1, computation->root_instruction()->operand(0)); + EXPECT_THAT(computation->root_instruction(), + op::Reshape(op::Reshape(param0))); AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, non_bitcasting_callback()); ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); - EXPECT_EQ(HloOpcode::kReshape, computation->root_instruction()->opcode()); - EXPECT_EQ(HloOpcode::kParameter, - computation->root_instruction()->operand(0)->opcode()); + EXPECT_THAT(computation->root_instruction(), op::Reshape(param0)); +} + +TEST_F(AlgebraicSimplifierTest, CopiesMerged) { + HloComputation::Builder builder(TestName()); + HloInstruction* param0 = + builder.AddInstruction(HloInstruction::CreateParameter( + 0, ShapeUtil::MakeShapeWithMonotonicDim0MajorLayout(F32, {2, 2, 2}), + "param0")); + + HloInstruction* copy1 = builder.AddInstruction(HloInstruction::CreateUnary( + ShapeUtil::MakeShapeWithLayout(F32, {2, 2, 2}, {0, 1, 2}), + HloOpcode::kCopy, param0)); + + builder.AddInstruction(HloInstruction::CreateUnary( + ShapeUtil::MakeShapeWithLayout(F32, {2, 2, 2}, {0, 2, 1}), + HloOpcode::kCopy, copy1)); + + auto module = CreateNewModule(); + auto computation = module->AddEntryComputation(builder.Build()); + + EXPECT_THAT(computation->root_instruction(), op::Copy(op::Copy(param0))); + + AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/true, + non_bitcasting_callback()); + ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); + + EXPECT_THAT(computation->root_instruction(), op::Copy(param0)); } TEST_F(AlgebraicSimplifierTest, TransposesMerged) { @@ -746,25 +1131,21 @@ TEST_F(AlgebraicSimplifierTest, TransposesMerged) { builder.AddInstruction(HloInstruction::CreateTranspose( ShapeUtil::MakeShape(F32, {3, 4, 2}), param0, {1, 2, 0})); - HloInstruction* transpose2 = - builder.AddInstruction(HloInstruction::CreateTranspose( - ShapeUtil::MakeShape(F32, {4, 3, 2}), transpose1, {1, 0, 2})); + builder.AddInstruction(HloInstruction::CreateTranspose( + ShapeUtil::MakeShape(F32, {4, 3, 2}), transpose1, {1, 0, 2})); - auto module = MakeUnique(TestName()); + auto module = CreateNewModule(); auto computation = module->AddEntryComputation(builder.Build()); - EXPECT_EQ(transpose2, computation->root_instruction()); - EXPECT_EQ(transpose1, computation->root_instruction()->operand(0)); + EXPECT_THAT(computation->root_instruction(), op::Transpose(transpose1)); AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, non_bitcasting_callback()); ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); - EXPECT_EQ(HloOpcode::kTranspose, computation->root_instruction()->opcode()); + EXPECT_THAT(computation->root_instruction(), op::Transpose(param0)); EXPECT_EQ(std::vector({2, 1, 0}), computation->root_instruction()->dimensions()); - EXPECT_EQ(HloOpcode::kParameter, - computation->root_instruction()->operand(0)->opcode()); } // Test merging reshape and broadcast. @@ -777,16 +1158,17 @@ TEST_F(AlgebraicSimplifierTest, ReshapeAndBroadcastMerged) { builder.AddInstruction(HloInstruction::CreateBroadcast( ShapeUtil::MakeShape(F32, {1, 2, 3, 5, 1}), reshape1, {0, 2, 3})); - auto module = MakeUnique(TestName()); + auto module = CreateNewModule(); auto computation = module->AddEntryComputation(builder.Build()); + EXPECT_THAT(computation->root_instruction(), + op::Broadcast(op::Reshape(param0))); + AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, non_bitcasting_callback()); ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); - EXPECT_EQ(HloOpcode::kBroadcast, computation->root_instruction()->opcode()); - EXPECT_EQ(HloOpcode::kParameter, - computation->root_instruction()->operand(0)->opcode()); + EXPECT_THAT(computation->root_instruction(), op::Broadcast(param0)); } // Test merging broadcast and reshape. @@ -799,16 +1181,17 @@ TEST_F(AlgebraicSimplifierTest, BroadcastAndReshapeMerged) { builder.AddInstruction(HloInstruction::CreateReshape( ShapeUtil::MakeShape(F32, {2, 3, 7, 2, 1, 3, 2}), broadcast1)); - auto module = MakeUnique(TestName()); + auto module = CreateNewModule(); auto computation = module->AddEntryComputation(builder.Build()); + EXPECT_THAT(computation->root_instruction(), + op::Reshape(op::Broadcast(param0))); + AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, non_bitcasting_callback()); ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); - EXPECT_EQ(HloOpcode::kBroadcast, computation->root_instruction()->opcode()); - EXPECT_EQ(HloOpcode::kParameter, - computation->root_instruction()->operand(0)->opcode()); + EXPECT_THAT(computation->root_instruction(), op::Broadcast(param0)); } TEST_F(AlgebraicSimplifierTest, BroadcastAndReshape_1_3x1_3) { @@ -820,12 +1203,18 @@ TEST_F(AlgebraicSimplifierTest, BroadcastAndReshape_1_3x1_3) { builder.AddInstruction( HloInstruction::CreateReshape(ShapeUtil::MakeShape(F32, {3}), broadcast)); - auto module = MakeUnique(TestName()); - module->AddEntryComputation(builder.Build()); + auto module = CreateNewModule(); + auto computation = module->AddEntryComputation(builder.Build()); + + EXPECT_THAT(computation->root_instruction(), + op::Reshape(op::Broadcast(param))); AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, non_bitcasting_callback()); EXPECT_FALSE(simplifier.Run(module.get()).ValueOrDie()); + + EXPECT_THAT(computation->root_instruction(), + op::Reshape(op::Broadcast(param))); } TEST_F(AlgebraicSimplifierTest, BroadcastAndReshape_4_3x2x4_6x1x1x4) { @@ -837,15 +1226,19 @@ TEST_F(AlgebraicSimplifierTest, BroadcastAndReshape_4_3x2x4_6x1x1x4) { builder.AddInstruction(HloInstruction::CreateReshape( ShapeUtil::MakeShape(F32, {6, 1, 1, 4}), broadcast)); - auto module = MakeUnique(TestName()); + auto module = CreateNewModule(); HloComputation* computation = module->AddEntryComputation(builder.Build()); + EXPECT_THAT(computation->root_instruction(), + op::Reshape(op::Broadcast(param))); + AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, non_bitcasting_callback()); ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); - EXPECT_EQ(HloOpcode::kBroadcast, computation->root_instruction()->opcode()); - EXPECT_MATCH(computation->root_instruction()->dimensions(), - testing::VectorMatcher({3})); + + EXPECT_THAT(computation->root_instruction(), op::Broadcast(param)); + EXPECT_THAT(computation->root_instruction()->dimensions(), + ::testing::ElementsAre(3)); } TEST_F(AlgebraicSimplifierTest, BroadcastAndReshape_1_3x2x1_6x1x1x1) { @@ -857,18 +1250,21 @@ TEST_F(AlgebraicSimplifierTest, BroadcastAndReshape_1_3x2x1_6x1x1x1) { builder.AddInstruction(HloInstruction::CreateReshape( ShapeUtil::MakeShape(F32, {6, 1, 1, 1}), broadcast)); - auto module = MakeUnique(TestName()); + auto module = CreateNewModule(); HloComputation* computation = module->AddEntryComputation(builder.Build()); + EXPECT_THAT(computation->root_instruction(), + op::Reshape(op::Broadcast(param))); + AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, non_bitcasting_callback()); ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); - EXPECT_EQ(HloOpcode::kBroadcast, computation->root_instruction()->opcode()); + + EXPECT_THAT(computation->root_instruction(), op::Broadcast(param)); const std::vector broadcast_dims = computation->root_instruction()->dimensions(); EXPECT_EQ(1, broadcast_dims.size()); - EXPECT_TRUE(broadcast_dims[0] == 1 || broadcast_dims[0] == 2 || - broadcast_dims[3] == 3); + EXPECT_THAT(broadcast_dims[0], ::testing::AnyOf(1, 2, 3)); } TEST_F(AlgebraicSimplifierTest, BroadcastAndReshape_4_3x2x4x2_6x8) { @@ -880,12 +1276,18 @@ TEST_F(AlgebraicSimplifierTest, BroadcastAndReshape_4_3x2x4x2_6x8) { builder.AddInstruction(HloInstruction::CreateReshape( ShapeUtil::MakeShape(F32, {6, 8}), broadcast)); - auto module = MakeUnique(TestName()); - module->AddEntryComputation(builder.Build()); + auto module = CreateNewModule(); + HloComputation* computation = module->AddEntryComputation(builder.Build()); + + EXPECT_THAT(computation->root_instruction(), + op::Reshape(op::Broadcast(param))); AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, non_bitcasting_callback()); EXPECT_FALSE(simplifier.Run(module.get()).ValueOrDie()); + + EXPECT_THAT(computation->root_instruction(), + op::Reshape(op::Broadcast(param))); } TEST_F(AlgebraicSimplifierTest, RemoveNoopPad) { @@ -894,7 +1296,7 @@ TEST_F(AlgebraicSimplifierTest, RemoveNoopPad) { builder.AddInstruction(HloInstruction::CreateParameter( 0, ShapeUtil::MakeShape(F32, {2, 2}), "param")); HloInstruction* zero = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(0.0f))); + HloInstruction::CreateConstant(Literal::CreateR0(0.0f))); PaddingConfig no_padding; for (int i = 0; i < 2; ++i) { auto dimension = no_padding.add_dimensions(); @@ -908,10 +1310,13 @@ TEST_F(AlgebraicSimplifierTest, RemoveNoopPad) { HloModule module(TestName()); HloComputation* computation = module.AddEntryComputation(builder.Build()); + EXPECT_THAT(computation->root_instruction(), op::Pad(param, zero)); + AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, non_bitcasting_callback()); ASSERT_TRUE(simplifier.Run(&module).ValueOrDie()); - EXPECT_EQ(1, computation->instruction_count()); + + EXPECT_THAT(computation->root_instruction(), param); } TEST_F(AlgebraicSimplifierTest, NegativePadding) { @@ -922,7 +1327,7 @@ TEST_F(AlgebraicSimplifierTest, NegativePadding) { builder.AddInstruction(HloInstruction::CreateParameter( 0, ShapeUtil::MakeShape(F32, {10, 10}), "param")); HloInstruction* zero = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(0.0f))); + HloInstruction::CreateConstant(Literal::CreateR0(0.0f))); PaddingConfig padding; int64 low_padding[2] = {-1, -2}; int64 high_padding[2] = {2, -3}; @@ -951,18 +1356,14 @@ TEST_F(AlgebraicSimplifierTest, NegativePadding) { return false; }; - EXPECT_EQ(3, computation->instruction_count()); - EXPECT_EQ(computation->root_instruction(), pad); + EXPECT_THAT(computation->root_instruction(), op::Pad(param, zero)); EXPECT_TRUE(has_negative_padding(pad)); ASSERT_TRUE(simplifier.Run(&module).ValueOrDie()); - EXPECT_EQ(4, computation->instruction_count()); - EXPECT_EQ(computation->root_instruction()->opcode(), HloOpcode::kSlice); - const HloInstruction* root_operand = - computation->root_instruction()->operand(0); - EXPECT_EQ(root_operand->opcode(), HloOpcode::kPad); - EXPECT_FALSE(has_negative_padding(root_operand)); + EXPECT_THAT(computation->root_instruction(), op::Slice(op::Pad(param, zero))); + EXPECT_FALSE( + has_negative_padding(computation->root_instruction()->operand(0))); } TEST_F(AlgebraicSimplifierTest, RemoveNoopReshape) { @@ -976,10 +1377,13 @@ TEST_F(AlgebraicSimplifierTest, RemoveNoopReshape) { HloModule module(TestName()); HloComputation* computation = module.AddEntryComputation(builder.Build()); + EXPECT_THAT(computation->root_instruction(), op::Reshape(param)); + AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, non_bitcasting_callback()); ASSERT_TRUE(simplifier.Run(&module).ValueOrDie()); - EXPECT_EQ(1, computation->instruction_count()); + + EXPECT_THAT(computation->root_instruction(), param); } TEST_F(AlgebraicSimplifierTest, RemoveNoopSlice) { @@ -991,15 +1395,18 @@ TEST_F(AlgebraicSimplifierTest, RemoveNoopSlice) { 0, ShapeUtil::MakeShape(F32, {dim0, dim1}), "param")); builder.AddInstruction(HloInstruction::CreateSlice( ShapeUtil::MakeShape(F32, {dim0, dim1}), param, /*start_indices=*/{0, 0}, - /*limit_indices=*/{dim0, dim1})); + /*limit_indices=*/{dim0, dim1}, /*strides=*/{1, 1})); HloModule module(TestName()); HloComputation* computation = module.AddEntryComputation(builder.Build()); + EXPECT_THAT(computation->root_instruction(), op::Slice(param)); + AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, non_bitcasting_callback()); ASSERT_TRUE(simplifier.Run(&module).ValueOrDie()); - EXPECT_EQ(1, computation->instruction_count()); + + EXPECT_THAT(computation->root_instruction(), param); } TEST_F(AlgebraicSimplifierTest, ConvertConvToMatmul) { @@ -1230,26 +1637,26 @@ TEST_F(AlgebraicSimplifierTest, MaxMinToClamp) { HloInstruction* param0 = builder.AddInstruction( HloInstruction::CreateParameter(0, r0f32, "param0")); HloInstruction* min_value = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(0.0f))); + HloInstruction::CreateConstant(Literal::CreateR0(0.0f))); HloInstruction* max_value = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(1.0f))); + HloInstruction::CreateConstant(Literal::CreateR0(1.0f))); HloInstruction* min = builder.AddInstruction(HloInstruction::CreateBinary( r0f32, HloOpcode::kMinimum, param0, min_value)); - HloInstruction* max = builder.AddInstruction( + builder.AddInstruction( HloInstruction::CreateBinary(r0f32, HloOpcode::kMaximum, min, max_value)); HloModule module(TestName()); auto computation = module.AddEntryComputation(builder.Build()); - HloInstruction* root = computation->root_instruction(); - EXPECT_EQ(root, max); + + EXPECT_THAT(computation->root_instruction(), + op::Maximum(op::Minimum(param0, min_value), max_value)); + AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, non_bitcasting_callback()); ASSERT_TRUE(simplifier.Run(&module).ValueOrDie()); - root = computation->root_instruction(); - ASSERT_EQ(root->opcode(), HloOpcode::kClamp); - EXPECT_EQ(root->operand(0), max_value); - EXPECT_EQ(root->operand(1), param0); - EXPECT_EQ(root->operand(2), min_value); + + EXPECT_THAT(computation->root_instruction(), + op::Clamp(max_value, param0, min_value)); } // Test that min(max(A, x), y) is transformed to clamp(x, A, y) for scalar @@ -1260,26 +1667,26 @@ TEST_F(AlgebraicSimplifierTest, MinMaxToClamp) { HloInstruction* param0 = builder.AddInstruction( HloInstruction::CreateParameter(0, r0f32, "param0")); HloInstruction* min_value = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(0.0f))); + HloInstruction::CreateConstant(Literal::CreateR0(0.0f))); HloInstruction* max_value = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(1.0f))); + HloInstruction::CreateConstant(Literal::CreateR0(1.0f))); HloInstruction* max = builder.AddInstruction(HloInstruction::CreateBinary( r0f32, HloOpcode::kMaximum, param0, max_value)); - HloInstruction* min = builder.AddInstruction( + builder.AddInstruction( HloInstruction::CreateBinary(r0f32, HloOpcode::kMinimum, max, min_value)); HloModule module(TestName()); auto computation = module.AddEntryComputation(builder.Build()); - HloInstruction* root = computation->root_instruction(); - EXPECT_EQ(root, min); + + EXPECT_THAT(computation->root_instruction(), + op::Minimum(op::Maximum(param0, max_value), min_value)); + AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, non_bitcasting_callback()); ASSERT_TRUE(simplifier.Run(&module).ValueOrDie()); - root = computation->root_instruction(); - EXPECT_EQ(root->opcode(), HloOpcode::kClamp); - EXPECT_EQ(root->operand(0), max_value); - EXPECT_EQ(root->operand(1), param0); - EXPECT_EQ(root->operand(2), min_value); + + EXPECT_THAT(computation->root_instruction(), + op::Clamp(max_value, param0, min_value)); } // Test that min(max(A, x), y) is transformed to clamp(x, A, y) for @@ -1291,26 +1698,26 @@ TEST_F(AlgebraicSimplifierTest, MinMaxWithBroadcastToClamp) { HloInstruction* param0 = builder.AddInstruction( HloInstruction::CreateParameter(0, r1f32, "param0")); HloInstruction* min_value = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(0.0f))); + HloInstruction::CreateConstant(Literal::CreateR0(0.0f))); HloInstruction* max_value = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(1.0f))); + HloInstruction::CreateConstant(Literal::CreateR0(1.0f))); HloInstruction* max = builder.AddInstruction(HloInstruction::CreateBinary( r1f32, HloOpcode::kMaximum, param0, max_value)); - HloInstruction* min = builder.AddInstruction( + builder.AddInstruction( HloInstruction::CreateBinary(r1f32, HloOpcode::kMinimum, max, min_value)); HloModule module(TestName()); auto computation = module.AddEntryComputation(builder.Build()); - HloInstruction* root = computation->root_instruction(); - EXPECT_EQ(root, min); + + EXPECT_THAT(computation->root_instruction(), + op::Minimum(op::Maximum(param0, max_value), min_value)); + AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, non_bitcasting_callback()); ASSERT_TRUE(simplifier.Run(&module).ValueOrDie()); - root = computation->root_instruction(); - EXPECT_EQ(root->opcode(), HloOpcode::kClamp); - EXPECT_EQ(root->operand(0), max_value); - EXPECT_EQ(root->operand(1), param0); - EXPECT_EQ(root->operand(2), min_value); + + EXPECT_THAT(computation->root_instruction(), + op::Clamp(max_value, param0, min_value)); } // Test that min(max(A, non-constant1), non-constant2) is not canonicalized to @@ -1326,17 +1733,21 @@ TEST_F(AlgebraicSimplifierTest, MinMaxNotToClamp) { HloInstruction::CreateParameter(2, r0f32, "param2")); HloInstruction* max = builder.AddInstruction(HloInstruction::CreateBinary( r0f32, HloOpcode::kMaximum, param0, max_value)); - HloInstruction* min = builder.AddInstruction( + builder.AddInstruction( HloInstruction::CreateBinary(r0f32, HloOpcode::kMinimum, max, min_value)); HloModule module(TestName()); auto computation = module.AddEntryComputation(builder.Build()); - HloInstruction* root = computation->root_instruction(); + + EXPECT_THAT(computation->root_instruction(), + op::Minimum(op::Maximum(param0, max_value), min_value)); + AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, non_bitcasting_callback()); EXPECT_FALSE(simplifier.Run(&module).ValueOrDie()); - root = computation->root_instruction(); - EXPECT_EQ(root, min); + + EXPECT_THAT(computation->root_instruction(), + op::Minimum(op::Maximum(param0, max_value), min_value)); } // Test that min(f(max(A, constant1)), constant2) is not transformed to @@ -1347,25 +1758,30 @@ TEST_F(AlgebraicSimplifierTest, MinEquationWithMaxNotToClamp) { HloInstruction* param0 = builder.AddInstruction( HloInstruction::CreateParameter(0, r0f32, "param0")); HloInstruction* min_value = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(0.0f))); + HloInstruction::CreateConstant(Literal::CreateR0(0.0f))); HloInstruction* max_value = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(1.0f))); + HloInstruction::CreateConstant(Literal::CreateR0(1.0f))); HloInstruction* max = builder.AddInstruction(HloInstruction::CreateBinary( r0f32, HloOpcode::kMaximum, param0, max_value)); HloInstruction* fmax = builder.AddInstruction( HloInstruction::CreateBinary(r0f32, HloOpcode::kAdd, max, max_value)); - HloInstruction* min = builder.AddInstruction(HloInstruction::CreateBinary( + builder.AddInstruction(HloInstruction::CreateBinary( r0f32, HloOpcode::kMinimum, fmax, min_value)); HloModule module(TestName()); auto computation = module.AddEntryComputation(builder.Build()); - HloInstruction* root = computation->root_instruction(); - EXPECT_EQ(root, min); + + EXPECT_THAT(computation->root_instruction(), + op::Minimum(op::Add(op::Maximum(param0, max_value), max_value), + min_value)); + AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, non_bitcasting_callback()); EXPECT_FALSE(simplifier.Run(&module).ValueOrDie()); - root = computation->root_instruction(); - EXPECT_EQ(root, min); + + EXPECT_THAT(computation->root_instruction(), + op::Minimum(op::Add(op::Maximum(param0, max_value), max_value), + min_value)); } // Test that slice(broadcast(/*scalar value*/)) simplifies to a single @@ -1384,7 +1800,7 @@ TEST_F(AlgebraicSimplifierTest, ScalarBroadcastToSlice) { Shape slice_shape = ShapeUtil::MakeShape(F32, {2, 2, 3, 3}); HloInstruction* slice = builder.AddInstruction(HloInstruction::CreateSlice( - slice_shape, broadcast, {0, 1, 2, 3}, {2, 3, 5, 6})); + slice_shape, broadcast, {0, 1, 2, 3}, {2, 3, 5, 6}, {1, 1, 1, 1})); HloModule module(TestName()); auto computation = module.AddEntryComputation(builder.Build()); @@ -1402,8 +1818,7 @@ TEST_F(AlgebraicSimplifierTest, ScalarBroadcastToSlice) { ASSERT_FALSE(simplifier.Run(&module).ValueOrDie()); root = computation->root_instruction(); - EXPECT_EQ(root->opcode(), HloOpcode::kBroadcast); - EXPECT_EQ(scalar_param, root->operand(0)); + EXPECT_THAT(root, op::Broadcast(scalar_param)); EXPECT_TRUE(ShapeUtil::Equal(root->shape(), slice_shape)); } @@ -1412,7 +1827,7 @@ TEST_F(AlgebraicSimplifierTest, ScalarBroadcastToSlice) { TEST_F(AlgebraicSimplifierTest, ScalarBroadcastToTransposeReshape) { HloComputation::Builder builder(TestName()); HloInstruction* forty_two = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(42.0f))); + HloInstruction::CreateConstant(Literal::CreateR0(42.0f))); Shape broadcast_shape = ShapeUtil::MakeShape(F32, {4, 5, 6}); HloInstruction* broadcast = @@ -1440,11 +1855,90 @@ TEST_F(AlgebraicSimplifierTest, ScalarBroadcastToTransposeReshape) { ASSERT_TRUE(simplifier.Run(&module).ValueOrDie()); root = computation->root_instruction(); - EXPECT_EQ(root->opcode(), HloOpcode::kBroadcast); - EXPECT_EQ(forty_two, root->operand(0)); + EXPECT_THAT(root, op::Broadcast(forty_two)); EXPECT_TRUE(ShapeUtil::Equal(root->shape(), reshape_shape)); } +// Test that ReduceWindow(Pad(op, x), y) can simplify to ReduceWindow(op, x). +TEST_F(AlgebraicSimplifierTest, FoldPadIntoReduceWindow) { + HloModule module(TestName()); + HloComputation::Builder builder(TestName()); + + // Create operand to the pad. + HloInstruction* operand = + builder.AddInstruction(HloInstruction::CreateParameter( + 0, ShapeUtil::MakeShape(F32, {1, 2, 3, 4}), "p0")); + + // Create the pad. + PaddingConfig padding = MakeNoPaddingConfig(4); + padding.mutable_dimensions(1)->set_edge_padding_low(1); + padding.mutable_dimensions(3)->set_edge_padding_high(2); + + HloInstruction* pad_value = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(5.0f))); + HloInstruction* pad = builder.AddInstruction(HloInstruction::CreatePad( + ShapeUtil::MakeShape(F32, {1, 3, 3, 5}), operand, pad_value, padding)); + + // 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 = module.AddEmbeddedComputation(builder.Build()); + } + + // Create the reduce-window. + Window window; + for (int64 i = 0; i < ShapeUtil::Rank(pad->shape()); ++i) { + auto* dim = window.add_dimensions(); + dim->set_size(1); + dim->set_padding_low(10); + dim->set_padding_high(100); + dim->set_window_dilation(1); + dim->set_base_dilation(1); + } + const Shape reduce_window_shape = + ShapeUtil::MakeShape(F32, {111, 113, 113, 115}); + HloInstruction* reduce_init_value = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(5.0f))); + HloInstruction* reduce_window = + builder.AddInstruction(HloInstruction::CreateReduceWindow( + reduce_window_shape, pad, reduce_init_value, window, + add_computation)); + + // Build the computation and run the simplifier. + 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()); + ASSERT_TRUE(simplifier.Run(&module).ValueOrDie()); + + // Running simplification again should not result in any further changes. + ASSERT_FALSE(simplifier.Run(&module).ValueOrDie()); + + // Verify the result + root = computation->root_instruction(); + EXPECT_THAT(root, op::ReduceWindow(operand, op::Constant())); + EXPECT_TRUE(ShapeUtil::Equal(root->shape(), reduce_window_shape)) + << ShapeUtil::HumanString(root->shape()) << " vs " + << ShapeUtil::HumanString(reduce_window_shape); + EXPECT_EQ(root->window().dimensions(0).padding_low(), 10); + EXPECT_EQ(root->window().dimensions(1).padding_low(), 11); + EXPECT_EQ(root->window().dimensions(2).padding_low(), 10); + EXPECT_EQ(root->window().dimensions(3).padding_low(), 10); + EXPECT_EQ(root->window().dimensions(0).padding_high(), 100); + EXPECT_EQ(root->window().dimensions(1).padding_high(), 100); + EXPECT_EQ(root->window().dimensions(2).padding_high(), 100); + EXPECT_EQ(root->window().dimensions(3).padding_high(), 102); +} + TEST_F(AlgebraicSimplifierTest, ReversalOfTrivialDimensionsToBitcast) { HloComputation::Builder builder(TestName()); const Shape shape = ShapeUtil::MakeShape(F32, {448, 2048, 1, 1}); @@ -1461,10 +1955,63 @@ TEST_F(AlgebraicSimplifierTest, ReversalOfTrivialDimensionsToBitcast) { ASSERT_TRUE(simplifier.Run(&module).ValueOrDie()); HloInstruction* root = computation->root_instruction(); - EXPECT_EQ(root->opcode(), HloOpcode::kParameter); EXPECT_EQ(a, root); EXPECT_TRUE(ShapeUtil::Equal(root->shape(), shape)); } +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. + Shape r1f32 = ShapeUtil::MakeShape(F32, {1}); + HloComputation::Builder builder(TestName() + ".Dot"); + HloInstruction* x = + builder.AddInstruction(HloInstruction::CreateParameter(0, r1f32, "x")); + HloInstruction* y = + builder.AddInstruction(HloInstruction::CreateParameter(1, r1f32, "y")); + builder.AddInstruction( + HloInstruction::CreateBinary(r1f32, HloOpcode::kDot, x, y)); + std::unique_ptr dot_computation(builder.Build()); + + HloComputation::Builder call_builder(TestName() + ".Call"); + HloInstruction* zero = call_builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR1({0.0f}))); + HloInstruction* one = call_builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR1({1.0f}))); + builder.AddInstruction( + HloInstruction::CreateCall(r1f32, {zero, one}, dot_computation.get())); + + auto module = CreateNewModule(); + 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.get()).ValueOrDie()); +} + +// Test that a constant with tuple shape becomes a tuple of constants. +TEST_F(AlgebraicSimplifierTest, ConstantTupleBecomesTupleOfConstants) { + HloComputation::Builder builder(TestName()); + const float constant_scalar = 7.3f; + std::initializer_list constant_vector = {1.1f, 2.0f, 3.3f}; + std::unique_ptr value = + Literal::MakeTuple({Literal::CreateR0(constant_scalar).get(), + Literal::CreateR1(constant_vector).get()}); + builder.AddInstruction(HloInstruction::CreateConstant(std::move(value))); + + auto module = CreateNewModule(); + auto computation = module->AddEntryComputation(builder.Build()); + + AlgebraicSimplifier simplifier(/*is_layout_sensitive=*/false, + non_bitcasting_callback()); + ASSERT_TRUE(simplifier.Run(module.get()).ValueOrDie()); + EXPECT_THAT(computation->root_instruction(), + op::Tuple(op::Constant(), op::Constant())); +} + } // namespace } // namespace xla + +int main(int argc, char** argv) { + return xla::ParseDebugOptionsFlagsAndRunTests(argc, argv); +} diff --git a/tensorflow/compiler/xla/service/allocation_tracker.cc b/tensorflow/compiler/xla/service/allocation_tracker.cc index e59fad4e05252ebd54b3a7cecbdf990127a5264c..ad2fee2d39a8ca183b87212bdeea22c351aaa88a 100644 --- a/tensorflow/compiler/xla/service/allocation_tracker.cc +++ b/tensorflow/compiler/xla/service/allocation_tracker.cc @@ -64,8 +64,9 @@ GlobalDataHandle AllocationTracker::RegisterInternal( auto& allocation = FindOrDie(handle_to_allocation_, handle); int ref_count = allocation->ref_count(); CHECK_GT(ref_count, 0); - VLOG(2) << "ref_count: " << ref_count << " -> " << ref_count + 1; - allocation->increment_ref_count(); + VLOG(2) << "ref_count: " << ref_count << " -> " << + (ref_count + initial_ref_count); + allocation->increment_ref_count(initial_ref_count); } else { handle = next_handle_++; VLOG(2) << "ref_count: " << initial_ref_count; @@ -125,9 +126,7 @@ tensorflow::Status AllocationTracker::DeallocateShape( handle_map.erase(device_memory->opaque()); } - // TODO(b/36256956) Ideally tuple elements could always be distinct buffers. - if (ShapeUtil::IsTuple(shape) && - backend->transfer_manager()->TupleElementsAreDistinctBuffers()) { + if (ShapeUtil::IsTuple(shape)) { // Traverse into tuple recursively deallocating buffers. TF_ASSIGN_OR_RETURN(se::StreamExecutor * executor, backend->stream_executor(device_ordinal)); @@ -172,6 +171,7 @@ StatusOr> AllocationTracker::DeconstructTuple( executor, allocation->device_memory(), allocation->shape())); std::vector element_handles; + element_handles.reserve(element_bases.size()); for (int i = 0; i < element_bases.size(); ++i) { element_handles.push_back(RegisterInternal( allocation->backend(), allocation->device_ordinal(), element_bases[i], diff --git a/tensorflow/compiler/xla/service/allocation_tracker.h b/tensorflow/compiler/xla/service/allocation_tracker.h index e00768001620275d702c2f96a89d981526ea81a7..ebbf35b6fe87bc7322ccb99cfe8f8eed56de06b3 100644 --- a/tensorflow/compiler/xla/service/allocation_tracker.h +++ b/tensorflow/compiler/xla/service/allocation_tracker.h @@ -63,10 +63,10 @@ class Allocation { CHECK_GE(ref_count_, 0); return ref_count_; } - void increment_ref_count() { + void increment_ref_count(int inc) { CHECK_GT(ref_count_, 0); - CHECK_LT(ref_count_, INT_MAX); - ++ref_count_; + CHECK_LE(ref_count_, INT_MAX - inc); + ref_count_ += inc; } void decrement_ref_count() { CHECK_GT(ref_count_, 0); diff --git a/tensorflow/compiler/xla/service/backend.cc b/tensorflow/compiler/xla/service/backend.cc index 5c05417c6dcb887b5352d1270c24a4eae62149e3..9abe30e3f371cc294c36c1dcd743224b11b0c4f5 100644 --- a/tensorflow/compiler/xla/service/backend.cc +++ b/tensorflow/compiler/xla/service/backend.cc @@ -22,7 +22,6 @@ limitations under the License. #define EIGEN_USE_THREADS #include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor" -#include "tensorflow/compiler/xla/legacy_flags/backend_flags.h" #include "tensorflow/compiler/xla/service/compiler.h" #include "tensorflow/compiler/xla/service/platform_util.h" #include "tensorflow/compiler/xla/status_macros.h" @@ -41,13 +40,32 @@ namespace se = ::perftools::gputools; namespace xla { +BackendOptions& BackendOptions::set_platform( + perftools::gputools::Platform* platform) { + platform_ = platform; + return *this; +} + +perftools::gputools::Platform* BackendOptions::platform() const { + return platform_; +} + +BackendOptions& BackendOptions::set_intra_op_parallelism_threads( + int num_threads) { + intra_op_parallelism_threads_ = num_threads; + return *this; +} + +int BackendOptions::intra_op_parallelism_threads() const { + return intra_op_parallelism_threads_; +} + // Define this in .cc file to avoid having to include eigen or forward declare // these types in the header. struct Backend::EigenThreadPoolWrapper { - explicit EigenThreadPoolWrapper() - : pool(new tensorflow::thread::ThreadPool( - tensorflow::Env::Default(), "XLAEigen", - tensorflow::port::NumSchedulableCPUs())), + explicit EigenThreadPoolWrapper(const int num_threads) + : pool(new tensorflow::thread::ThreadPool(tensorflow::Env::Default(), + "XLAEigen", num_threads)), wrapper(new tensorflow::EigenThreadPoolWrapper(pool.get())), device(new Eigen::ThreadPoolDevice(wrapper.get(), wrapper->NumThreads())) {} @@ -58,20 +76,18 @@ struct Backend::EigenThreadPoolWrapper { }; /* static */ StatusOr> Backend::CreateBackend( - perftools::gputools::Platform* platform, int64 replica_count) { - if (replica_count == -1) { - legacy_flags::BackendFlags* flags = legacy_flags::GetBackendFlags(); - replica_count = flags->xla_replicas; - } + const BackendOptions& options) { + perftools::gputools::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 transfer_manager, TransferManager::GetForPlatform(platform)); - std::unique_ptr backend(new Backend( - replica_count, platform, compiler, stream_executors, transfer_manager)); - TF_RETURN_IF_ERROR(backend->PoolStreams(kInitialStreamsToPool, - backend->default_stream_executor())); + TF_ASSIGN_OR_RETURN(auto computation_placer, + ComputationPlacer::GetForPlatform(platform)); + std::unique_ptr backend( + new Backend(platform, compiler, stream_executors, transfer_manager, + computation_placer, options.intra_op_parallelism_threads())); return std::move(backend); } @@ -79,16 +95,9 @@ struct Backend::EigenThreadPoolWrapper { Backend::CreateDefaultBackend() { TF_ASSIGN_OR_RETURN(se::Platform * platform, PlatformUtil::GetDefaultPlatform()); - return CreateBackend(platform); -} - -tensorflow::Status Backend::PoolStreams(int n, se::StreamExecutor* executor) { - std::vector primed; - for (int i = 0; i < n; ++i) { - TF_ASSIGN_OR_RETURN(auto stream, BorrowStream(executor)); - primed.emplace_back(std::move(stream)); - } - return tensorflow::Status::OK(); + BackendOptions backend_options; + backend_options.set_platform(platform); + return CreateBackend(backend_options); } StatusOr Backend::BorrowStream(int device_ordinal) { @@ -98,6 +107,7 @@ 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), @@ -111,40 +121,35 @@ StatusOr Backend::BorrowStream( } Backend::Backend( - int64 replica_count, perftools::gputools::Platform* platform, - Compiler* compiler, + perftools::gputools::Platform* platform, Compiler* compiler, tensorflow::gtl::ArraySlice stream_executors, - TransferManager* transfer_manager) + TransferManager* transfer_manager, ComputationPlacer* computation_placer, + int intra_op_parallelism_threads) : platform_(platform), compiler_(compiler), transfer_manager_(transfer_manager), - replica_count_(replica_count) { + computation_placer_(computation_placer) { // The given set of stream executors set may include invalid executors. for (se::StreamExecutor* exec : stream_executors) { if (exec != nullptr) { stream_executors_.push_back(exec); } } - CHECK_GE(replica_count, 1) << "Must request at least 1 replica."; - // Create a memory allocator for the valid stream executors. memory_allocator_ = MakeUnique(platform, stream_executors); - - // First check that there are some non-null stream executors to avoid issuing - // an error mentioning replicas in the common case of requesting just 1 - // replica, which means no replication. CHECK(!stream_executors_.empty()) << "Service found no devices for backend " << platform_->Name() << '.'; - CHECK_GE(stream_executors_.size(), replica_count) - << "Requested more replicas than there are devices for backend " - << platform_->Name() << '.'; if (platform->id() == se::host::kHostPlatformId) { inter_op_thread_pool_.reset(new tensorflow::thread::ThreadPool( tensorflow::Env::Default(), "xla_inter_op", tensorflow::port::NumSchedulableCPUs())); - intra_op_thread_pool_wrapper_.reset(new EigenThreadPoolWrapper()); + const int num_threads = intra_op_parallelism_threads > 0 + ? intra_op_parallelism_threads + : tensorflow::port::NumSchedulableCPUs(); + intra_op_thread_pool_wrapper_.reset( + new EigenThreadPoolWrapper(num_threads)); } } @@ -154,46 +159,25 @@ int Backend::default_device_ordinal() const { return default_stream_executor()->device_ordinal(); } -StatusOr> Backend::Replicas( - int device_ordinal) const { - if (stream_executors_[device_ordinal] == nullptr) { - return InvalidArgument("device %s not supported by XLA service", - device_name(device_ordinal).c_str()); - } - - // Find replica_count_ stream executors starting from the given device - // ordinal. - std::vector replicas; - for (se::StreamExecutor* exec : stream_executors_) { - CHECK(exec != nullptr); - if (exec->device_ordinal() >= device_ordinal) { - replicas.push_back(exec); - if (replicas.size() >= replica_count_) { - return replicas; - } - } - } - - return InvalidArgument( - "Not enough devices for replicas for the device ordinal %d", - device_ordinal); -} - -std::vector Backend::Replicas() const { - CHECK_GE(stream_executors_.size(), replica_count_); - return Replicas(default_device_ordinal()).ValueOrDie(); -} - tensorflow::thread::ThreadPool* Backend::inter_op_thread_pool() const { return inter_op_thread_pool_.get(); } const Eigen::ThreadPoolDevice* Backend::eigen_intra_op_thread_pool_device() const { - if (intra_op_thread_pool_wrapper_ == nullptr) return nullptr; + if (intra_op_thread_pool_wrapper_ == nullptr) { + return nullptr; + } return intra_op_thread_pool_wrapper_->device.get(); } +tensorflow::thread::ThreadPool* Backend::eigen_intra_op_thread_pool() const { + if (intra_op_thread_pool_wrapper_ == nullptr) { + return nullptr; + } + return intra_op_thread_pool_wrapper_->pool.get(); +} + StatusOr Backend::stream_executor( int device_ordinal) const { if (device_ordinal < 0 || diff --git a/tensorflow/compiler/xla/service/backend.h b/tensorflow/compiler/xla/service/backend.h index 9f6829b7d937cec6a67d4016a40506de5df8572d..b5ca483b7274d20c31e932d748b6a4c9dea926f9 100644 --- a/tensorflow/compiler/xla/service/backend.h +++ b/tensorflow/compiler/xla/service/backend.h @@ -22,6 +22,7 @@ limitations under the License. #include #include "tensorflow/compiler/xla/service/compiler.h" +#include "tensorflow/compiler/xla/service/computation_placer.h" #include "tensorflow/compiler/xla/service/device_memory_allocator.h" #include "tensorflow/compiler/xla/service/pool.h" #include "tensorflow/compiler/xla/service/transfer_manager.h" @@ -39,6 +40,24 @@ struct ThreadPoolDevice; namespace xla { +// Options to configure the backend when it is created. +class BackendOptions { + public: + // Set the platform backing the backend, or nullptr for the default platform. + BackendOptions& set_platform(perftools::gputools::Platform* platform); + perftools::gputools::Platform* platform() const; + + // Sets the thread pool size for parallel execution of an individual operator. + // The default value of -1 will result in initializing the thread pool with + // the number of threads equal to the number of cores in the system. + BackendOptions& set_intra_op_parallelism_threads(int num_threads); + int intra_op_parallelism_threads() const; + + private: + perftools::gputools::Platform* platform_ = nullptr; + int intra_op_parallelism_threads_ = -1; +}; + // Class which encapsulates an XLA backend. It includes everything necessary // to compile and execute computations on a particular platform. // @@ -49,13 +68,9 @@ class Backend { public: using StreamPtr = Pool::SmartPtr; - // The number of streams we create for the pool at initialization time. - static constexpr int kInitialStreamsToPool = 8; - - // Creates a new backend for the given platform with the given number of - // replicas. A value of -1 means to use the flag value. + // Creates a new backend. static StatusOr> CreateBackend( - perftools::gputools::Platform* platform, int64 replica_count = -1); + const BackendOptions& options); // Creates a backend for the default platform. The default platform is defined // in PlatformUtil. @@ -70,6 +85,7 @@ class Backend { return memory_allocator_.get(); } TransferManager* transfer_manager() const { return transfer_manager_; } + ComputationPlacer* computation_placer() const { return computation_placer_; } // Returns the number of devices of the platform type which are visible. Not // all of these devices may be usable by XLA. @@ -85,34 +101,18 @@ class Backend { return stream_executors_; } - // Returns the replicas for the default stream executor. - // - // When the number of replicas is R, the first R stream executors are assigned - // to the replicas of the default stream executor. - std::vector Replicas() const; - - // Returns the replicas for the given device_ordinal. The given device ordinal - // is considered to be the first device ordinal among the replicas. Returns an - // error status if the stream executor for the given given device ordinal does - // not exist or if there are not enough stream executors for the replicas. - StatusOr> Replicas( - int device_ordinal) const; - - // Return the stream executor for the given device ordinal. + // Returns the stream executor for the given device ordinal. StatusOr stream_executor( int device_ordinal) const; - // Return the stream executor for the default device ordinal. + // Returns the stream executor for the default device ordinal. This stream + // executor can only be used when the number of computations is 1 (replication + // can be > 1). perftools::gputools::StreamExecutor* default_stream_executor() const { CHECK(!stream_executors_.empty()); return stream_executors_[0]; } - // Primes the internal pool of streams for BorrowStream with n initialized - // stream instances. - tensorflow::Status PoolStreams(int n, - perftools::gputools::StreamExecutor* executor); - // Borrows a stream for use by the caller, either by grabbing it from an // internal pool, or by constructing/initializating it, and returns the result // to the caller. @@ -150,32 +150,36 @@ class Backend { // For the host platform, returns the configured eigen threadpool device to be // used for scheduling work. For other platforms, returns NULL. const Eigen::ThreadPoolDevice* eigen_intra_op_thread_pool_device() const; + tensorflow::thread::ThreadPool* eigen_intra_op_thread_pool() const; // Resets the devices associated with this backend. Status ResetDevices(); private: struct EigenThreadPoolWrapper; - Backend(int64 replica_count, perftools::gputools::Platform* platform, - Compiler* compiler, + Backend(perftools::gputools::Platform* platform, Compiler* compiler, tensorflow::gtl::ArraySlice stream_executors, - TransferManager* transfer_manager); + TransferManager* transfer_manager, + ComputationPlacer* computation_placer, + int intra_op_parallelism_threads); Backend(const Backend&) = delete; Backend& operator=(const Backend&) = delete; perftools::gputools::Platform* platform_; Compiler* compiler_; TransferManager* transfer_manager_; - int64 replica_count_ = -1; + ComputationPlacer* computation_placer_; // Vector of stream executors. stream_executors_[0] is the default executor. std::vector stream_executors_; + tensorflow::mutex mu_; + // Mapping from stream executor to stream pools, used by `BorrowStream` above. std::map> - stream_pools_; + stream_pools_ GUARDED_BY(mu_); // The default memory allocator to use. std::unique_ptr memory_allocator_; diff --git a/tensorflow/compiler/xla/service/batchnorm_rewriter.cc b/tensorflow/compiler/xla/service/batchnorm_rewriter.cc new file mode 100644 index 0000000000000000000000000000000000000000..41d32d0c8b1cc31522f2c8012fb5350816cadbec --- /dev/null +++ b/tensorflow/compiler/xla/service/batchnorm_rewriter.cc @@ -0,0 +1,554 @@ +/* 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/batchnorm_rewriter.h" + +#include +#include +#include +#include +#include +#include +#include + +#include "tensorflow/compiler/xla/layout_util.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/service/hlo_query.h" +#include "tensorflow/compiler/xla/service/shape_inference.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/errors.h" +#include "tensorflow/core/lib/core/status.h" +#include "tensorflow/core/lib/gtl/array_slice.h" +#include "tensorflow/core/platform/logging.h" +#include "tensorflow/core/platform/types.h" + +namespace xla { + +// BatchNormRewriterVisitor traverses the HLO computation and rewrites BatchNorm +// operations into smaller operations. +class BatchNormRewriterVisitor : public DfsHloVisitorWithDefault { + public: + // Default visitor action is to do nothing and return OK. + Status DefaultAction(HloInstruction* /*hlo_instruction*/) override { + return Status::OK(); + } + + Status HandleBatchNormTraining(HloInstruction* batch_norm) override; + + Status HandleBatchNormInference(HloInstruction* batch_norm) override; + + Status HandleBatchNormGrad(HloInstruction* batch_norm) override; + + // Runs the visitor on a computation. + static bool Run(HloComputation* computation, bool rewrite_training_op, + bool rewrite_inference_op, bool rewrite_grad_op, + bool use_fusion); + + // Returns whether any batch norm ops were rewritten. + const bool changed() const { return changed_; } + + ~BatchNormRewriterVisitor() override = default; + + private: + explicit BatchNormRewriterVisitor(HloComputation* computation, + bool rewrite_training_op, + bool rewrite_inference_op, + bool rewrite_grad_op, bool use_fusion) + : computation_(computation), + rewrite_training_op_(rewrite_training_op), + rewrite_inference_op_(rewrite_inference_op), + rewrite_grad_op_(rewrite_grad_op), + use_fusion_(use_fusion) {} + + HloComputation* GetScalarBinaryComputation(PrimitiveType primitive_type, + HloOpcode opcode) { + HloComputation::Builder b("scalar computation"); + auto scalar_lhs = b.AddInstruction(HloInstruction::CreateParameter( + 0, ShapeUtil::MakeShape(F32, {}), "scalar lhs")); + auto scalar_rhs = b.AddInstruction(HloInstruction::CreateParameter( + 1, ShapeUtil::MakeShape(F32, {}), "scalar rhs")); + auto scalar_op = b.AddInstruction( + HloInstruction::CreateBinary(ShapeUtil::MakeShape(primitive_type, {}), + opcode, scalar_lhs, scalar_rhs)); + return computation_->parent()->AddEmbeddedComputation(b.Build(scalar_op)); + } + + // Current HloComputation instance the BatchNormRewriter is + // traversing. + HloComputation* computation_; + + bool rewrite_training_op_; + bool rewrite_inference_op_; + bool rewrite_grad_op_; + bool use_fusion_; + + // Whether rewrite has occurred. + bool changed_ = false; + + // Replaces the existing HLO instruction old_instruction, with + // new_instruction, and marks the optimizer status as changed. + // Returns the Status representing the result of the replace operation. + Status ReplaceWithNewInstruction( + HloInstruction* old_instruction, + std::unique_ptr new_instruction) { + TF_RETURN_IF_ERROR(computation_->ReplaceWithNewInstruction( + old_instruction, std::move(new_instruction))); + changed_ = true; + return Status::OK(); + } + + // Replaces the existing HLO instruction old_instruction, with + // new_instruction, and marks the optimizer status as changed. + // Returns the Status representing the result of the replace operation. + Status ReplaceInstruction(HloInstruction* old_instruction, + HloInstruction* new_instruction) { + TF_RETURN_IF_ERROR( + computation_->ReplaceInstruction(old_instruction, new_instruction)); + changed_ = true; + return Status::OK(); + } +}; + +bool BatchNormRewriterVisitor::Run(HloComputation* computation, + bool rewrite_training_op, + bool rewrite_inference_op, + bool rewrite_grad_op, bool use_fusion) { + BatchNormRewriterVisitor visitor( + computation, + /*rewrite_training_op=*/rewrite_training_op, + /*rewrite_inference_op=*/rewrite_inference_op, + /*rewrite_grad_op=*/rewrite_grad_op, + /*use_fusion=*/use_fusion); + TF_CHECK_OK(computation->Accept(&visitor)); + return visitor.changed_; +} + +Status BatchNormRewriterVisitor::HandleBatchNormTraining( + HloInstruction* batch_norm) { + if (!rewrite_training_op_) { + return Status::OK(); + } + // Expand batch norm training into smaller HLO ops. + HloInstruction* operand = batch_norm->mutable_operand(0); + const Shape operand_shape = operand->shape(); + 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); + auto elements_per_feature = + computation_->AddInstruction(HloInstruction::CreateConstant( + Literal::CreateR0(size_in_elements / feature_count))); + + HloInstruction* scale = batch_norm->mutable_operand(1); + HloInstruction* offset = batch_norm->mutable_operand(2); + const Shape feature_shape = scale->shape(); + + auto zero = computation_->AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(0.0f))); + + auto epsilon = computation_->AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(batch_norm->epsilon()))); + + std::vector dimensions_without_feature; + + for (int64 i = 0; i < ShapeUtil::Rank(operand_shape); ++i) { + if (i != feature_index) { + dimensions_without_feature.push_back(i); + } + } + + auto scale_broadcasted = computation_->AddInstruction( + HloInstruction::CreateBroadcast(operand_shape, scale, {feature_index})); + + auto offset_broadcasted = computation_->AddInstruction( + HloInstruction::CreateBroadcast(operand_shape, offset, {feature_index})); + + HloComputation* add_reduce_computation = + GetScalarBinaryComputation(F32, HloOpcode::kAdd); + + // X^2. + auto operand_squared = + computation_->AddInstruction(HloInstruction::CreateBinary( + operand_shape, HloOpcode::kMultiply, operand, operand)); + // Sum[X]. + auto sum = computation_->AddInstruction(HloInstruction::CreateReduce( + feature_shape, operand, zero, dimensions_without_feature, + add_reduce_computation)); + + // Sum[X^2]. + auto squared_sum = computation_->AddInstruction(HloInstruction::CreateReduce( + feature_shape, operand_squared, zero, dimensions_without_feature, + add_reduce_computation)); + + // Fuse two parallel reduces together to improve performance. + if (use_fusion_) { + auto tuple = computation_->AddInstruction( + HloInstruction::CreateTuple({sum, squared_sum})); + + auto fused = computation_->CreateFusionInstruction( + {tuple, sum, squared_sum, operand_squared}, + HloInstruction::FusionKind::kInput); + + sum = computation_->AddInstruction( + HloInstruction::CreateGetTupleElement(feature_shape, fused, 0)); + + squared_sum = computation_->AddInstruction( + HloInstruction::CreateGetTupleElement(feature_shape, fused, 1)); + } + + // E[X]. + auto mean = computation_->AddInstruction(HloInstruction::CreateBinary( + feature_shape, HloOpcode::kDivide, sum, elements_per_feature)); + + auto mean_broadcasted = computation_->AddInstruction( + HloInstruction::CreateBroadcast(operand_shape, mean, {feature_index})); + + // E[X^2]. + auto square_mean = computation_->AddInstruction(HloInstruction::CreateBinary( + feature_shape, HloOpcode::kDivide, squared_sum, elements_per_feature)); + + // E^2[X]. + auto mean_square = computation_->AddInstruction(HloInstruction::CreateBinary( + feature_shape, HloOpcode::kMultiply, mean, mean)); + + // Var[X]. + auto var = computation_->AddInstruction(HloInstruction::CreateBinary( + feature_shape, HloOpcode::kSubtract, square_mean, mean_square)); + + auto var_broadcasted = computation_->AddInstruction( + HloInstruction::CreateBroadcast(operand_shape, var, {feature_index})); + + // Var[X] + epsilon. + auto var_add_epsilon = + computation_->AddInstruction(HloInstruction::CreateBinary( + operand_shape, HloOpcode::kAdd, var_broadcasted, epsilon)); + + auto neg_half = computation_->AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(-0.5f))); + + // 1 / Sqrt[Var[X] + epsilon]. + auto rsqrt_var_add_epsilon = + computation_->AddInstruction(HloInstruction::CreateBinary( + operand_shape, HloOpcode::kPower, var_add_epsilon, neg_half)); + + // X - E[X]. + auto operand_minus_mean = + computation_->AddInstruction(HloInstruction::CreateBinary( + operand_shape, HloOpcode::kSubtract, operand, mean_broadcasted)); + + // (X - E[X]) / Sqrt[Var[X] + epsilon]. + auto normalized = computation_->AddInstruction( + HloInstruction::CreateBinary(operand_shape, HloOpcode::kMultiply, + operand_minus_mean, rsqrt_var_add_epsilon)); + + // (X - E[X]) / Sqrt[Var[X] + epsilon] * scale. + auto scaled_normalized = + computation_->AddInstruction(HloInstruction::CreateBinary( + operand_shape, HloOpcode::kMultiply, normalized, scale_broadcasted)); + + // (X - E[X]) / Sqrt[Var[X] + epsilon] * scale + offset. + auto shifted_normalized = computation_->AddInstruction( + HloInstruction::CreateBinary(operand_shape, HloOpcode::kAdd, + scaled_normalized, offset_broadcasted)); + + TF_CHECK_OK(ReplaceWithNewInstruction( + batch_norm, + HloInstruction::CreateTuple({shifted_normalized, mean, var}))); + return Status::OK(); +} + +Status BatchNormRewriterVisitor::HandleBatchNormInference( + HloInstruction* batch_norm) { + if (!rewrite_inference_op_) { + return Status::OK(); + } + // Expand batch norm inference into smaller HLO ops. + HloInstruction* operand = batch_norm->mutable_operand(0); + const Shape operand_shape = operand->shape(); + int64 feature_index = batch_norm->feature_index(); + + HloInstruction* scale = batch_norm->mutable_operand(1); + HloInstruction* offset = batch_norm->mutable_operand(2); + HloInstruction* mean = batch_norm->mutable_operand(3); + HloInstruction* var = batch_norm->mutable_operand(4); + const Shape feature_shape = scale->shape(); + + auto epsilon = computation_->AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(batch_norm->epsilon()))); + + std::vector dimensions_without_feature; + + for (int64 i = 0; i < ShapeUtil::Rank(operand_shape); ++i) { + if (i != feature_index) { + dimensions_without_feature.push_back(i); + } + } + + auto scale_broadcasted = computation_->AddInstruction( + HloInstruction::CreateBroadcast(operand_shape, scale, {feature_index})); + + auto offset_broadcasted = computation_->AddInstruction( + HloInstruction::CreateBroadcast(operand_shape, offset, {feature_index})); + + auto mean_broadcasted = computation_->AddInstruction( + HloInstruction::CreateBroadcast(operand_shape, mean, {feature_index})); + + auto var_broadcasted = computation_->AddInstruction( + HloInstruction::CreateBroadcast(operand_shape, var, {feature_index})); + + // Var[X] + epsilon. + auto var_add_epsilon = + computation_->AddInstruction(HloInstruction::CreateBinary( + operand_shape, HloOpcode::kAdd, var_broadcasted, epsilon)); + + auto neg_half = computation_->AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(-0.5f))); + + // 1 / Sqrt[Var[X] + epsilon]. + auto rsqrt_var_add_epsilon = + computation_->AddInstruction(HloInstruction::CreateBinary( + operand_shape, HloOpcode::kPower, var_add_epsilon, neg_half)); + + // X - E[X]. + auto operand_minus_mean = + computation_->AddInstruction(HloInstruction::CreateBinary( + operand_shape, HloOpcode::kSubtract, operand, mean_broadcasted)); + + // (X - E[X]) / Sqrt[Var[X] + epsilon]. + auto normalized = computation_->AddInstruction( + HloInstruction::CreateBinary(operand_shape, HloOpcode::kMultiply, + operand_minus_mean, rsqrt_var_add_epsilon)); + + // (X - E[X]) / Sqrt[Var[X] + epsilon] * scale. + auto scaled_normalized = + computation_->AddInstruction(HloInstruction::CreateBinary( + operand_shape, HloOpcode::kMultiply, normalized, scale_broadcasted)); + + // (X - E[X]) / Sqrt[Var[X] + epsilon] * scale + offset. + auto shifted_normalized = HloInstruction::CreateBinary( + operand_shape, HloOpcode::kAdd, scaled_normalized, offset_broadcasted); + + TF_CHECK_OK( + ReplaceWithNewInstruction(batch_norm, std::move(shifted_normalized))); + return Status::OK(); +} + +Status BatchNormRewriterVisitor::HandleBatchNormGrad( + HloInstruction* batch_norm) { + // Use the following formulas to calculate gradients: + // scale_grad = + // sum(output_grad * (activation - mean(activation))) * rsqrt(var + epsilon) + // + // offset_grad = + // sum(output_grad) + // + // activation_grad = + // 1/N * scale * rsqrt(var + epsilon) * + // (N * output_grad - sum(output_grad) - (activation - mean(activation)) * + // sum(output_grad * (activation - mean(activation))) / (variance + + // epsilon)) + if (!rewrite_grad_op_) { + return Status::OK(); + } + + HloInstruction* activation = batch_norm->mutable_operand(0); + const Shape activation_shape = activation->shape(); + HloInstruction* scale = batch_norm->mutable_operand(1); + const Shape feature_shape = scale->shape(); + HloInstruction* mean = batch_norm->mutable_operand(2); + HloInstruction* variance = batch_norm->mutable_operand(3); + HloInstruction* grad_output = batch_norm->mutable_operand(4); + + 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); + auto elements_per_feature = + computation_->AddInstruction(HloInstruction::CreateConstant( + Literal::CreateR0(size_in_elements / feature_count))); + + auto zero = computation_->AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(0.0f))); + + auto neg_half = computation_->AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(-0.5f))); + + auto epsilon = computation_->AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(batch_norm->epsilon()))); + + std::vector dimensions_without_feature; + + for (int64 i = 0; i < ShapeUtil::Rank(activation_shape); ++i) { + if (i != feature_index) { + dimensions_without_feature.push_back(i); + } + } + + auto scale_broadcasted = + computation_->AddInstruction(HloInstruction::CreateBroadcast( + activation_shape, scale, {feature_index})); + auto variance_broadcasted = + computation_->AddInstruction(HloInstruction::CreateBroadcast( + activation_shape, variance, {feature_index})); + + // E[X]. + auto mean_broadcasted = computation_->AddInstruction( + HloInstruction::CreateBroadcast(activation_shape, mean, {feature_index})); + + // rsqrt[Var[X] + epsilon]. + auto rsqrt_var_add_epsilon_broadcasted = + computation_->AddInstruction(HloInstruction::CreateBinary( + activation_shape, HloOpcode::kPower, + computation_->AddInstruction( + HloInstruction::CreateBinary(activation_shape, HloOpcode::kAdd, + variance_broadcasted, epsilon)), + neg_half)); + + auto rsqrt_var_add_epsilon = + computation_->AddInstruction(HloInstruction::CreateBinary( + feature_shape, HloOpcode::kPower, + computation_->AddInstruction(HloInstruction::CreateBinary( + feature_shape, HloOpcode::kAdd, variance, epsilon)), + neg_half)); + + // X - E[X]. + auto activation_minus_mean = computation_->AddInstruction( + HloInstruction::CreateBinary(activation_shape, HloOpcode::kSubtract, + activation, mean_broadcasted)); + + // Grad[Y] * (X - E[X]). + auto grad_output_times_activiation_minus_mean = computation_->AddInstruction( + HloInstruction::CreateBinary(activation_shape, HloOpcode::kMultiply, + grad_output, activation_minus_mean)); + + HloComputation* add_reduce_computation = + GetScalarBinaryComputation(F32, HloOpcode::kAdd); + + // sum(Grad[Y] * (X - E[X])). + auto sum_grad_output_times_activiation_minus_mean = + computation_->AddInstruction(HloInstruction::CreateReduce( + feature_shape, grad_output_times_activiation_minus_mean, zero, + dimensions_without_feature, add_reduce_computation)); + + // Grad[beta] = Sum(Grad[Y]). + auto grad_beta = computation_->AddInstruction(HloInstruction::CreateReduce( + feature_shape, grad_output, zero, dimensions_without_feature, + add_reduce_computation)); + + if (use_fusion_) { + auto tuple = computation_->AddInstruction(HloInstruction::CreateTuple( + {sum_grad_output_times_activiation_minus_mean, grad_beta})); + + auto fused = computation_->CreateFusionInstruction( + {tuple, sum_grad_output_times_activiation_minus_mean, grad_beta}, + HloInstruction::FusionKind::kInput); + + sum_grad_output_times_activiation_minus_mean = computation_->AddInstruction( + HloInstruction::CreateGetTupleElement(feature_shape, fused, 0)); + + grad_beta = computation_->AddInstruction( + HloInstruction::CreateGetTupleElement(feature_shape, fused, 1)); + } + + // Grad[scale] = Sum(Grad[Y] * (X - E[X]) * rsqrt[Var[X] + epsilon]). + auto grad_scale = computation_->AddInstruction(HloInstruction::CreateBinary( + feature_shape, HloOpcode::kMultiply, + sum_grad_output_times_activiation_minus_mean, rsqrt_var_add_epsilon)); + + // I2 = Sum(Grad[Y]) + auto I2 = computation_->AddInstruction(HloInstruction::CreateBroadcast( + activation_shape, grad_beta, {feature_index})); + + // I3 = Sum(Grad[Y] * (X - E[X])) + auto I3 = computation_->AddInstruction(HloInstruction::CreateBroadcast( + activation_shape, sum_grad_output_times_activiation_minus_mean, + {feature_index})); + + // I4 = (X - E[X]) * I3 + auto I4 = computation_->AddInstruction(HloInstruction::CreateBinary( + activation_shape, HloOpcode::kMultiply, I3, activation_minus_mean)); + + // I5 = I4 / (Var[X] + epsilon) + auto I5 = computation_->AddInstruction(HloInstruction::CreateBinary( + activation_shape, HloOpcode::kDivide, I4, + computation_->AddInstruction(HloInstruction::CreateBinary( + activation_shape, HloOpcode::kAdd, variance_broadcasted, epsilon)))); + + // scale * rsqrt[Var[X] + epsilon] * 1/N + auto scale_times_rsqrt_var_add_epsilon = + computation_->AddInstruction(HloInstruction::CreateBinary( + activation_shape, HloOpcode::kMultiply, scale_broadcasted, + rsqrt_var_add_epsilon_broadcasted)); + + scale_times_rsqrt_var_add_epsilon = + computation_->AddInstruction(HloInstruction::CreateBinary( + activation_shape, HloOpcode::kDivide, + scale_times_rsqrt_var_add_epsilon, elements_per_feature)); + + auto I1 = computation_->AddInstruction( + HloInstruction::CreateBinary(activation_shape, HloOpcode::kMultiply, + grad_output, elements_per_feature)); + + // I6 = I1 - I2 - I5 + auto I6 = computation_->AddInstruction(HloInstruction::CreateBinary( + activation_shape, HloOpcode::kSubtract, + computation_->AddInstruction(HloInstruction::CreateBinary( + activation_shape, HloOpcode::kSubtract, I1, I2)), + I5)); + + // Grad[X] = scale * rsqrt[Var[X] + epsilon] * 1/N * I6. + auto grad_activation = computation_->AddInstruction( + HloInstruction::CreateBinary(activation_shape, HloOpcode::kMultiply, + scale_times_rsqrt_var_add_epsilon, I6)); + + TF_CHECK_OK(ReplaceWithNewInstruction( + batch_norm, + HloInstruction::CreateTuple({grad_activation, grad_scale, grad_beta}))); + + return Status::OK(); +} + +StatusOr BatchNormRewriter::Run(HloModule* module) { + XLA_VLOG_LINES(2, "BatchNormRewriter::Run(), before:\n" + module->ToString()); + bool changed = false; + // Make a copy of the computations because we may add computations to the + // module, invalidating iteration. + std::vector computations; + for (auto& comp : module->computations()) { + if (comp->IsFusionComputation()) { + continue; + } + computations.push_back(comp.get()); + } + for (auto& comp : computations) { + if (BatchNormRewriterVisitor::Run(comp, rewrite_training_op_, + rewrite_inference_op_, rewrite_grad_op_, + use_fusion_)) { + changed = true; + } + } + XLA_VLOG_LINES(2, "BatchNormRewriter::Run(), after:\n" + module->ToString()); + return changed; +} + +} // namespace xla diff --git a/tensorflow/compiler/xla/service/batchnorm_rewriter.h b/tensorflow/compiler/xla/service/batchnorm_rewriter.h new file mode 100644 index 0000000000000000000000000000000000000000..f601741d964376058a2bafade311ede4c8567fd2 --- /dev/null +++ b/tensorflow/compiler/xla/service/batchnorm_rewriter.h @@ -0,0 +1,55 @@ +/* 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_BATCHNORM_REWRITER_H_ +#define TENSORFLOW_COMPILER_XLA_SERVICE_BATCHNORM_REWRITER_H_ + +#include + +#include "tensorflow/compiler/xla/service/hlo_module.h" +#include "tensorflow/compiler/xla/service/hlo_pass_interface.h" + +namespace xla { + +// A pass which rewrites batch norm operations into more operations. Breaking a +// big operation into smaller operations helps leverage our generic fusion +// logic. +class BatchNormRewriter : public HloPassInterface { + public: + // When use_fusion is set, a multi-output fusion node is created. + BatchNormRewriter(bool rewrite_training_op = false, + bool rewrite_inference_op = false, + bool rewrite_grad_op = false, bool use_fusion = true) + : rewrite_training_op_(rewrite_training_op), + rewrite_inference_op_(rewrite_inference_op), + rewrite_grad_op_(rewrite_grad_op), + use_fusion_(use_fusion) {} + ~BatchNormRewriter() = default; + tensorflow::StringPiece name() const override { return "batchnorm_rewriter"; } + + // Run operation expander on the given computation. Returns whether the + // computation was changed. + StatusOr Run(HloModule* module) override; + + private: + bool rewrite_training_op_; + bool rewrite_inference_op_; + bool rewrite_grad_op_; + bool use_fusion_; +}; + +} // namespace xla + +#endif // TENSORFLOW_COMPILER_XLA_SERVICE_BATCHNORM_REWRITER_H_ diff --git a/tensorflow/compiler/xla/service/batchnorm_rewriter_test.cc b/tensorflow/compiler/xla/service/batchnorm_rewriter_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..07775623e75f379696624d1fe233eaafc80a8994 --- /dev/null +++ b/tensorflow/compiler/xla/service/batchnorm_rewriter_test.cc @@ -0,0 +1,122 @@ +/* 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/batchnorm_rewriter.h" + +#include +#include + +#include "tensorflow/compiler/xla/layout_util.h" +#include "tensorflow/compiler/xla/literal_util.h" +#include "tensorflow/compiler/xla/ptr_util.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_pass_fix.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/types.h" +#include "tensorflow/compiler/xla/xla_data.pb.h" +#include "tensorflow/core/lib/strings/str_util.h" + +namespace xla { +namespace { + +using BatchNormRewriterTest = HloTestBase; + +// Test that we expand BatchNormTraining. +TEST_F(BatchNormRewriterTest, BatchNormTraining) { + Shape input_shape = ShapeUtil::MakeShape(F32, {2, 2, 2, 2}); + Shape scale_shape = ShapeUtil::MakeShape(F32, {2}); + Shape offset_shape = ShapeUtil::MakeShape(F32, {2}); + + HloComputation::Builder builder(TestName()); + HloInstruction* param0 = builder.AddInstruction( + HloInstruction::CreateParameter(0, input_shape, "activiation")); + + HloInstruction* param1 = builder.AddInstruction( + HloInstruction::CreateParameter(1, scale_shape, "scale")); + + HloInstruction* param2 = builder.AddInstruction( + HloInstruction::CreateParameter(2, offset_shape, "offset")); + + builder.AddInstruction(HloInstruction::CreateBatchNormTraining( + ShapeUtil::MakeTupleShape({input_shape, scale_shape, offset_shape}), + param0, param1, param2, + /*epsilon=*/0.001, /*feature_index=*/3)); + + auto module = CreateNewModule(); + auto computation = module->AddEntryComputation(builder.Build()); + HloInstruction* root = computation->root_instruction(); + EXPECT_EQ(root->opcode(), HloOpcode::kBatchNormTraining); + BatchNormRewriter rewriter(/*rewrite_training_op=*/true, + /*rewrite_inference_op=*/true, + /*rewrite_grad_op=*/true); + ASSERT_TRUE(rewriter.Run(module.get()).ValueOrDie()); + root = computation->root_instruction(); + // Make sure this operation is expanded. + EXPECT_EQ(root->opcode(), HloOpcode::kTuple); +} + +// Test that we expand BatchNormGrad. +TEST_F(BatchNormRewriterTest, BatchNormGrad) { + Shape input_shape = ShapeUtil::MakeShape(F32, {2, 2, 2, 2}); + Shape scale_shape = ShapeUtil::MakeShape(F32, {2}); + Shape mean_shape = ShapeUtil::MakeShape(F32, {2}); + Shape var_shape = ShapeUtil::MakeShape(F32, {2}); + Shape grad_output_shape = ShapeUtil::MakeShape(F32, {2, 2, 2, 2}); + + HloComputation::Builder builder(TestName()); + HloInstruction* param0 = builder.AddInstruction( + HloInstruction::CreateParameter(0, input_shape, "activation")); + + HloInstruction* param1 = builder.AddInstruction( + HloInstruction::CreateParameter(1, scale_shape, "scale")); + + HloInstruction* param2 = builder.AddInstruction( + HloInstruction::CreateParameter(2, mean_shape, "mean")); + + HloInstruction* param3 = builder.AddInstruction( + HloInstruction::CreateParameter(3, var_shape, "var")); + + HloInstruction* param4 = builder.AddInstruction( + HloInstruction::CreateParameter(4, grad_output_shape, "grad_output")); + + builder.AddInstruction(HloInstruction::CreateBatchNormGrad( + ShapeUtil::MakeTupleShape({input_shape, scale_shape, mean_shape}), param0, + param1, param2, param3, param4, + /*epsilon=*/0.001, /*feature_index=*/3)); + + auto module = CreateNewModule(); + auto computation = module->AddEntryComputation(builder.Build()); + HloInstruction* root = computation->root_instruction(); + EXPECT_EQ(root->opcode(), HloOpcode::kBatchNormGrad); + BatchNormRewriter rewriter(/*rewrite_training_op=*/true, + /*rewrite_inference_op=*/true, + /*rewrite_grad_op=*/true); + ASSERT_TRUE(rewriter.Run(module.get()).ValueOrDie()); + root = computation->root_instruction(); + // Make sure this operation is expanded. + EXPECT_EQ(root->opcode(), HloOpcode::kTuple); +} + +} // namespace +} // namespace xla + +int main(int argc, char** argv) { + return xla::ParseDebugOptionsFlagsAndRunTests(argc, argv); +} diff --git a/tensorflow/compiler/xla/service/buffer_assignment.cc b/tensorflow/compiler/xla/service/buffer_assignment.cc index e2b550fc022610c72aa312281727c9c2aea66388..6bc0ca4f827b44c78336100b5380ac4c86e8df01 100644 --- a/tensorflow/compiler/xla/service/buffer_assignment.cc +++ b/tensorflow/compiler/xla/service/buffer_assignment.cc @@ -22,12 +22,12 @@ limitations under the License. #include #include -#include "tensorflow/compiler/xla/legacy_flags/buffer_assignment_flags.h" #include "tensorflow/compiler/xla/map_util.h" #include "tensorflow/compiler/xla/ptr_util.h" #include "tensorflow/compiler/xla/service/heap_simulator.h" +#include "tensorflow/compiler/xla/service/hlo.pb.h" #include "tensorflow/compiler/xla/service/hlo_opcode.h" -#include "tensorflow/compiler/xla/service/tuple_points_to_analysis.h" +#include "tensorflow/compiler/xla/service/hlo_scheduling.h" #include "tensorflow/compiler/xla/shape_util.h" #include "tensorflow/compiler/xla/status_macros.h" #include "tensorflow/compiler/xla/types.h" @@ -41,6 +41,8 @@ limitations under the License. namespace xla { +using ::tensorflow::gtl::FlatMap; +using ::tensorflow::gtl::FlatSet; using ::tensorflow::strings::Appendf; using ::tensorflow::strings::HumanReadableNumBytes; @@ -64,6 +66,7 @@ BufferAllocation::Slice BufferAllocation::GetSlice( void BufferAllocation::AddAssignment(const LogicalBuffer& buffer, int64 offset, int64 size) { + VLOG(4) << "Trying to add " << buffer << " to " << this; CHECK(assigned_buffers_.count(&buffer) == 0) << "LogicalBuffer " << buffer << " already assigned to allocation " << index_; @@ -71,17 +74,44 @@ void BufferAllocation::AddAssignment(const LogicalBuffer& buffer, int64 offset, << " offset out of range"; CHECK_LE(offset + size, size_) << "LogicalBuffer " << buffer << " size out of range"; + CHECK_EQ(buffer.color(), color()) + << "Buffer color " << buffer.color() + << " does not match allocation color " << color() << "."; OffsetSize offset_size; offset_size.offset = offset; offset_size.size = size; assigned_buffers_.emplace(&buffer, offset_size); } +BufferAllocationProto BufferAllocation::ToProto() const { + BufferAllocationProto proto; + proto.set_index(index_); + proto.set_size(size_); + proto.set_is_thread_local(is_thread_local_); + proto.set_is_reusable(is_reusable_); + proto.set_color(color_.value()); + if (is_entry_computation_parameter_) { + proto.set_is_entry_computation_parameter(true); + proto.set_parameter_number(parameter_number_); + } + proto.set_maybe_live_out(maybe_live_out_); + for (const auto& buffer_offset_size : assigned_buffers_) { + BufferAllocationProto::Assigned* proto_assigned = proto.add_assigned(); + proto_assigned->set_logical_buffer_id(buffer_offset_size.first->id()); + proto_assigned->set_offset(buffer_offset_size.second.offset); + proto_assigned->set_size(buffer_offset_size.second.size); + } + return proto; +} + string BufferAllocation::ToString() const { string output; tensorflow::strings::StrAppend( &output, tensorflow::strings::Printf("allocation %lld: %p, size %lld", index_, this, size())); + if (color().value() != 0) { + tensorflow::strings::StrAppend(&output, ", color ", color().value()); + } if (is_entry_computation_parameter()) { tensorflow::strings::StrAppend(&output, ", parameter ", parameter_number()); } @@ -170,10 +200,10 @@ BufferAllocation* BufferAssignment::GetMutableAllocation( return const_cast(&GetAllocation(index)); } -bool BufferAssignment::HasTopLevelAllocation( - const HloInstruction* instruction) const { +bool BufferAssignment::HasAllocationAt(const HloInstruction* instruction, + const ShapeIndex& index) const { for (const LogicalBuffer* buffer : - GetPointsToSet(instruction).element(/*index=*/{})) { + GetPointsToSet(instruction).element(index)) { if (allocation_index_for_buffer_.count(buffer) > 0) { return true; } @@ -181,12 +211,21 @@ bool BufferAssignment::HasTopLevelAllocation( return false; } +bool BufferAssignment::HasTopLevelAllocation( + const HloInstruction* instruction) const { + return HasAllocationAt(instruction, /*index=*/{}); +} + StatusOr BufferAssignment::GetUniqueSlice( const HloInstruction* instruction, const ShapeIndex& index) const { + VLOG(3) << "Trying to find unique slice for " << instruction->name() << " [" + << index << "]"; BufferAllocation::Slice result; for (const LogicalBuffer* buffer : GetPointsToSet(instruction).element(index)) { + VLOG(3) << "Examining buffer " << *buffer; if (HasAllocation(*buffer)) { + VLOG(3) << "Has allocation"; const BufferAllocation::Slice slice = GetAssignedAllocation(*buffer).GetSlice(*buffer); if (result.allocation() == nullptr) { @@ -197,6 +236,8 @@ StatusOr BufferAssignment::GetUniqueSlice( "be determined at compile-time.", instruction->name().c_str(), index.ToString().c_str()); } + } else { + VLOG(3) << "No allocation"; } } if (result.allocation() == nullptr) { @@ -225,11 +266,11 @@ BufferAssignment::GetUniqueTopLevelOutputSlice() const { module_->entry_computation()->root_instruction()); } -BufferAllocation* BufferAssignment::NewEmptyAllocation(int64 size, - bool is_thread_local, - bool is_reusable) { +BufferAllocation* BufferAssignment::NewEmptyAllocation( + int64 size, bool is_thread_local, bool is_reusable, + LogicalBuffer::Color color) { BufferAllocation::Index index = allocations_.size(); - allocations_.emplace_back(index, size, is_thread_local, is_reusable); + allocations_.emplace_back(index, size, is_thread_local, is_reusable, color); BufferAllocation* allocation = &allocations_.back(); return allocation; } @@ -239,7 +280,7 @@ BufferAllocation* BufferAssignment::NewAllocation(const LogicalBuffer& buffer, bool is_thread_local, bool is_reusable) { BufferAllocation* allocation = - NewEmptyAllocation(size, is_thread_local, is_reusable); + NewEmptyAllocation(size, is_thread_local, is_reusable, buffer.color()); AddAssignment(allocation, buffer, /*offset=*/0, size); return allocation; } @@ -259,33 +300,56 @@ void BufferAssignment::AddAssignment(BufferAllocation* allocation, allocation_index_for_buffer_[&buffer] = allocation->index(); } -// Combines allocations of temporary buffers into one big BufferAllocation. +// Combines allocations of temporary buffers of the same color into one big +// BufferAllocation. void BufferAssignment::CombineTempAllocations() { + FlatMap + combined_allocation_map; + // Move all temp allocations into a single run at the end of the allocations - // vector, and combine them into the first allocation of the run. + // vector. const auto first_temp_it = std::partition(allocations_.begin(), allocations_.end(), [](const BufferAllocation& allocation) { return !allocation.IsPreallocatedTempBuffer(); }); + + // Walk over the run of temp allocations, collecting the allocations belonging + // to the same color. if (first_temp_it != allocations_.end()) { - BufferAllocation* combined = &*first_temp_it; - const auto second_temp_it = std::next(first_temp_it); - for (auto it = second_temp_it; it != allocations_.end(); ++it) { + for (auto it = first_temp_it; it != allocations_.end(); ++it) { + const BufferAllocation& temp_allocation = *it; + LogicalBuffer::Color color = temp_allocation.color(); + auto combined_it = combined_allocation_map.find(color); + if (combined_it == combined_allocation_map.end()) { + // We have found the first temp allocation of this color. Collect + // the other temp allocations of the same color into it. + combined_allocation_map.emplace(color, temp_allocation); + continue; + } + + auto* combined_allocation = &combined_it->second; // Each temp allocation is placed end-to-end, accounting for alignment. // The offset of each buffer in the combined allocation is computed from // the base offset of the allocation. - const int64 base = RoundUpToNearest(combined->size(), alignment_); - combined->set_size(base + it->size()); - for (const auto& buffer_offset_size : it->assigned_buffers_) { + int64 alignment = color_alignment_(color); + const int64 base = + RoundUpToNearest(combined_allocation->size(), alignment); + combined_allocation->set_size(base + temp_allocation.size()); + for (const auto& buffer_offset_size : temp_allocation.assigned_buffers_) { const LogicalBuffer* buffer = buffer_offset_size.first; const int64 offset = buffer_offset_size.second.offset; const int64 size = buffer_offset_size.second.size; - combined->AddAssignment(*buffer, base + offset, size); + combined_allocation->AddAssignment(*buffer, base + offset, size); } } - allocations_.erase(second_temp_it, allocations_.end()); - temp_allocation_ = combined; + // Replace all existing temporary allocations with the new combined + // allocations. + allocations_.erase(first_temp_it, allocations_.end()); + for (auto& combined : combined_allocation_map) { + allocations_.push_back(combined.second); + temp_allocation_total_size_ += combined.second.size(); + } } // Update allocation indices to their new positions. @@ -300,8 +364,7 @@ void BufferAssignment::CombineTempAllocations() { } } -Status BufferAssignment::ComputeSummaryStats( - const LogicalBuffer::SizeFunction& buffer_size) { +Status BufferAssignment::ComputeSummaryStats() { for (auto& allocation : Allocations()) { if (allocation.is_entry_computation_parameter()) { stats_.parameter_allocation_count++; @@ -329,8 +392,9 @@ Status BufferAssignment::ComputeSummaryStats( } } if (module_sequence.size() == module_->computations().size()) { - TF_ASSIGN_OR_RETURN(const int64 min_size, - MinimumMemoryForSequence(module_sequence, buffer_size)); + TF_ASSIGN_OR_RETURN( + const int64 min_size, + MinimumMemoryForSequence(module_sequence, buffer_size_)); stats_.total_fragmentation_bytes = stats_.total_allocation_bytes - min_size; } @@ -374,6 +438,46 @@ string BufferAssignment::ToString() const { return output; } +BufferAssignmentProto BufferAssignment::ToProto() const { + BufferAssignmentProto proto; + // NOTE: TuplePointsToAnalysis state is serialized here in BufferAssigment, + // because we need to do the HasAllocation check for each buffer. Otherwise + // the buffer_size_ call might fail for some backends. + const TuplePointsToAnalysis& points_to_analysis = + liveness_->points_to_analysis(); + for (LogicalBuffer::Id id = 0; id < points_to_analysis.num_logical_buffers(); + id++) { + auto& buffer = points_to_analysis.logical_buffer(id); + if (HasAllocation(buffer)) { + LogicalBufferProto proto_buffer = buffer.ToProto(buffer_size_); + proto.add_logical_buffers()->Swap(&proto_buffer); + + // Fill buffer aliases. + for (const BufferAlias& alias : + points_to_analysis.GetBufferAliases(buffer)) { + if (alias.instruction() == buffer.instruction() && + alias.index() == buffer.index()) { + continue; // skip self-aliases + } + BufferAssignmentProto::BufferAlias* proto_alias = + proto.add_buffer_aliases(); + LogicalBufferProto::Location proto_alias_location = + LogicalBuffer::ToLocationProto(*alias.instruction(), alias.index()); + proto_alias->set_source_buffer_id(buffer.id()); + proto_alias->mutable_location()->Swap(&proto_alias_location); + } + } + } + for (const BufferAllocation& allocation : Allocations()) { + BufferAllocationProto proto_allocation = allocation.ToProto(); + proto.add_buffer_allocations()->Swap(&proto_allocation); + } + for (const HeapSimulatorTrace& trace : heap_simulator_traces_) { + *proto.add_heap_simulator_traces() = trace; + } + return proto; +} + namespace { // Walk the call graph of the HLO module and place each computation into either @@ -394,8 +498,8 @@ Status GatherComputationsByAllocationType( // Sets for quickly checking membership. Computations are returned in vectors // for stable iteration. - tensorflow::gtl::FlatSet thread_local_set; - tensorflow::gtl::FlatSet global_set; + FlatSet thread_local_set; + FlatSet global_set; while (!worklist.empty()) { auto worklist_front = worklist.front(); @@ -486,35 +590,34 @@ Status GatherComputationsByAllocationType( /* static */ StatusOr> BufferAssigner::Run( const HloModule* module, std::unique_ptr hlo_ordering, - LogicalBuffer::SizeFunction buffer_size, int64 alignment, - bool colocate_related_buffers, - const std::vector* hlos_to_allocate) { - BufferAssigner assigner(std::move(buffer_size), alignment, - colocate_related_buffers); + LogicalBuffer::SizeFunction buffer_size, + LogicalBuffer::AlignmentFunction color_alignment, + bool allow_input_output_aliasing, BufferLiveness::Colorer colorer) { + BufferAssigner assigner(allow_input_output_aliasing, std::move(colorer)); return assigner.CreateAssignment(module, std::move(hlo_ordering), - hlos_to_allocate); -} - -/* static */ -StatusOr> BufferAssigner::Run( - const HloModule* module, std::unique_ptr hlo_ordering, - LogicalBuffer::SizeFunction buffer_size, int64 alignment) { - return BufferAssigner::Run(module, std::move(hlo_ordering), - std::move(buffer_size), alignment, - /*colocate_related_buffers=*/true); + std::move(buffer_size), + std::move(color_alignment)); } bool BufferAssigner::MaybeAssignBuffer(BufferAllocation* allocation, const LogicalBuffer& buffer, BufferAssignment* assignment) { + const LogicalBuffer::SizeFunction& buffer_size = assignment->buffer_size_; + CHECK(!assignment->HasAllocation(buffer)) << "buffer " << buffer << " already has an allocation assigned."; VLOG(4) << "Trying to assign " << buffer << " to allocation: " << *allocation; - if (buffer_size_(buffer) > allocation->size()) { + if (buffer.color() != allocation->color()) { + VLOG(4) << "Can't assign: buffer has color" << buffer.color() + << " and allocation has color " << allocation->color() << "."; + return false; + } + + if (buffer_size(buffer) > allocation->size()) { VLOG(4) << "Can't assign: buffer is larger than allocation (" - << buffer_size_(buffer) << " > " << allocation->size() << ")"; + << buffer_size(buffer) << " > " << allocation->size() << ")"; return false; } @@ -535,50 +638,70 @@ bool BufferAssigner::MaybeAssignBuffer(BufferAllocation* allocation, << " may interfere with " << buffer; return false; } + // Copy instruction don't share a buffer with their input operand. + if (buffer.instruction()->IsUserOf(assigned_buffer.instruction()) && + buffer.instruction()->opcode() == HloOpcode::kCopy) { + VLOG(4) << "Can't assign: assignee " << assigned_buffer + << " is used at copy instruction " << buffer; + return false; + } + } + + if (allow_input_output_aliasing_ && allocation->maybe_live_out()) { + HloComputation* entry_computation = + assignment->module_->entry_computation(); + for (auto param : entry_computation->parameter_instructions()) { + for (auto& param_buffer : + assignment->points_to_analysis().GetBuffersDefinedByInstruction( + param)) { + if (assignment->liveness().MayInterfere(*param_buffer, buffer)) { + VLOG(4) << "Can't assign: Parameter interference with result"; + return false; + } + } + } } // If the buffer is live out of the computation then it should only be // assigned a buffer which exactly fits the result to avoid wasting memory // (result buffers can have arbitrary lifetimes). if (assignment->liveness().MaybeLiveOut(buffer) && - allocation->size() != buffer_size_(buffer)) { + allocation->size() != buffer_size(buffer)) { VLOG(4) << "Can't assign: buffer " << buffer << "is live out and size not the same as allocation"; return false; } assignment->AddAssignment(allocation, buffer, /*offset=*/0, - buffer_size_(buffer)); + buffer_size(buffer)); return true; } Status BufferAssigner::AssignBuffersForComputation( - const HloComputation* computation, bool is_thread_local, - const tensorflow::gtl::FlatSet* hlos_to_allocate, - const tensorflow::gtl::FlatSet& colocated_buffers, - const tensorflow::gtl::FlatSet& - colocated_allocations, + const HloComputation* computation, const DebugOptions& debug_options, + bool is_thread_local, + const FlatSet& colocated_buffers, + const FlatSet& colocated_allocations, + FlatMap>* + buffers_to_assign_sequentially, BufferAssignment* assignment) { // Buffers are sorted and assigned to BufferAllocations in decreasing order of // size. std::vector sorted_buffers; for (auto& instruction : computation->instructions()) { - if (hlos_to_allocate == nullptr || - hlos_to_allocate->count(instruction.get()) > 0) { - // Add all buffers which this instruction defines. Instruction which don't - // define buffers (eg, bitcast which just forwards a pointer) don't need - // any allocations. - for (const LogicalBuffer* buffer : - assignment->points_to_analysis().GetBuffersDefinedByInstruction( - instruction.get())) { - sorted_buffers.push_back(buffer); - } + // Add all buffers which this instruction defines. Instruction which don't + // define buffers (eg, bitcast which just forwards a pointer) don't need + // any allocations. + for (const LogicalBuffer* buffer : + assignment->points_to_analysis().GetBuffersDefinedByInstruction( + instruction.get())) { + sorted_buffers.push_back(buffer); } } // Generate a post order sort of instructions for sorting of the // LogicalBuffers. - tensorflow::gtl::FlatMap post_order_position; + FlatMap post_order_position; int position = 0; for (auto* instruction : computation->MakeInstructionPostOrder()) { post_order_position.emplace(instruction, position); @@ -588,9 +711,16 @@ Status BufferAssigner::AssignBuffersForComputation( // If there is a sequential instruction ordering, we'll delay assignment of // temp buffers until after the main assignment loop. const BufferLiveness& liveness = assignment->liveness(); - const std::vector* sequential_order = - liveness.hlo_ordering().SequentialOrder(*computation); - tensorflow::gtl::FlatSet unassigned_temp_buffers; + const bool has_sequential_order = + liveness.hlo_ordering().SequentialOrder(*computation) != nullptr; + if (has_sequential_order && buffers_to_assign_sequentially != nullptr) { + // Every sequential computation must get an entry in the + // buffers_to_assign_sequentially map, even if we end up with an empty set + // of buffers. This ensures we can correctly determine whether to run + // whole-module heap simulation. + buffers_to_assign_sequentially->emplace(computation, + FlatSet()); + } // Sort the LogicalBuffers first by size. We assign the larger LogicalBuffers // first for simplicity. This means any previously created BufferAllocation is @@ -609,17 +739,17 @@ Status BufferAssigner::AssignBuffersForComputation( // 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(), - [this, sequential_order, &liveness, &post_order_position]( - const LogicalBuffer* a, const LogicalBuffer* b) { + [this, 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 = buffer_size_(*a); - const int64 b_size = buffer_size_(*b); + 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 (sequential_order != nullptr) { + 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) { @@ -652,7 +782,7 @@ Status BufferAssigner::AssignBuffersForComputation( continue; } - const int64 buffer_size = buffer_size_(*buffer); + const int64 buffer_size = assignment->buffer_size_(*buffer); const bool is_entry_parameter = instruction->opcode() == HloOpcode::kParameter && @@ -673,10 +803,7 @@ Status BufferAssigner::AssignBuffersForComputation( continue; } - legacy_flags::BufferAssignmentFlags* flags = - legacy_flags::GetBufferAssignmentFlags(); - if (!flags->xla_enable_buffer_reuse || is_thread_local || - instruction->opcode() == HloOpcode::kCustomCall) { + if (is_thread_local || instruction->opcode() == HloOpcode::kCustomCall) { // Custom call operations never have reusable buffers. Also we do not // reuse thread-local buffers for now, because they are dynamically // allocated and their lifetimes are hard to compute. @@ -687,6 +814,17 @@ Status BufferAssigner::AssignBuffersForComputation( continue; } + if (instruction->opcode() == HloOpcode::kRecv) { + // Make sure that recv operations get a new unique allocation so that + // don't share their buffer with any other operations. + BufferAllocation* allocation = assignment->NewAllocation( + *buffer, buffer_size, is_thread_local, /*is_reusable=*/false); + allocation_indices.push_back(allocation->index()); + VLOG(3) << "New allocation #" << allocation->index() + << " for recv: " << *buffer; + continue; + } + if (ShapeUtil::IsTuple(buffer->shape())) { // TODO(b/34669761): Don't reuse tuple buffers because the GPU backend // assumes longer buffer liveness than indicated by the analysis. @@ -756,7 +894,7 @@ Status BufferAssigner::AssignBuffersForComputation( } } - if (!assignment->HasAllocation(*buffer) && sequential_order != nullptr && + if (!assignment->HasAllocation(*buffer) && has_sequential_order && !liveness.MaybeLiveOut(*buffer)) { // There is a sequential instruction ordering, so we delay assignment of // temp buffers until after the loop. We do this right before we decide to @@ -768,7 +906,7 @@ Status BufferAssigner::AssignBuffersForComputation( // for the definition of temp buffers. CHECK(!is_entry_parameter) << *buffer; CHECK(!is_thread_local) << *buffer; - unassigned_temp_buffers.insert(buffer); + (*buffers_to_assign_sequentially)[computation].insert(buffer); VLOG(3) << "Delaying assignment of temp buffer: " << *buffer; continue; } @@ -782,27 +920,98 @@ Status BufferAssigner::AssignBuffersForComputation( } } - if (!unassigned_temp_buffers.empty()) { - TF_RETURN_IF_ERROR(AssignBuffersWithSequentialOrdering( - *sequential_order, unassigned_temp_buffers, *computation, assignment)); - } return Status::OK(); } +FlatMap, + LogicalBuffer::Color::Hasher> +BufferAssigner::SplitBuffersByColor( + const FlatSet& buffers) { + FlatMap, + LogicalBuffer::Color::Hasher> + color_map; + for (auto buffer : buffers) { + color_map[buffer->color()].insert(buffer); + } + return color_map; +} + Status BufferAssigner::AssignBuffersWithSequentialOrdering( - const std::vector& sequence, - const tensorflow::gtl::FlatSet& buffers_to_assign, - const HloComputation& computation, BufferAssignment* assignment) { + const FlatMap>& + buffers_to_assign_sequentially, + bool run_whole_module_heap_simulation, BufferAssignment* assignment) { // Run the sequence of instructions through the heap simulator. The heuristic // that seems to give the best results is lazy-best-fit, with all runs of // alloc / free calls sorted in decreasing size order. - TF_ASSIGN_OR_RETURN( - HeapSimulator::Result result, - HeapSimulator::Run(MakeUnique( - MakeUnique(alignment_)), - sequence, computation, - assignment->points_to_analysis(), buffer_size_, - &buffers_to_assign)); + const HloOrdering& hlo_ordering = assignment->liveness().hlo_ordering(); + if (run_whole_module_heap_simulation) { + // Run the heap simulation over the whole module. This reduces memory usage, + // since buffers for kCall and kWhile sub-computations are only live for the + // duration of their calling instructions. + VLOG(1) << "Running whole-module heap simulation"; + SequentialHloOrdering::HloModuleSequence module_sequence; + FlatSet all_buffers_to_assign; + for (const auto& pair : buffers_to_assign_sequentially) { + const HloComputation* computation = pair.first; + const FlatSet& buffers_to_assign = pair.second; + const std::vector* instruction_sequence = + hlo_ordering.SequentialOrder(*computation); + CHECK(instruction_sequence != nullptr) << computation->name(); + module_sequence[computation] = *instruction_sequence; + all_buffers_to_assign.insert(buffers_to_assign.begin(), + buffers_to_assign.end()); + } + auto color_map = SplitBuffersByColor(all_buffers_to_assign); + for (auto& single_colored_set : color_map) { + auto color = single_colored_set.first; + VLOG(2) << "Simulating heap for color " << color; + int64 alignment = assignment->color_alignment_(color); + TF_ASSIGN_OR_RETURN( + const HeapSimulator::Result result, + HeapSimulator::Run(MakeUnique( + MakeUnique(alignment)), + assignment->module(), module_sequence, + assignment->points_to_analysis(), + assignment->buffer_size_, + &single_colored_set.second)); + AssignBuffersFromHeapSimulator(result, assignment, + single_colored_set.first); + } + } else { + // Run the heap-simulation on a per-computation basis. Buffers for + // sub-computations are assigned disjoint BufferAllocations, assuming the + // worst-case that they may all be live concurrently. + VLOG(1) << "Running per-computation heap simulation"; + for (const auto& pair : buffers_to_assign_sequentially) { + const HloComputation* computation = pair.first; + const FlatSet& buffers_to_assign = pair.second; + const std::vector* instruction_sequence = + hlo_ordering.SequentialOrder(*computation); + CHECK(instruction_sequence != nullptr) << computation->name(); + auto color_map = SplitBuffersByColor(buffers_to_assign); + for (auto& single_colored_set : color_map) { + auto color = single_colored_set.first; + VLOG(2) << "Simulating heap for color " << color; + int64 alignment = assignment->color_alignment_(color); + TF_ASSIGN_OR_RETURN( + const HeapSimulator::Result result, + HeapSimulator::Run(MakeUnique( + MakeUnique(alignment)), + *computation, *instruction_sequence, + assignment->points_to_analysis(), + assignment->buffer_size_, + &single_colored_set.second)); + AssignBuffersFromHeapSimulator(result, assignment, + single_colored_set.first); + } + } + } + return Status::OK(); +} + +void BufferAssigner::AssignBuffersFromHeapSimulator( + const HeapSimulator::Result& result, BufferAssignment* assignment, + LogicalBuffer::Color color) { if (assignment->stats_.preallocated_temp_fragmentation_bytes == -1) { assignment->stats_.preallocated_temp_fragmentation_bytes = result.fragmentation_size; @@ -811,16 +1020,15 @@ Status BufferAssigner::AssignBuffersWithSequentialOrdering( result.fragmentation_size; } - // Use the results of the heap simulator to create one allocation per - // computation, with LogicalBuffers packed to specific offsets. BufferAllocation* allocation = assignment->NewEmptyAllocation( - result.heap_size, /*is_thread_local=*/false, /*is_reusable=*/true); + result.heap_size, /*is_thread_local=*/false, /*is_reusable=*/true, color); for (const auto& buffer_chunk : result.chunk_map) { const LogicalBuffer& buffer = *buffer_chunk.first; const HeapSimulator::Chunk& chunk = buffer_chunk.second; assignment->AddAssignment(allocation, buffer, chunk.offset, chunk.size); } - return Status::OK(); + + assignment->heap_simulator_traces_.push_back(result.debug_trace); } // Adds the 'colocated_set' of buffers to 'colocated_buffer_sets', maintaining @@ -881,43 +1089,174 @@ void BufferAssigner::AddSetToColocatedBufferSets( } } +// Conceptually the same as AddSetToColocatedBufferSets, but specific to the +// colocated buffers for while instructions. 'colocated_set' contains the +// buffers for a single while instruction that must be colocated. The idea here +// is to apply a memory-saving heuristic for separate while instructions whose +// buffers are disjoint in liveness, by using the colocation mechanism to force +// buffer sharing. This often reduces memory for multi-layer RNNs. +// +// TODO(b/32491382): We should be able to remove this heuristic after we +// implement module-level liveness analysis, which would let us directly detect +// buffer sharing opportunities between the while instruction buffer and the +// buffers from the predicate and body computation, as well as sharing across +// different while instructions. +void BufferAssigner::AddWhileSetToColocatedBufferSets( + const std::vector& colocated_set, + const LogicalBuffer* while_init_buffer, + const LogicalBuffer* while_result_buffer, const HloInstruction* while_hlo, + const HloComputation& computation, const BufferLiveness& buffer_liveness, + const LogicalBuffer::SizeFunction& buffer_size, + std::vector* colocated_buffer_sets) { + CHECK(!colocated_set.empty()); + const TuplePointsToAnalysis& points_to_analysis = + buffer_liveness.points_to_analysis(); + + // Parallel while loops cannot safely share colocated buffer sets. + if (buffer_liveness.hlo_ordering().SequentialOrder(computation) == nullptr) { + AddSetToColocatedBufferSets(colocated_set, colocated_buffer_sets); + return; + } + + // Scan 'colocated_buffer_sets' in reverse order for locality; colocated sets + // are added in postorder over computations and instructions. + const int64 init_buffer_size = buffer_size(*while_init_buffer); + for (int i = colocated_buffer_sets->size() - 1; i >= 0; --i) { + const ColocatedBufferSet& predecessor_set = (*colocated_buffer_sets)[i]; + + // Skip predecessor sets not associated with while loops. + if (std::all_of(predecessor_set.begin(), predecessor_set.end(), + [](const LogicalBuffer* buffer) { + return buffer->instruction()->opcode() != + HloOpcode::kWhile; + })) { + continue; + } + + // Skip predecessor sets already associated with 'while_hlo'. + if (std::any_of(predecessor_set.begin(), predecessor_set.end(), + [&while_hlo](const LogicalBuffer* buffer) { + return buffer->instruction() == while_hlo; + })) { + continue; + } + + // Build vector of predecessor while result and init buffers, which are + // checked for liveness interference below. We must check both the result + // and init buffers because they're aliased together, but + // TuplePointsToAnalysis is unaware of this aliasing. + std::vector predecessor_while_buffers; + for (const LogicalBuffer* buffer : predecessor_set) { + const HloInstruction* instruction = buffer->instruction(); + if (instruction->opcode() == HloOpcode::kWhile && + buffer_size(*buffer) == init_buffer_size && + instruction->parent() == &computation) { + predecessor_while_buffers.push_back(buffer); + // Add the init buffer at the same index, which must also exist in the + // predecessor set, and must be unambiguous. + const PointsToSet& init_points_to = + points_to_analysis.GetPointsToSet(instruction->operand(0)); + const auto& init_buffers = init_points_to.element(buffer->index()); + CHECK_EQ(init_buffers.size(), 1); + CHECK_GT(predecessor_set.count(init_buffers[0]), 0); + predecessor_while_buffers.push_back(init_buffers[0]); + } + } + if (predecessor_while_buffers.empty()) { + continue; + } + + // Skip predecessor set if the live range of any predecessor + // buffers overlaps with 'while_init_buffer' or + // 'while_result_buffer' (we need to check both since they're + // aliased together, but the points-to analysis is unaware of this + // aliasing). Note that tuple element buffer forwarding can cause + // the same buffer to appear on both sides of the interference + // comparison below. + auto may_interfere_with_init_or_result = [&](const LogicalBuffer* buffer) { + if (while_init_buffer->id() != buffer->id() && + buffer_liveness.MayInterfere(*while_init_buffer, *buffer)) { + return true; + } + + if (while_result_buffer->id() != buffer->id() && + buffer_liveness.MayInterfere(*while_result_buffer, *buffer)) { + return true; + } + + return false; + }; + + if (std::any_of(predecessor_while_buffers.begin(), + predecessor_while_buffers.end(), + may_interfere_with_init_or_result)) { + continue; + } + + // All our checks have passed; merge 'predecessor_set' with 'colocated_set', + // and add the merged set to 'colocated_buffer_sets'. This forces the + // colocation of buffers across different while instructions. + FlatSet unique; + unique.insert(predecessor_set.begin(), predecessor_set.end()); + unique.insert(colocated_set.begin(), colocated_set.end()); + std::vector merged_set(unique.begin(), unique.end()); + AddSetToColocatedBufferSets(merged_set, colocated_buffer_sets); + return; + } + + // Failed to merge into predecessor set; add 'colocated_set' as-is. + AddSetToColocatedBufferSets(colocated_set, colocated_buffer_sets); +} + namespace { + // Checks that points-to set of 'instruction' is unambiguous and distinct // (ensured by CopyInsertion), then adds the buffer from the points-to set at // 'index' to 'colocated_set'. -void AddBufferToColocatedSet(const HloInstruction* instruction, - const ShapeIndex& index, - const TuplePointsToAnalysis& points_to_analysis, - std::vector* colocated_set) { +const LogicalBuffer* AddBufferToColocatedSet( + const HloInstruction* instruction, const ShapeIndex& index, + const TuplePointsToAnalysis& points_to_analysis, + std::vector* colocated_set) { // CopyInsertion ensures root points-to set is unambiguous and distinct. const auto& points_to = points_to_analysis.GetPointsToSet(instruction); - CHECK(!points_to.IsAmbiguous()); - CHECK(points_to.IsDistinct()); + DCHECK(!points_to.IsAmbiguous()); + DCHECK(points_to.IsDistinct()); colocated_set->push_back(points_to.element(index)[0]); + return colocated_set->back(); } + } // namespace // Builds sets of buffers in 'colocated_buffer_sets' which should be colocated // in the same allocation (currently just supports kWhile and kCall). void BufferAssigner::BuildColocatedBufferSets( - const HloModule* module, const TuplePointsToAnalysis& points_to_analysis, + const HloModule* module, const BufferLiveness& buffer_liveness, + const LogicalBuffer::SizeFunction& buffer_size, std::vector* colocated_buffer_sets) { - for (auto& computation : module->computations()) { - for (auto& instruction : computation->instructions()) { + const TuplePointsToAnalysis& points_to_analysis = + buffer_liveness.points_to_analysis(); + for (const HloComputation* computation : module->MakeComputationPostOrder()) { + if (computation->IsFusionComputation()) { + continue; + } + for (const HloInstruction* instruction : + computation->MakeInstructionPostOrder()) { const HloOpcode opcode = instruction->opcode(); if (opcode == HloOpcode::kWhile) { - HloInstruction* while_hlo = instruction.get(); - TF_CHECK_OK(ShapeUtil::ForEachSubshape( + const HloInstruction* while_hlo = instruction; + ShapeUtil::ForEachSubshape( while_hlo->shape(), - [this, while_hlo, &points_to_analysis, colocated_buffer_sets]( + [this, while_hlo, &points_to_analysis, &buffer_liveness, + buffer_size, computation, colocated_buffer_sets]( const Shape& /*subshape*/, const ShapeIndex& index) { std::vector colocated_set; // Add while.init. - AddBufferToColocatedSet(while_hlo->operand(0), index, - points_to_analysis, &colocated_set); + auto* init_buffer = + AddBufferToColocatedSet(while_hlo->operand(0), index, + points_to_analysis, &colocated_set); // Add while.result. - AddBufferToColocatedSet(while_hlo, index, points_to_analysis, - &colocated_set); + auto* result_buffer = AddBufferToColocatedSet( + while_hlo, index, points_to_analysis, &colocated_set); // Add while.cond.parameter. AddBufferToColocatedSet( while_hlo->while_condition()->parameter_instruction(0), index, @@ -930,13 +1269,16 @@ void BufferAssigner::BuildColocatedBufferSets( AddBufferToColocatedSet( while_hlo->while_body()->root_instruction(), index, points_to_analysis, &colocated_set); - AddSetToColocatedBufferSets(colocated_set, colocated_buffer_sets); - return Status::OK(); - })); + AddWhileSetToColocatedBufferSets( + colocated_set, init_buffer, result_buffer, while_hlo, + *computation, buffer_liveness, buffer_size, + colocated_buffer_sets); + }); } else if (opcode == HloOpcode::kCall) { - HloInstruction* call_hlo = instruction.get(); - HloInstruction* root_hlo = call_hlo->to_apply()->root_instruction(); - TF_CHECK_OK(ShapeUtil::ForEachSubshape( + const HloInstruction* call_hlo = instruction; + const HloInstruction* root_hlo = + call_hlo->to_apply()->root_instruction(); + ShapeUtil::ForEachSubshape( call_hlo->shape(), [this, call_hlo, root_hlo, &points_to_analysis, colocated_buffer_sets](const Shape& /*subshape*/, @@ -949,8 +1291,7 @@ void BufferAssigner::BuildColocatedBufferSets( AddBufferToColocatedSet(root_hlo, index, points_to_analysis, &colocated_set); AddSetToColocatedBufferSets(colocated_set, colocated_buffer_sets); - return Status::OK(); - })); + }); } } } @@ -961,23 +1302,43 @@ void BufferAssigner::BuildColocatedBufferSets( void BufferAssigner::AssignColocatedBufferSets( const std::vector& colocated_buffer_sets, BufferAssignment* assignment, - tensorflow::gtl::FlatSet* colocated_buffers, - tensorflow::gtl::FlatSet* colocated_allocations) { + FlatSet* colocated_buffers, + FlatSet* colocated_allocations) { for (const ColocatedBufferSet& colocated_buffer_set : colocated_buffer_sets) { BufferAllocation* allocation = nullptr; + // Set 'entry_parameter_number' if entry param in 'colocated_buffer_set'. + int64 entry_parameter_number = -1; + for (const LogicalBuffer* buffer : colocated_buffer_set) { + const HloInstruction* instruction = buffer->instruction(); + const HloComputation* computation = instruction->parent(); + if (instruction->opcode() == HloOpcode::kParameter && + computation == computation->parent()->entry_computation()) { + entry_parameter_number = instruction->parameter_number(); + break; + } + } + for (const LogicalBuffer* buffer : colocated_buffer_set) { if (allocation == nullptr) { // TODO(b/32491382) Avoid current trivial solution of using new // allocations for each colocated buffer set. When liveness has // module-level scope, we can allow buffers to be shared across // computations (in some cases). - allocation = assignment->NewAllocation(*buffer, buffer_size_(*buffer), - /*is_thread_local=*/false, - /*is_reusable=*/true); + allocation = assignment->NewAllocation( + *buffer, assignment->buffer_size_(*buffer), + /*is_thread_local=*/false, /*is_reusable=*/true); + if (entry_parameter_number >= 0) { + // This colocated buffer set contains an entry parameter and other + // logical buffers which use the parameter as read-only in a while + // body computation (which updates in place). + // Set 'entry_computation_parameter' to indicate that it contains + // an entry parameter, and to prevent reuse in MaybeAssignBuffer. + allocation->set_entry_computation_parameter(entry_parameter_number); + } colocated_allocations->insert(allocation->index()); } else { assignment->AddAssignment(allocation, *buffer, /*offset=*/0, - buffer_size_(*buffer)); + assignment->buffer_size_(*buffer)); } colocated_buffers->insert(buffer); } @@ -986,62 +1347,72 @@ void BufferAssigner::AssignColocatedBufferSets( StatusOr> BufferAssigner::CreateAssignment( const HloModule* module, std::unique_ptr hlo_ordering, - const std::vector* hlos_to_allocate) { + LogicalBuffer::SizeFunction buffer_size, + LogicalBuffer::AlignmentFunction color_alignment) { TF_ASSIGN_OR_RETURN(std::unique_ptr liveness, BufferLiveness::Run(module, std::move(hlo_ordering))); - std::vector thread_local_computations; - std::vector global_computations; VLOG(1) << "Assigning buffers to module " << module->name(); - if (hlos_to_allocate != nullptr) { - VLOG(3) << "LogicalBuffer assignment restricted to hlos: "; - for (auto hlo : *hlos_to_allocate) { - VLOG(3) << " " << hlo->parent()->name() << "::" << hlo->name(); - } - } - XLA_VLOG_LINES(3, module->ToString()); + XLA_VLOG_LINES(2, module->ToString()); XLA_VLOG_LINES(3, liveness->ToString()); XLA_VLOG_LINES(3, liveness->points_to_analysis().ToString()); - TF_RETURN_IF_ERROR(GatherComputationsByAllocationType( - module, &thread_local_computations, &global_computations)); - - // Set of HLO's to allocate if hlos_to_allocate is given. Passed as a set to - // AssignBuffersForComputation for fast membership testing. - std::unique_ptr> hlo_set; - if (hlos_to_allocate != nullptr) { - hlo_set = MakeUnique>( - hlos_to_allocate->begin(), hlos_to_allocate->end()); - } - // Can't use MakeUnique because BufferAssignment constructor is private. std::unique_ptr assignment( - new BufferAssignment(module, std::move(liveness), alignment_)); + new BufferAssignment(module, std::move(liveness), std::move(buffer_size), + std::move(color_alignment))); // Assign buffers with the tightest constraints first (colocated buffer sets). // Once b/32491382 enables module-level liveness analysis, we may be able // to assign colocated buffers (or at least reuse their allocation for // buffers outside of the set) in AssignBuffersForComputation. - tensorflow::gtl::FlatSet colocated_buffers; - tensorflow::gtl::FlatSet colocated_allocations; - if (colocate_related_buffers_) { - std::vector colocated_buffer_sets; - BuildColocatedBufferSets(module, assignment->points_to_analysis(), - &colocated_buffer_sets); - AssignColocatedBufferSets(colocated_buffer_sets, assignment.get(), - &colocated_buffers, &colocated_allocations); - } + FlatSet colocated_buffers; + FlatSet colocated_allocations; + std::vector colocated_buffer_sets; + BuildColocatedBufferSets(module, assignment->liveness(), + assignment->buffer_size_, &colocated_buffer_sets); + TF_RETURN_IF_ERROR(colorer_(assignment->liveness())); + VLOG(3) << "After coloring:"; + XLA_VLOG_LINES(3, assignment->points_to_analysis().ToString()); + + AssignColocatedBufferSets(colocated_buffer_sets, assignment.get(), + &colocated_buffers, &colocated_allocations); + + std::vector thread_local_computations; + std::vector global_computations; + TF_RETURN_IF_ERROR(GatherComputationsByAllocationType( + module, &thread_local_computations, &global_computations)); + // First assign buffers for global computatations. Temporary buffers for + // sequential computations are collected in 'buffers_to_assign_sequentially'. + FlatMap> + buffers_to_assign_sequentially; for (auto* computation : global_computations) { TF_RETURN_IF_ERROR(AssignBuffersForComputation( - computation, /*is_thread_local=*/false, hlo_set.get(), - colocated_buffers, colocated_allocations, assignment.get())); + computation, module->config().debug_options(), + /*is_thread_local=*/false, colocated_buffers, colocated_allocations, + &buffers_to_assign_sequentially, assignment.get())); } + // Assign buffers with sequential ordering, if any. If all global computations + // are sequential, we can run heap simuation on the whole module, which + // reduces memory usage. + const bool run_whole_module_heap_simulation = + buffers_to_assign_sequentially.size() == global_computations.size(); + TF_RETURN_IF_ERROR(AssignBuffersWithSequentialOrdering( + buffers_to_assign_sequentially, run_whole_module_heap_simulation, + assignment.get())); + + // Now assign buffers for thread-local computations. All LogicalBuffers get + // their own BufferAllocation. for (auto* computation : thread_local_computations) { TF_RET_CHECK(computation != module->entry_computation()); + if (computation->IsFusionComputation()) { + continue; + } TF_RETURN_IF_ERROR(AssignBuffersForComputation( - computation, /*is_thread_local=*/true, hlo_set.get(), colocated_buffers, - colocated_allocations, assignment.get())); + computation, module->config().debug_options(), + /*is_thread_local=*/true, colocated_buffers, colocated_allocations, + /*buffers_to_assign_sequentially=*/nullptr, assignment.get())); } // Mark all buffers which may be live out of the entry computation as @@ -1064,7 +1435,7 @@ StatusOr> BufferAssigner::CreateAssignment( assignment->CombineTempAllocations(); XLA_VLOG_LINES(2, assignment->ToString()); - TF_RETURN_IF_ERROR(assignment->ComputeSummaryStats(buffer_size_)); + TF_RETURN_IF_ERROR(assignment->ComputeSummaryStats()); XLA_VLOG_LINES(1, assignment->GetStats().ToString()); return std::move(assignment); } diff --git a/tensorflow/compiler/xla/service/buffer_assignment.h b/tensorflow/compiler/xla/service/buffer_assignment.h index b82acb19b3488884bdc8d2d5c4a1524ac165676a..688aff89125ce3e30be8918a9dfe9f17e22e6243 100644 --- a/tensorflow/compiler/xla/service/buffer_assignment.h +++ b/tensorflow/compiler/xla/service/buffer_assignment.h @@ -23,10 +23,13 @@ limitations under the License. #include #include "tensorflow/compiler/xla/service/buffer_liveness.h" +#include "tensorflow/compiler/xla/service/heap_simulator.h" +#include "tensorflow/compiler/xla/service/hlo.pb.h" #include "tensorflow/compiler/xla/service/hlo_computation.h" #include "tensorflow/compiler/xla/service/hlo_instruction.h" #include "tensorflow/compiler/xla/service/hlo_module.h" #include "tensorflow/compiler/xla/service/logical_buffer.h" +#include "tensorflow/compiler/xla/service/tuple_points_to_analysis.h" #include "tensorflow/compiler/xla/statusor.h" #include "tensorflow/compiler/xla/types.h" #include "tensorflow/compiler/xla/xla_data.pb.h" @@ -56,11 +59,12 @@ class BufferAllocation { using Index = int64; BufferAllocation(Index index, int64 size, bool is_thread_local, - bool is_reusable) + bool is_reusable, LogicalBuffer::Color color) : index_(index), size_(size), is_thread_local_(is_thread_local), - is_reusable_(is_reusable) {} + is_reusable_(is_reusable), + color_(color) {} ~BufferAllocation() {} // Returns the index of this allocation. @@ -95,6 +99,10 @@ class BufferAllocation { // large as any LogicalBuffer assigned to this allocation. int64 size() const { return size_; } + // Returns the color of the allocation. Only logical buffers with a matching + // color can reside in this allocation. + LogicalBuffer::Color color() const { return color_; } + struct OffsetSize { int64 offset = 0; int64 size = 0; @@ -158,6 +166,7 @@ class BufferAllocation { Slice GetSlice(const LogicalBuffer& buffer) const; string ToString() const; + BufferAllocationProto ToProto() const; // Whether the buffer is a parameter to or live out of the entry computation. bool IsInputOrOutput() const { @@ -214,6 +223,9 @@ class BufferAllocation { // Whether this buffer is usable by more than one logical buffer. bool is_reusable_; + // Color of the allocation. + LogicalBuffer::Color color_; + // Whether this allocation holds an entry computation parameter. Entry // computation parameters are special be cause they have lifetimes which may // outlast the computation. @@ -247,10 +259,10 @@ class BufferAssignment { return allocations_; } - // Returns the single allocation holding all temporary buffers. Returns - // nullptr if there are no temporary buffers, or if the assignment uses more - // than one allocation to hold temporary buffers. - const BufferAllocation* GetTempAllocation() const { return temp_allocation_; } + // Returns the total size allocation holding all temporary buffers. + int64 temp_allocation_total_size() const { + return temp_allocation_total_size_; + } // Returns whether the given buffer has been assigned an allocation. bool HasAllocation(const LogicalBuffer& buffer) const; @@ -269,6 +281,11 @@ class BufferAssignment { std::set GetAllSlices( const HloInstruction* instruction, const ShapeIndex& index) const; + // Convenience function which returns whether the buffer of the + // instruction at the given index is assigned an allocation. + bool HasAllocationAt(const HloInstruction* instruction, + const ShapeIndex& index) const; + // Convenience function which returns whether the top-level buffer of the // instruction (index == {}) is assigned an allocation. bool HasTopLevelAllocation(const HloInstruction* instruction) const; @@ -289,7 +306,7 @@ class BufferAssignment { // Returns the set LogicalBuffers which may be the source of the value at the // given index and instruction. - const std::vector& GetSourceBuffers( + const PointsToSet::BufferList& GetSourceBuffers( const HloInstruction* instruction, const ShapeIndex& index) const { return GetPointsToSet(instruction).element(index); } @@ -308,7 +325,11 @@ class BufferAssignment { return liveness_->points_to_analysis(); } + // Returns the BufferLiveness object used to construct this assignment. + const BufferLiveness& liveness() const { return *liveness_; } + string ToString() const; + BufferAssignmentProto ToProto() const; // Statistics for the assignment. Values initialized to -1 are not always // collected; fragmentation is only collected for instructions that have a @@ -335,15 +356,18 @@ class BufferAssignment { explicit BufferAssignment(const HloModule* module, std::unique_ptr liveness, - int64 alignment) + LogicalBuffer::SizeFunction buffer_size, + LogicalBuffer::AlignmentFunction color_alignment) : module_(module), liveness_(std::move(liveness)), - alignment_(alignment) {} + buffer_size_(std::move(buffer_size)), + color_alignment_(std::move(color_alignment)) {} // Creates and returns a new BufferAllocation, with no assigned // LogicalBuffers. Ownership is maintained internally. BufferAllocation* NewEmptyAllocation(int64 size, bool is_thread_local, - bool is_reusable); + bool is_reusable, + LogicalBuffer::Color color); // Helper that calls NewEmptyAllocation and AddAssignment in one call, // creating an allocation containing a single LogicalBuffer. @@ -354,8 +378,8 @@ class BufferAssignment { void AddAssignment(BufferAllocation* allocation, const LogicalBuffer& buffer, int64 offset, int64 size); - // Returns the BufferLiveness object used to construct this assignment. - const BufferLiveness& liveness() { return *liveness_; } + // Returns the HloModule used to construct this assignment. + const HloModule& module() const { return *module_; } // Convenience function which returns the PointsToSet for the given // instruction. Extracted from the liveness object. @@ -369,13 +393,13 @@ class BufferAssignment { void CombineTempAllocations(); // Computes stats for the assignment, to be retrieved by GetStats. - Status ComputeSummaryStats(const LogicalBuffer::SizeFunction& buffer_size); + Status ComputeSummaryStats(); // The vector of buffer allocations. Indexed by BufferAllocation::Index. std::vector allocations_; - // The single allocation holding all temporary buffers. - BufferAllocation* temp_allocation_ = nullptr; + // The total size of all temporary buffers. + int64 temp_allocation_total_size_ = 0; // Maps Buffers to the index of the BufferAllocation which holds the buffer. tensorflow::gtl::FlatMap @@ -383,9 +407,15 @@ class BufferAssignment { const HloModule* module_; const std::unique_ptr liveness_; - const int64 alignment_; + + // Function which returns the buffer size for a given logical buffer (shape). + LogicalBuffer::SizeFunction buffer_size_; + + // Function which returns the alignment for a given logical buffer color. + LogicalBuffer::AlignmentFunction color_alignment_; Stats stats_; + std::vector heap_simulator_traces_; TF_DISALLOW_COPY_AND_ASSIGN(BufferAssignment); }; @@ -394,60 +424,62 @@ class BufferAssignment { class BufferAssigner { public: // Build and return a BufferAssignment for the given module. The given - // HloOrdering is used to determine buffer liveness. buffer_size is a function - // which returns the size of a LogicalBuffer. Alignment is the the minimum - // alignment of any buffer. If hlos_to_allocate is not null then only - // instructions in this vector are considered for buffer assignment. If - // hlos_to_allocate is null then all instructions are considered. If - // 'colocate_related_buffers' is true, related LogicalBuffers will be - // colocated in the same allocation (i.e buffers for while result will share - // an allocation with buffers related to that same while instruction: init - // operand, condition/body parameter and body result). + // HloOrdering is used to determine buffer liveness. buffer_size and + // color_alignment are functions which returns the size and alignment of a + // LogicalBuffer. allow_input_output_aliasing specifies whether input buffer + // are allowed to be reused as outbut buffers by the client code. static StatusOr> Run( const HloModule* module, std::unique_ptr hlo_ordering, - LogicalBuffer::SizeFunction buffer_size, int64 alignment, - bool colocate_related_buffers, - const std::vector* hlos_to_allocate = nullptr); - - // Overload of Run which uses ShapeUtil::ByteSizeOf to determine buffer size - // and assigns buffers to all HLO instructions in the module. - static StatusOr> Run( - const HloModule* module, std::unique_ptr hlo_ordering, - LogicalBuffer::SizeFunction buffer_size, int64 alignment); + LogicalBuffer::SizeFunction buffer_size, + LogicalBuffer::AlignmentFunction color_alignment, + bool allow_input_output_aliasing = false, + BufferLiveness::Colorer colorer = BufferLiveness::DefaultColorer()); private: - explicit BufferAssigner(LogicalBuffer::SizeFunction buffer_size, - int64 alignment, bool colocate_related_buffers) - : buffer_size_(std::move(buffer_size)), - alignment_(alignment), - colocate_related_buffers_(colocate_related_buffers) {} + BufferAssigner(bool allow_input_output_aliasing, + BufferLiveness::Colorer colorer) + : allow_input_output_aliasing_(allow_input_output_aliasing), + colorer_(colorer) {} virtual ~BufferAssigner() = default; // Create a buffer assignment. StatusOr> CreateAssignment( const HloModule* module, std::unique_ptr hlo_ordering, - const std::vector* hlos_to_allocate = nullptr); + LogicalBuffer::SizeFunction buffer_size, + LogicalBuffer::AlignmentFunction color_alignment); // Assigns buffers to the instructions in the given computation. "assignment" // is modified to reflect the new buffer assignments. If is_thread_local is // true, then all assigned buffers have the is_thread_local flag set to - // true. If hlos_to_allocate is not null it indicates which HLOs to include in - // buffer assignment. If null, all instructions in the computation are - // included. + // true. Status AssignBuffersForComputation( - const HloComputation* computation, bool is_thread_local, - const tensorflow::gtl::FlatSet* hlos_to_allocate, + const HloComputation* computation, const DebugOptions& debug_options, + bool is_thread_local, const tensorflow::gtl::FlatSet& colocated_buffers, const tensorflow::gtl::FlatSet& colocated_allocations, + tensorflow::gtl::FlatMap>* + buffers_to_assign_sequentially, BufferAssignment* assignment); - // Assigns 'buffers_to_assign' assuming the HLO instructions will be executed - // in the given 'sequential_order'. + // Assigns 'buffers_to_assign_sequentially' using heap simulation, assuming + // the HLO instructions will be executed in the sequential order given by + // assignment->liveness().hlo_ordering().SequentialOrder. If + // 'run_whole_module_heap_simulation' is true, the heap simulation will be run + // assuming all global computations are sequentially ordered. Status AssignBuffersWithSequentialOrdering( - const std::vector& sequential_order, - const tensorflow::gtl::FlatSet& buffers_to_assign, - const HloComputation& computation, BufferAssignment* assignment); + const tensorflow::gtl::FlatMap< + const HloComputation*, + tensorflow::gtl::FlatSet>& + buffers_to_assign_sequentially, + bool run_whole_module_heap_simulation, BufferAssignment* assignment); + + // Uses the results of the heap simulator to create a single allocation, with + // LogicalBuffers packed to specific offsets. + void AssignBuffersFromHeapSimulator(const HeapSimulator::Result& result, + BufferAssignment* assignment, + LogicalBuffer::Color color); // Tries to assign the given instruction to the given buffer. Returns if the // assignment was successful. @@ -465,7 +497,8 @@ class BufferAssigner { // ColocatedBufferSet aggregates a set of related LogicalBuffers from 'module' // which should be colocated in the same buffer allocation. void BuildColocatedBufferSets( - const HloModule* module, const TuplePointsToAnalysis& points_to_analysis, + const HloModule* module, const BufferLiveness& buffer_liveness, + const LogicalBuffer::SizeFunction& buffer_size, std::vector* colocated_buffer_sets); // For each buffer set in 'colocated_buffer_sets', assigns all buffers in the @@ -482,16 +515,30 @@ class BufferAssigner { const std::vector& colocated_set, std::vector* colocated_buffer_sets); - const HloModule* module_; + // Conceptually the same as AddSetToColocatedBufferSets, but specific to the + // colocated buffers for while instructions. + void AddWhileSetToColocatedBufferSets( + const std::vector& colocated_set, + const LogicalBuffer* while_init_buffer, + const LogicalBuffer* while_result_buffer, const HloInstruction* while_hlo, + const HloComputation& computation, const BufferLiveness& buffer_liveness, + const LogicalBuffer::SizeFunction& buffer_size, + std::vector* colocated_buffer_sets); - // Function which returns the buffer size for a given logical buffer (shape). - LogicalBuffer::SizeFunction buffer_size_; + // Split a set of buffers into several sets, each of which contains buffers + // colored with the same color. + tensorflow::gtl::FlatMap, + LogicalBuffer::Color::Hasher> + SplitBuffersByColor( + const tensorflow::gtl::FlatSet& buffers); - // Minimum alignment of any buffer. - int64 alignment_; + // If true, buffer assignments assumes that input parameter buffers and output + // buffers can be shared if their sizes match. + bool allow_input_output_aliasing_; - // Indicates whether related buffers should share the same buffer allocation. - const bool colocate_related_buffers_; + // Functor used to assign colors to newly allocated logical buffers. + BufferLiveness::Colorer colorer_; 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 bb7342d5081af32c9882311af8dddf08c115becc..f2e922672b2c579ccbae110ef51fa2f58c4f0732 100644 --- a/tensorflow/compiler/xla/service/buffer_assignment_test.cc +++ b/tensorflow/compiler/xla/service/buffer_assignment_test.cc @@ -18,16 +18,24 @@ limitations under the License. #include #include #include +#include #include #include "tensorflow/compiler/xla/literal_util.h" #include "tensorflow/compiler/xla/ptr_util.h" +#include "tensorflow/compiler/xla/service/call_graph.h" #include "tensorflow/compiler/xla/service/computation_tracker.h" +#include "tensorflow/compiler/xla/service/copy_insertion.h" #include "tensorflow/compiler/xla/service/dfs_hlo_visitor_with_default.h" +#include "tensorflow/compiler/xla/service/flatten_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/hlo_ordering.h" +#include "tensorflow/compiler/xla/service/hlo_scheduling.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/types.h" #include "tensorflow/compiler/xla/xla_data.pb.h" @@ -78,10 +86,18 @@ class BufferAssignmentTest : public HloTestBase { int64 alignment = 1) { return BufferAssigner::Run( module, MakeUnique(module), - [this](const LogicalBuffer& buffer) { - return backend_->compiler()->ShapeSizeBytes(buffer.shape()); - }, - alignment) + backend_->compiler()->BufferSizeBytesFunction(), + [alignment](LogicalBuffer::Color) { return alignment; }) + .ConsumeValueOrDie(); + } + + std::unique_ptr RunColoredBufferAssignment( + HloModule* module, BufferLiveness::Colorer colorer, int64 alignment = 1) { + return BufferAssigner::Run( + module, MakeUnique(module), + backend_->compiler()->BufferSizeBytesFunction(), + [alignment](LogicalBuffer::Color) { return alignment; }, false, + std::move(colorer)) .ConsumeValueOrDie(); } @@ -91,7 +107,7 @@ class BufferAssignmentTest : public HloTestBase { auto param = builder.AddInstruction(HloInstruction::CreateParameter(0, r0f32_, "x")); auto value = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(1.0))); + HloInstruction::CreateConstant(Literal::CreateR0(1.0))); builder.AddInstruction( HloInstruction::CreateBinary(r0f32_, HloOpcode::kAdd, param, value)); return builder.Build(); @@ -108,7 +124,7 @@ class BufferAssignmentTest : public HloTestBase { const string& name) { auto builder = HloComputation::Builder(name); auto const4 = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(4))); + HloInstruction::CreateConstant(Literal::CreateR0(4))); auto param = builder.AddInstruction( HloInstruction::CreateParameter(0, t_s32_f32v4_, "x")); auto index = builder.AddInstruction( @@ -133,9 +149,9 @@ class BufferAssignmentTest : public HloTestBase { const string& name) { auto builder = HloComputation::Builder(name); auto const1 = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(1))); + HloInstruction::CreateConstant(Literal::CreateR0(1))); auto constv = builder.AddInstruction(HloInstruction::CreateConstant( - LiteralUtil::CreateR1({1.1f, 2.2f, 3.3f, 4.4f}))); + Literal::CreateR1({1.1f, 2.2f, 3.3f, 4.4f}))); auto param = builder.AddInstruction( HloInstruction::CreateParameter(0, t_s32_f32v4_, "x")); auto indexc = builder.AddInstruction( @@ -208,30 +224,6 @@ class BufferAssignmentTest : public HloTestBase { return total_size; } - // Returns true if the buffers assigned to instructions in "a" are distinct - // from the buffers assigned to those in "b" (ie, intersection is empty). - bool BuffersDistinct(const std::vector& a, - const std::vector& b, - const BufferAssignment& assignment) { - std::set a_slices; - for (const HloInstruction* instruction : a) { - if (assignment.HasTopLevelAllocation(instruction)) { - a_slices.insert( - assignment.GetUniqueTopLevelSlice(instruction).ConsumeValueOrDie()); - } - } - - for (const HloInstruction* instruction : b) { - if (assignment.HasTopLevelAllocation(instruction)) { - if (a_slices.count(assignment.GetUniqueTopLevelSlice(instruction) - .ConsumeValueOrDie())) { - return false; - } - } - } - return true; - } - // Computation tracker for nested computations. ComputationTracker computation_tracker_; @@ -246,12 +238,36 @@ class BufferAssignmentTest : public HloTestBase { Shape t_s32_f32v10_ = ShapeUtil::MakeTupleShape({s32_, f32vec10_}); }; +// Returns true if the buffers assigned to instructions in "a" are distinct +// from the buffers assigned to those in "b" (ie, intersection is empty). +static bool BuffersDistinct(const std::vector& a, + const std::vector& b, + const BufferAssignment& assignment) { + std::set a_slices; + for (const HloInstruction* instruction : a) { + if (assignment.HasTopLevelAllocation(instruction)) { + a_slices.insert( + assignment.GetUniqueTopLevelSlice(instruction).ConsumeValueOrDie()); + } + } + + for (const HloInstruction* instruction : b) { + if (assignment.HasTopLevelAllocation(instruction)) { + if (a_slices.count(assignment.GetUniqueTopLevelSlice(instruction) + .ConsumeValueOrDie())) { + return false; + } + } + } + return true; +} + // Tests a computation consisting of a single scalar constant node. TEST_F(BufferAssignmentTest, ScalarConstant) { auto builder = HloComputation::Builder(TestName()); auto const0 = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(1.0))); - auto module = MakeUnique(TestName()); + HloInstruction::CreateConstant(Literal::CreateR0(1.0))); + auto module = CreateNewModule(); module->AddEntryComputation(builder.Build()); auto buffers = RunBufferAssignment(module.get()); @@ -264,12 +280,12 @@ TEST_F(BufferAssignmentTest, BufferForConst) { // no buffers assigned, and their consumer has a buffer. auto builder = HloComputation::Builder(TestName()); auto const0 = builder.AddInstruction(HloInstruction::CreateConstant( - LiteralUtil::CreateR1({1.1f, 2.2f, 3.3f, 4.4f}))); + Literal::CreateR1({1.1f, 2.2f, 3.3f, 4.4f}))); auto const1 = builder.AddInstruction(HloInstruction::CreateConstant( - LiteralUtil::CreateR1({4.1f, 4.2f, 4.3f, 4.4f}))); + Literal::CreateR1({4.1f, 4.2f, 4.3f, 4.4f}))); auto add = builder.AddInstruction( HloInstruction::CreateBinary(f32vec4_, HloOpcode::kAdd, const0, const1)); - auto module = MakeUnique(TestName()); + auto module = CreateNewModule(); module->AddEntryComputation(builder.Build()); auto buffers = RunBufferAssignment(module.get()); @@ -280,14 +296,42 @@ TEST_F(BufferAssignmentTest, BufferForConst) { GetAssignedOutputAllocation(*buffers, add); } +TEST_F(BufferAssignmentTest, HasAllocationAt) { + // Create a tuple with non-const and const elements and check that + // HasAllocationAt works correctly. + auto builder = HloComputation::Builder(TestName()); + auto param0 = builder.AddInstruction( + HloInstruction::CreateParameter(0, f32vec100_, "param0")); + auto constant = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(1))); + auto negate = builder.AddInstruction( + HloInstruction::CreateUnary(f32vec100_, HloOpcode::kNegate, param0)); + auto tuple = builder.AddInstruction( + HloInstruction::CreateTuple({negate, param0, constant})); + auto module = CreateNewModule(); + module->AddEntryComputation(builder.Build()); + + auto buffers = RunBufferAssignment(module.get()); + // Make sure that HasAllocationAt() agrees with what HasTopLevelAllocation() + // reports for the instruction directly. + EXPECT_EQ(buffers->HasTopLevelAllocation(tuple), + buffers->HasAllocationAt(tuple, /*index=*/{})); + EXPECT_EQ(buffers->HasTopLevelAllocation(negate), + buffers->HasAllocationAt(tuple, /*index=*/{0})); + EXPECT_EQ(buffers->HasTopLevelAllocation(param0), + buffers->HasAllocationAt(tuple, /*index=*/{1})); + EXPECT_EQ(buffers->HasTopLevelAllocation(constant), + buffers->HasAllocationAt(tuple, /*index=*/{2})); +} + TEST_F(BufferAssignmentTest, BufferForOutputConst) { // This computation copies a constant to output. auto builder = HloComputation::Builder(TestName()); auto const0 = builder.AddInstruction(HloInstruction::CreateConstant( - LiteralUtil::CreateR1({1.1f, 2.2f, 3.3f, 4.4f}))); + Literal::CreateR1({1.1f, 2.2f, 3.3f, 4.4f}))); auto copy = builder.AddInstruction( HloInstruction::CreateUnary(const0->shape(), HloOpcode::kCopy, const0)); - auto module = MakeUnique(TestName()); + auto module = CreateNewModule(); module->AddEntryComputation(builder.Build()); auto buffers = RunBufferAssignment(module.get()); @@ -314,7 +358,7 @@ TEST_F(BufferAssignmentTest, Basic) { HloInstruction::CreateBinary(f32vec100_, HloOpcode::kAdd, mul, param1)); auto sub = builder.AddInstruction(HloInstruction::CreateBinary( f32vec100_, HloOpcode::kSubtract, add, param1)); - auto module = MakeUnique(TestName()); + auto module = CreateNewModule(); module->AddEntryComputation(builder.Build()); auto buffers = RunBufferAssignment(module.get()); @@ -334,7 +378,133 @@ TEST_F(BufferAssignmentTest, Basic) { // The add node can reuse the mul node's buffer. const BufferAllocation& add_buffer = GetTopLevelAllocation(*buffers, add); - EXPECT_EQ(add_buffer.index(), add_buffer.index()); + EXPECT_EQ(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) + // / / / + // param0[100] -------/ / / + // / / + // param1[100] --------------/--------/ + // The output of each op is colored with a different color, so we can not + // share anything. + auto builder = HloComputation::Builder(TestName()); + auto paramscalar = + builder.AddInstruction(HloInstruction::CreateParameter(0, r0f32_, "")); + auto param0 = builder.AddInstruction( + HloInstruction::CreateParameter(1, f32vec100_, "")); + auto param1 = builder.AddInstruction( + HloInstruction::CreateParameter(2, f32vec100_, "")); + auto mul = builder.AddInstruction(HloInstruction::CreateBinary( + f32vec100_, HloOpcode::kMultiply, paramscalar, 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 = CreateNewModule(); + module->AddEntryComputation(builder.Build()); + + auto colorer = [](const BufferLiveness& buffer_liveness) { + int color = 0; + + for (LogicalBuffer::Id id = 0; + id < buffer_liveness.points_to_analysis().num_logical_buffers(); + id++) { + auto& buffer = buffer_liveness.points_to_analysis().logical_buffer(id); + buffer.set_color(LogicalBuffer::Color(color++)); + } + return Status::OK(); + }; + + auto buffers = RunColoredBufferAssignment(module.get(), colorer); + + // 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 can not reuse the mul node's buffer due to coloring. + 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, BasicPartiallyColored) { + // paramscalar ------- (mul) -- (add) -- (sub) + // / / / + // param0[100] -------/ / / + // / / + // param1[100] --------------/--------/ + // The output of the mul and the add have the color 1, and the other buffers + // have the color 0, which allows the mul and add to share buffers. + auto builder = HloComputation::Builder(TestName()); + auto paramscalar = + builder.AddInstruction(HloInstruction::CreateParameter(0, r0f32_, "")); + auto param0 = builder.AddInstruction( + HloInstruction::CreateParameter(1, f32vec100_, "")); + auto param1 = builder.AddInstruction( + HloInstruction::CreateParameter(2, f32vec100_, "")); + auto mul = builder.AddInstruction(HloInstruction::CreateBinary( + f32vec100_, HloOpcode::kMultiply, paramscalar, 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 = CreateNewModule(); + module->AddEntryComputation(builder.Build()); + + auto colorer = [](const BufferLiveness& buffer_liveness) { + for (LogicalBuffer::Id id = 0; + id < buffer_liveness.points_to_analysis().num_logical_buffers(); + id++) { + auto& buffer = buffer_liveness.points_to_analysis().logical_buffer(id); + const auto& aliases = + buffer_liveness.points_to_analysis().GetBufferAliases(buffer); + for (const auto& alias : aliases) { + if (alias.instruction()->opcode() == HloOpcode::kAdd || + alias.instruction()->opcode() == HloOpcode::kMultiply) { + buffer.set_color(LogicalBuffer::Color(1)); + } + } + if (!buffer.has_color()) { + buffer.set_color(LogicalBuffer::Color(0)); + } + } + return Status::OK(); + }; + + auto buffers = RunColoredBufferAssignment(module.get(), colorer); + + // 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 can reuse the mul node's buffer. + const BufferAllocation& add_buffer = GetTopLevelAllocation(*buffers, add); + EXPECT_EQ(add_buffer.index(), mul_buffer.index()); // The sub node has a valid output buffer assigned. GetAssignedOutputAllocation(*buffers, sub); @@ -364,7 +534,7 @@ TEST_F(BufferAssignmentTest, MultipleUsersForNode) { HloInstruction::CreateBinary(f32vec100_, HloOpcode::kAdd, mul, param1)); auto sub = builder.AddInstruction( HloInstruction::CreateBinary(f32vec100_, HloOpcode::kSubtract, add, mul)); - auto module = MakeUnique(TestName()); + auto module = CreateNewModule(); module->AddEntryComputation(builder.Build()); auto buffers = RunBufferAssignment(module.get()); @@ -399,7 +569,7 @@ TEST_F(BufferAssignmentTest, TrivialMap) { // param0[100x10] ---> (map x+1) // // Builds the map function. - auto module = MakeUnique(TestName()); + auto module = CreateNewModule(); auto map_computation = module->AddEmbeddedComputation(BuildMapComputationPlus1("f32+1")); auto inner_last = map_computation->root_instruction(); @@ -410,13 +580,13 @@ TEST_F(BufferAssignmentTest, TrivialMap) { HloInstruction::CreateParameter(0, f32a100x10_, "")); auto map = builder.AddInstruction( HloInstruction::CreateMap(f32a100x10_, {param0}, map_computation)); + module->AddEntryComputation(builder.Build()); + const std::vector level0 = GetInstructions(map); EXPECT_EQ(2, level0.size()) << "Invalid main kernel size"; const std::vector level1 = GetInstructions(inner_last); EXPECT_EQ(3, level1.size()) << "Invalid nested add+1 size"; - module->AddEntryComputation(builder.Build()); - // Assigns buffers and fetches sizes. auto buffers = RunBufferAssignment(module.get()); int64 size0 = ValidateBuffers(level0, *buffers); @@ -454,7 +624,7 @@ TEST_F(BufferAssignmentTest, CannotReuseInputBufferOfReduce) { // out-of-order reductions could overwrite an element before a use.) // // param0[100] --- (exp1) --- (exp2) --- (reduce x+1) --- (exp3) - auto module = MakeUnique(TestName()); + auto module = CreateNewModule(); auto reduce_computation = module->AddEmbeddedComputation(BuildMapComputationPlus1("f32+1")); @@ -466,7 +636,7 @@ TEST_F(BufferAssignmentTest, CannotReuseInputBufferOfReduce) { auto exp2 = builder.AddInstruction( HloInstruction::CreateUnary(f32a100x10_, HloOpcode::kExp, exp1)); auto const0 = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(0.0f))); + HloInstruction::CreateConstant(Literal::CreateR0(0.0f))); auto reduce = builder.AddInstruction(HloInstruction::CreateReduce( /*shape=*/f32vec10_, /*operand=*/exp2, @@ -505,7 +675,7 @@ TEST_F(BufferAssignmentTest, ExampleWhile) { // const4[f32[4]] --- tuple --- while[condition, body] // // Builds the nested condition and body. - auto module = MakeUnique(TestName()); + auto module = CreateNewModule(); auto condition_computation = module->AddEmbeddedComputation(BuildWhileConditionComputation("if<4")); auto body_computation = @@ -514,13 +684,14 @@ TEST_F(BufferAssignmentTest, ExampleWhile) { // Creates the main kernel and verifies instruction counts. auto builder = HloComputation::Builder(TestName()); auto const3 = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(0))); + HloInstruction::CreateConstant(Literal::CreateR0(0))); auto const4 = builder.AddInstruction(HloInstruction::CreateConstant( - LiteralUtil::CreateR1({1.1f, 2.2f, 3.3f, 4.4f}))); + Literal::CreateR1({1.1f, 2.2f, 3.3f, 4.4f}))); auto tuple = builder.AddInstruction(HloInstruction::CreateTuple({const3, const4})); auto while_op = builder.AddInstruction(HloInstruction::CreateWhile( t_s32_f32v4_, condition_computation, body_computation, tuple)); + module->AddEntryComputation(builder.Build()); const std::vector level0 = GetInstructions(while_op); EXPECT_EQ(4, level0.size()) << "Invalid while kernel size"; @@ -531,8 +702,6 @@ TEST_F(BufferAssignmentTest, ExampleWhile) { GetInstructions(body_computation->root_instruction()); EXPECT_EQ(8, levelb.size()) << "Invalid nested body size"; - module->AddEntryComputation(builder.Build()); - // Assigns buffers and fetches sizes. auto buffers = RunBufferAssignment(module.get()); int64 size0 = ValidateBuffers(level0, *buffers); @@ -556,15 +725,14 @@ TEST_F(BufferAssignmentTest, ExampleWhile) { // Check that buffer for each subshape of 'while_op' shares allocation with // corresponding buffer from while body computation at same index. - TF_CHECK_OK(ShapeUtil::ForEachSubshape( + ShapeUtil::ForEachSubshape( while_op->shape(), [this, &buffers, while_op, body_root](const Shape& /*subshape*/, const ShapeIndex& index) { auto while_op_allocation = GetAllocation(*buffers, while_op, index); auto body_root_allocation = GetAllocation(*buffers, body_root, index); EXPECT_EQ(while_op_allocation.index(), body_root_allocation.index()); - return Status::OK(); - })); + }); // Log size information for inspection. LOG(INFO) << "LogicalBuffer count " << buffers->Allocations().size() @@ -586,7 +754,7 @@ TEST_F(BufferAssignmentTest, UnaryOpReuseChain) { auto neg = builder.AddInstruction( HloInstruction::CreateUnary(f32vec100_, HloOpcode::kNegate, exp2)); - auto module = MakeUnique(TestName()); + auto module = CreateNewModule(); module->AddEntryComputation(builder.Build()); auto assignment = RunBufferAssignment(module.get()); @@ -611,11 +779,11 @@ TEST_F(BufferAssignmentTest, ReuseNonOperandBuffer) { auto negate = builder.AddInstruction( HloInstruction::CreateUnary(f32vec100_, HloOpcode::kNegate, param0)); auto slice = builder.AddInstruction( - HloInstruction::CreateSlice(f32vec10_, negate, {0}, {10})); + HloInstruction::CreateSlice(f32vec10_, negate, {0}, {10}, {1})); auto broadcast = builder.AddInstruction( HloInstruction::CreateBroadcast(f32a100x10_, slice, {1})); - auto module = MakeUnique(TestName()); + auto module = CreateNewModule(); module->AddEntryComputation(builder.Build()); auto assignment = RunBufferAssignment(module.get()); @@ -643,12 +811,12 @@ TEST_F(BufferAssignmentTest, NoReuseLiveBuffer) { auto negate = builder.AddInstruction( HloInstruction::CreateUnary(f32vec100_, HloOpcode::kNegate, param0)); auto slice = builder.AddInstruction( - HloInstruction::CreateSlice(f32vec10_, negate, {0}, {10})); + HloInstruction::CreateSlice(f32vec10_, negate, {0}, {10}, {1})); auto broadcast = builder.AddInstruction( HloInstruction::CreateBroadcast(f32a100x10_, slice, {1})); builder.AddInstruction(HloInstruction::CreateTuple({negate, broadcast})); - auto module = MakeUnique(TestName()); + auto module = CreateNewModule(); module->AddEntryComputation(builder.Build()); auto assignment = RunBufferAssignment(module.get()); @@ -680,12 +848,12 @@ TEST_F(BufferAssignmentTest, NoReuseAliasedBuffer) { auto tuple_element = builder.AddInstruction( HloInstruction::CreateGetTupleElement(f32vec100_, tuple, 0)); auto slice = builder.AddInstruction( - HloInstruction::CreateSlice(f32vec10_, tuple_element, {0}, {10})); + HloInstruction::CreateSlice(f32vec10_, tuple_element, {0}, {10}, {1})); auto broadcast = builder.AddInstruction( HloInstruction::CreateBroadcast(f32a100x10_, slice, {1})); builder.AddInstruction(HloInstruction::CreateTuple({tuple, broadcast})); - auto module = MakeUnique(TestName()); + auto module = CreateNewModule(); module->AddEntryComputation(builder.Build()); auto assignment = RunBufferAssignment(module.get()); @@ -715,12 +883,12 @@ TEST_F(BufferAssignmentTest, DoNotReuseOversizedOutputBuffer) { HloInstruction::CreateUnary(f32vec100_, HloOpcode::kNegate, param0)); // Slice output is 10 elements. auto slice = builder.AddInstruction( - HloInstruction::CreateSlice(f32vec10_, negate, {0}, {10})); + HloInstruction::CreateSlice(f32vec10_, negate, {0}, {10}, {1})); // Broadcast output is 40 elements. auto broadcast = builder.AddInstruction(HloInstruction::CreateBroadcast( ShapeUtil::MakeShape(F32, {10, 4}), slice, {0})); - auto module = MakeUnique(TestName()); + auto module = CreateNewModule(); module->AddEntryComputation(builder.Build()); auto assignment = RunBufferAssignment(module.get()); @@ -747,12 +915,12 @@ TEST_F(BufferAssignmentTest, ReuseOutputBufferIfExactlySized) { auto negate = builder.AddInstruction( HloInstruction::CreateUnary(f32vec100_, HloOpcode::kNegate, param0)); auto slice = builder.AddInstruction( - HloInstruction::CreateSlice(f32vec10_, negate, {0}, {10})); + HloInstruction::CreateSlice(f32vec10_, negate, {0}, {10}, {1})); // Broadcast output is 40 elements. auto broadcast = builder.AddInstruction(HloInstruction::CreateBroadcast( ShapeUtil::MakeShape(F32, {10, 10}), slice, {0})); - auto module = MakeUnique(TestName()); + auto module = CreateNewModule(); module->AddEntryComputation(builder.Build()); auto assignment = RunBufferAssignment(module.get()); @@ -784,13 +952,13 @@ TEST_F(BufferAssignmentTest, DoNotReuseOversizedOutputBufferInTuple) { HloInstruction::CreateUnary(f32vec100_, HloOpcode::kNegate, param0)); // Slice output is 10 elements. auto slice = builder.AddInstruction( - HloInstruction::CreateSlice(f32vec10_, negate, {0}, {10})); + HloInstruction::CreateSlice(f32vec10_, negate, {0}, {10}, {1})); // Broadcast output is 40 elements. auto broadcast = builder.AddInstruction(HloInstruction::CreateBroadcast( ShapeUtil::MakeShape(F32, {10, 4}), slice, {0})); builder.AddInstruction(HloInstruction::CreateTuple({broadcast})); - auto module = MakeUnique(TestName()); + auto module = CreateNewModule(); module->AddEntryComputation(builder.Build()); auto assignment = RunBufferAssignment(module.get()); @@ -805,7 +973,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 = MakeUnique(TestName()); + auto module = CreateNewModule(); auto vec_shape = ShapeUtil::MakeShape(F32, {42}); auto scalar_shape = ShapeUtil::MakeShape(F32, {}); @@ -850,8 +1018,7 @@ TEST_F(BufferAssignmentTest, EmbeddedComputationBuffers) { EXPECT_FALSE(map_root_alloc.maybe_live_out()); EXPECT_TRUE(map_root_alloc.is_thread_local()); - // Allocations for the call computation should not be thread-local and not - // live-out. + // Allocations for the call computation should not be thread-local. auto& call_param_alloc = GetTopLevelAllocation(*assignment, call_param); EXPECT_FALSE(call_param_alloc.is_entry_computation_parameter()); EXPECT_FALSE(call_param_alloc.maybe_live_out()); @@ -859,7 +1026,6 @@ TEST_F(BufferAssignmentTest, EmbeddedComputationBuffers) { auto& call_root_alloc = GetTopLevelAllocation(*assignment, call_root); EXPECT_FALSE(call_root_alloc.is_entry_computation_parameter()); - EXPECT_FALSE(call_root_alloc.maybe_live_out()); EXPECT_FALSE(call_root_alloc.is_thread_local()); // Entry computation allocations can be marked liveout and @@ -879,12 +1045,13 @@ TEST_F(BufferAssignmentTest, TupleParameterAsOutput) { // Test a computation that returns a tuple parameter. auto builder = HloComputation::Builder(TestName()); auto tuple_param = builder.AddInstruction(HloInstruction::CreateParameter( - 0, ShapeUtil::MakeTupleShape({ShapeUtil::MakeShape(PRED, {1, 2, 3, 4}), - ShapeUtil::MakeShape(F32, {}), - ShapeUtil::MakeShape(S32, {42})}), + 0, + ShapeUtil::MakeTupleShape({ShapeUtil::MakeShape(PRED, {1, 2, 3, 4}), + ShapeUtil::MakeShape(F32, {}), + ShapeUtil::MakeShape(S32, {42})}), "param0")); - auto module = MakeUnique(TestName()); + auto module = CreateNewModule(); module->AddEntryComputation(builder.Build()); auto assignment = RunBufferAssignment(module.get()); @@ -894,7 +1061,7 @@ TEST_F(BufferAssignmentTest, TupleParameterAsOutput) { // Verify each buffer allocation is marked as an entry computation parameter // and is liveout. - TF_CHECK_OK(ShapeUtil::ForEachSubshape( + ShapeUtil::ForEachSubshape( tuple_param->shape(), [this, &assignment, tuple_param](const Shape& /*subshape*/, const ShapeIndex& index) { @@ -902,8 +1069,7 @@ TEST_F(BufferAssignmentTest, TupleParameterAsOutput) { EXPECT_TRUE(allocation.is_entry_computation_parameter()); EXPECT_EQ(0, allocation.parameter_number()); EXPECT_TRUE(allocation.maybe_live_out()); - return Status::OK(); - })); + }); } TEST_F(BufferAssignmentTest, ElementOfNestedTupleParameterAsOutput) { @@ -911,16 +1077,17 @@ TEST_F(BufferAssignmentTest, ElementOfNestedTupleParameterAsOutput) { // parameter. auto builder = HloComputation::Builder(TestName()); auto tuple_param = builder.AddInstruction(HloInstruction::CreateParameter( - 0, ShapeUtil::MakeTupleShape( - {ShapeUtil::MakeShape(PRED, {1, 2, 3, 4}), - ShapeUtil::MakeTupleShape({ShapeUtil::MakeShape(S32, {42}), - ShapeUtil::MakeShape(S32, {101})})}), + 0, + ShapeUtil::MakeTupleShape( + {ShapeUtil::MakeShape(PRED, {1, 2, 3, 4}), + ShapeUtil::MakeTupleShape({ShapeUtil::MakeShape(S32, {42}), + ShapeUtil::MakeShape(S32, {101})})}), "param0")); auto tuple_element = builder.AddInstruction(HloInstruction::CreateGetTupleElement( ShapeUtil::GetSubshape(tuple_param->shape(), {1}), tuple_param, 1)); - auto module = MakeUnique(TestName()); + auto module = CreateNewModule(); module->AddEntryComputation(builder.Build()); auto assignment = RunBufferAssignment(module.get()); @@ -956,14 +1123,13 @@ TEST_F(BufferAssignmentTest, ElementOfNestedTupleParameterAsOutput) { // TODO(b/32248867): Enable when buffer assignment gives allocations to // constants. TEST_F(BufferAssignmentTest, DISABLED_TupleConstantAsOutput) { - // Test that a tuple constant which is forwarded to the computation output is - // properly handled. + // Test that a tuple constant which is forwarded to the computation output + // is properly handled. auto builder = HloComputation::Builder(TestName()); - builder.AddInstruction(HloInstruction::CreateConstant( - LiteralUtil::MakeTuple({LiteralUtil::CreateR0(0).get(), - LiteralUtil::CreateR0(1).get()}))); + builder.AddInstruction(HloInstruction::CreateConstant(Literal::MakeTuple( + {Literal::CreateR0(0).get(), Literal::CreateR0(1).get()}))); - auto module = MakeUnique(TestName()); + auto module = CreateNewModule(); module->AddEntryComputation(builder.Build()); auto assignment = RunBufferAssignment(module.get()); @@ -977,7 +1143,7 @@ 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 = MakeUnique(TestName()); + auto module = CreateNewModule(); module->AddEntryComputation(builder.Build()); auto assignment = RunBufferAssignment(module.get()); @@ -992,7 +1158,7 @@ TEST_F(BufferAssignmentTest, TupleCustomCallAsOutput) { TEST_F(BufferAssignmentTest, TupleCallAsOutput) { // Test a computation which returns a tuple call value. - auto module = MakeUnique(TestName()); + auto module = CreateNewModule(); auto elem_shape = f32vec4_; auto tuple_shape = ShapeUtil::MakeTupleShape({elem_shape}); @@ -1031,7 +1197,7 @@ TEST_F(BufferAssignmentTest, TupleChainedCallAsOutput) { // B: call(C, param) // C: call(D, param) // D: param - auto module = MakeUnique(TestName()); + auto module = CreateNewModule(); auto elem_shape = f32vec4_; auto tuple_shape = ShapeUtil::MakeTupleShape({elem_shape}); @@ -1102,7 +1268,7 @@ TEST_F(BufferAssignmentTest, BitcastAsOutput) { auto bitcast = builder.AddInstruction( HloInstruction::CreateUnary(param->shape(), HloOpcode::kBitcast, param)); - auto module = MakeUnique(TestName()); + auto module = CreateNewModule(); module->AddEntryComputation(builder.Build()); auto assignment = RunBufferAssignment(module.get()); @@ -1113,8 +1279,8 @@ TEST_F(BufferAssignmentTest, BitcastAsOutput) { } TEST_F(BufferAssignmentTest, AmbiguousBufferAsOutput) { - // Test a computation with an output that has an ambiguous points-to set. This - // is constructed using a select among tuple shapes. + // Test a computation with an output that has an ambiguous points-to set. + // This is constructed using a select among tuple shapes. auto builder = HloComputation::Builder(TestName()); auto tuple_shape = ShapeUtil::MakeTupleShape({ShapeUtil::MakeShape(PRED, {1, 2, 3, 4})}); @@ -1128,7 +1294,7 @@ TEST_F(BufferAssignmentTest, AmbiguousBufferAsOutput) { auto select = builder.AddInstruction(HloInstruction::CreateTernary( tuple_shape, HloOpcode::kSelect, pred_param, tuple_param0, tuple_param1)); - auto module = MakeUnique(TestName()); + auto module = CreateNewModule(); module->AddEntryComputation(builder.Build()); auto assignment = RunBufferAssignment(module.get()); @@ -1144,12 +1310,12 @@ TEST_F(BufferAssignmentTest, AmbiguousBufferAsOutput) { // should include the slices of both of the elements in the parameters. auto element_slices = assignment->GetAllSlices(select, /*index=*/{0}); EXPECT_EQ(2, element_slices.size()); - EXPECT_MATCH(testing::SetToVec(element_slices), - testing::UnorderedMatcher( - assignment->GetUniqueSlice(tuple_param0, /*index=*/{0}) - .ConsumeValueOrDie(), - assignment->GetUniqueSlice(tuple_param1, /*index=*/{0}) - .ConsumeValueOrDie())); + EXPECT_THAT(element_slices, + ::testing::UnorderedElementsAre( + assignment->GetUniqueSlice(tuple_param0, /*index=*/{0}) + .ConsumeValueOrDie(), + assignment->GetUniqueSlice(tuple_param1, /*index=*/{0}) + .ConsumeValueOrDie())); } // TODO(b/34669761): Remove this test when buffers are allowed to share @@ -1166,7 +1332,7 @@ TEST_F(BufferAssignmentTest, TupleBufferNotReused) { auto copy = builder.AddInstruction(HloInstruction::CreateUnary( scalar_shape, HloOpcode::kCopy, tuple_element)); - auto module = MakeUnique(TestName()); + auto module = CreateNewModule(); module->AddEntryComputation(builder.Build()); auto assignment = RunBufferAssignment(module.get()); @@ -1178,8 +1344,8 @@ TEST_F(BufferAssignmentTest, TupleBufferNotReused) { } TEST_F(BufferAssignmentTest, OneTempAllocation) { - // Test a computation that requires multiple temp buffers, and ensure they are - // combined into a single allocation. + // Test a computation that requires multiple temp buffers, and ensure they + // are combined into a single allocation. auto builder = HloComputation::Builder(TestName()); Shape shape_2x3 = ShapeUtil::MakeShape(F32, {2, 3}); Shape shape_2x4 = ShapeUtil::MakeShape(F32, {2, 4}); @@ -1202,7 +1368,7 @@ TEST_F(BufferAssignmentTest, OneTempAllocation) { HloInstruction::CreateConcatenate(shape_5x4, {dot_ab, dot_bc}, 1)); // Run buffer assignment with alignment=1. - auto module = MakeUnique(TestName()); + auto module = CreateNewModule(); module->AddEntryComputation(builder.Build()); auto assignment = RunBufferAssignment(module.get(), /*alignment=*/1); @@ -1245,6 +1411,362 @@ TEST_F(BufferAssignmentTest, OneTempAllocation) { } } -} // namespace +class WhileBufferAssignmentTest : public HloTestBase { + protected: + std::unique_ptr BuildWhileConditionComputation( + const string& name) { + auto builder = HloComputation::Builder(name); + builder.AddInstruction( + HloInstruction::CreateParameter(0, loop_state_shape_, "loop_state")); + auto zero = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(0))); + auto ten = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(10))); + builder.AddInstruction(HloInstruction::CreateBinary( + ShapeUtil::MakeShape(PRED, {}), HloOpcode::kLt, zero, ten)); + return builder.Build(); + } + + std::unique_ptr BuildWhileBodyComputation( + const string& name) { + auto builder = HloComputation::Builder(name); + auto loop_state = builder.AddInstruction( + HloInstruction::CreateParameter(0, loop_state_shape_, "loop_state")); + auto input = builder.AddInstruction( + HloInstruction::CreateGetTupleElement(data_shape_, loop_state, 0)); + auto weights = builder.AddInstruction( + HloInstruction::CreateGetTupleElement(data_shape_, loop_state, 1)); + auto output = builder.AddInstruction(HloInstruction::CreateBinary( + data_shape_, HloOpcode::kMultiply, input, weights)); + builder.AddInstruction( + HloInstruction::CreateTuple({input, weights, output})); + return builder.Build(); + } + + std::unique_ptr RunBufferAssignment(HloModule* module, + int64 alignment = 1) { + auto sequence = + CreateMemoryMinimizingSequence(*module, ByteSizeOf).ConsumeValueOrDie(); + return BufferAssigner::Run( + module, MakeUnique(module, sequence), + ByteSizeOf, + [alignment](LogicalBuffer::Color) { return alignment; }) + .ConsumeValueOrDie(); + } + + static int64 ByteSizeOf(const LogicalBuffer& buffer) { + return ShapeUtil::ByteSizeOf(buffer.shape(), sizeof(void*)); + } + + Shape data_shape_ = ShapeUtil::MakeShape(F32, {4}); + Shape loop_state_shape_ = + ShapeUtil::MakeTupleShape({data_shape_, data_shape_, data_shape_}); +}; + +static void RunCopyInsertion(HloModule* module) { + CopyInsertion copy_insertion; + EXPECT_IS_OK(copy_insertion.Run(module).status()); +} +TEST_F(WhileBufferAssignmentTest, TwoForwardWhileLoops) { + auto module = MakeUnique(TestName()); + auto builder = HloComputation::Builder("entry"); + + auto input0 = builder.AddInstruction( + HloInstruction::CreateParameter(0, data_shape_, "input0")); + auto weights0 = builder.AddInstruction( + HloInstruction::CreateParameter(1, data_shape_, "weights0")); + auto weights1 = builder.AddInstruction( + HloInstruction::CreateParameter(2, data_shape_, "weights1")); + + auto zero = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(0.0))); + auto output0 = builder.AddInstruction( + HloInstruction::CreateBroadcast(data_shape_, zero, {1})); + auto output1 = builder.AddInstruction( + HloInstruction::CreateBroadcast(data_shape_, zero, {1})); + + auto cond0 = + module->AddEmbeddedComputation(BuildWhileConditionComputation("cond")); + auto body0 = + module->AddEmbeddedComputation(BuildWhileBodyComputation("body")); + + auto tuple0 = builder.AddInstruction( + HloInstruction::CreateTuple({input0, weights0, output0})); + auto while0 = builder.AddInstruction( + HloInstruction::CreateWhile(loop_state_shape_, cond0, body0, tuple0)); + + auto cond1 = + module->AddEmbeddedComputation(BuildWhileConditionComputation("cond")); + auto body1 = + module->AddEmbeddedComputation(BuildWhileBodyComputation("body")); + auto input1 = builder.AddInstruction( + HloInstruction::CreateGetTupleElement(data_shape_, while0, 2)); + auto tuple1 = builder.AddInstruction( + HloInstruction::CreateTuple({input1, weights1, output1})); + auto while1 = builder.AddInstruction( + HloInstruction::CreateWhile(loop_state_shape_, cond1, body1, tuple1)); + + module->AddEntryComputation(builder.Build()); + RunCopyInsertion(module.get()); + auto assignment = RunBufferAssignment(module.get()); + + // Verify 'input0' and read-only use while0{0} alias. + EXPECT_EQ(assignment->GetUniqueSlice(input0, {}).ConsumeValueOrDie(), + assignment->GetUniqueSlice(while0, {0}).ConsumeValueOrDie()); + // Verify 'weights0' and read-only use while0{1} alias. + EXPECT_EQ(assignment->GetUniqueSlice(weights0, {}).ConsumeValueOrDie(), + assignment->GetUniqueSlice(while0, {1}).ConsumeValueOrDie()); + // Verify 'while0{2}' and read-only use while1{0} alias. + EXPECT_EQ(assignment->GetUniqueSlice(while0, {2}).ConsumeValueOrDie(), + assignment->GetUniqueSlice(while1, {0}).ConsumeValueOrDie()); + // Verify 'weights1' and read-only use while1{1} alias. + EXPECT_EQ(assignment->GetUniqueSlice(weights1, {}).ConsumeValueOrDie(), + assignment->GetUniqueSlice(while1, {1}).ConsumeValueOrDie()); +} + +TEST_F(WhileBufferAssignmentTest, OneForwardBackwardWhileLoopSet) { + auto module = MakeUnique(TestName()); + auto builder = HloComputation::Builder("entry"); + + auto input0 = builder.AddInstruction( + HloInstruction::CreateParameter(0, data_shape_, "input0")); + auto weights0 = builder.AddInstruction( + HloInstruction::CreateParameter(1, data_shape_, "weights0")); + + auto zero = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(0.0))); + auto output0 = builder.AddInstruction( + HloInstruction::CreateBroadcast(data_shape_, zero, {1})); + auto output1 = builder.AddInstruction( + HloInstruction::CreateBroadcast(data_shape_, zero, {1})); + + auto cond0 = + module->AddEmbeddedComputation(BuildWhileConditionComputation("cond")); + auto body0 = + module->AddEmbeddedComputation(BuildWhileBodyComputation("body")); + + auto tuple0 = builder.AddInstruction( + HloInstruction::CreateTuple({input0, weights0, output0})); + auto while0 = builder.AddInstruction( + HloInstruction::CreateWhile(loop_state_shape_, cond0, body0, tuple0)); + + auto cond1 = + module->AddEmbeddedComputation(BuildWhileConditionComputation("cond")); + auto body1 = + module->AddEmbeddedComputation(BuildWhileBodyComputation("body")); + + auto tuple1 = builder.AddInstruction( + HloInstruction::CreateTuple({input0, weights0, output1})); + auto while1 = builder.AddInstruction( + HloInstruction::CreateWhile(loop_state_shape_, cond1, body1, tuple1)); + + module->AddEntryComputation(builder.Build()); + RunCopyInsertion(module.get()); + auto assignment = RunBufferAssignment(module.get()); + + // while0 and while1 buffers should be completely aligned. + EXPECT_EQ(assignment->GetUniqueSlice(while0, {0}).ConsumeValueOrDie(), + assignment->GetUniqueSlice(while1, {0}).ConsumeValueOrDie()); + EXPECT_EQ(assignment->GetUniqueSlice(while0, {1}).ConsumeValueOrDie(), + assignment->GetUniqueSlice(while1, {1}).ConsumeValueOrDie()); + EXPECT_EQ(assignment->GetUniqueSlice(while0, {2}).ConsumeValueOrDie(), + assignment->GetUniqueSlice(while1, {2}).ConsumeValueOrDie()); +} + +TEST_F(BufferAssignmentTest, TwoCalls) { + auto module = MakeUnique(TestName()); + Shape r0f32 = ShapeUtil::MakeShape(xla::F32, {}); + HloComputation* sub_computation; + { + auto builder = HloComputation::Builder(TestName() + "_sub_comp"); + auto param = builder.AddInstruction( + HloInstruction::CreateParameter(0, r0f32, "param")); + auto constant1 = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(1.0))); + auto add = builder.AddInstruction( + HloInstruction::CreateBinary(r0f32, HloOpcode::kAdd, param, constant1)); + sub_computation = module->AddEmbeddedComputation(builder.Build(add)); + } + auto builder = HloComputation::Builder(TestName()); + auto constant2 = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(2.0))); + auto constant3 = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(3.0))); + auto call1 = builder.AddInstruction( + HloInstruction::CreateCall(r0f32, {constant2}, sub_computation)); + auto call2 = builder.AddInstruction( + HloInstruction::CreateCall(r0f32, {constant3}, sub_computation)); + auto add1 = builder.AddInstruction( + HloInstruction::CreateBinary(r0f32, HloOpcode::kAdd, call1, constant2)); + auto add2 = builder.AddInstruction( + HloInstruction::CreateBinary(r0f32, HloOpcode::kAdd, call2, add1)); + module->AddEntryComputation(builder.Build(add2)); + + { + FlattenCallGraph flatten; + TF_ASSERT_OK_AND_ASSIGN(bool result, flatten.Run(module.get())); + EXPECT_TRUE(result); + std::unique_ptr call_graph = CallGraph::Build(module.get()); + } + + RunCopyInsertion(module.get()); + auto assignment = RunBufferAssignment(module.get()); + + EXPECT_TRUE(BuffersDistinct({call1}, {call2}, *assignment)); +} + +static bool IsPostOrderTraversal( + const std::vector& sequence) { + tensorflow::gtl::FlatSet seen_so_far; + auto has_not_been_seen_yet = [&](const HloInstruction* instruction) { + return seen_so_far.count(instruction) == 0; + }; + + for (auto instruction : sequence) { + if (std::any_of(instruction->operands().begin(), + instruction->operands().end(), has_not_been_seen_yet) || + std::any_of(instruction->control_predecessors().begin(), + instruction->control_predecessors().end(), + has_not_been_seen_yet)) { + return false; // Not a post order. + } + if (!seen_so_far.insert(instruction).second) { + return false; // Not a "traversal". + } + } + + return true; +} + +TEST_F(WhileBufferAssignmentTest, WhileLoopsInterferingResultRange) { + auto module = MakeUnique(TestName()); + auto builder = HloComputation::Builder(TestName()); + + auto zero = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(0.0))); + auto one = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(1.0))); + + auto input0 = builder.AddInstruction( + HloInstruction::CreateParameter(0, data_shape_, "input0")); + auto weights0 = builder.AddInstruction( + HloInstruction::CreateParameter(1, data_shape_, "weights0")); + auto output0 = builder.AddInstruction( + HloInstruction::CreateBroadcast(data_shape_, zero, {1})); + + auto input1 = builder.AddInstruction( + HloInstruction::CreateParameter(2, data_shape_, "input1")); + auto weights1 = builder.AddInstruction( + HloInstruction::CreateParameter(3, data_shape_, "weights1")); + auto output1 = builder.AddInstruction( + HloInstruction::CreateBroadcast(data_shape_, one, {1})); + + auto cond = + module->AddEmbeddedComputation(BuildWhileConditionComputation("cond")); + auto body = module->AddEmbeddedComputation(BuildWhileBodyComputation("body")); + + auto tuple0 = builder.AddInstruction( + HloInstruction::CreateTuple({input0, weights0, output0})); + auto tuple1 = builder.AddInstruction( + HloInstruction::CreateTuple({input1, weights1, output1})); + + auto while0 = builder.AddInstruction( + HloInstruction::CreateWhile(loop_state_shape_, cond, body, tuple0)); + auto while1 = builder.AddInstruction( + HloInstruction::CreateWhile(loop_state_shape_, cond, body, tuple1)); + + auto root_add = builder.AddInstruction(HloInstruction::CreateBinary( + while0->shape(), HloOpcode::kAdd, while0, while1)); + module->AddEntryComputation(builder.Build()); + + RunCopyInsertion(module.get()); + + { + FlattenCallGraph flatten; + TF_ASSERT_OK_AND_ASSIGN(bool result, flatten.Run(module.get())); + EXPECT_TRUE(result); + } + + auto sequence = + CreateMemoryMinimizingSequence(*module, ByteSizeOf).ConsumeValueOrDie(); + + // To trigger b/38494731, we want a specific Hlo sequence for the + // root computation, so we overwrite that entry with a manually + // crafted sequence. + std::vector sequence_for_buffer_assigment = { + input1, weights1, one, output1, tuple1, while1, input0, + weights0, zero, output0, tuple0, while0, root_add}; + + // If this ASSERT_TRUE fails, we constructed a bogus sequence above + // and this test itself is buggy. + ASSERT_TRUE(IsPostOrderTraversal(sequence_for_buffer_assigment)); + + sequence[module->entry_computation()] = + std::move(sequence_for_buffer_assigment); + + auto assignment = + BufferAssigner::Run( + module.get(), + MakeUnique(module.get(), sequence), ByteSizeOf, + [](LogicalBuffer::Color) { return 1; }) + .ConsumeValueOrDie(); + + EXPECT_TRUE(BuffersDistinct({while0}, {while1}, *assignment)); +} + +// Test buffer assignment for while nodes with multiple uses. +// TODO(b/37245345): Fix buffer assignment for this case. +TEST_F(WhileBufferAssignmentTest, DISABLED_TwoWhiles) { + auto module = MakeUnique(TestName()); + auto builder = HloComputation::Builder(TestName()); + + auto input0 = builder.AddInstruction( + HloInstruction::CreateParameter(0, data_shape_, "input0")); + auto weights0 = builder.AddInstruction( + HloInstruction::CreateParameter(1, data_shape_, "weights0")); + + auto zero = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(0.0))); + auto output0 = builder.AddInstruction( + HloInstruction::CreateBroadcast(data_shape_, zero, {1})); + + auto cond0 = + module->AddEmbeddedComputation(BuildWhileConditionComputation("cond")); + auto body0 = + module->AddEmbeddedComputation(BuildWhileBodyComputation("body")); + + auto tuple0 = builder.AddInstruction( + HloInstruction::CreateTuple({input0, weights0, output0})); + auto while0 = builder.AddInstruction( + HloInstruction::CreateWhile(loop_state_shape_, cond0, body0, tuple0)); + auto while1 = builder.AddInstruction( + HloInstruction::CreateWhile(loop_state_shape_, cond0, body0, while0)); + + auto get0 = builder.AddInstruction( + HloInstruction::CreateGetTupleElement(data_shape_, while0, 2)); + auto get1 = builder.AddInstruction( + HloInstruction::CreateGetTupleElement(data_shape_, while1, 2)); + builder.AddInstruction( + HloInstruction::CreateBinary(data_shape_, HloOpcode::kAdd, get0, get1)); + module->AddEntryComputation(builder.Build()); + + RunCopyInsertion(module.get()); + + { + FlattenCallGraph flatten; + TF_ASSERT_OK_AND_ASSIGN(bool result, flatten.Run(module.get())); + EXPECT_TRUE(result); + } + + auto assignment = RunBufferAssignment(module.get()); + + EXPECT_TRUE(BuffersDistinct({while0}, {while1}, *assignment)); +} + +} // namespace } // namespace xla + +int main(int argc, char** argv) { + return xla::ParseDebugOptionsFlagsAndRunTests(argc, argv); +} diff --git a/tensorflow/compiler/xla/service/buffer_liveness.cc b/tensorflow/compiler/xla/service/buffer_liveness.cc index 736f227aa423120ecb4a5e82824defac2d345b2e..f085ffa6bc40b212339a97604455a07c1e662952 100644 --- a/tensorflow/compiler/xla/service/buffer_liveness.cc +++ b/tensorflow/compiler/xla/service/buffer_liveness.cc @@ -45,10 +45,11 @@ StatusOr> BufferLiveness::Run( } tensorflow::Status BufferLiveness::Analyze() { - TF_ASSIGN_OR_RETURN(points_to_analysis_, - TuplePointsToAnalysis::Run( - module_, /*include_loop_fusion_instructions=*/true)); + TF_ASSIGN_OR_RETURN(points_to_analysis_, TuplePointsToAnalysis::Run(module_)); for (auto& computation : module_->computations()) { + if (computation->IsFusionComputation()) { + continue; + } // Gather all instructions whose buffers might alias other instructions into // the set aliased_buffers_. This includes those contained as a tuple // element in other instruction's output. @@ -117,26 +118,42 @@ bool BufferLiveness::live_range_strictly_before(const LogicalBuffer& a, // If 'b' is a user of 'a' then the buffers interfere unless 'a.instruction' // and 'b.instruction' emit the same shape/layout, and 'b.instruction' meets - // one of following qualifications: - // *) Is element-wise. - // *) Is a loop fusion instruction (with DynamicUpdateSlice fused root) where - // the singleton use of 'a' at 'a.index' is the fused root at operand 0. - // *) Use of 'operand' is DynamicUpdateSlice at operand index 0. + // the qualifications specified in CanShareOperandBufferWithUser. for (const BufferAlias& alias : points_to_analysis_->GetBufferAliases(a)) { if (b.instruction()->IsUserOf(alias.instruction()) && !CanShareOperandBufferWithUser(alias.instruction(), alias.index(), b.instruction(), b.index(), - points_to_analysis())) { + &points_to_analysis())) { return false; } } return true; } +namespace { +bool IsEntryParameter(const HloInstruction* instruction) { + const HloComputation* computation = instruction->parent(); + return instruction->opcode() == HloOpcode::kParameter && + computation == computation->parent()->entry_computation(); +} +} // namespace + bool BufferLiveness::MayInterfere(const LogicalBuffer& a, const LogicalBuffer& b) const { - return (!live_range_strictly_before(a, b) && - !live_range_strictly_before(b, a)); + // Entry parameters live at the entry of the execution, thus always interfere + // with all other instructions executing before them in the ordering. + const HloInstruction* a_instruction = a.instruction(); + const HloInstruction* b_instruction = b.instruction(); + if (IsEntryParameter(a_instruction) && + hlo_ordering_->ExecutesBefore(b_instruction, a_instruction)) { + return true; + } + if (IsEntryParameter(b_instruction) && + hlo_ordering_->ExecutesBefore(a_instruction, b_instruction)) { + return true; + } + // Buffers without disjoint liveness may interfere. + return !live_range_strictly_before(a, b) && !live_range_strictly_before(b, a); } bool BufferLiveness::MaybeLiveOut(const LogicalBuffer& buffer) const { diff --git a/tensorflow/compiler/xla/service/buffer_liveness.h b/tensorflow/compiler/xla/service/buffer_liveness.h index 4c94d1a27d7a7d196b3662b53f5c3fa2d013b11e..11834a5127e383cc2ec2ab3fe1bb82ba86e4abed 100644 --- a/tensorflow/compiler/xla/service/buffer_liveness.h +++ b/tensorflow/compiler/xla/service/buffer_liveness.h @@ -36,6 +36,8 @@ namespace xla { // interference. class BufferLiveness { public: + using Colorer = std::function; + // Constructs a buffer liveness object for the given module assuming the given // HLO instruction ordering. static StatusOr> Run( @@ -52,8 +54,7 @@ class BufferLiveness { bool MaybeLiveOut(const LogicalBuffer& buffer) const; // Returns the complete set of buffers that may be live out of the module. - const tensorflow::gtl::FlatSet& maybe_live_out_buffers() - const { + const PointsToSet::BufferSet& maybe_live_out_buffers() const { return maybe_live_out_buffers_; } @@ -65,8 +66,22 @@ class BufferLiveness { // Returns the underlying hlo ordering used for this liveness analysis. const HloOrdering& hlo_ordering() const { return *hlo_ordering_; } + const HloModule& module() const { return *module_; } + string ToString() const; + static Colorer DefaultColorer() { + return [](const BufferLiveness& buffer_liveness) { + for (LogicalBuffer::Id id = 0; + id < buffer_liveness.points_to_analysis().num_logical_buffers(); + id++) { + auto& buffer = buffer_liveness.points_to_analysis().logical_buffer(id); + buffer.set_color(LogicalBuffer::Color(0)); + } + return Status::OK(); + }; + } + private: explicit BufferLiveness(const HloModule* module, std::unique_ptr hlo_ordering) @@ -90,7 +105,7 @@ class BufferLiveness { tensorflow::gtl::FlatSet aliased_buffers_; // LogicalBuffers that may be live out of the entry computation. - tensorflow::gtl::FlatSet maybe_live_out_buffers_; + PointsToSet::BufferSet maybe_live_out_buffers_; std::unique_ptr points_to_analysis_; }; diff --git a/tensorflow/compiler/xla/service/buffer_liveness_test.cc b/tensorflow/compiler/xla/service/buffer_liveness_test.cc index e71b98298b344b5689785bfa67a8bea54e0248e3..7a102d65ce3d755bad9272fcd55f6f18e2cb7d67 100644 --- a/tensorflow/compiler/xla/service/buffer_liveness_test.cc +++ b/tensorflow/compiler/xla/service/buffer_liveness_test.cc @@ -37,10 +37,9 @@ class BufferLivenessTest : public HloTestBase { const LogicalBuffer& GetBuffer(const BufferLiveness& liveness, const HloInstruction* instruction, const ShapeIndex& index) { - const std::vector& pointed_to = - liveness.points_to_analysis() - .GetPointsToSet(instruction) - .element(index); + const auto& pointed_to = liveness.points_to_analysis() + .GetPointsToSet(instruction) + .element(index); CHECK_EQ(1, pointed_to.size()); CHECK_EQ(instruction, pointed_to[0]->instruction()); CHECK(index == pointed_to[0]->index()); @@ -72,9 +71,9 @@ class BufferLivenessTest : public HloTestBase { ShapeUtil::GetSubshape(b->shape(), index))); // Lookup PointsTo set for instructions 'a' and 'b'. auto& points_to_analysis = liveness.points_to_analysis(); - const std::vector& points_to_a = + const auto& points_to_a = points_to_analysis.GetPointsToSet(a).element(index); - const std::vector& points_to_b = + const auto& points_to_b = points_to_analysis.GetPointsToSet(b).element(index); // Make sure PointsTo sets for 'a' and 'b' are unambiguous. EXPECT_EQ(1, points_to_a.size()); @@ -92,6 +91,12 @@ class BufferLivenessTest : public HloTestBase { GetBuffer(liveness, instruction, /*index=*/{})); } + std::unique_ptr BuildDummyComputation() { + auto builder = HloComputation::Builder(TestName() + "_dummy"); + builder.AddInstruction(HloInstruction::CreateParameter(0, vec_, "param")); + return builder.Build(); + } + const Shape vec_ = ShapeUtil::MakeShape(xla::F32, {42}); }; @@ -110,7 +115,7 @@ TEST_F(BufferLivenessTest, ElementwiseChain) { auto log = builder.AddInstruction( HloInstruction::CreateUnary(vec_, HloOpcode::kLog, exp)); - auto module = MakeUnique(TestName()); + auto module = CreateNewModule(); module->AddEntryComputation(builder.Build()); auto liveness = @@ -118,12 +123,17 @@ TEST_F(BufferLivenessTest, ElementwiseChain) { MakeUnique(module.get())) .ConsumeValueOrDie(); - // No buffers should interfere. EXPECT_FALSE(InstructionsMayInterfere(*liveness, param, negate)); + EXPECT_FALSE(InstructionsMayInterfere(*liveness, param, exp)); + EXPECT_FALSE(InstructionsMayInterfere(*liveness, param, log)); + + // No buffers should interfere. EXPECT_FALSE(InstructionsMayInterfere(*liveness, negate, exp)); + EXPECT_FALSE(InstructionsMayInterfere(*liveness, negate, log)); EXPECT_FALSE(InstructionsMayInterfere(*liveness, exp, negate)); EXPECT_FALSE(InstructionsMayInterfere(*liveness, exp, log)); - EXPECT_FALSE(InstructionsMayInterfere(*liveness, param, log)); + EXPECT_FALSE(InstructionsMayInterfere(*liveness, log, negate)); + EXPECT_FALSE(InstructionsMayInterfere(*liveness, log, exp)); // Buffers should interfere with itself. EXPECT_TRUE(InstructionsMayInterfere(*liveness, exp, exp)); @@ -135,22 +145,73 @@ TEST_F(BufferLivenessTest, ElementwiseChain) { EXPECT_TRUE(InstructionMaybeLiveOut(*liveness, log)); } +TEST_F(BufferLivenessTest, MultipleEntryParameters_Sequential) { + // Two entry params, which interfere with each other. + // + // param0 --> negate ---------------\ + // param1 --> exp --> add + auto builder = HloComputation::Builder(TestName()); + auto param0 = builder.AddInstruction( + HloInstruction::CreateParameter(0, vec_, "param0")); + auto param1 = builder.AddInstruction( + HloInstruction::CreateParameter(1, vec_, "param1")); + auto negate = builder.AddInstruction( + HloInstruction::CreateUnary(vec_, HloOpcode::kNegate, param0)); + auto exp = builder.AddInstruction( + HloInstruction::CreateUnary(vec_, HloOpcode::kExp, param1)); + auto add = builder.AddInstruction( + HloInstruction::CreateBinary(vec_, HloOpcode::kAdd, negate, exp)); + + auto module = CreateNewModule(); + HloComputation* entry = module->AddEntryComputation(builder.Build()); + + SequentialHloOrdering::HloModuleSequence sequence; + sequence.insert({entry, {param0, negate, param1, exp, add}}); + auto liveness = BufferLiveness::Run( + module.get(), + MakeUnique(module.get(), sequence)) + .ConsumeValueOrDie(); + + // Entry parameters interfere as if they are defined simultaneously at + // the very beginning. + EXPECT_TRUE(InstructionsMayInterfere(*liveness, param0, param1)); + EXPECT_FALSE(InstructionsMayInterfere(*liveness, param0, negate)); + EXPECT_FALSE(InstructionsMayInterfere(*liveness, param0, exp)); + EXPECT_FALSE(InstructionsMayInterfere(*liveness, param0, add)); + EXPECT_TRUE(InstructionsMayInterfere(*liveness, param1, param0)); + EXPECT_TRUE(InstructionsMayInterfere(*liveness, param1, negate)); + EXPECT_FALSE(InstructionsMayInterfere(*liveness, param1, exp)); + EXPECT_FALSE(InstructionsMayInterfere(*liveness, param1, add)); + + // Negate and exp still interfere. + EXPECT_TRUE(InstructionsMayInterfere(*liveness, negate, exp)); + EXPECT_TRUE(InstructionsMayInterfere(*liveness, exp, negate)); + + // But {negate, add} and {exp, add} don't interfere. + EXPECT_FALSE(InstructionsMayInterfere(*liveness, negate, add)); + EXPECT_FALSE(InstructionsMayInterfere(*liveness, add, negate)); + EXPECT_FALSE(InstructionsMayInterfere(*liveness, exp, add)); + EXPECT_FALSE(InstructionsMayInterfere(*liveness, add, exp)); +} + TEST_F(BufferLivenessTest, NonElementwiseOperand) { - // A chain of operations with one elementwise and one non-elementwise. The + // A chain of operations with two elementwise and one non-elementwise. The // elementwise op should not interfere with its operand, while the - // non-elementwise op should interfere. + // non-elementwise op should interfere. Entry params always interfere. // - // param --> negate -> reverse + // param --> exp -> negate -> reverse // auto builder = HloComputation::Builder(TestName()); auto param = builder.AddInstruction(HloInstruction::CreateParameter(0, vec_, "param")); + auto exp = builder.AddInstruction( + HloInstruction::CreateUnary(vec_, HloOpcode::kExp, param)); auto negate = builder.AddInstruction( - HloInstruction::CreateUnary(vec_, HloOpcode::kNegate, param)); + HloInstruction::CreateUnary(vec_, HloOpcode::kNegate, exp)); auto reverse = builder.AddInstruction(HloInstruction::CreateReverse(vec_, negate, {0})); - auto module = MakeUnique(TestName()); + auto module = CreateNewModule(); module->AddEntryComputation(builder.Build()); auto liveness = @@ -158,10 +219,14 @@ TEST_F(BufferLivenessTest, NonElementwiseOperand) { MakeUnique(module.get())) .ConsumeValueOrDie(); - // No buffers should interfere. + EXPECT_FALSE(InstructionsMayInterfere(*liveness, param, exp)); EXPECT_FALSE(InstructionsMayInterfere(*liveness, param, negate)); + EXPECT_FALSE(InstructionsMayInterfere(*liveness, param, reverse)); + + // Negate is elementwise, so doesn't interfere with its operand. + // Reverse is non-elementwise, so does interfere with its operand. + EXPECT_FALSE(InstructionsMayInterfere(*liveness, exp, negate)); EXPECT_TRUE(InstructionsMayInterfere(*liveness, negate, reverse)); - EXPECT_FALSE(InstructionsMayInterfere(*liveness, param, negate)); } TEST_F(BufferLivenessTest, OverlappedBuffers) { @@ -180,7 +245,7 @@ TEST_F(BufferLivenessTest, OverlappedBuffers) { auto add = builder.AddInstruction( HloInstruction::CreateBinary(vec_, HloOpcode::kAdd, negate, exp)); - auto module = MakeUnique(TestName()); + auto module = CreateNewModule(); module->AddEntryComputation(builder.Build()); auto liveness = @@ -190,8 +255,15 @@ TEST_F(BufferLivenessTest, OverlappedBuffers) { EXPECT_TRUE(InstructionsMayInterfere(*liveness, param, negate)); EXPECT_TRUE(InstructionsMayInterfere(*liveness, param, exp)); - EXPECT_TRUE(InstructionsMayInterfere(*liveness, negate, exp)); EXPECT_FALSE(InstructionsMayInterfere(*liveness, param, add)); + + // Negate and exp interfere with each other, but not with add. + EXPECT_TRUE(InstructionsMayInterfere(*liveness, negate, exp)); + EXPECT_TRUE(InstructionsMayInterfere(*liveness, exp, negate)); + EXPECT_FALSE(InstructionsMayInterfere(*liveness, negate, add)); + EXPECT_FALSE(InstructionsMayInterfere(*liveness, add, negate)); + EXPECT_FALSE(InstructionsMayInterfere(*liveness, exp, add)); + EXPECT_FALSE(InstructionsMayInterfere(*liveness, add, exp)); } TEST_F(BufferLivenessTest, OverlappedBuffersSequentialOrder) { @@ -204,8 +276,7 @@ TEST_F(BufferLivenessTest, OverlappedBuffersSequentialOrder) { // Sequential order: // param, negate, exp, add // - // Liveness is identical to the DependencyHloOrdering except that 'param' and - // exp no longer interfere. + // Liveness is identical to the DependencyHloOrdering. auto builder = HloComputation::Builder(TestName()); auto param = builder.AddInstruction(HloInstruction::CreateParameter(0, vec_, "param")); @@ -216,7 +287,7 @@ TEST_F(BufferLivenessTest, OverlappedBuffersSequentialOrder) { auto add = builder.AddInstruction( HloInstruction::CreateBinary(vec_, HloOpcode::kAdd, negate, exp)); - auto module = MakeUnique(TestName()); + auto module = CreateNewModule(); auto computation = module->AddEntryComputation(builder.Build()); SequentialHloOrdering::HloModuleSequence module_sequence; @@ -229,8 +300,15 @@ TEST_F(BufferLivenessTest, OverlappedBuffersSequentialOrder) { EXPECT_TRUE(InstructionsMayInterfere(*liveness, param, negate)); EXPECT_FALSE(InstructionsMayInterfere(*liveness, param, exp)); - EXPECT_TRUE(InstructionsMayInterfere(*liveness, negate, exp)); EXPECT_FALSE(InstructionsMayInterfere(*liveness, param, add)); + + // Negate and exp interfere with each other, but not with add. + EXPECT_TRUE(InstructionsMayInterfere(*liveness, negate, exp)); + EXPECT_TRUE(InstructionsMayInterfere(*liveness, exp, negate)); + EXPECT_FALSE(InstructionsMayInterfere(*liveness, negate, add)); + EXPECT_FALSE(InstructionsMayInterfere(*liveness, add, negate)); + EXPECT_FALSE(InstructionsMayInterfere(*liveness, exp, add)); + EXPECT_FALSE(InstructionsMayInterfere(*liveness, add, exp)); } TEST_F(BufferLivenessTest, TupleLiveOut) { @@ -251,7 +329,7 @@ TEST_F(BufferLivenessTest, TupleLiveOut) { auto outer_tuple = builder.AddInstruction(HloInstruction::CreateTuple({inner_tuple, exp})); - auto module = MakeUnique(TestName()); + auto module = CreateNewModule(); module->AddEntryComputation(builder.Build()); auto liveness = @@ -271,7 +349,7 @@ TEST_F(BufferLivenessTest, TupleLiveOut) { TEST_F(BufferLivenessTest, EmbeddedComputation) { // Test MaybeLiveOut and MayInterfere for embedded computation. - auto module = MakeUnique(TestName()); + auto module = CreateNewModule(); auto embedded_builder = HloComputation::Builder(TestName() + "_embedded"); auto embedded_param = embedded_builder.AddInstruction( @@ -318,17 +396,15 @@ TEST_F(BufferLivenessTest, TupleConstantLiveOut) { // computation. The buffer containing {0, 1} is copied by GetTupleElement, and // the buffers containing {3} and 3 are dead. auto builder = HloComputation::Builder(TestName()); - auto inner_tuple0 = - LiteralUtil::MakeTuple({LiteralUtil::CreateR0(0).get(), - LiteralUtil::CreateR0(1).get()}); - auto inner_tuple1 = - LiteralUtil::MakeTuple({LiteralUtil::CreateR0(3).get()}); + auto inner_tuple0 = Literal::MakeTuple( + {Literal::CreateR0(0).get(), Literal::CreateR0(1).get()}); + auto inner_tuple1 = Literal::MakeTuple({Literal::CreateR0(3).get()}); auto tuple_constant = builder.AddInstruction(HloInstruction::CreateConstant( - LiteralUtil::MakeTuple({inner_tuple0.get(), inner_tuple1.get()}))); + Literal::MakeTuple({inner_tuple0.get(), inner_tuple1.get()}))); builder.AddInstruction(HloInstruction::CreateGetTupleElement( inner_tuple0->shape(), tuple_constant, 0)); - auto module = MakeUnique(TestName()); + auto module = CreateNewModule(); module->AddEntryComputation(builder.Build()); auto liveness = @@ -358,8 +434,9 @@ TEST_F(BufferLivenessTest, IndependentTupleElements) { auto builder = HloComputation::Builder(TestName()); // Create param0 Tuple. auto tuple_param0 = builder.AddInstruction(HloInstruction::CreateParameter( - 0, ShapeUtil::MakeTupleShape( - {ShapeUtil::MakeShape(F32, {8}), ShapeUtil::MakeShape(S32, {4})}), + 0, + ShapeUtil::MakeTupleShape( + {ShapeUtil::MakeShape(F32, {8}), ShapeUtil::MakeShape(S32, {4})}), "param0")); // Create independent computations for each tuple elememt. @@ -371,7 +448,7 @@ TEST_F(BufferLivenessTest, IndependentTupleElements) { builder.AddInstruction(HloInstruction::CreateGetTupleElement( tuple_element0_shape, tuple_param0, 0)); auto const0 = builder.AddInstruction(HloInstruction::CreateConstant( - LiteralUtil::CreateR1({1.f, 1.f, 1.f, 1.f, 1.f, 1.f, 1.f, 1.f}))); + Literal::CreateR1({1.f, 1.f, 1.f, 1.f, 1.f, 1.f, 1.f, 1.f}))); auto add0 = builder.AddInstruction(HloInstruction::CreateBinary( tuple_element0_shape, HloOpcode::kAdd, tuple_element0, const0)); @@ -383,7 +460,7 @@ TEST_F(BufferLivenessTest, IndependentTupleElements) { builder.AddInstruction(HloInstruction::CreateGetTupleElement( tuple_element1_shape, tuple_param0, 1)); auto const1 = builder.AddInstruction(HloInstruction::CreateConstant( - LiteralUtil::CreateR1({2.f, 2.f, 2.f, 2.f, 2.f, 2.f, 2.f, 2.f}))); + Literal::CreateR1({2.f, 2.f, 2.f, 2.f, 2.f, 2.f, 2.f, 2.f}))); auto add1 = builder.AddInstruction(HloInstruction::CreateBinary( tuple_element1_shape, HloOpcode::kAdd, tuple_element1, const1)); @@ -391,8 +468,9 @@ TEST_F(BufferLivenessTest, IndependentTupleElements) { auto tuple_root = builder.AddInstruction(HloInstruction::CreateTuple({add0, add1})); - auto module = MakeUnique(TestName()); - module->AddEntryComputation(builder.Build()); + auto module = CreateNewModule(); + module->AddEntryComputation(BuildDummyComputation()); + module->AddEmbeddedComputation(builder.Build()); auto liveness = BufferLiveness::Run(module.get(), @@ -420,8 +498,9 @@ TEST_F(BufferLivenessTest, DependentTupleElements) { auto builder = HloComputation::Builder(TestName()); // Create param0 Tuple. auto tuple_param0 = builder.AddInstruction(HloInstruction::CreateParameter( - 0, ShapeUtil::MakeTupleShape( - {ShapeUtil::MakeShape(F32, {8}), ShapeUtil::MakeShape(F32, {8})}), + 0, + ShapeUtil::MakeTupleShape( + {ShapeUtil::MakeShape(F32, {8}), ShapeUtil::MakeShape(F32, {8})}), "param0")); // Create dependent computations for each tuple elememt. @@ -433,7 +512,7 @@ TEST_F(BufferLivenessTest, DependentTupleElements) { builder.AddInstruction(HloInstruction::CreateGetTupleElement( tuple_element0_shape, tuple_param0, 0)); auto const0 = builder.AddInstruction(HloInstruction::CreateConstant( - LiteralUtil::CreateR1({1.f, 1.f, 1.f, 1.f, 1.f, 1.f, 1.f, 1.f}))); + Literal::CreateR1({1.f, 1.f, 1.f, 1.f, 1.f, 1.f, 1.f, 1.f}))); auto add0 = builder.AddInstruction(HloInstruction::CreateBinary( tuple_element0_shape, HloOpcode::kAdd, tuple_element0, const0)); @@ -451,8 +530,9 @@ TEST_F(BufferLivenessTest, DependentTupleElements) { auto tuple_root = builder.AddInstruction(HloInstruction::CreateTuple({add0, add1})); - auto module = MakeUnique(TestName()); - module->AddEntryComputation(builder.Build()); + auto module = CreateNewModule(); + module->AddEntryComputation(BuildDummyComputation()); + module->AddEmbeddedComputation(builder.Build()); auto liveness = BufferLiveness::Run(module.get(), @@ -504,18 +584,18 @@ class FusedDynamicUpdateSliceLivenessTest : public BufferLivenessTest { HloInstruction::CreateGetTupleElement(data_shape, tuple_param0, 1)); auto update = builder.AddInstruction(HloInstruction::CreateConstant( - LiteralUtil::CreateR1({2.f, 2.f, 2.f}))); + Literal::CreateR1({2.f, 2.f, 2.f}))); HloInstruction* slice = nullptr; if (update_uses_tuple_element1) { // Create a slice instruction as an additional user of 'gte1'. slice = builder.AddInstruction( - HloInstruction::CreateSlice(update_shape, gte1, {0}, {3})); + HloInstruction::CreateSlice(update_shape, gte1, {0}, {3}, {1})); update = builder.AddInstruction(HloInstruction::CreateBinary( update_shape, HloOpcode::kAdd, update, slice)); } // Create a DynamicUpdateSlice instruction of tuple element 1 with 'update'. auto starts = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR1({2}))); + HloInstruction::CreateConstant(Literal::CreateR1({2}))); auto dynamic_update_slice = builder.AddInstruction(HloInstruction::CreateDynamicUpdateSlice( data_shape, gte1, update, starts)); @@ -523,8 +603,9 @@ class FusedDynamicUpdateSliceLivenessTest : public BufferLivenessTest { auto tuple_root = builder.AddInstruction( HloInstruction::CreateTuple({gte0, dynamic_update_slice})); // Build module and get reference to entry computation. - auto module = MakeUnique(TestName()); - auto* computation = module->AddEntryComputation(builder.Build()); + auto module = CreateNewModule(); + module->AddEntryComputation(BuildDummyComputation()); + auto* computation = module->AddEmbeddedComputation(builder.Build()); // Create fusion instruction based on number of tuple element 1 users. if (update_uses_tuple_element1) { computation->CreateFusionInstruction( @@ -546,7 +627,7 @@ class FusedDynamicUpdateSliceLivenessTest : public BufferLivenessTest { BufferLiveness::Run(module.get(), MakeUnique(module.get())) .ConsumeValueOrDie(); - // Return whether or not buffers interfernce is detected between + // Return whether or not buffers interference is detected between // 'tuple_param0' and 'tuple_root' at shape index '{1}'. return TupleElementsMayInterfere(*liveness, tuple_param0, tuple_root, {1}); } @@ -633,7 +714,7 @@ class DynamicUpdateSliceLivenessTest : public BufferLivenessTest { HloInstruction::CreateGetTupleElement(data_shape, tuple_param0, 1)); auto update = builder.AddInstruction(HloInstruction::CreateConstant( - LiteralUtil::CreateR1({2.f, 2.f, 2.f}))); + Literal::CreateR1({2.f, 2.f, 2.f}))); if (tuple_element1_has_two_uses) { // Add 'gte0' and 'gte1' to create another user of 'gte1'. @@ -642,7 +723,7 @@ 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(Literal::CreateR1({2}))); auto dynamic_update_slice = builder.AddInstruction(HloInstruction::CreateDynamicUpdateSlice( data_shape, gte1, update, starts)); @@ -650,14 +731,15 @@ class DynamicUpdateSliceLivenessTest : public BufferLivenessTest { auto tuple_root = builder.AddInstruction( HloInstruction::CreateTuple({gte0, dynamic_update_slice})); // Build module and get reference to entry computation. - auto module = MakeUnique(TestName()); - module->AddEntryComputation(builder.Build()); + auto module = CreateNewModule(); + module->AddEntryComputation(BuildDummyComputation()); + module->AddEmbeddedComputation(builder.Build()); // Run BufferLiveness on 'module'. auto liveness = BufferLiveness::Run(module.get(), MakeUnique(module.get())) .ConsumeValueOrDie(); - // Return whether or not buffers interfernce is detected between + // Return whether or not buffers interference is detected between // 'tuple_param0' and 'tuple_root' at shape index '{1}'. return TupleElementsMayInterfere(*liveness, tuple_param0, tuple_root, {1}); } @@ -702,3 +784,7 @@ TEST_F(DynamicUpdateSliceLivenessTest, WithInterference) { } // namespace } // namespace xla + +int main(int argc, char** argv) { + return xla::ParseDebugOptionsFlagsAndRunTests(argc, argv); +} diff --git a/tensorflow/compiler/xla/service/call_graph.cc b/tensorflow/compiler/xla/service/call_graph.cc index ab3eae2495ec55e8667db86b025f980157517ccc..c0f3bcdc2218199288eaa3d0010ee70632c8f959 100644 --- a/tensorflow/compiler/xla/service/call_graph.cc +++ b/tensorflow/compiler/xla/service/call_graph.cc @@ -51,6 +51,22 @@ std::ostream& operator<<(std::ostream& out, const CallContext& context) { return out; } +CallContext GetInstructionCallContext(const HloInstruction* instruction) { + switch (instruction->opcode()) { + case HloOpcode::kCall: + case HloOpcode::kWhile: + return CallContext::kSequential; + case HloOpcode::kMap: + case HloOpcode::kReduce: + case HloOpcode::kReduceWindow: + case HloOpcode::kSelectAndScatter: + case HloOpcode::kFusion: + return CallContext::kParallel; + default: + return CallContext::kNone; + } +} + string CallSite::ToString() const { return StrCat(instruction()->name(), " calls in context ", CallContextToString(context()), ": ", @@ -82,32 +98,12 @@ void CallGraphNode::AddCallerCallSite(const CallSite& caller_callsite) { } } -namespace { - -CallContext GetInstructionCallContext(const HloInstruction* instruction) { - switch (instruction->opcode()) { - case HloOpcode::kCall: - case HloOpcode::kWhile: - return CallContext::kSequential; - case HloOpcode::kMap: - case HloOpcode::kReduce: - case HloOpcode::kReduceWindow: - case HloOpcode::kSelectAndScatter: - case HloOpcode::kFusion: - return CallContext::kParallel; - default: - return CallContext::kNone; - } -} - -} // namespace - -Status CallGraphNode::AddCallSiteForInstruction(HloInstruction* instruction) { - TF_RET_CHECK(instruction->parent() == computation()); +void CallGraphNode::AddCallSiteForInstruction(HloInstruction* instruction) { + CHECK_EQ(instruction->parent(), computation()); const CallContext context = GetInstructionCallContext(instruction); if (!instruction->called_computations().empty()) { - TF_RET_CHECK(context == CallContext::kSequential || - context == CallContext::kParallel); + CHECK(context == CallContext::kSequential || + context == CallContext::kParallel); callsite_instructions_.insert({instruction, callsites_.size()}); callsites_.push_back( CallSite(instruction, instruction->called_computations(), context)); @@ -120,22 +116,52 @@ Status CallGraphNode::AddCallSiteForInstruction(HloInstruction* instruction) { } } } - return Status::OK(); } CallGraph::CallGraph(const HloModule* module) : module_(module) {} -StatusOr CallGraph::GetNode( +const CallGraphNode& CallGraph::GetNode( const HloComputation* computation) const { auto it = node_indices_.find(computation); - TF_RET_CHECK(it != node_indices_.end()); - return &nodes_[it->second]; + CHECK(it != node_indices_.end()); + return nodes_[it->second]; } -StatusOr CallGraph::GetNode(const HloComputation* computation) { +CallGraphNode& CallGraph::GetNode(const HloComputation* computation) { auto it = node_indices_.find(computation); - TF_RET_CHECK(it != node_indices_.end()); - return &nodes_[it->second]; + CHECK(it != node_indices_.end()); + return nodes_[it->second]; +} + +bool CallGraph::DominatesHelper( + const HloComputation* a, const HloComputation* b, + tensorflow::gtl::FlatSet* visited) const { + if (a == b || ContainsKey(*visited, b)) { + // The call graph is guaranteed to be acyclic so any previously visited node + // we encounter was already determined to be dominated. + return true; + } + + const CallGraphNode& b_node = GetNode(b); + if (b_node.callers().empty()) { + // We reached a root node without hitting 'a'. 'a' does not dominate 'b'. + return false; + } + + // Walk up the callers of 'b' until we hit 'a' or a root node (no callers). + visited->insert(b); + for (const HloComputation* b_caller : b_node.callers()) { + if (!DominatesHelper(a, b_caller, visited)) { + return false; + } + } + return true; +} + +bool CallGraph::Dominates(const HloComputation* a, + const HloComputation* b) const { + tensorflow::gtl::FlatSet visited; + return DominatesHelper(a, b, &visited); } namespace { @@ -158,17 +184,17 @@ CallContext UnionContexts(CallContext a, CallContext b) { } // namespace -Status CallGraph::SetCallContexts() { +void CallGraph::SetCallContexts() { std::queue worklist; // Initialize worklist with all roots of the call graph (computations without // callers). for (const std::unique_ptr& computation : module_->computations()) { - TF_ASSIGN_OR_RETURN(CallGraphNode * node, GetNode(computation.get())); - if (node->callers().empty()) { - node->set_context(CallContext::kSequential); - worklist.push(node); + CallGraphNode& node = GetNode(computation.get()); + if (node.callers().empty()) { + node.set_context(CallContext::kSequential); + worklist.push(&node); } } @@ -178,7 +204,7 @@ Status CallGraph::SetCallContexts() { for (const CallSite& callsite : node->callsites()) { for (const HloComputation* callee : callsite.called_computations()) { - TF_ASSIGN_OR_RETURN(CallGraphNode * callee_node, GetNode(callee)); + CallGraphNode& callee_node = GetNode(callee); // Update context of callee computation based on the callsite and its // current context. @@ -186,16 +212,16 @@ Status CallGraph::SetCallContexts() { if (callsite.context() == CallContext::kParallel) { context_to_add = CallContext::kParallel; } else { - TF_RET_CHECK(callsite.context() == CallContext::kSequential); + CHECK_EQ(callsite.context(), CallContext::kSequential); context_to_add = node->context(); } CallContext new_context = - UnionContexts(context_to_add, callee_node->context()); + UnionContexts(context_to_add, callee_node.context()); - if (new_context != callee_node->context()) { + if (new_context != callee_node.context()) { // Context of computation has been changed so add node to worklist. - callee_node->set_context(new_context); - worklist.push(callee_node); + callee_node.set_context(new_context); + worklist.push(&callee_node); } } } @@ -204,14 +230,12 @@ Status CallGraph::SetCallContexts() { // No node should have a kNone calling context. for (const std::unique_ptr& computation : module_->computations()) { - TF_ASSIGN_OR_RETURN(CallGraphNode * node, GetNode(computation.get())); - TF_RET_CHECK(node->context() != CallContext::kNone); + CHECK_NE(GetNode(computation.get()).context(), CallContext::kNone); } - return Status::OK(); } /* static */ -StatusOr> CallGraph::Build(const HloModule* module) { +std::unique_ptr CallGraph::Build(const HloModule* module) { // Constructor for CallGraph is private so MakeUnique can't be used. auto call_graph = WrapUnique(new CallGraph(module)); @@ -223,56 +247,51 @@ StatusOr> CallGraph::Build(const HloModule* module) { module->computations()) { auto it_added = call_graph->node_indices_.insert( {computation.get(), call_graph->nodes_.size()}); - // All computation should be unique, so the computation should not already + // All computations should be unique, so the computation should not already // exist in the map. - TF_RET_CHECK(it_added.second); + CHECK(it_added.second); call_graph->nodes_.emplace_back(computation.get()); // Add all callsites in this computation. for (const std::unique_ptr& instruction : computation->instructions()) { - TF_RETURN_IF_ERROR(call_graph->nodes_.back().AddCallSiteForInstruction( - instruction.get())); + call_graph->nodes_.back().AddCallSiteForInstruction(instruction.get()); } } // Add caller callsites to each node. for (const std::unique_ptr& computation : module->computations()) { - TF_ASSIGN_OR_RETURN(CallGraphNode * caller_node, - call_graph->GetNode(computation.get())); - for (const CallSite& callsite : caller_node->callsites()) { + for (const CallSite& callsite : + call_graph->GetNode(computation.get()).callsites()) { for (auto* callee : callsite.called_computations()) { // Add caller callsites. - TF_ASSIGN_OR_RETURN(CallGraphNode * callee_node, - call_graph->GetNode(callee)); - callee_node->AddCallerCallSite(callsite); + call_graph->GetNode(callee).AddCallerCallSite(callsite); } } } - TF_RETURN_IF_ERROR(call_graph->SetCallContexts()); - + call_graph->SetCallContexts(); XLA_VLOG_LINES(1, call_graph->ToString()); - return std::move(call_graph); + return call_graph; } Status CallGraph::VisitNodesInternal( - const VisitorFunction& visitor_func, const CallGraphNode* node, + const VisitorFunction& visitor_func, const CallGraphNode& node, tensorflow::gtl::FlatSet* visited) const { - auto pair = visited->insert(node); + auto pair = visited->insert(&node); if (!pair.second) { // Node was not inserted. Node has already been visited. return Status::OK(); } - for (const HloComputation* computation : node->callees()) { - TF_ASSIGN_OR_RETURN(const CallGraphNode* callee_node, GetNode(computation)); - TF_RETURN_IF_ERROR(VisitNodesInternal(visitor_func, callee_node, visited)); + for (const HloComputation* computation : node.callees()) { + TF_RETURN_IF_ERROR( + VisitNodesInternal(visitor_func, GetNode(computation), visited)); } - return visitor_func(*node); + return visitor_func(node); } Status CallGraph::VisitNodes(const VisitorFunction& visitor_func, @@ -282,19 +301,61 @@ Status CallGraph::VisitNodes(const VisitorFunction& visitor_func, // Traverse from all roots in the call graph. for (const CallGraphNode& node : nodes()) { if (node.callers().empty()) { - TF_RETURN_IF_ERROR(VisitNodesInternal(visitor_func, &node, &visited)); + TF_RETURN_IF_ERROR(VisitNodesInternal(visitor_func, node, &visited)); } } } else { // Traverse only from the entry computation. - TF_ASSIGN_OR_RETURN(const CallGraphNode* entry_node, - GetNode(module_->entry_computation())); - TF_RETURN_IF_ERROR(VisitNodesInternal(visitor_func, entry_node, &visited)); + TF_RETURN_IF_ERROR(VisitNodesInternal( + visitor_func, GetNode(module_->entry_computation()), &visited)); } return Status::OK(); } +bool CallGraph::IsFlattened() const { + for (const CallGraphNode& node : nodes_) { + if (node.context() == CallContext::kBoth) { + return false; + } + if (node.context() == CallContext::kSequential && + node.caller_callsites().size() > 1) { + return false; + } + } + return true; +} + +std::pair +CallGraph::NearestAncestorsInSameComputation(HloInstruction* a, + HloInstruction* b) const { + // Lambda which returns the next instruction in the callee->caller chain in + // the call graph. This is the unique instruction which calls the computation + // containing 'instruction'. If more than one instruction calls the + // computation containing 'instruction' or no instructions call the + // computation then nullptr is returned. + auto next_caller = [this](HloInstruction* instruction) -> HloInstruction* { + const CallGraphNode& node = GetNode(instruction->parent()); + if (node.caller_callsites().size() != 1) { + return nullptr; + } + return node.caller_callsites()[0].instruction(); + }; + + // 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}; + } + } + } + return {nullptr, nullptr}; +} + string CallGraph::ToString() const { string out; Appendf(&out, "Call graph for module %s:\n", module_->name().c_str()); diff --git a/tensorflow/compiler/xla/service/call_graph.h b/tensorflow/compiler/xla/service/call_graph.h index e2fed044c88008d0a7e43f0166d397627ed72267..688c4085dfb4f47d3e08a4abee5e7b645f595b11 100644 --- a/tensorflow/compiler/xla/service/call_graph.h +++ b/tensorflow/compiler/xla/service/call_graph.h @@ -23,7 +23,6 @@ limitations under the License. #include "tensorflow/compiler/xla/service/hlo_computation.h" #include "tensorflow/compiler/xla/service/hlo_instruction.h" #include "tensorflow/compiler/xla/service/hlo_module.h" -#include "tensorflow/compiler/xla/statusor.h" #include "tensorflow/core/lib/gtl/flatmap.h" #include "tensorflow/core/lib/gtl/flatset.h" @@ -54,6 +53,8 @@ enum class CallContext { string CallContextToString(CallContext context); std::ostream& operator<<(std::ostream& out, const CallContext& context); +CallContext GetInstructionCallContext(const HloInstruction* instruction); + // Represents an HLO instruction which calls one or more computations. class CallSite { public: @@ -136,7 +137,7 @@ class CallGraphNode { // If instruction calls any computations adds a call site for this instruction // to the call graph node. If the instruction calls no computations then no // call site is added. - Status AddCallSiteForInstruction(HloInstruction* instruction); + void AddCallSiteForInstruction(HloInstruction* instruction); // Computation represented by this call graph node. HloComputation* computation_; @@ -172,12 +173,11 @@ class CallGraph { using VisitorFunction = std::function; // Builds and returns a call graph for the given HLO module. - static StatusOr> Build(const HloModule* module); + static std::unique_ptr Build(const HloModule* module); // Returns the node associated with the given computation. - StatusOr GetNode( - const HloComputation* computation) const; - StatusOr GetNode(const HloComputation* computation); + const CallGraphNode& GetNode(const HloComputation* computation) const; + CallGraphNode& GetNode(const HloComputation* computation); // Returns the vector of all nodes in the call graph. const std::vector& nodes() const { return nodes_; } @@ -189,22 +189,76 @@ class CallGraph { Status VisitNodes(const VisitorFunction& visitor_func, bool visit_unreachable_nodes = true) const; + // Returns true if 'a' dominates 'b' in the call graph. Computation 'a' + // dominates computation 'b' iff all callgraph paths in the caller-to-callee + // direction from a root computation to 'b' pass through computation + // 'a'. Trivially, a computation dominates itself. + bool Dominates(const HloComputation* a, const HloComputation* b) const; + + // Returns whether 'instruction' is contained in 'computation' either directly + // ('instruction->parent' is 'computation') or indirectly ('computation' + // dominates 'instruction->parent' in the call graph). + bool InstructionIsNestedIn(const HloInstruction* instruction, + const HloComputation* computation) const { + return Dominates(computation, instruction->parent()); + } + + // Returns the nearest call graph ancestors of instructions 'a' and 'b' for + // which the ancestors are in the same computation. An instruction is an call + // graph ancestor of 'a' if the instruction calls the computation containing + // 'a' either directly or transitively. Degeneratively an instruction is an + // ancestor of itself. nullptr is returned if there is no common ancestor or + // if the caller chain of 'a' or 'b' diverges (has multiple callers) before + // the nearest common ancestor. + // + // Example: + // + // Entry computation: + // %x = Call(A, {Constant(42.0)}) + // %y = Call(B, {%x}) + // + // Computation A: + // %a = Negate(Param()) + // + // Computation B: + // %b = Exp(Param()); + // + // If called with %a and %b, this function would return (%x, %y). %x is an + // ancestor of %a, and %y is an ancestor of %b, and %x and %y are in the same + // computation. + std::pair NearestAncestorsInSameComputation( + HloInstruction* a, HloInstruction* b) const; + + // Returns whether the call graph is flattened. A call graph is flattened if + // every computation called in a sequential context (eg, kWhile or kCall) has + // zero or one callsite, and no computation is called from both a parallel and + // sequential context. The call graph of a module can be flattened with + // FlattenCallGraph. + bool IsFlattened() const; + string ToString() const; private: CallGraph(const HloModule* module); // Sets the call contexts for every node in the graph. - Status SetCallContexts(); + void SetCallContexts(); // 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 // 'visited'. Status VisitNodesInternal( - const VisitorFunction& visitor_func, const CallGraphNode* node, + const VisitorFunction& visitor_func, const CallGraphNode& node, tensorflow::gtl::FlatSet* visited) const; + // Recursive helper for computing whether 'a' dominates 'b' in the call + // graph. 'b_ancestor' is the currently visited node (which starts at 'b'), + // and 'visited' is the set of computations which have been visited. + bool DominatesHelper( + const HloComputation* a, const HloComputation* b, + tensorflow::gtl::FlatSet* visited) const; + // The HLO module represented by this call graph. const HloModule* module_ = nullptr; diff --git a/tensorflow/compiler/xla/service/call_graph_test.cc b/tensorflow/compiler/xla/service/call_graph_test.cc index 65900fd4f86cd07d5d956da0df429d30fcdf7561..4243d37a77e10dce950d421f87a16d56e4829e4c 100644 --- a/tensorflow/compiler/xla/service/call_graph_test.cc +++ b/tensorflow/compiler/xla/service/call_graph_test.cc @@ -19,6 +19,7 @@ limitations under the License. #include "tensorflow/compiler/xla/service/hlo_computation.h" #include "tensorflow/compiler/xla/shape_util.h" #include "tensorflow/compiler/xla/status_macros.h" +#include "tensorflow/compiler/xla/test.h" #include "tensorflow/compiler/xla/test_helpers.h" #include "tensorflow/compiler/xla/tests/hlo_test_base.h" #include "tensorflow/compiler/xla/util.h" @@ -28,7 +29,7 @@ limitations under the License. namespace xla { namespace { -using testing::UnorderedMatcher; +using ::testing::UnorderedElementsAre; class CallGraphTest : public HloTestBase { protected: @@ -60,14 +61,15 @@ class CallGraphTest : public HloTestBase { // Build and return a computation which takes a scalar and calls (kCall) the // given computation with value 'callsites' number of times. std::unique_ptr MakeCallingComputation( - HloComputation* map_computation, int64 callsites) { - HloComputation::Builder builder(TestName() + ".CallingComputation"); + HloComputation* callee_computation, int64 callsites, + const string& suffix = ".CallingComputation") { + HloComputation::Builder builder(TestName() + suffix); HloInstruction* param0 = builder.AddInstruction( HloInstruction::CreateParameter(0, kScalarShape, "param0")); HloInstruction* last_value = param0; for (int64 i = 0; i < callsites; ++i) { last_value = builder.AddInstruction(HloInstruction::CreateCall( - kScalarShape, {last_value}, map_computation)); + kScalarShape, {last_value}, callee_computation)); } return builder.Build(); } @@ -79,7 +81,7 @@ class CallGraphTest : public HloTestBase { HloInstruction* param0 = builder.AddInstruction( HloInstruction::CreateParameter(0, kScalarShape, "param0")); HloInstruction* zero = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(0.0f))); + HloInstruction::CreateConstant(Literal::CreateR0(0.0f))); builder.AddInstruction(HloInstruction::CreateBinary( ShapeUtil::MakeShape(PRED, {}), HloOpcode::kGt, param0, zero)); return builder.Build(); @@ -90,116 +92,112 @@ class CallGraphTest : public HloTestBase { TEST_F(CallGraphTest, SingletonComputation) { // Test the call graph of a module with a single computation. - HloModule module(TestName()); + auto module = CreateNewModule(); HloComputation* computation = - module.AddEntryComputation(MakeScalarComputation()); - TF_ASSIGN_OR_ASSERT_OK(std::unique_ptr call_graph, - CallGraph::Build(&module)); + module->AddEntryComputation(MakeScalarComputation()); + std::unique_ptr call_graph = CallGraph::Build(module.get()); EXPECT_EQ(1, call_graph->nodes().size()); - TF_ASSIGN_OR_ASSERT_OK(const CallGraphNode* node, - call_graph->GetNode(computation)); - EXPECT_EQ(computation, node->computation()); - EXPECT_TRUE(node->callsites().empty()); - EXPECT_TRUE(node->callees().empty()); - EXPECT_TRUE(node->caller_callsites().empty()); - EXPECT_TRUE(node->callers().empty()); - EXPECT_EQ(CallContext::kSequential, node->context()); + EXPECT_TRUE(call_graph->IsFlattened()); + + const CallGraphNode& node = call_graph->GetNode(computation); + EXPECT_EQ(computation, node.computation()); + EXPECT_TRUE(node.callsites().empty()); + EXPECT_TRUE(node.callees().empty()); + EXPECT_TRUE(node.caller_callsites().empty()); + EXPECT_TRUE(node.callers().empty()); + EXPECT_EQ(CallContext::kSequential, node.context()); } TEST_F(CallGraphTest, UnreachableComputation) { // Test the call graph of a module with an entry computation and an // unreachable computation. - HloModule module(TestName()); + auto module = CreateNewModule(); HloComputation* entry_computation = - module.AddEntryComputation(MakeScalarComputation()); + module->AddEntryComputation(MakeScalarComputation()); HloComputation* unreachable_computation = - module.AddEmbeddedComputation(MakeScalarComputation()); + module->AddEmbeddedComputation(MakeScalarComputation()); - TF_ASSIGN_OR_ASSERT_OK(std::unique_ptr call_graph, - CallGraph::Build(&module)); + std::unique_ptr call_graph = CallGraph::Build(module.get()); EXPECT_EQ(2, call_graph->nodes().size()); - TF_ASSIGN_OR_ASSERT_OK(const CallGraphNode* entry_node, - call_graph->GetNode(entry_computation)); - EXPECT_EQ(entry_computation, entry_node->computation()); - EXPECT_EQ(CallContext::kSequential, entry_node->context()); + const CallGraphNode& entry_node = call_graph->GetNode(entry_computation); + EXPECT_EQ(entry_computation, entry_node.computation()); + EXPECT_EQ(CallContext::kSequential, entry_node.context()); - TF_ASSIGN_OR_ASSERT_OK(const CallGraphNode* unreachable_node, - call_graph->GetNode(unreachable_computation)); - EXPECT_EQ(unreachable_computation, unreachable_node->computation()); - EXPECT_EQ(CallContext::kSequential, unreachable_node->context()); + const CallGraphNode& unreachable_node = + call_graph->GetNode(unreachable_computation); + EXPECT_EQ(unreachable_computation, unreachable_node.computation()); + EXPECT_EQ(CallContext::kSequential, unreachable_node.context()); } 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. - HloModule module(TestName()); + auto module = CreateNewModule(); HloComputation* map_computation = - module.AddEmbeddedComputation(MakeScalarComputation()); - HloComputation* entry_computation = module.AddEmbeddedComputation( + module->AddEmbeddedComputation(MakeScalarComputation()); + HloComputation* entry_computation = module->AddEntryComputation( MakeMappingComputation(map_computation, /*callsites=*/5)); - TF_ASSIGN_OR_ASSERT_OK(std::unique_ptr call_graph, - CallGraph::Build(&module)); + std::unique_ptr call_graph = CallGraph::Build(module.get()); EXPECT_EQ(2, call_graph->nodes().size()); - TF_ASSIGN_OR_ASSERT_OK(const CallGraphNode* entry_node, - call_graph->GetNode(entry_computation)); - EXPECT_EQ(entry_computation, entry_node->computation()); - EXPECT_EQ(CallContext::kSequential, entry_node->context()); - EXPECT_EQ(5, entry_node->callsites().size()); - EXPECT_EQ(1, entry_node->callees().size()); - EXPECT_TRUE(entry_node->caller_callsites().empty()); - EXPECT_TRUE(entry_node->callers().empty()); - - TF_ASSIGN_OR_ASSERT_OK(const CallGraphNode* map_node, - call_graph->GetNode(map_computation)); - EXPECT_EQ(map_computation, map_node->computation()); - EXPECT_EQ(CallContext::kParallel, map_node->context()); - EXPECT_TRUE(map_node->callsites().empty()); - EXPECT_TRUE(map_node->callees().empty()); - EXPECT_EQ(5, map_node->caller_callsites().size()); - EXPECT_EQ(1, map_node->callers().size()); + const CallGraphNode& entry_node = call_graph->GetNode(entry_computation); + EXPECT_EQ(entry_computation, entry_node.computation()); + EXPECT_EQ(CallContext::kSequential, entry_node.context()); + EXPECT_EQ(5, entry_node.callsites().size()); + EXPECT_EQ(1, entry_node.callees().size()); + EXPECT_TRUE(entry_node.caller_callsites().empty()); + EXPECT_TRUE(entry_node.callers().empty()); + + const CallGraphNode& map_node = call_graph->GetNode(map_computation); + EXPECT_EQ(map_computation, map_node.computation()); + EXPECT_EQ(CallContext::kParallel, map_node.context()); + EXPECT_TRUE(map_node.callsites().empty()); + EXPECT_TRUE(map_node.callees().empty()); + EXPECT_EQ(5, map_node.caller_callsites().size()); + EXPECT_EQ(1, map_node.callers().size()); } 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. - HloModule module(TestName()); + auto module = CreateNewModule(); HloComputation* called_computation = - module.AddEmbeddedComputation(MakeScalarComputation()); - HloComputation* entry_computation = module.AddEmbeddedComputation( + module->AddEmbeddedComputation(MakeScalarComputation()); + HloComputation* entry_computation = module->AddEntryComputation( MakeCallingComputation(called_computation, /*callsites=*/3)); - TF_ASSIGN_OR_ASSERT_OK(std::unique_ptr call_graph, - CallGraph::Build(&module)); + std::unique_ptr call_graph = CallGraph::Build(module.get()); EXPECT_EQ(2, call_graph->nodes().size()); - TF_ASSIGN_OR_ASSERT_OK(const CallGraphNode* entry_node, - call_graph->GetNode(entry_computation)); - EXPECT_EQ(entry_computation, entry_node->computation()); - EXPECT_EQ(CallContext::kSequential, entry_node->context()); - EXPECT_EQ(3, entry_node->callsites().size()); - EXPECT_EQ(1, entry_node->callees().size()); - EXPECT_TRUE(entry_node->caller_callsites().empty()); - EXPECT_TRUE(entry_node->callers().empty()); - - TF_ASSIGN_OR_ASSERT_OK(const CallGraphNode* called_node, - call_graph->GetNode(called_computation)); - EXPECT_EQ(called_computation, called_node->computation()); - EXPECT_EQ(CallContext::kSequential, called_node->context()); - EXPECT_TRUE(called_node->callsites().empty()); - EXPECT_TRUE(called_node->callees().empty()); - EXPECT_EQ(3, called_node->caller_callsites().size()); - EXPECT_EQ(1, called_node->callers().size()); + // The called computation is only called from one other computation, but there + // are multiple callsites. + EXPECT_FALSE(call_graph->IsFlattened()); + + const CallGraphNode& entry_node = call_graph->GetNode(entry_computation); + EXPECT_EQ(entry_computation, entry_node.computation()); + EXPECT_EQ(CallContext::kSequential, entry_node.context()); + EXPECT_EQ(3, entry_node.callsites().size()); + EXPECT_EQ(1, entry_node.callees().size()); + EXPECT_TRUE(entry_node.caller_callsites().empty()); + EXPECT_TRUE(entry_node.callers().empty()); + + const CallGraphNode& called_node = call_graph->GetNode(called_computation); + EXPECT_EQ(called_computation, called_node.computation()); + EXPECT_EQ(CallContext::kSequential, called_node.context()); + EXPECT_TRUE(called_node.callsites().empty()); + EXPECT_TRUE(called_node.callees().empty()); + EXPECT_EQ(3, called_node.caller_callsites().size()); + EXPECT_EQ(1, called_node.callers().size()); } 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. - HloModule module(TestName()); + auto module = CreateNewModule(); HloComputation* subcomputation = - module.AddEmbeddedComputation(MakeScalarComputation()); + module->AddEmbeddedComputation(MakeScalarComputation()); HloComputation::Builder builder(TestName()); HloInstruction* param0 = builder.AddInstruction( @@ -209,34 +207,33 @@ TEST_F(CallGraphTest, ContextBothComputations) { HloInstruction* map = builder.AddInstruction( HloInstruction::CreateMap(kScalarShape, {call}, subcomputation)); HloComputation* entry_computation = - module.AddEmbeddedComputation(builder.Build()); + module->AddEntryComputation(builder.Build()); - TF_ASSIGN_OR_ASSERT_OK(std::unique_ptr call_graph, - CallGraph::Build(&module)); + std::unique_ptr call_graph = CallGraph::Build(module.get()); EXPECT_EQ(2, call_graph->nodes().size()); - TF_ASSIGN_OR_ASSERT_OK(const CallGraphNode* entry_node, - call_graph->GetNode(entry_computation)); - EXPECT_EQ(entry_computation, entry_node->computation()); - EXPECT_EQ(2, entry_node->callsites().size()); + EXPECT_FALSE(call_graph->IsFlattened()); + + const CallGraphNode& entry_node = call_graph->GetNode(entry_computation); + EXPECT_EQ(entry_computation, entry_node.computation()); + EXPECT_EQ(2, entry_node.callsites().size()); - const CallSite& call_callsite = entry_node->callsites()[0]; + const CallSite& call_callsite = entry_node.callsites()[0]; EXPECT_EQ(call, call_callsite.instruction()); - EXPECT_MATCH(call_callsite.called_computations(), - UnorderedMatcher(subcomputation)); + EXPECT_THAT(call_callsite.called_computations(), + UnorderedElementsAre(subcomputation)); EXPECT_EQ(CallContext::kSequential, call_callsite.context()); - EXPECT_EQ(entry_node->GetCallSite(call), &call_callsite); + EXPECT_EQ(entry_node.GetCallSite(call), &call_callsite); - const CallSite& map_callsite = entry_node->callsites()[1]; + const CallSite& map_callsite = entry_node.callsites()[1]; EXPECT_EQ(map, map_callsite.instruction()); - EXPECT_MATCH(map_callsite.called_computations(), - UnorderedMatcher(subcomputation)); + EXPECT_THAT(map_callsite.called_computations(), + UnorderedElementsAre(subcomputation)); EXPECT_EQ(CallContext::kParallel, map_callsite.context()); - EXPECT_EQ(entry_node->GetCallSite(map), &map_callsite); + EXPECT_EQ(entry_node.GetCallSite(map), &map_callsite); - TF_ASSIGN_OR_ASSERT_OK(const CallGraphNode* sub_node, - call_graph->GetNode(subcomputation)); - EXPECT_EQ(CallContext::kBoth, sub_node->context()); + const CallGraphNode& sub_node = call_graph->GetNode(subcomputation); + EXPECT_EQ(CallContext::kBoth, sub_node.context()); } TEST_F(CallGraphTest, ComplexGraph) { @@ -252,12 +249,12 @@ TEST_F(CallGraphTest, ComplexGraph) { // c // // Calls are made via kCall, kWhile, and kMap instructions. - HloModule module(TestName()); + auto module = CreateNewModule(); HloComputation* cond_computation = - module.AddEmbeddedComputation(MakeConditionComputation()); + module->AddEmbeddedComputation(MakeConditionComputation()); HloComputation* c_computation = - module.AddEmbeddedComputation(MakeScalarComputation()); - HloComputation* b_computation = module.AddEmbeddedComputation( + module->AddEmbeddedComputation(MakeScalarComputation()); + HloComputation* b_computation = module->AddEmbeddedComputation( MakeMappingComputation(c_computation, /*callsites=*/1)); HloComputation* a_computation; @@ -269,7 +266,7 @@ TEST_F(CallGraphTest, ComplexGraph) { HloInstruction::CreateCall(kScalarShape, {param0}, c_computation)); builder.AddInstruction(HloInstruction::CreateWhile( kScalarShape, cond_computation, b_computation, call)); - a_computation = module.AddEmbeddedComputation(builder.Build()); + a_computation = module->AddEmbeddedComputation(builder.Build()); } HloComputation* entry_computation; @@ -279,30 +276,28 @@ TEST_F(CallGraphTest, ComplexGraph) { HloInstruction::CreateParameter(0, kScalarShape, "param0")); builder.AddInstruction(HloInstruction::CreateWhile( kScalarShape, cond_computation, a_computation, param0)); - entry_computation = module.AddEntryComputation(builder.Build()); + entry_computation = module->AddEntryComputation(builder.Build()); } - TF_ASSIGN_OR_ASSERT_OK(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()); // Entry computation has one while instruction calling two computations // (cond_computation and a_computation). - TF_ASSIGN_OR_ASSERT_OK(const CallGraphNode* entry_node, - call_graph->GetNode(entry_computation)); - ASSERT_EQ(1, entry_node->callsites().size()); + 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(); - EXPECT_MATCH(called_computations, - UnorderedMatcher(cond_computation, a_computation)); - EXPECT_EQ(CallContext::kSequential, entry_node->context()); - - TF_ASSIGN_OR_ASSERT_OK(const CallGraphNode* c_node, - call_graph->GetNode(c_computation)); - EXPECT_TRUE(c_node->callsites().empty()); - EXPECT_MATCH(c_node->callers(), - UnorderedMatcher(a_computation, b_computation)); - EXPECT_EQ(CallContext::kBoth, c_node->context()); + entry_node.callsites()[0].called_computations(); + EXPECT_THAT(called_computations, + 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)); + EXPECT_EQ(CallContext::kBoth, c_node.context()); // Visit the graph and verify nodes were visited in callee-before-caller // order. @@ -328,33 +323,134 @@ TEST_F(CallGraphTest, ComplexGraph) { EXPECT_LT(index_of(cond_computation), index_of(a_computation)); EXPECT_LT(index_of(c_computation), index_of(b_computation)); EXPECT_LT(index_of(b_computation), index_of(a_computation)); + + // Verify dominance relations between computation in the graph. + + // Entry dominates everybody, and is dominated by no one except itself. + EXPECT_TRUE(call_graph->Dominates(entry_computation, entry_computation)); + EXPECT_TRUE(call_graph->Dominates(entry_computation, a_computation)); + EXPECT_TRUE(call_graph->Dominates(entry_computation, b_computation)); + EXPECT_TRUE(call_graph->Dominates(entry_computation, c_computation)); + EXPECT_TRUE(call_graph->Dominates(entry_computation, cond_computation)); + EXPECT_FALSE(call_graph->Dominates(a_computation, entry_computation)); + EXPECT_FALSE(call_graph->Dominates(b_computation, entry_computation)); + EXPECT_FALSE(call_graph->Dominates(c_computation, entry_computation)); + EXPECT_FALSE(call_graph->Dominates(cond_computation, entry_computation)); + + // 'a' only dominates 'b' and 'c'. + EXPECT_TRUE(call_graph->Dominates(a_computation, a_computation)); + EXPECT_TRUE(call_graph->Dominates(a_computation, b_computation)); + EXPECT_TRUE(call_graph->Dominates(a_computation, c_computation)); + EXPECT_FALSE(call_graph->Dominates(b_computation, a_computation)); + EXPECT_FALSE(call_graph->Dominates(c_computation, a_computation)); + EXPECT_FALSE(call_graph->Dominates(a_computation, cond_computation)); + + EXPECT_TRUE(call_graph->Dominates(b_computation, b_computation)); + EXPECT_FALSE(call_graph->Dominates(b_computation, c_computation)); + EXPECT_FALSE(call_graph->Dominates(b_computation, cond_computation)); + + EXPECT_TRUE(call_graph->Dominates(c_computation, c_computation)); + EXPECT_FALSE(call_graph->Dominates(c_computation, cond_computation)); + EXPECT_FALSE(call_graph->Dominates(cond_computation, c_computation)); + + EXPECT_TRUE(call_graph->Dominates(cond_computation, cond_computation)); +} + +TEST_F(CallGraphTest, ComplexGraphNearestAncestors) { + // Test NearestAncestorsInSameComputation on a call graph of a module with + // several computation called in various contexts. The call graph looks like: + // + // entry + // / | + // a | + // / | \ | + // b | cond + // \ | + // c + // + // Calls are made via kCall, kWhile, and kMap instructions. + auto module = CreateNewModule(); + HloComputation* cond_computation = + module->AddEmbeddedComputation(MakeConditionComputation()); + HloComputation* c_computation = + module->AddEmbeddedComputation(MakeScalarComputation()); + HloComputation* b_computation = module->AddEmbeddedComputation( + MakeMappingComputation(c_computation, /*callsites=*/1)); + HloInstruction* b_map = b_computation->root_instruction(); + + HloComputation* a_computation; + HloInstruction* a_call; + HloInstruction* a_while; + { + HloComputation::Builder builder(TestName() + ".a"); + HloInstruction* param0 = builder.AddInstruction( + HloInstruction::CreateParameter(0, kScalarShape, "param0")); + a_call = builder.AddInstruction( + HloInstruction::CreateCall(kScalarShape, {param0}, c_computation)); + a_while = builder.AddInstruction(HloInstruction::CreateWhile( + kScalarShape, cond_computation, b_computation, a_call)); + a_computation = module->AddEmbeddedComputation(builder.Build()); + } + + HloComputation* entry_computation; + HloInstruction* entry_while; + { + HloComputation::Builder builder(TestName() + ".entry"); + HloInstruction* param0 = builder.AddInstruction( + HloInstruction::CreateParameter(0, kScalarShape, "param0")); + entry_while = builder.AddInstruction(HloInstruction::CreateWhile( + kScalarShape, cond_computation, a_computation, param0)); + entry_computation = module->AddEntryComputation(builder.Build()); + } + + std::unique_ptr call_graph = CallGraph::Build(module.get()); + EXPECT_EQ(5, call_graph->nodes().size()); + + // Verify NearestAncestorsInSameComputation for various instructions in the + // module. + EXPECT_EQ(call_graph->NearestAncestorsInSameComputation(a_call, a_call), + std::make_pair(a_call, a_call)); + + // c_computation is called from more than one site, so + // NearestAncestorsInSameComputation bails and returns nullptrs. + std::pair null_pair = {nullptr, nullptr}; + EXPECT_EQ(call_graph->NearestAncestorsInSameComputation( + b_map, c_computation->root_instruction()), + null_pair); + + EXPECT_EQ(call_graph->NearestAncestorsInSameComputation(b_map, entry_while), + std::make_pair(entry_while, entry_while)); + EXPECT_EQ(call_graph->NearestAncestorsInSameComputation(b_map, a_call), + std::make_pair(a_while, a_call)); + EXPECT_EQ(call_graph->NearestAncestorsInSameComputation(a_while, a_call), + std::make_pair(a_while, a_call)); + EXPECT_EQ(call_graph->NearestAncestorsInSameComputation(a_while, b_map), + std::make_pair(a_while, a_while)); } TEST_F(CallGraphTest, VisitSingletonComputation) { // Test the call graph visitor with a call graph with a single node. - HloModule module(TestName()); + auto module = CreateNewModule(); HloComputation* computation = - module.AddEntryComputation(MakeScalarComputation()); - TF_ASSIGN_OR_ASSERT_OK(std::unique_ptr call_graph, - CallGraph::Build(&module)); + module->AddEntryComputation(MakeScalarComputation()); + std::unique_ptr call_graph = CallGraph::Build(module.get()); std::vector visited; TF_ASSERT_OK(call_graph->VisitNodes([&visited](const CallGraphNode& node) { visited.push_back(node.computation()); return Status::OK(); })); - EXPECT_MATCH(visited, UnorderedMatcher(computation)); + EXPECT_THAT(visited, UnorderedElementsAre(computation)); } TEST_F(CallGraphTest, VisitUnreachableComputation) { // Test the call graph visitor with a call graph with an unreachable node. - HloModule module(TestName()); + auto module = CreateNewModule(); HloComputation* entry_computation = - module.AddEntryComputation(MakeScalarComputation()); + module->AddEntryComputation(MakeScalarComputation()); HloComputation* unreachable_computation = - module.AddEmbeddedComputation(MakeScalarComputation()); - TF_ASSIGN_OR_ASSERT_OK(std::unique_ptr call_graph, - CallGraph::Build(&module)); + module->AddEmbeddedComputation(MakeScalarComputation()); + std::unique_ptr call_graph = CallGraph::Build(module.get()); // Test visitation of only reachable nodes. { @@ -379,25 +475,29 @@ TEST_F(CallGraphTest, VisitUnreachableComputation) { }, /*visit_unreachable_nodes=*/true)); EXPECT_EQ(visited.size(), 2); - EXPECT_MATCH(visited, - UnorderedMatcher(entry_computation, unreachable_computation)); + EXPECT_THAT(visited, UnorderedElementsAre(entry_computation, + unreachable_computation)); } } TEST_F(CallGraphTest, VisitWithError) { // Test that the call graph visitor properly propagates errors. - HloModule module(TestName()); - module.AddEntryComputation(MakeScalarComputation()); - TF_ASSIGN_OR_ASSERT_OK(std::unique_ptr call_graph, - CallGraph::Build(&module)); + auto module = CreateNewModule(); + module->AddEntryComputation(MakeScalarComputation()); + std::unique_ptr call_graph = CallGraph::Build(module.get()); Status status = call_graph->VisitNodes( [](const CallGraphNode&) { return InternalError("Visitation failed"); }); ASSERT_FALSE(status.ok()); ASSERT_EQ(status.code(), tensorflow::error::INTERNAL); - ASSERT_MATCH(status.error_message(), testing::HasSubstr("Visitation failed")); + ASSERT_THAT(status.error_message(), + ::testing::HasSubstr("Visitation failed")); } } // namespace } // namespace xla + +int main(int argc, char** argv) { + return xla::ParseDebugOptionsFlagsAndRunTests(argc, argv); +} diff --git a/tensorflow/compiler/xla/service/call_inliner.cc b/tensorflow/compiler/xla/service/call_inliner.cc new file mode 100644 index 0000000000000000000000000000000000000000..817b59f7627fa53f74b90e9a33994688ac2ac8c9 --- /dev/null +++ b/tensorflow/compiler/xla/service/call_inliner.cc @@ -0,0 +1,156 @@ +/* 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/call_inliner.h" + +#include + +#include "tensorflow/core/lib/core/errors.h" + +namespace xla { + +StatusOr CallInliner::Run(HloModule* module) { + std::deque work_queue; + + // Seed the work queue with call instructions from the main computation. + TF_RETURN_IF_ERROR( + module->entry_computation()->Accept([&](HloInstruction* hlo) { + if (hlo->opcode() == HloOpcode::kCall) { + work_queue.push_back(hlo); + } + return Status::OK(); + })); + + VLOG(1) << "Work queue seeded with " << work_queue.size() << " entries."; + + bool mutated = false; + while (!work_queue.empty()) { + mutated = true; + HloInstruction* call = work_queue.front(); + work_queue.pop_front(); + TF_RETURN_IF_ERROR(ReplaceWithInlinedBody(call, &work_queue)); + } + return mutated; +} + +// Traverses the callee computation, inlining cloned nodes into the caller +// computation and connecting them to producers/consumers appropriately. +// When the traversal has completed, the provided call instruction is entriely +// replaced in the caller's graph, and any calls encountered in the callee +// computation have been added to the work_queue. +class SubcomputationInsertionVisitor : public DfsHloVisitorWithDefault { + public: + SubcomputationInsertionVisitor(HloInstruction* call, + std::deque* work_queue) + : call_(call), outer_(call->parent()), work_queue_(work_queue) {} + + // Resolves the operands to the HLO instruction in the inlined (caller) graph, + // and clones the HLO instruction into that graph with the new operands. + // If the instruction is a call, it is added to the work queue. + Status DefaultAction(HloInstruction* hlo) override { + std::vector new_operands; + for (HloInstruction* operand : hlo->operands()) { + TF_ASSIGN_OR_RETURN(HloInstruction * new_operand, Resolve(operand)); + new_operands.push_back(new_operand); + } + VLOG(1) << "Cloning HLO and adding to caller: " << hlo->ToString(); + auto new_hlo = hlo->CloneWithNewOperands(hlo->shape(), new_operands); + HloInstruction* new_hlo_pointer = + outer_->AddInstruction(std::move(new_hlo)); + TF_RETURN_IF_ERROR(NoteMapping(hlo, new_hlo_pointer)); + + // Account for control edges. + for (HloInstruction* control_predecessor : hlo->control_predecessors()) { + TF_ASSIGN_OR_RETURN(HloInstruction * new_control_predecessor, + Resolve(control_predecessor)); + TF_RETURN_IF_ERROR( + new_control_predecessor->AddControlDependencyTo(new_hlo_pointer)); + } + + if (new_hlo_pointer->opcode() == HloOpcode::kCall) { + VLOG(1) << "Adding new call HLO to work queue."; + // Call instructions we observe in the subcomputation are added to the + // inliner work queue. + work_queue_->push_back(new_hlo_pointer); + } + return Status::OK(); + } + + // Does not create new nodes for the parameter; rather, notes the mapping from + // the subcomputation parameter node to the call operands in the caller + // computation. + Status HandleParameter(HloInstruction* parameter) override { + TF_RETURN_IF_ERROR(NoteMapping( + parameter, call_->mutable_operand(parameter->parameter_number()))); + return Status::OK(); + } + + // Wires the consumers of the call to instead point at the newly created root, + // replacing the call operation in the caller computation. + Status FinishVisit(HloInstruction* root) override { + TF_ASSIGN_OR_RETURN(HloInstruction * new_root, Resolve(root)); + VLOG(1) << "Replacing all uses of " << call_->ToString() + << " with new root " << new_root->ToString(); + return outer_->ReplaceInstruction(call_, new_root); + } + + private: + // Resolves the callee subcomputation_hlo to the new (inline) HLO in the + // caller computation, or returns a NotFound error if that subcomputation HLO + // has not been mapped. + StatusOr Resolve(HloInstruction* subcomputation_hlo) { + auto it = subcomputation_hlo_to_new_hlo_.find(subcomputation_hlo); + if (it == subcomputation_hlo_to_new_hlo_.end()) { + return NotFound( + "Could not find mapping from subcomputation HLO %s to a cloned HLO.", + subcomputation_hlo->ToString().c_str()); + } + return it->second; + } + + // Notes that the given subcomputation_hlo in the callee has been mapped to + // the (inline) new_hlo in the caller computation. + // + // Returns an error status if the subcomputation_hlo is mapped more than + // once. + Status NoteMapping(HloInstruction* subcomputation_hlo, + HloInstruction* new_hlo) { + auto result = subcomputation_hlo_to_new_hlo_.insert( + std::make_pair(subcomputation_hlo, new_hlo)); + TF_RET_CHECK(result.second) + << "A mapping for the subcomputation HLO is already present."; + return Status::OK(); + } + + HloInstruction* call_; + HloComputation* outer_; + std::unordered_map + subcomputation_hlo_to_new_hlo_; + std::deque* work_queue_; +}; + +Status CallInliner::ReplaceWithInlinedBody( + HloInstruction* call, std::deque* work_queue) { + TF_RET_CHECK(call->opcode() == HloOpcode::kCall); + TF_RET_CHECK(call->called_computations().size() == 1); + HloComputation* called = call->called_computations()[0]; + VLOG(1) << "Replacing call " << call->ToString() << " with inlined body of " + << called->name(); + + SubcomputationInsertionVisitor visitor(call, work_queue); + return called->Accept(&visitor); +} + +} // namespace xla diff --git a/tensorflow/compiler/xla/service/call_inliner.h b/tensorflow/compiler/xla/service/call_inliner.h new file mode 100644 index 0000000000000000000000000000000000000000..8647edffa7f04c93fa7c393ead12d7d5a2b93955 --- /dev/null +++ b/tensorflow/compiler/xla/service/call_inliner.h @@ -0,0 +1,48 @@ +/* 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__CALL_INLINER_H_ +#define TENSORFLOW_COMPILER_XLA_SERVICE__CALL_INLINER_H_ + +#include + +#include "tensorflow/compiler/xla/service/hlo_pass_interface.h" +#include "tensorflow/compiler/xla/statusor.h" + +namespace xla { + +// For every kCall operation in the main computation, we inline the body of the +// called function, and proceed recursively. +class CallInliner : public HloPassInterface { + public: + ~CallInliner() override = default; + tensorflow::StringPiece name() const override { return "CallInliner"; } + + StatusOr Run(HloModule* module) override; + + private: + // Replaces the given call operation -- which must be an operation inside the + // entry computation with opcode kCall -- with the called computation's body, + // such that the called computation is inline in the entry computation. + // + // On successful inlining, the inlined computation may have itself contained + // calls; if so, they are added to the work_queue. + Status ReplaceWithInlinedBody(HloInstruction* call, + std::deque* work_queue); +}; + +} // namespace xla + +#endif // TENSORFLOW_COMPILER_XLA_SERVICE__CALL_INLINER_H_ diff --git a/tensorflow/compiler/xla/service/call_inliner_test.cc b/tensorflow/compiler/xla/service/call_inliner_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..77528d0b75f85f262141d98517db74bd29e1500d --- /dev/null +++ b/tensorflow/compiler/xla/service/call_inliner_test.cc @@ -0,0 +1,77 @@ +/* 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/call_inliner.h" + +#include +#include + +#include "tensorflow/compiler/xla/layout_util.h" +#include "tensorflow/compiler/xla/literal_util.h" +#include "tensorflow/compiler/xla/ptr_util.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_pass_fix.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/types.h" +#include "tensorflow/compiler/xla/xla_data.pb.h" +#include "tensorflow/core/lib/core/status_test_util.h" +#include "tensorflow/core/lib/strings/str_util.h" + +namespace op = xla::testing::opcode_matchers; + +namespace xla { +namespace { + +// Tests for call inlining that are most tractable at the HLO level (vs +// ComputationBuilder API in call_test.cc). +using CallInlinerTest = HloTestBase; + +TEST_F(CallInlinerTest, ControlDependenciesAreCarriedToCaller) { + HloComputation::Builder inner(TestName() + ".inner"); + HloInstruction* zero = inner.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(24.0f))); + HloInstruction* one = inner.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(42.0f))); + TF_ASSERT_OK(zero->AddControlDependencyTo(one)); + auto module = CreateNewModule(); + HloComputation* inner_computation = + module->AddEmbeddedComputation(inner.Build()); + + HloComputation::Builder outer(TestName() + ".outer"); + Shape r0f32 = ShapeUtil::MakeShape(F32, {}); + outer.AddInstruction( + HloInstruction::CreateCall(r0f32, {}, inner_computation)); + + auto computation = module->AddEntryComputation(outer.Build()); + + CallInliner call_inliner; + 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(), + 42); + ASSERT_EQ(1, computation->root_instruction()->control_predecessors().size()); + auto prior = computation->root_instruction()->control_predecessors()[0]; + EXPECT_THAT(prior, op::Constant()); + EXPECT_EQ(prior->literal().GetFirstElement(), 24); +} + +} // namespace +} // namespace xla diff --git a/tensorflow/compiler/xla/service/compile_only_service.cc b/tensorflow/compiler/xla/service/compile_only_service.cc new file mode 100644 index 0000000000000000000000000000000000000000..62dab56a71cfd172aa73a781d1879ed4f0f4b66b --- /dev/null +++ b/tensorflow/compiler/xla/service/compile_only_service.cc @@ -0,0 +1,112 @@ +/* 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/compile_only_service.h" + +#include +#include +#include + +#include "tensorflow/compiler/xla/legacy_flags/debug_options_flags.h" +#include "tensorflow/compiler/xla/service/backend.h" +#include "tensorflow/compiler/xla/service/computation_layout.h" +#include "tensorflow/compiler/xla/service/computation_tracker.h" +#include "tensorflow/compiler/xla/service/platform_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/gtl/cleanup.h" +#include "tensorflow/core/lib/strings/strcat.h" +#include "tensorflow/core/platform/logging.h" +#include "tensorflow/core/platform/stream_executor_no_cuda.h" + +namespace xla { + +/* static */ StatusOr> +CompileOnlyService::NewService(perftools::gputools::Platform* platform) { + ServiceOptions default_options; + default_options.set_platform(platform); + return NewService(default_options); +} + +/* static */ StatusOr> +CompileOnlyService::NewService(const ServiceOptions& options) { + perftools::gputools::Platform* platform = options.platform(); + if (platform == nullptr) { + TF_ASSIGN_OR_RETURN(platform, PlatformUtil::GetDefaultPlatform()); + } + + TF_ASSIGN_OR_RETURN(auto compiler, Compiler::GetForPlatform(platform)); + + std::unique_ptr service( + new CompileOnlyService(options, compiler)); + return std::move(service); +} + +CompileOnlyService::CompileOnlyService(const ServiceOptions& options, + Compiler* compiler) + : Service(options, /*execute_backend=*/nullptr), compiler_(compiler) {} + +StatusOr>> +CompileOnlyService::CompileAheadOfTime( + const tensorflow::gtl::ArraySlice computations, + const AotCompilationOptions& options) { + std::vector> hlo_modules; + for (const AotComputationInstance& instance : computations) { + TF_ASSIGN_OR_RETURN(UserComputation * user_computation, + computation_tracker_.Resolve(instance.computation)); + VersionedComputationHandle versioned_handle = + user_computation->GetVersionedHandle(); + + // TODO(b/63773457): Track DebugOptions in AotCompilationOptions. + DebugOptions debug_options = legacy_flags::GetDebugOptionsFromFlags(); + + // Dump computation proto state if flag is set. + const string& directory_path = debug_options.xla_dump_computations_to(); + if (!directory_path.empty()) { + TF_ASSIGN_OR_RETURN( + std::unique_ptr session_module, + computation_tracker_.SnapshotComputation(versioned_handle.handle)); + string filename = tensorflow::strings::StrCat( + "computation_", versioned_handle.handle.handle(), "__", + session_module->entry().name(), "__version_", + versioned_handle.version); + TF_RETURN_IF_ERROR(Executable::DumpToDirectory(directory_path, filename, + *session_module)); + } + + TF_ASSIGN_OR_RETURN( + std::shared_ptr program_shape, + user_computation->ComputeProgramShape(versioned_handle.version)); + + ExecutionOptions execution_options; + *execution_options.mutable_debug_options() = debug_options; + TF_ASSIGN_OR_RETURN( + std::unique_ptr module_config, + CreateModuleConfig(*program_shape, instance.argument_layouts, + &execution_options, + /*has_hybrid_result=*/false)); + + TF_ASSIGN_OR_RETURN(std::unique_ptr hlo_module, + computation_tracker_.BuildHloModule( + versioned_handle, *module_config, + /*include_unreachable_instructions=*/true)); + hlo_modules.push_back(std::move(hlo_module)); + } + + return compiler_->CompileAheadOfTime(std::move(hlo_modules), options); +} + +} // namespace xla diff --git a/tensorflow/compiler/xla/service/compile_only_service.h b/tensorflow/compiler/xla/service/compile_only_service.h new file mode 100644 index 0000000000000000000000000000000000000000..9859941c6c17460939e5b6817f1c7c415e63443c --- /dev/null +++ b/tensorflow/compiler/xla/service/compile_only_service.h @@ -0,0 +1,118 @@ +/* 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_COMPILE_ONLY_SERVICE_H_ +#define TENSORFLOW_COMPILER_XLA_SERVICE_COMPILE_ONLY_SERVICE_H_ + +#include "tensorflow/compiler/xla/service/backend.h" +#include "tensorflow/compiler/xla/service/compiler.h" +#include "tensorflow/compiler/xla/service/service.h" +#include "tensorflow/compiler/xla/statusor.h" +#include "tensorflow/compiler/xla/xla_data.pb.h" +#include "tensorflow/core/platform/stream_executor_no_cuda.h" + +namespace xla { + +// An XLA Service specialization for ahead-of-time compilation. This only +// instantiates a Compiler object for the relevant platform; it does not +// instantiate or require an execution backend. +class CompileOnlyService : public Service { + public: + // Factory for creating a CompileOnlyService. The parameter platform is the + // platform that the service should target. If platform is null then the + // default platform is used. + static StatusOr> NewService( + perftools::gputools::Platform* platform); + static StatusOr> NewService( + const ServiceOptions& options); + + // A description of a computation to compile using CompileAheadOfTime. + struct AotComputationInstance { + ComputationHandle computation; + std::vector argument_layouts; + const Shape* result_layout = nullptr; + }; + + // Compiles a list of computations for ahead-of-time execution. This is + // intended for use in static compilation. See + // |CompileOnlyClient::CompileAheadOfTime| for additional details. + StatusOr>> + CompileAheadOfTime( + const tensorflow::gtl::ArraySlice computations, + const AotCompilationOptions& Options); + + // Override Service methods that require or imply the existence of an + // execute backend. Note that this does not include TransferToClient, as + // computing constants produces global data that we may wish to transfer. + tensorflow::Status Execute(const ExecuteRequest* arg, + ExecuteResponse* result) override { + return Unimplemented("CompileOnlyService does not support execution."); + } + tensorflow::Status ExecuteParallel(const ExecuteParallelRequest* arg, + ExecuteParallelResponse* result) override { + return Unimplemented("CompileOnlyService does not support execution."); + } + tensorflow::Status GetDeviceHandles( + const GetDeviceHandlesRequest* arg, + GetDeviceHandlesResponse* result) override { + return Unimplemented("CompileOnlyService does not support devices."); + } + tensorflow::Status ExecuteAsync(const ExecuteAsyncRequest* arg, + ExecuteAsyncResponse* result) override { + return Unimplemented("CompileOnlyService does not support execution."); + } + tensorflow::Status WaitForExecution( + const WaitForExecutionRequest* arg, + WaitForExecutionResponse* result) override { + return Unimplemented("CompileOnlyService does not support execution."); + } + tensorflow::Status TransferToServer( + const TransferToServerRequest* arg, + TransferToServerResponse* result) override { + return Unimplemented( + "CompileOnlyService does not support device data transfers."); + } + tensorflow::Status TransferToInfeed( + const TransferToInfeedRequest* arg, + TransferToInfeedResponse* result) override { + return Unimplemented( + "CompileOnlyService does not support device data transfers."); + } + tensorflow::Status TransferFromOutfeed( + const TransferFromOutfeedRequest* arg, + TransferFromOutfeedResponse* result) override { + return Unimplemented( + "CompileOnlyService does not support device data transfers."); + } + tensorflow::Status ResetDevice(const ResetDeviceRequest* arg, + ResetDeviceResponse* result) override { + return Unimplemented("CompileOnlyService does not support devices."); + } + + private: + explicit CompileOnlyService(const ServiceOptions& options, + Compiler* compiler); + CompileOnlyService(const CompileOnlyService&) = delete; + void operator=(const CompileOnlyService&) = delete; + + // The compiler for the target platform. This is included in place of + // the Service::execute_backend_'s compiler, since execute_backend_ is a + // nullptr in CompileOnlyService. + Compiler* compiler_; +}; + +} // namespace xla + +#endif // TENSORFLOW_COMPILER_XLA_SERVICE_COMPILE_ONLY_SERVICE_H_ diff --git a/tensorflow/compiler/xla/service/compiler.h b/tensorflow/compiler/xla/service/compiler.h index 6f43c9b8040e9b21e7c0fcf86e2dc5b8ff8c6475..d5bd9214be44f4abd5f672168335ae1a259c9118 100644 --- a/tensorflow/compiler/xla/service/compiler.h +++ b/tensorflow/compiler/xla/service/compiler.h @@ -28,6 +28,7 @@ limitations under the License. #include "tensorflow/compiler/xla/service/executable.h" #include "tensorflow/compiler/xla/service/hlo_module.h" #include "tensorflow/compiler/xla/service/hlo_module_config.h" +#include "tensorflow/compiler/xla/service/logical_buffer.h" #include "tensorflow/compiler/xla/statusor.h" #include "tensorflow/compiler/xla/types.h" #include "tensorflow/core/lib/gtl/array_slice.h" @@ -91,13 +92,6 @@ class AotCompilationOptions { // platform. class Compiler { public: - // Callback signature used to dump the HLO graph during compilation. - // Different compiler backends will call this as they please, providing - // a view of the HLO at different points in compilation -- context for the - // dump is indicated by the label string. - using HloDumper = - std::function; - virtual ~Compiler() {} // Returns the ID of the platform that this compiler targets. @@ -113,25 +107,20 @@ class Compiler { // Use the overload below to compile computations that run in parallel. virtual StatusOr> Compile( std::unique_ptr module, - std::unique_ptr module_config, HloDumper dump_hlo, perftools::gputools::StreamExecutor* executor) = 0; // Compiles a set of HLO modules that can run in parallel, potentially // communicating data between the modules, and returns a corresponding // sequence of executable objects. virtual StatusOr>> Compile( - std::vector> hlo_module, - std::vector> module_config, - HloDumper dump_hlo, + std::vector> modules, std::vector stream_exec) = 0; // Compiles the HLO module for ahead-of-time execution. This is intended for // use in static compilation. virtual StatusOr>> - CompileAheadOfTime( - std::vector> module, - std::vector> module_config, - HloDumper dump_hlo, const AotCompilationOptions& options) = 0; + CompileAheadOfTime(std::vector> modules, + const AotCompilationOptions& options) = 0; ///// // The Compiler class also serves as a point to register compiler objects @@ -152,8 +141,18 @@ class Compiler { static StatusOr GetForPlatform( const perftools::gputools::Platform* platform); - // Returns the size in bytes of the top-level buffer of a shape. - virtual int64 ShapeSizeBytes(const Shape& shape) const = 0; + // Returns a function that computes the size in bytes of the logical + // buffer that contains a shape. + virtual HloCostAnalysis::ShapeSizeFunction ShapeSizeBytesFunction() const = 0; + + // Returns a function that computes the size in bytes of a given + // logical buffer. + std::function BufferSizeBytesFunction() { + HloCostAnalysis::ShapeSizeFunction shape_size = ShapeSizeBytesFunction(); + return [shape_size](const LogicalBuffer& buffer) { + return shape_size(buffer.shape()); + }; + } private: // Mutex that guards the platform-compiler map. diff --git a/tensorflow/compiler/xla/service/computation_placer.cc b/tensorflow/compiler/xla/service/computation_placer.cc new file mode 100644 index 0000000000000000000000000000000000000000..cdfa30dd9a7b6a5b9e58087491a9d99caaa1b998 --- /dev/null +++ b/tensorflow/compiler/xla/service/computation_placer.cc @@ -0,0 +1,152 @@ +/* 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/computation_placer.h" + +#include +#include +#include + +#include "tensorflow/compiler/xla/literal_util.h" +#include "tensorflow/compiler/xla/ptr_util.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/types.h" +#include "tensorflow/compiler/xla/util.h" +#include "tensorflow/core/lib/core/errors.h" +#include "tensorflow/core/lib/core/status.h" +#include "tensorflow/core/platform/logging.h" +#include "tensorflow/core/platform/stream_executor_no_cuda.h" + +namespace se = ::perftools::gputools; + +namespace xla { + +Status DeviceAssignment::Serialize(DeviceAssignmentProto* proto) const { + proto->set_replica_count(replica_count()); + proto->set_computation_count(computation_count()); + for (int computation = 0; computation < computation_count(); ++computation) { + DeviceAssignmentProto::ComputationDevice* computation_device = + proto->add_computation_devices(); + for (int replica = 0; replica < replica_count(); ++replica) { + computation_device->add_replica_device_ids((*this)(replica, computation)); + } + } + return Status::OK(); +} + +/* static */ StatusOr> +DeviceAssignment::Deserialize(const DeviceAssignmentProto& proto) { + TF_RET_CHECK(proto.computation_devices_size() == proto.computation_count()); + auto assignment = MakeUnique(proto.replica_count(), + proto.computation_count()); + for (int computation = 0; computation < proto.computation_count(); + ++computation) { + const auto& computation_device = proto.computation_devices(computation); + TF_RET_CHECK(computation_device.replica_device_ids_size() == + proto.replica_count()); + for (int replica = 0; replica < proto.replica_count(); ++replica) { + (*assignment)(replica, computation) = + computation_device.replica_device_ids(replica); + } + } + return std::move(assignment); +} + +StatusOr ComputationPlacer::DeviceId(int replica, int computation, + int replica_count, + int computation_count) { + TF_RET_CHECK(replica < replica_count); + TF_RET_CHECK(computation < computation_count); + + return computation * replica_count + replica; +} + +StatusOr ComputationPlacer::AssignDevices( + int replica_count, int computation_count) { + DeviceAssignment assignment(replica_count, computation_count); + for (int replica = 0; replica < replica_count; ++replica) { + for (int computation = 0; computation < computation_count; ++computation) { + TF_ASSIGN_OR_RETURN( + int device_id, + DeviceId(replica, computation, replica_count, computation_count)); + assignment(replica, computation) = device_id; + } + } + return std::move(assignment); +} + +/* static */ void ComputationPlacer::RegisterComputationPlacer( + se::Platform::Id platform_id, + ComputationPlacerCreationFunction creation_function) { + tensorflow::mutex_lock lock( + *ComputationPlacer::platform_computation_placer_mutex()); + auto* computation_placers = GetPlatformComputationPlacers(); + CHECK(computation_placers->find(platform_id) == computation_placers->end()); + (*computation_placers)[platform_id].creation_function = creation_function; +} + +/* static */ StatusOr ComputationPlacer::GetForPlatform( + const se::Platform* platform) { + tensorflow::mutex_lock lock( + *ComputationPlacer::platform_computation_placer_mutex()); + auto* computation_placers = GetPlatformComputationPlacers(); + + auto it = computation_placers->find(platform->id()); + if (it == computation_placers->end()) { + return NotFound( + "could not find registered computation placer for platform %s -- check " + "target linkage", + platform->Name().c_str()); + } + + if (it->second.placer == nullptr) { + // Lazily create the computation placer the first time it is needed. + it->second.placer = (*it->second.creation_function)(); + } + + return it->second.placer.get(); +} + +/* static */ tensorflow::mutex* +ComputationPlacer::platform_computation_placer_mutex() { + static tensorflow::mutex* m = new tensorflow::mutex; + return m; +} + +/* static */ std::map* +ComputationPlacer::GetPlatformComputationPlacers() { + static auto* r = + new std::map; + return r; +} + +} // namespace xla + +static std::unique_ptr CreateComputationPlacer() { + return xla::MakeUnique(); +} + +static bool InitModule() { + xla::ComputationPlacer::RegisterComputationPlacer(se::host::kHostPlatformId, + &CreateComputationPlacer); + xla::ComputationPlacer::RegisterComputationPlacer(se::cuda::kCudaPlatformId, + &CreateComputationPlacer); + return true; +} +static bool module_initialized = InitModule(); diff --git a/tensorflow/compiler/xla/service/computation_placer.h b/tensorflow/compiler/xla/service/computation_placer.h new file mode 100644 index 0000000000000000000000000000000000000000..7d9abcd100dd9e878da885110bc1bd1ac65e3f84 --- /dev/null +++ b/tensorflow/compiler/xla/service/computation_placer.h @@ -0,0 +1,117 @@ +/* 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_COMPUTATION_PLACER_H_ +#define TENSORFLOW_COMPILER_XLA_SERVICE_COMPUTATION_PLACER_H_ + +#include +#include +#include + +#include "tensorflow/compiler/xla/array2d.h" +#include "tensorflow/compiler/xla/status.h" +#include "tensorflow/compiler/xla/statusor.h" +#include "tensorflow/compiler/xla/xla_data.pb.h" +#include "tensorflow/core/lib/core/status.h" +#include "tensorflow/core/platform/macros.h" +#include "tensorflow/core/platform/stream_executor_no_cuda.h" +#include "tensorflow/core/platform/types.h" + +namespace xla { + +// Class that represents the device assignment for a set of XLA replicated +// computations. For R replicas and C computations, R * C devices are required +// execute the computation in parallel. The assigned device ids can be accessed +// by assignment(replica, computation). +class DeviceAssignment : public Array2D { + public: + DeviceAssignment() {} + DeviceAssignment(int replica_count, int computation_count) + : Array2D(replica_count, computation_count, -1) { + CHECK_GT(replica_count, 0); + CHECK_GT(computation_count, 0); + } + + int replica_count() const { return height(); } + int computation_count() const { return width(); } + + // Protocol buffer serialization and deserialization. + Status Serialize(DeviceAssignmentProto* proto) const; + + // Return a std::unique_ptr instead of a DeviceAssignment + // directly because one of the supported TF platforms (mac) does not compile + // due to a StatusOr of an incomplete type (DeviceAssignment). + static StatusOr> Deserialize( + const DeviceAssignmentProto& proto); +}; + +// A generic implementation of the XLA computation placer, which assigns device +// ids to a set of replicated computations. +class ComputationPlacer { + public: + ComputationPlacer() {} + virtual ~ComputationPlacer() {} + + // Returns the device id assigned to the given replica and computation + // instance for [replica_count x computation_count] setup. The returned device + // id must match the assignement from PlaceReplicatedComputation(). + virtual StatusOr DeviceId(int replica, int computation, + int replica_count, int computation_count); + + // Returns the device ids assigned to a set of replicated computations, given + // the number of replicas and the number of computations. + virtual StatusOr AssignDevices(int replica_count, + int computation_count); + + using ComputationPlacerCreationFunction = + std::unique_ptr (*)(); + + // Registers a computation placer creation function for a particular platform. + static void RegisterComputationPlacer( + perftools::gputools::Platform::Id platform_id, + ComputationPlacerCreationFunction creation_function); + + // Returns the computation placer singleton pointer if it is available for the + // given platform, or an error status if it is not. + static StatusOr GetForPlatform( + const perftools::gputools::Platform* platform); + + private: + // Routine that returns the mutex that guards the platform-to-computation + // placer map. Done as a routine to ensure correct initialization ordering, + // since RegisterComputationPlacer can be called during program initialization + // time. + static tensorflow::mutex* platform_computation_placer_mutex(); + + // State kept for each kind of ComputationPlacer. Registration functions set + // up creation_function, and then we use that to lazily create "placer" the + // first time GetForPlatform is invoked for a particular id. + struct State { + std::unique_ptr placer; + ComputationPlacerCreationFunction creation_function = nullptr; + }; + + // Map from platform kind to computation placer singleton. + static std::map* + GetPlatformComputationPlacers(); + + perftools::gputools::Platform::Id platform_id_; + + TF_DISALLOW_COPY_AND_ASSIGN(ComputationPlacer); +}; + +} // namespace xla + +#endif // TENSORFLOW_COMPILER_XLA_SERVICE_COMPUTATION_PLACER_H_ diff --git a/tensorflow/compiler/xla/service/computation_tracker.cc b/tensorflow/compiler/xla/service/computation_tracker.cc index f78806bce82f7f524ba2bf80fbf602ad49e103c7..70e25eebdb068db893e24aec0f72d09090ac7027 100644 --- a/tensorflow/compiler/xla/service/computation_tracker.cc +++ b/tensorflow/compiler/xla/service/computation_tracker.cc @@ -164,11 +164,11 @@ void ComputationTracker::ComputeComputationPostOrder( visited->insert(versioned_handle); post_order->push_back(versioned_handle); - return; } StatusOr> ComputationTracker::BuildHloModule( const VersionedComputationHandle& entry_handle, + const HloModuleConfig& config, bool include_unreachable_instructions) const { tensorflow::mutex_lock lock(computation_mutex_); @@ -208,7 +208,7 @@ StatusOr> ComputationTracker::BuildHloModule( string module_name = tensorflow::strings::StrCat(entry_computation->name(), "_module"); - auto module = MakeUnique(module_name, entry_handle); + auto module = MakeUnique(module_name, entry_handle, config); for (auto versioned_handle : post_order) { UserComputation* computation = ResolveInternal(versioned_handle.handle).ValueOrDie(); @@ -216,6 +216,7 @@ StatusOr> ComputationTracker::BuildHloModule( TF_ASSIGN_OR_RETURN( std::unique_ptr hlo_computation, computation->BuildHloComputation(versioned_handle.version, resolver, + config.debug_options(), include_unreachable_instructions)); // Add the newly created computation to VersionedHandle-to-HloComputation diff --git a/tensorflow/compiler/xla/service/computation_tracker.h b/tensorflow/compiler/xla/service/computation_tracker.h index 1922908747c6ef3b74c5b87d3c3924e5ffb38fc5..d42d66adefe7faa2751da4cd80b392a38917ce70 100644 --- a/tensorflow/compiler/xla/service/computation_tracker.h +++ b/tensorflow/compiler/xla/service/computation_tracker.h @@ -23,6 +23,7 @@ limitations under the License. #include #include "tensorflow/compiler/xla/service/hlo_module.h" +#include "tensorflow/compiler/xla/service/hlo_module_config.h" #include "tensorflow/compiler/xla/service/session.pb.h" #include "tensorflow/compiler/xla/service/user_computation.h" #include "tensorflow/compiler/xla/service/versioned_computation_handle.h" @@ -72,12 +73,15 @@ class ComputationTracker { // Builds an HLO module using the specified computation as the entry. The // module will include the entry computation as well as all computations which // are called directly or indirectly from the entry computation via operations - // like "map". If include_unreachable_instructions is true, then instructions + // like "map". config is the HLO module configuration to use for the + // constructed module. + // If include_unreachable_instructions is true, then instructions // which are not reachable from the root are lowered into HloInstructions // including unreachable parameters. This ensures the entry HloComputation has // the same program shape (ProgramShape) as the entry UserComputation. StatusOr> BuildHloModule( const VersionedComputationHandle& entry_handle, + const HloModuleConfig& config, bool include_unreachable_instructions = true) const; string ToString() const; diff --git a/tensorflow/compiler/xla/service/copy_insertion.cc b/tensorflow/compiler/xla/service/copy_insertion.cc index 7dae49acad388e6d18a8cb1e4ea70244616978bb..628f729e0b4388cf258ec7f393f14c48042c1e3e 100644 --- a/tensorflow/compiler/xla/service/copy_insertion.cc +++ b/tensorflow/compiler/xla/service/copy_insertion.cc @@ -16,19 +16,20 @@ limitations under the License. #include "tensorflow/compiler/xla/service/copy_insertion.h" #include -#include -#include #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_opcode.h" +#include "tensorflow/compiler/xla/service/liveness_util.h" #include "tensorflow/compiler/xla/service/logical_buffer.h" #include "tensorflow/compiler/xla/service/tuple_points_to_analysis.h" #include "tensorflow/compiler/xla/status_macros.h" #include "tensorflow/compiler/xla/statusor.h" #include "tensorflow/compiler/xla/types.h" #include "tensorflow/compiler/xla/util.h" +#include "tensorflow/core/lib/gtl/flatmap.h" +#include "tensorflow/core/lib/gtl/flatset.h" #include "tensorflow/core/lib/strings/str_util.h" #include "tensorflow/core/platform/logging.h" @@ -36,6 +37,9 @@ namespace xla { namespace { +using tensorflow::gtl::FlatMap; +using tensorflow::gtl::FlatSet; + // InstructionCopier encapsulates indices at which to copy 'instruction'. // All 'instruction' users in 'copy_users' are updated to use the copy. // @@ -52,7 +56,7 @@ namespace { // // Example two-element tuple with one element that needs a copy: // -// Tuple // instruction +// original-instruction // / \ // GTE(0) GTE(1) // | | @@ -60,23 +64,54 @@ namespace { // \ / // Tuple // copied-instruction // +// As an optimization, if the original instruction is itself a Tuple +// instruction, we elide the unnecessary extra GTE and Tuple instructions, +// and just insert the copy into a new Tuple instruction, with control +// dependencies to ensure the copy occurs after any possible interference. class InstructionCopier { public: - InstructionCopier(const bool init_value, HloInstruction* instruction, - const std::vector& copy_users); + InstructionCopier(HloInstruction* instruction, + const std::vector& copy_users) + : instruction_(instruction), + copy_users_(copy_users), + indices_to_copy_(instruction->shape()), + control_predecessors_(instruction->shape()) {} + + // Sets indices that are read-only, and thus do not need to be copied. + void SetReadOnlyIndices(const ShapeTree& read_only_indices) { + read_only_indices_ = read_only_indices; + } + + // Sets copy overrides, which are copy instructions to use at each index. This + // is used to share a single copy of read-only entry parameters and constants + // between multiple While loops. + void SetCopyOverrides(const ShapeTree& copy_overrides) { + copy_overrides_ = copy_overrides; + } // Returns true if all recorded indices are false (returns true otherwise). bool HasAllIndicesFalse() const; // Records instruction buffer indices which point-to a Parameter or Constant. - tensorflow::Status RecordIndicesWhichPointToParamOrConstant( + Status RecordIndicesWhichPointToParamOrConstant( const TuplePointsToAnalysis& points_to_analysis); // Records instruction buffer indices to copy which are necessary to ensure: // *) PointsToSet of 'instruction_' is unambiguous and distinct. // *) No liveness interference between 'instruction_' and 'other_instruction'. - tensorflow::Status RecordIndicesToCopyForColocatingBuffers( - BufferLiveness* liveness, HloInstruction* other_instruction); + // + // If 'read_only_indices_out' is non-null, read-only indices are set to true. + Status RecordIndicesToCopyForColocatingBuffers( + const BufferLiveness& liveness, const HloInstruction* other_instruction, + ShapeTree* read_only_indices_out); + + // Records control predecessors to add for inserted copy instructions. + // 'parameter' must have the same shape as the instruction that will be + // copied, and must define all buffers in the shape. Control predecessors are + // only recorded for indices that have already been marked for copying. + Status RecordControlPredecessors( + const TuplePointsToAnalysis& points_to_analysis, + HloInstruction* parameter); // Inserts copies of 'instruction' buffers at indices in 'indices_to_copy', // and replaces all uses for instructions in 'copy_users_' with copy. @@ -88,15 +123,29 @@ class InstructionCopier { const std::vector& copy_users() const { return copy_users_; } private: + // Does the given index represent a read-only buffer? + bool IsReadOnlyIndex(const ShapeIndex& index) const { + return !ShapeUtil::IsNil(read_only_indices_.shape()) && + read_only_indices_.element(index); + } + + // Returns the copy override at the given index, or nullptr. + HloInstruction* GetCopyOverride(const ShapeIndex& index) const { + return ShapeUtil::IsNil(copy_overrides_.shape()) + ? nullptr + : copy_overrides_.element(index); + } + // Records instruction buffer indices which have ambiguous or non-distinct // points-to sets. - tensorflow::Status RecordAmbiguousOrNonDistinctIndices( + Status RecordAmbiguousOrNonDistinctIndices( const TuplePointsToAnalysis& points_to_analysis); - // Records instruction buffer indices which have interferring live ranges + // Records instruction buffer indices which have interfering live ranges // with 'other_instruction' buffers at same index. - tensorflow::Status RecordIndicesWhichInterfereWithOtherInstruction( - BufferLiveness* liveness, HloInstruction* other_instruction); + Status RecordIndicesWhichInterfereWithOtherInstruction( + const BufferLiveness& liveness, const HloInstruction* other_instruction, + ShapeTree* read_only_indices_out); // Recursively inserts copies of 'instruction' tuple elements at indices // specified in 'indices_to_copy', and returns the copy of 'instruction'. @@ -107,28 +156,25 @@ class InstructionCopier { } HloInstruction* instruction_; - std::vector copy_users_; + const std::vector copy_users_; ShapeTree indices_to_copy_; + ShapeTree> control_predecessors_; + ShapeTree read_only_indices_; + ShapeTree copy_overrides_; }; -InstructionCopier::InstructionCopier( - const bool init_value, HloInstruction* instruction, - const std::vector& copy_users) - : instruction_(instruction), - copy_users_(copy_users), - indices_to_copy_(instruction->shape(), init_value) {} - bool InstructionCopier::HasAllIndicesFalse() const { bool all_indices_false = true; - TF_CHECK_OK(indices_to_copy_.ForEachElement([&all_indices_false]( - const ShapeIndex& /*index*/, bool /*is_leaf*/, const bool& data) { - if (data) all_indices_false = false; - return tensorflow::Status::OK(); - })); + indices_to_copy_.ForEachElement( + [&all_indices_false](const ShapeIndex& /*index*/, bool data) { + if (data) { + all_indices_false = false; + } + }); return all_indices_false; } -tensorflow::Status InstructionCopier::RecordIndicesWhichPointToParamOrConstant( +Status InstructionCopier::RecordIndicesWhichPointToParamOrConstant( const TuplePointsToAnalysis& points_to_analysis) { const PointsToSet& points_to = points_to_analysis.GetPointsToSet(instruction_); @@ -141,10 +187,12 @@ tensorflow::Status InstructionCopier::RecordIndicesWhichPointToParamOrConstant( // Multiple buffers within a parameter/constant may be live out, so collect // a set of indices at which to copy first. - TF_RETURN_IF_ERROR(points_to.ForEachElement([this]( - const ShapeIndex& index, bool /*is_leaf*/, - const std::vector& buffers) { - for (auto buffer : buffers) { + points_to.ForEachElement([this](const ShapeIndex& index, + const PointsToSet::BufferList& buffers) { + if (IsReadOnlyIndex(index)) { + return; + } + for (const LogicalBuffer* buffer : buffers) { // pointee is the HloInstruction producing the buffer which may be // liveout. HloInstruction* pointee = buffer->instruction(); @@ -154,59 +202,55 @@ tensorflow::Status InstructionCopier::RecordIndicesWhichPointToParamOrConstant( << " index: " << tensorflow::str_util::Join(index, ",") << " may be live out of computation: " << pointee->ToString(); RecordIndex(index); + break; } } - return tensorflow::Status::OK(); - })); - return tensorflow::Status::OK(); + }); + return Status::OK(); } -tensorflow::Status InstructionCopier::RecordIndicesToCopyForColocatingBuffers( - BufferLiveness* liveness, HloInstruction* other_instruction) { +Status InstructionCopier::RecordIndicesToCopyForColocatingBuffers( + const BufferLiveness& liveness, const HloInstruction* other_instruction, + ShapeTree* read_only_indices_out) { TF_RETURN_IF_ERROR( - RecordAmbiguousOrNonDistinctIndices(liveness->points_to_analysis())); + RecordAmbiguousOrNonDistinctIndices(liveness.points_to_analysis())); TF_RETURN_IF_ERROR(RecordIndicesWhichInterfereWithOtherInstruction( - liveness, other_instruction)); - return tensorflow::Status::OK(); + liveness, other_instruction, read_only_indices_out)); + return Status::OK(); } -tensorflow::Status InstructionCopier::RecordAmbiguousOrNonDistinctIndices( +Status InstructionCopier::RecordAmbiguousOrNonDistinctIndices( const TuplePointsToAnalysis& points_to_analysis) { const PointsToSet& points_to = points_to_analysis.GetPointsToSet(instruction_); // Mapping from LogicalBuffer to index (used to detect non-distinct indices). - std::unordered_map> + FlatMap> buffer_to_source_indices; - TF_RETURN_IF_ERROR(points_to.ForEachElement([this, &buffer_to_source_indices]( - const ShapeIndex& index, bool /*is_leaf*/, - const std::vector& buffers) { - if (buffers.size() > 1) { - // Record ambiguous points-to set at 'index'. - if (!indices_to_copy_.element(index)) { - VLOG(2) << "Adding copy of buffer for instruction: " - << instruction_->name() - << " at index: " << tensorflow::str_util::Join(index, ",") - << " with ambiguous points-to set."; - RecordIndex(index); - } - } - // For each 'buffer': record a mapping from 'buffer' to 'index'. - for (auto& buffer : buffers) { - auto it = buffer_to_source_indices.find(buffer); - if (it == buffer_to_source_indices.end()) { - buffer_to_source_indices.insert({buffer, std::vector()}); - } - buffer_to_source_indices[buffer].push_back(index); - } - return tensorflow::Status::OK(); - })); + points_to.ForEachElement( + [this, &buffer_to_source_indices]( + const ShapeIndex& index, const PointsToSet::BufferList& buffers) { + if (buffers.size() > 1) { + // Record ambiguous points-to set at 'index'. + if (!indices_to_copy_.element(index)) { + VLOG(2) << "Adding copy of buffer for instruction: " + << instruction_->name() + << " at index: " << tensorflow::str_util::Join(index, ",") + << " with ambiguous points-to set."; + RecordIndex(index); + } + } + // For each 'buffer': record a mapping from 'buffer' to 'index'. + for (const LogicalBuffer* buffer : buffers) { + buffer_to_source_indices[buffer].push_back(index); + } + }); // Record all non-distinct indices detected in 'buffer_to_source_indices'. - for (auto& buff_to_src : buffer_to_source_indices) { + for (const auto& buff_to_src : buffer_to_source_indices) { if (buff_to_src.second.size() == 1) { continue; } - for (auto& src_index : buff_to_src.second) { + for (const ShapeIndex& src_index : buff_to_src.second) { // Record non-distinct points-to set at 'src_index'. if (!indices_to_copy_.element(src_index)) { VLOG(2) << "Adding copy of buffer for instruction: " @@ -217,25 +261,28 @@ tensorflow::Status InstructionCopier::RecordAmbiguousOrNonDistinctIndices( } } } - return tensorflow::Status::OK(); + return Status::OK(); } -tensorflow::Status -InstructionCopier::RecordIndicesWhichInterfereWithOtherInstruction( - BufferLiveness* liveness, HloInstruction* other_instruction) { +Status InstructionCopier::RecordIndicesWhichInterfereWithOtherInstruction( + const BufferLiveness& liveness, const HloInstruction* other_instruction, + ShapeTree* read_only_indices_out) { // Record all buffer indices for 'instruction_', which interfere with // 'other_instruction' at the same index. - TF_RETURN_IF_ERROR(ShapeUtil::ForEachSubshape( + ShapeUtil::ForEachSubshape( instruction_->shape(), - [this, &liveness, &other_instruction](const Shape& /*subshape*/, - const ShapeIndex& index) { + [this, &liveness, other_instruction, read_only_indices_out]( + const Shape& /*subshape*/, const ShapeIndex& index) { + if (IsReadOnlyIndex(index)) { + return; + } if (indices_to_copy_.element(index)) { // Return if previous pass already set index. - return tensorflow::Status::OK(); + return; } - auto& points_to_analysis = liveness->points_to_analysis(); + const auto& points_to_analysis = liveness.points_to_analysis(); // Lookup buffers for 'instruction_' and 'other_instruction'. - const std::vector instruction_buffers = + const auto instruction_buffers = points_to_analysis.GetPointsToSet(instruction_).element(index); // If 'instruction_' has ambiguous points-to-set at 'index', it would // have been recorded in a previous pass (and we would have returned @@ -244,7 +291,7 @@ InstructionCopier::RecordIndicesWhichInterfereWithOtherInstruction( CHECK_EQ(1, instruction_buffers.size()); const LogicalBuffer* instruction_buffer = instruction_buffers[0]; - const std::vector other_instruction_buffers = + const auto other_instruction_buffers = points_to_analysis.GetPointsToSet(other_instruction).element(index); // Do not insert a copy if both instructions point at the same buffer. // This eliminates unnecessary copies of read-only tuple elements. @@ -252,20 +299,24 @@ InstructionCopier::RecordIndicesWhichInterfereWithOtherInstruction( // then that buffer is not updated on the path between the two // instructions. Therefore, any other (possibly interference-causing) // users of that buffer from 'other_instruction' will see the same data, - // irrespecive of whether we insert a copy of this buffer at + // irrespective of whether we insert a copy of this buffer at // 'instruction_' or not. if (other_instruction_buffers.size() == 1 && other_instruction_buffers[0]->id() == instruction_buffer->id()) { - return tensorflow::Status::OK(); + if (read_only_indices_out != nullptr) { + *read_only_indices_out->mutable_element(index) = true; + } + return; } - // We cant say anything about the ambiguity of 'other_instruction' at + // We can't say anything about the ambiguity of 'other_instruction' at // this point, so we need to check interference between the single // buffer in the points-to set of 'instruction_' and all buffers in // 'other_instruction_buffers'. - for (auto& other_buffer : other_instruction_buffers) { - if (liveness->MayInterfere(*instruction_buffer, *other_buffer)) { + for (const LogicalBuffer* other_buffer : other_instruction_buffers) { + if (liveness.MayInterfere(*instruction_buffer, *other_buffer)) { VLOG(2) << "Adding copy of buffer for instruction: " << instruction_->name() + << " instruction_buffer: " << instruction_buffer->ToString() << " at index: " << tensorflow::str_util::Join(index, ",") << " because of interference with buffer: " << other_buffer->ToString(); @@ -273,40 +324,88 @@ InstructionCopier::RecordIndicesWhichInterfereWithOtherInstruction( break; } } - return tensorflow::Status::OK(); - })); - return tensorflow::Status::OK(); + }); + return Status::OK(); +} + +// This is called when 'instruction_' is a while body root, and 'parameter' is +// the while body parameter. We record all users of all aliases of 'parameter' +// as control predecessors, so that when we add a copy of 'instruction_', we can +// mark the control dependencies. This is necessary because points-to and +// liveness analysis doesn't know about the aliasing between the while body root +// and param. Without these control dependencies, the copy might get scheduled +// to run at a point that interferes with users of the buffer. +Status InstructionCopier::RecordControlPredecessors( + const TuplePointsToAnalysis& points_to_analysis, + HloInstruction* parameter) { + return indices_to_copy_.ForEachElementWithStatus( + [this, &points_to_analysis, parameter](const ShapeIndex& index, + bool will_copy) { + if (will_copy) { + TF_ASSIGN_OR_RETURN( + const LogicalBuffer* buffer, + points_to_analysis.GetBufferDefinedAt(parameter, index)); + for (const BufferAlias& alias : + points_to_analysis.GetBufferAliases(*buffer)) { + for (HloInstruction* user : alias.instruction()->users()) { + if (DoesNotUseOperandBuffer(alias.instruction(), alias.index(), + user, points_to_analysis)) { + continue; + } + + if (user != instruction_) { + control_predecessors_.mutable_element(index)->push_back(user); + } + } + } + } + return Status::OK(); + }); } // Recursively inserts copies of 'instruction' tuple element buffers at // indices in 'indices_to_copy_', expanding tuples as needed. -// TODO(b/31159897) Remove superfluous Tuple->GTE->Tuple expressions. HloInstruction* InstructionCopier::CopyTuple(HloInstruction* instruction, ShapeIndex* index) { - std::vector element_copies; const int64 num_tuple_elements = ShapeUtil::TupleElementCount(instruction->shape()); + std::vector elem_copies(num_tuple_elements); for (int64 i = 0; i < num_tuple_elements; ++i) { - HloInstruction* gte = instruction->parent()->AddInstruction( - HloInstruction::CreateGetTupleElement( - ShapeUtil::GetSubshape(instruction->shape(), {i}), instruction, i)); - HloInstruction* element_copy; + HloInstruction* elem; + if (instruction->opcode() == HloOpcode::kTuple) { + // If the instruction is already a Tuple instruction, we know that the + // element buffers are aliased, so we can just grab the operand directly. + elem = instruction->mutable_operand(i); + } else { + // Otherwise we need to add a GTE to unpack the element out of the tuple. + elem = instruction->parent()->AddInstruction( + HloInstruction::CreateGetTupleElement( + ShapeUtil::GetSubshape(instruction->shape(), {i}), instruction, + i)); + } index->push_back(i); - if (ShapeUtil::IsTuple(gte->shape())) { - element_copy = CopyTuple(gte, index); + if (ShapeUtil::IsTuple(elem->shape())) { + elem_copies[i] = CopyTuple(elem, index); + } else if (!indices_to_copy_.element(*index)) { + elem_copies[i] = elem; + } else if (HloInstruction* copy_override = GetCopyOverride(*index)) { + elem_copies[i] = copy_override; } else { - if (indices_to_copy_.element(*index)) { - element_copy = gte->parent()->AddInstruction( - HloInstruction::CreateUnary(gte->shape(), HloOpcode::kCopy, gte)); - } else { - element_copy = gte; + HloInstruction* elem_copy = elem->parent()->AddInstruction( + HloInstruction::CreateUnary(elem->shape(), HloOpcode::kCopy, elem)); + for (HloInstruction* control_predecessor : + control_predecessors_.element(*index)) { + VLOG(2) << "Adding control dependency from " + << control_predecessor->ToString() << " to " + << elem_copy->ToString(); + TF_CHECK_OK(control_predecessor->AddControlDependencyTo(elem_copy)); } + elem_copies[i] = elem_copy; } index->pop_back(); - element_copies.push_back(element_copy); } return instruction->parent()->AddInstruction( - HloInstruction::CreateTuple(element_copies)); + HloInstruction::CreateTuple(elem_copies)); } // Inserts copies of 'instruction_' buffers at indices in 'indices_to_copy_'. @@ -327,8 +426,87 @@ HloInstruction* InstructionCopier::Copy() { return copy; } +// The 'read_only_indices' are initialized based on points-to analysis on the +// while body corresponding to 'while_hlo'. If the init buffer corresponding to +// a read-only index aliases with a constant, it cannot be considered read-only, +// and must be copied. This is necessary because BufferAssignment does not +// currently assign an allocation for constants (b/32248867). +// This function performs this fix-up of 'read_only_indices'. +// +// Returns a ShapeTree of copy_overrides, which implements an optimization to +// allow multiple while loops that share the same read-only constants to +// share a single copy. +StatusOr> RevertReadOnlyIndicesForConstants( + const HloInstruction* while_hlo, + const TuplePointsToAnalysis& points_to_analysis, + ShapeTree* read_only_indices, + FlatMap* shared_copies) { + const HloInstruction* init_hlo = while_hlo->operand(0); + const PointsToSet& points_to = points_to_analysis.GetPointsToSet(init_hlo); + + // Mapping from LogicalBuffer to index (used to detect non-distinct indices). + FlatSet buffer_set; + + ShapeTree copy_overrides(init_hlo->shape()); + points_to.ForEachElement([init_hlo, read_only_indices, shared_copies, + &buffer_set, ©_overrides]( + const ShapeIndex& index, + const PointsToSet::BufferList& buffers) { + // Look for read-only entry parameters. + if (!read_only_indices->element(index)) { + return; + } + for (const LogicalBuffer* buffer : buffers) { + HloInstruction* pointee = buffer->instruction(); + const bool is_constant = pointee->opcode() == HloOpcode::kConstant; + if (!is_constant) { + continue; + } + + // We have found an constant that is read-only in + // the while body. These buffers are managed by the caller, and cannot + // be aliased with HLO buffers. Revert this read-only index, + // to allow it to be copied. + *read_only_indices->mutable_element(index) = false; + + // Optimization to allow multiple while loops that share the same + // read-only entry constants to share a single copy. + // Only unambiguous and distinct array-shaped buffers are allowed, to + // reduce code complexity. The shape of the entry parameter must be + // identical to the shape of the init_hlo at this index, to ensure + // there were no intervening bitcast or GTE instructions, which are + // also hard to handle. + const Shape& pointee_shape = pointee->shape(); + const Shape& init_shape = + ShapeUtil::GetSubshape(init_hlo->shape(), index); + if (buffers.size() == 1 && ShapeUtil::IsArray(pointee_shape) && + ShapeUtil::Equal(pointee_shape, init_shape) && + buffer_set.count(buffer) < 1) { + HloInstruction** copy = &(*shared_copies)[pointee]; + if (*copy == nullptr) { + *copy = pointee->parent()->AddInstruction(HloInstruction::CreateUnary( + pointee_shape, HloOpcode::kCopy, pointee)); + } + // Add the copy as an override. + *copy_overrides.mutable_element(index) = *copy; + } + + // Tracks whether this current buffer is distinct. + buffer_set.insert(buffer); + + // We've already reverted the read-only index and handled the + // single-copy optimization above, so there's nothing more to do. + break; + } + }); + return copy_overrides; +} + } // anonymous namespace +// NOTE: This is only called by gpu::CopyInsertion. It's not called here in the +// base class, since the regular CopyInsertion logic above selectively copies +// tuple elements, while this method assumes all buffers need to be deep copied. StatusOr CopyInsertion::FindOrInsertCopy(HloInstruction* hlo) { auto copy_it = inserted_copies_.find(hlo); if (copy_it == inserted_copies_.end()) { @@ -347,85 +525,99 @@ StatusOr CopyInsertion::Run(HloModule* module) { TF_ASSIGN_OR_RETURN( std::unique_ptr liveness, BufferLiveness::Run(module, MakeUnique(module))); - auto& points_to_analysis = liveness->points_to_analysis(); + const auto& points_to_analysis = liveness->points_to_analysis(); XLA_VLOG_LINES(2, points_to_analysis.ToString()); XLA_VLOG_LINES(2, module->ToString()); - // Gather references to all while body computations in 'module'. - std::unordered_set while_body_computations; - // Gather references to all while instructions in 'module' by computation. - std::unordered_map> - while_instructions; + // Gather all while body computations and while instructions. + FlatSet while_body_computations; + std::vector while_instructions; for (auto& computation : module->computations()) { for (auto& instruction : computation->instructions()) { - if (instruction->opcode() != HloOpcode::kWhile) { - continue; + if (instruction->opcode() == HloOpcode::kWhile) { + while_body_computations.insert(instruction->while_body()); + while_instructions.push_back(instruction.get()); } - while_body_computations.insert(instruction->while_body()); - auto it = while_instructions.find(computation.get()); - if (it == while_instructions.end()) { - while_instructions.insert( - {computation.get(), std::vector()}); - } - while_instructions[computation.get()].emplace_back(instruction.get()); } } - for (auto& computation : module->computations()) { - VLOG(2) << "computation " << computation->name(); + // Collect instruction buffer indices to copy in 'instructions_to_copy'. + std::vector instructions_to_copy; - // Collect instruction buffer indices to copy in 'instructions_to_copy'. - std::vector instructions_to_copy; - - // Add copies of while 'init' operand instructions (if needed). - // TODO(b/33301720) Remove redundant while instruction copies. - auto it = while_instructions.find(computation.get()); - if (it != while_instructions.end()) { - for (auto& while_hlo : it->second) { - // Create InstructionCopier for init operand of while instruction. - HloInstruction* init_hlo = while_hlo->mutable_operand(0); - instructions_to_copy.push_back( - InstructionCopier(/*init_value=*/false, init_hlo, {while_hlo})); - InstructionCopier& init_copier = instructions_to_copy.back(); - // Record 'init' buffer indices which point-to a Constant or Parameter. - TF_RETURN_IF_ERROR(init_copier.RecordIndicesWhichPointToParamOrConstant( - liveness->points_to_analysis())); - // Record indices necessary to colocate while and init operand buffers. - TF_RETURN_IF_ERROR(init_copier.RecordIndicesToCopyForColocatingBuffers( - liveness.get(), while_hlo)); - } + // Add copies of computation root instructions, if needed. + FlatMap> while_body_read_only_indices; + for (auto& computation : module->computations()) { + if (computation->IsFusionComputation()) { + continue; } - - // Create InstructionCopier for computation root instruction. - instructions_to_copy.push_back(InstructionCopier( - /*init_value=*/false, computation->root_instruction(), {})); - InstructionCopier& root_copier = instructions_to_copy.back(); - + VLOG(2) << "computation " << computation->name(); + InstructionCopier root_copier(computation->root_instruction(), + /*copy_users=*/{}); if (while_body_computations.count(computation.get()) > 0) { - // Record root indices to copy for while body sub-computations. - // We do not need to call RecordIndicesWhichPointToParamOrConstant for - // the while root instruction here, because any neccessary copies needed - // to avoid constant or parameters in the output are handled by while.init - // operand copy insertion above (which will share an allocation). + // Record root indices to copy for while body sub-computations. We do not + // need to call RecordIndicesWhichPointToParamOrConstant for the while + // body root instruction here, because any necessary copies needed to + // avoid constants or parameters in the output are handled by while.init + // operand copy insertion below (which will share an allocation). + HloInstruction* while_body_param = computation->parameter_instruction(0); + ShapeTree read_only_indices(while_body_param->shape()); TF_RETURN_IF_ERROR(root_copier.RecordIndicesToCopyForColocatingBuffers( - liveness.get(), computation->parameter_instruction(0))); - } else if (copy_param_and_const_) { + *liveness, while_body_param, &read_only_indices)); + while_body_read_only_indices[computation.get()] = read_only_indices; + + // Mark control predecessors, based on the body param, for any copies + // we'll be inserting. This ensures the copy doesn't run too early. + TF_RETURN_IF_ERROR(root_copier.RecordControlPredecessors( + points_to_analysis, while_body_param)); + } else { // Record root indices to copy for general computations. TF_RETURN_IF_ERROR(root_copier.RecordIndicesWhichPointToParamOrConstant( - liveness->points_to_analysis())); + points_to_analysis)); } + instructions_to_copy.push_back(root_copier); + } - for (auto& to_copy : instructions_to_copy) { - if (to_copy.HasAllIndicesFalse()) { - continue; - } - changed = true; + // Add copies of while 'init' operand instructions, if needed. 'shared_copies' + // is used to ensure that multiple while loops can share a single copy of the + // same entry parameter or constant, if all loops use it read-only. + // + // TODO(b/33301720) Remove redundant while instruction copies. + FlatMap shared_copies; + for (HloInstruction* while_hlo : while_instructions) { + // Fix read_only_indices to account for entry constants. Also + // initialize copy_overrides, which ensures a single copy for each read-only + // constant that is used in multiple while loops. + ShapeTree* read_only_indices = + &while_body_read_only_indices[while_hlo->while_body()]; + TF_ASSIGN_OR_RETURN( + const ShapeTree copy_overrides, + RevertReadOnlyIndicesForConstants(while_hlo, points_to_analysis, + read_only_indices, &shared_copies)); + // Create InstructionCopier for init operand of while instruction. + HloInstruction* init_hlo = while_hlo->mutable_operand(0); + InstructionCopier init_copier(init_hlo, {while_hlo}); + init_copier.SetReadOnlyIndices(*read_only_indices); + init_copier.SetCopyOverrides(copy_overrides); + // Record 'init' buffer indices which point-to a Constant or Parameter. + TF_RETURN_IF_ERROR(init_copier.RecordIndicesWhichPointToParamOrConstant( + points_to_analysis)); + // Record indices necessary to colocate while and init operand buffers. + TF_RETURN_IF_ERROR(init_copier.RecordIndicesToCopyForColocatingBuffers( + *liveness, while_hlo, /*read_only_indices_out=*/nullptr)); + instructions_to_copy.push_back(init_copier); + } - // Copy instruction at recorded buffer indices. - HloInstruction* copy = to_copy.Copy(); - if (to_copy.instruction() == computation->root_instruction()) { - computation->set_root_instruction(copy); - } + for (InstructionCopier& to_copy : instructions_to_copy) { + if (to_copy.HasAllIndicesFalse()) { + continue; + } + changed = true; + + // Copy instruction at recorded buffer indices. + HloComputation* computation = to_copy.instruction()->parent(); + HloInstruction* copy = to_copy.Copy(); + if (to_copy.instruction() == computation->root_instruction()) { + computation->set_root_instruction(copy); } } diff --git a/tensorflow/compiler/xla/service/copy_insertion.h b/tensorflow/compiler/xla/service/copy_insertion.h index ce91ac0de56f3fc1101c38cee838c0b0593214ad..28bb62e40c7674960dbb1bb63dc8967b06956028 100644 --- a/tensorflow/compiler/xla/service/copy_insertion.h +++ b/tensorflow/compiler/xla/service/copy_insertion.h @@ -21,6 +21,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/core/lib/gtl/flatmap.h" namespace xla { @@ -32,9 +33,6 @@ namespace xla { // different lifetimes than computation results. class CopyInsertion : public HloPassInterface { public: - explicit CopyInsertion(bool copy_param_and_const = true) - : copy_param_and_const_(copy_param_and_const) {} - ~CopyInsertion() override {} tensorflow::StringPiece name() const override { return "copy-insertion"; } // Run the pass on the given module. Returns whether the module was changed @@ -46,13 +44,9 @@ class CopyInsertion : public HloPassInterface { // duplicate copies. StatusOr FindOrInsertCopy(HloInstruction* hlo); - // Determines whether to insert copies if the root instruction is, or - // points-to, any constant or parameter instruction. - const bool copy_param_and_const_; - // A map containing all copies inserted during the copy insertion pass. The // key is the copied instruction and the value is the copy. - std::unordered_map inserted_copies_; + tensorflow::gtl::FlatMap inserted_copies_; }; } // namespace xla diff --git a/tensorflow/compiler/xla/service/copy_insertion_test.cc b/tensorflow/compiler/xla/service/copy_insertion_test.cc index 4c26b2de124b0b42f6de1ebdf82d4584f2904cab..daaf8d10bb828fb118ea575d005ff1274ed193fb 100644 --- a/tensorflow/compiler/xla/service/copy_insertion_test.cc +++ b/tensorflow/compiler/xla/service/copy_insertion_test.cc @@ -20,18 +20,23 @@ limitations under the License. #include "tensorflow/compiler/xla/literal_util.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/tuple_points_to_analysis.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/compiler/xla/test_helpers.h" +namespace op = xla::testing::opcode_matchers; namespace xla { namespace { +using ::testing::UnorderedElementsAre; + class CopyInsertionTest : public HloTestBase { protected: void InsertCopies(HloModule* module) { @@ -39,55 +44,25 @@ class CopyInsertionTest : public HloTestBase { EXPECT_IS_OK(copy_insertion.Run(module).status()); // Verify the points to set of the root of the computation after copy - // insertion contains no constants or parameters. + // insertion contains no constants or parameters, and is distinct and + // non-ambiguous. auto points_to_analysis = TuplePointsToAnalysis::Run(module).ConsumeValueOrDie(); - tensorflow::gtl::FlatSet maybe_live_out_buffers = + const auto& points_to = points_to_analysis->GetPointsToSet( + module->entry_computation()->root_instruction()); + EXPECT_TRUE(points_to.IsDistinct()); + EXPECT_TRUE(!points_to.IsAmbiguous()); + + auto maybe_live_out_buffers = points_to_analysis ->GetPointsToSet(module->entry_computation()->root_instruction()) .CreateFlattenedSet(); + for (const LogicalBuffer* buffer : maybe_live_out_buffers) { EXPECT_NE(buffer->instruction()->opcode(), HloOpcode::kConstant); EXPECT_NE(buffer->instruction()->opcode(), HloOpcode::kParameter); } } - - // OperandTree is a test helper class that simplifies the expression of - // an expected tree of operands (starting at some root instruction) in a - // unit test. - // Each HLO instruction is represented as a node in the OperandTree. - struct OperandTree { - // The expected opcode for this OperandTree node. - HloOpcode opcode; - // The set of operands expected for this OperandTree node. - std::vector operands; - // If non-null, a pointer to the expected HloInstruction at this node. - const HloInstruction* instruction = nullptr; - - // Returns a mutable reference to operand 'i' of this node. - OperandTree& op(int i) { - if (i >= operands.size()) { - operands.resize(i + 1); - } - return operands[i]; - } - - // Check that 'instruction' and its operands match expected values recorded - // in OperandTree. - void Check(const HloInstruction* instruction) { - EXPECT_EQ(opcode, instruction->opcode()); - if (instruction != nullptr) { - EXPECT_EQ(instruction, instruction); - } - if (operands.empty()) { - return; - } - EXPECT_EQ(operands.size(), instruction->operand_count()); - for (int i = 0; i < instruction->operand_count(); ++i) { - operands[i].Check(instruction->operand(i)); - } - } - }; }; TEST_F(CopyInsertionTest, SingleParameter) { @@ -97,53 +72,35 @@ TEST_F(CopyInsertionTest, SingleParameter) { HloInstruction* tuple = builder.AddInstruction(HloInstruction::CreateTuple({x})); - ExpectEqUnordered(x->users(), {tuple}); + EXPECT_THAT(x->users(), UnorderedElementsAre(tuple)); - HloModule module(TestName()); - module.AddEntryComputation(builder.Build()); + auto module = CreateNewModule(); + module->AddEntryComputation(builder.Build()); - HloInstruction* old_root = module.entry_computation()->root_instruction(); - InsertCopies(&module); - HloInstruction* new_root = module.entry_computation()->root_instruction(); + HloInstruction* old_root = module->entry_computation()->root_instruction(); + InsertCopies(module.get()); - // Check path from 'new_root' to 'old_root'. - OperandTree op_tree; - op_tree.opcode = HloOpcode::kTuple; - - op_tree.op(0).opcode = HloOpcode::kCopy; - op_tree.op(0).op(0).opcode = HloOpcode::kGetTupleElement; - op_tree.op(0).op(0).op(0).opcode = HloOpcode::kTuple; - op_tree.op(0).op(0).op(0).instruction = old_root; - - op_tree.Check(new_root); + EXPECT_THAT(module->entry_computation()->root_instruction(), + op::Tuple(op::Copy(old_root->operand(0)))); } TEST_F(CopyInsertionTest, SingleConstant) { auto builder = HloComputation::Builder(TestName()); HloInstruction* constant = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(1.0))); + HloInstruction::CreateConstant(Literal::CreateR0(1.0))); HloInstruction* tuple = builder.AddInstruction(HloInstruction::CreateTuple({constant})); - ExpectEqUnordered(constant->users(), {tuple}); - - HloModule module(TestName()); - module.AddEntryComputation(builder.Build()); - - HloInstruction* old_root = module.entry_computation()->root_instruction(); - InsertCopies(&module); - HloInstruction* new_root = module.entry_computation()->root_instruction(); + EXPECT_THAT(constant->users(), UnorderedElementsAre(tuple)); - // Check path from 'new_root' to 'old_root'. - OperandTree op_tree; - op_tree.opcode = HloOpcode::kTuple; + auto module = CreateNewModule(); + module->AddEntryComputation(builder.Build()); - op_tree.op(0).opcode = HloOpcode::kCopy; - op_tree.op(0).op(0).opcode = HloOpcode::kGetTupleElement; - op_tree.op(0).op(0).op(0).opcode = HloOpcode::kTuple; - op_tree.op(0).op(0).op(0).instruction = old_root; + HloInstruction* old_root = module->entry_computation()->root_instruction(); + InsertCopies(module.get()); - op_tree.Check(new_root); + EXPECT_THAT(module->entry_computation()->root_instruction(), + op::Tuple(op::Copy(old_root->operand(0)))); } TEST_F(CopyInsertionTest, MultipleConstantsAndParameters) { @@ -153,9 +110,9 @@ TEST_F(CopyInsertionTest, MultipleConstantsAndParameters) { auto builder = HloComputation::Builder(TestName()); HloInstruction* constant1 = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(1.0))); + HloInstruction::CreateConstant(Literal::CreateR0(1.0))); HloInstruction* constant2 = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(2.0))); + HloInstruction::CreateConstant(Literal::CreateR0(2.0))); HloInstruction* x = builder.AddInstruction( HloInstruction::CreateParameter(0, ShapeUtil::MakeShape(F32, {}), "x")); @@ -167,35 +124,15 @@ TEST_F(CopyInsertionTest, MultipleConstantsAndParameters) { builder.AddInstruction(HloInstruction::CreateTuple({constant2, x, add})); - HloModule module(TestName()); - module.AddEntryComputation(builder.Build()); - - HloInstruction* old_root = module.entry_computation()->root_instruction(); - InsertCopies(&module); - HloInstruction* new_root = module.entry_computation()->root_instruction(); - - // "constant2" and parameter "x" are pointed to by the tuple and should be - // copied. - - // Check all paths from 'new_root' to 'old_root'. - OperandTree op_tree; - op_tree.opcode = HloOpcode::kTuple; - - op_tree.op(0).opcode = HloOpcode::kCopy; - op_tree.op(0).op(0).opcode = HloOpcode::kGetTupleElement; - op_tree.op(0).op(0).op(0).opcode = HloOpcode::kTuple; - op_tree.op(0).op(0).op(0).instruction = old_root; + auto module = CreateNewModule(); + module->AddEntryComputation(builder.Build()); - op_tree.op(1).opcode = HloOpcode::kCopy; - op_tree.op(1).op(0).opcode = HloOpcode::kGetTupleElement; - op_tree.op(1).op(0).op(0).opcode = HloOpcode::kTuple; - op_tree.op(1).op(0).op(0).instruction = old_root; + HloInstruction* old_root = module->entry_computation()->root_instruction(); + InsertCopies(module.get()); - op_tree.op(2).opcode = HloOpcode::kGetTupleElement; - op_tree.op(2).op(0).opcode = HloOpcode::kTuple; - op_tree.op(2).op(0).instruction = old_root; - - op_tree.Check(new_root); + EXPECT_THAT(module->entry_computation()->root_instruction(), + op::Tuple(op::Copy(old_root->operand(0)), + op::Copy(old_root->operand(1)), old_root->operand(2))); } TEST_F(CopyInsertionTest, AmbiguousPointsToSet) { @@ -203,11 +140,11 @@ TEST_F(CopyInsertionTest, AmbiguousPointsToSet) { // the computation result. Verify that copies are added properly. auto builder = HloComputation::Builder(TestName()); HloInstruction* constant1 = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(1.0))); + HloInstruction::CreateConstant(Literal::CreateR0(1.0))); HloInstruction* constant2 = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(2.0))); + HloInstruction::CreateConstant(Literal::CreateR0(2.0))); HloInstruction* constant3 = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(3.0))); + HloInstruction::CreateConstant(Literal::CreateR0(3.0))); HloInstruction* tuple1 = builder.AddInstruction( HloInstruction::CreateTuple({constant1, constant2})); @@ -215,36 +152,23 @@ TEST_F(CopyInsertionTest, AmbiguousPointsToSet) { HloInstruction::CreateTuple({constant3, constant2})); HloInstruction* pred = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(false))); + HloInstruction::CreateConstant(Literal::CreateR0(false))); builder.AddInstruction(HloInstruction::CreateTernary( tuple1->shape(), HloOpcode::kSelect, pred, tuple1, tuple2)); - ExpectEqUnordered(constant1->users(), {tuple1}); - ExpectEqUnordered(constant2->users(), {tuple1, tuple2}); - ExpectEqUnordered(constant3->users(), {tuple2}); - - HloModule module(TestName()); - module.AddEntryComputation(builder.Build()); - - HloInstruction* old_root = module.entry_computation()->root_instruction(); - InsertCopies(&module); - HloInstruction* new_root = module.entry_computation()->root_instruction(); - - // Check all paths from 'new_root' to 'old_root'. - OperandTree op_tree; - op_tree.opcode = HloOpcode::kTuple; + EXPECT_THAT(constant1->users(), UnorderedElementsAre(tuple1)); + EXPECT_THAT(constant2->users(), UnorderedElementsAre(tuple1, tuple2)); + EXPECT_THAT(constant3->users(), UnorderedElementsAre(tuple2)); - op_tree.op(0).opcode = HloOpcode::kCopy; - op_tree.op(0).op(0).opcode = HloOpcode::kGetTupleElement; - op_tree.op(0).op(0).op(0).opcode = HloOpcode::kSelect; - op_tree.op(0).op(0).op(0).instruction = old_root; + auto module = CreateNewModule(); + module->AddEntryComputation(builder.Build()); - op_tree.op(1).opcode = HloOpcode::kCopy; - op_tree.op(1).op(0).opcode = HloOpcode::kGetTupleElement; - op_tree.op(1).op(0).op(0).opcode = HloOpcode::kSelect; - op_tree.op(1).op(0).op(0).instruction = old_root; + HloInstruction* old_root = module->entry_computation()->root_instruction(); + InsertCopies(module.get()); - op_tree.Check(new_root); + EXPECT_THAT(module->entry_computation()->root_instruction(), + op::Tuple(op::Copy(op::GetTupleElement(old_root)), + op::Copy(op::GetTupleElement(old_root)))); } TEST_F(CopyInsertionTest, BitcastParameter) { @@ -256,50 +180,37 @@ TEST_F(CopyInsertionTest, BitcastParameter) { HloInstruction* bitcast = builder.AddInstruction(HloInstruction::CreateUnary( ShapeUtil::MakeShape(F32, {2, 2}), HloOpcode::kBitcast, x)); - HloModule module(TestName()); - module.AddEntryComputation(builder.Build()); + auto module = CreateNewModule(); + module->AddEntryComputation(builder.Build()); - ExpectEqUnordered(x->users(), {bitcast}); + EXPECT_THAT(x->users(), UnorderedElementsAre(bitcast)); - HloInstruction* old_root = module.entry_computation()->root_instruction(); - InsertCopies(&module); - HloInstruction* new_root = module.entry_computation()->root_instruction(); + HloInstruction* old_root = module->entry_computation()->root_instruction(); + InsertCopies(module.get()); - // Check path from 'new_root' to 'old_root'. - OperandTree op_tree; - op_tree.opcode = HloOpcode::kCopy; - op_tree.op(0).opcode = HloOpcode::kBitcast; - op_tree.op(0).instruction = old_root; - - op_tree.Check(new_root); + EXPECT_THAT(module->entry_computation()->root_instruction(), + op::Copy(old_root)); } TEST_F(CopyInsertionTest, BitcastConstant) { // The output of a bitcast is its operand (same buffer), so a bitcast // constant feeding the result must have a copy added. auto builder = HloComputation::Builder(TestName()); - HloInstruction* constant = - builder.AddInstruction(HloInstruction::CreateConstant( - LiteralUtil::CreateR1({1.0, 42.0}))); + HloInstruction* constant = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR1({1.0, 42.0}))); HloInstruction* bitcast = builder.AddInstruction(HloInstruction::CreateUnary( ShapeUtil::MakeShape(F32, {2, 2}), HloOpcode::kBitcast, constant)); - HloModule module(TestName()); - module.AddEntryComputation(builder.Build()); + auto module = CreateNewModule(); + module->AddEntryComputation(builder.Build()); - ExpectEqUnordered(constant->users(), {bitcast}); + EXPECT_THAT(constant->users(), UnorderedElementsAre(bitcast)); - HloInstruction* old_root = module.entry_computation()->root_instruction(); - InsertCopies(&module); - HloInstruction* new_root = module.entry_computation()->root_instruction(); + HloInstruction* old_root = module->entry_computation()->root_instruction(); + InsertCopies(module.get()); - // Check path from 'new_root' to 'old_root'. - OperandTree op_tree; - op_tree.opcode = HloOpcode::kCopy; - op_tree.op(0).opcode = HloOpcode::kBitcast; - op_tree.op(0).instruction = old_root; - - op_tree.Check(new_root); + EXPECT_THAT(module->entry_computation()->root_instruction(), + op::Copy(old_root)); } TEST_F(CopyInsertionTest, BitcastTupleElementParameter) { @@ -311,24 +222,16 @@ TEST_F(CopyInsertionTest, BitcastTupleElementParameter) { ShapeUtil::MakeShape(F32, {2, 2}), HloOpcode::kBitcast, x)); builder.AddInstruction(HloInstruction::CreateTuple({bitcast})); - HloModule module(TestName()); - module.AddEntryComputation(builder.Build()); - - ExpectEqUnordered(x->users(), {bitcast}); + auto module = CreateNewModule(); + module->AddEntryComputation(builder.Build()); - HloInstruction* old_root = module.entry_computation()->root_instruction(); - InsertCopies(&module); - HloInstruction* new_root = module.entry_computation()->root_instruction(); + EXPECT_THAT(x->users(), UnorderedElementsAre(bitcast)); - // Check path from 'new_root' to 'old_root'. - OperandTree op_tree; - op_tree.opcode = HloOpcode::kTuple; - op_tree.op(0).opcode = HloOpcode::kCopy; - op_tree.op(0).op(0).opcode = HloOpcode::kGetTupleElement; - op_tree.op(0).op(0).op(0).opcode = HloOpcode::kTuple; - op_tree.op(0).op(0).op(0).instruction = old_root; + HloInstruction* old_root = module->entry_computation()->root_instruction(); + InsertCopies(module.get()); - op_tree.Check(new_root); + EXPECT_THAT(module->entry_computation()->root_instruction(), + op::Tuple(op::Copy(old_root->operand(0)))); } TEST_F(CopyInsertionTest, NestedTupleParameter) { @@ -339,47 +242,31 @@ TEST_F(CopyInsertionTest, NestedTupleParameter) { // Param shape is: ((F32[], S32[1,2,3]), F32[42]) builder.AddInstruction(HloInstruction::CreateParameter( - 0, ShapeUtil::MakeTupleShape( - {ShapeUtil::MakeTupleShape({ShapeUtil::MakeShape(F32, {}), - ShapeUtil::MakeShape(S32, {1, 2, 3})}), - ShapeUtil::MakeShape(F32, {42})}), + 0, + ShapeUtil::MakeTupleShape( + {ShapeUtil::MakeTupleShape({ShapeUtil::MakeShape(F32, {}), + ShapeUtil::MakeShape(S32, {1, 2, 3})}), + ShapeUtil::MakeShape(F32, {42})}), "param0")); - HloModule module(TestName()); - module.AddEntryComputation(builder.Build()); + auto module = CreateNewModule(); + module->AddEntryComputation(builder.Build()); EXPECT_EQ(HloOpcode::kParameter, - module.entry_computation()->root_instruction()->opcode()); + module->entry_computation()->root_instruction()->opcode()); - HloInstruction* old_root = module.entry_computation()->root_instruction(); - InsertCopies(&module); - HloInstruction* new_root = module.entry_computation()->root_instruction(); + HloInstruction* old_root = module->entry_computation()->root_instruction(); + InsertCopies(module.get()); + HloInstruction* new_root = module->entry_computation()->root_instruction(); EXPECT_NE(old_root, new_root); - // Check all paths from 'new_root' to 'old_root'. - OperandTree op_tree; - op_tree.opcode = HloOpcode::kTuple; - - op_tree.op(0).opcode = HloOpcode::kTuple; - op_tree.op(0).op(0).opcode = HloOpcode::kCopy; - op_tree.op(0).op(0).op(0).opcode = HloOpcode::kGetTupleElement; - op_tree.op(0).op(0).op(0).op(0).opcode = HloOpcode::kGetTupleElement; - op_tree.op(0).op(0).op(0).op(0).op(0).opcode = HloOpcode::kParameter; - op_tree.op(0).op(0).op(0).op(0).op(0).instruction = old_root; - - op_tree.op(0).opcode = HloOpcode::kTuple; - op_tree.op(0).op(1).opcode = HloOpcode::kCopy; - op_tree.op(0).op(1).op(0).opcode = HloOpcode::kGetTupleElement; - op_tree.op(0).op(1).op(0).op(0).opcode = HloOpcode::kGetTupleElement; - op_tree.op(0).op(1).op(0).op(0).op(0).opcode = HloOpcode::kParameter; - op_tree.op(0).op(1).op(0).op(0).op(0).instruction = old_root; - - op_tree.op(1).opcode = HloOpcode::kCopy; - op_tree.op(1).op(0).opcode = HloOpcode::kGetTupleElement; - op_tree.op(1).op(0).op(0).opcode = HloOpcode::kParameter; - op_tree.op(1).op(0).op(0).instruction = old_root; - - op_tree.Check(new_root); + EXPECT_THAT( + new_root, + op::Tuple( + op::Tuple( + op::Copy(op::GetTupleElement(op::GetTupleElement(old_root))), + op::Copy(op::GetTupleElement(op::GetTupleElement(old_root)))), + op::Copy(op::GetTupleElement(old_root)))); } TEST_F(CopyInsertionTest, ElementOfNestedTupleParameter) { @@ -389,10 +276,11 @@ TEST_F(CopyInsertionTest, ElementOfNestedTupleParameter) { // Param shape is: ((F32[], S32[1,2,3]), F32[42]) auto param = builder.AddInstruction(HloInstruction::CreateParameter( - 0, ShapeUtil::MakeTupleShape( - {ShapeUtil::MakeTupleShape({ShapeUtil::MakeShape(F32, {}), - ShapeUtil::MakeShape(S32, {1, 2, 3})}), - ShapeUtil::MakeShape(F32, {42})}), + 0, + ShapeUtil::MakeTupleShape( + {ShapeUtil::MakeTupleShape({ShapeUtil::MakeShape(F32, {}), + ShapeUtil::MakeShape(S32, {1, 2, 3})}), + ShapeUtil::MakeShape(F32, {42})}), "param0")); // The return value of the computation is the zero-th elemnt of the nested @@ -400,30 +288,17 @@ TEST_F(CopyInsertionTest, ElementOfNestedTupleParameter) { auto gte = builder.AddInstruction(HloInstruction::CreateGetTupleElement( ShapeUtil::GetSubshape(param->shape(), {0}), param, 0)); - HloModule module(TestName()); - module.AddEntryComputation(builder.Build()); + auto module = CreateNewModule(); + module->AddEntryComputation(builder.Build()); - EXPECT_EQ(gte, module.entry_computation()->root_instruction()); + EXPECT_EQ(gte, module->entry_computation()->root_instruction()); - HloInstruction* old_root = module.entry_computation()->root_instruction(); - InsertCopies(&module); - HloInstruction* new_root = module.entry_computation()->root_instruction(); + HloInstruction* old_root = module->entry_computation()->root_instruction(); + InsertCopies(module.get()); - // Check all paths from 'new_root' to 'old_root'. - OperandTree op_tree; - op_tree.opcode = HloOpcode::kTuple; - - op_tree.op(0).opcode = HloOpcode::kCopy; - op_tree.op(0).op(0).opcode = HloOpcode::kGetTupleElement; - op_tree.op(0).op(0).op(0).opcode = HloOpcode::kGetTupleElement; - op_tree.op(0).op(0).op(0).instruction = old_root; - - op_tree.op(1).opcode = HloOpcode::kCopy; - op_tree.op(1).op(0).opcode = HloOpcode::kGetTupleElement; - op_tree.op(1).op(0).op(0).opcode = HloOpcode::kGetTupleElement; - op_tree.op(1).op(0).op(0).instruction = old_root; - - op_tree.Check(new_root); + EXPECT_THAT(module->entry_computation()->root_instruction(), + op::Tuple(op::Copy(op::GetTupleElement(old_root)), + op::Copy(op::GetTupleElement(old_root)))); } TEST_F(CopyInsertionTest, AmbiguousTopLevelRoot) { @@ -432,9 +307,9 @@ TEST_F(CopyInsertionTest, AmbiguousTopLevelRoot) { // copy is added. auto builder = HloComputation::Builder(TestName()); HloInstruction* constant1 = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(1.0))); + HloInstruction::CreateConstant(Literal::CreateR0(1.0))); HloInstruction* constant2 = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(2.0))); + HloInstruction::CreateConstant(Literal::CreateR0(2.0))); HloInstruction* tuple1 = builder.AddInstruction( HloInstruction::CreateTuple({constant1, constant2})); @@ -442,34 +317,28 @@ TEST_F(CopyInsertionTest, AmbiguousTopLevelRoot) { HloInstruction::CreateTuple({constant2, constant1})); HloInstruction* pred = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(false))); + HloInstruction::CreateConstant(Literal::CreateR0(false))); HloInstruction* select = builder.AddInstruction(HloInstruction::CreateTernary( tuple1->shape(), HloOpcode::kSelect, pred, tuple1, tuple2)); HloInstruction* gte = builder.AddInstruction(HloInstruction::CreateGetTupleElement( ShapeUtil::GetSubshape(select->shape(), {0}), select, 0)); - HloModule module(TestName()); - module.AddEntryComputation(builder.Build()); - - EXPECT_EQ(gte, module.entry_computation()->root_instruction()); + auto module = CreateNewModule(); + module->AddEntryComputation(builder.Build()); - HloInstruction* old_root = module.entry_computation()->root_instruction(); - InsertCopies(&module); - HloInstruction* new_root = module.entry_computation()->root_instruction(); + EXPECT_EQ(gte, module->entry_computation()->root_instruction()); - // Check path from 'new_root' to 'old_root'. - OperandTree op_tree; - op_tree.opcode = HloOpcode::kCopy; - op_tree.op(0).opcode = HloOpcode::kGetTupleElement; - op_tree.op(0).instruction = old_root; + HloInstruction* old_root = module->entry_computation()->root_instruction(); + InsertCopies(module.get()); - op_tree.Check(new_root); + EXPECT_THAT(module->entry_computation()->root_instruction(), + op::Copy(old_root)); } class WhileCopyInsertionTest : public CopyInsertionTest { protected: - WhileCopyInsertionTest() : module_(TestName()) {} + WhileCopyInsertionTest() : module_(CreateNewModule()) {} // Builds a While condition computation which reads the induction variable // from the tuple parameter, and returns a predicate indicating whether this @@ -480,7 +349,7 @@ class WhileCopyInsertionTest : public CopyInsertionTest { bool nested = false) { auto builder = HloComputation::Builder(TestName() + ".Condition"); auto limit_const = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(10))); + HloInstruction::CreateConstant(Literal::CreateR0(10))); const Shape& loop_state_shape = nested ? nested_loop_state_shape_ : loop_state_shape_; auto loop_state = builder.AddInstruction( @@ -511,7 +380,7 @@ class WhileCopyInsertionTest : public CopyInsertionTest { builder.AddInstruction(HloInstruction::CreateGetTupleElement( induction_variable_shape_, loop_state, 0)); auto inc = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(1))); + HloInstruction::CreateConstant(Literal::CreateR0(1))); auto add0 = builder.AddInstruction(HloInstruction::CreateBinary( induction_variable->shape(), HloOpcode::kAdd, induction_variable, inc)); // Update data GTE(1). @@ -527,8 +396,48 @@ class WhileCopyInsertionTest : public CopyInsertionTest { return builder.Build(); } - // Builds a While body computation with read-only tuple element 0. + // Builds a While body computation with two output tuple elements dependent on // both input tuple elements. + // + // EX: Body({in0, in1, in2}) + // out0 = Add(in0, 1) + // out1 = in1 + // out2 = in2 + // Tuple(out0, out1, out2) + std::unique_ptr BuildDependentBodyComputation2() { + auto builder = HloComputation::Builder(TestName() + ".Body"); + + const Shape& loop_state_shape = ShapeUtil::MakeTupleShape( + {induction_variable_shape_, data_shape_, data_shape_}); + + auto loop_state = builder.AddInstruction( + HloInstruction::CreateParameter(0, loop_state_shape, "loop_state")); + + // Update the induction variable GTE(0). + auto induction_variable = + builder.AddInstruction(HloInstruction::CreateGetTupleElement( + induction_variable_shape_, loop_state, 0)); + auto inc = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(1))); + + // add0 = Add(in0, 1) + auto add0 = builder.AddInstruction(HloInstruction::CreateBinary( + induction_variable->shape(), HloOpcode::kAdd, induction_variable, inc)); + // data1 = GTE(1). + HloInstruction* data1 = builder.AddInstruction( + HloInstruction::CreateGetTupleElement(data_shape_, loop_state, 1)); + + // data2 = GTE(2). + HloInstruction* data2 = builder.AddInstruction( + HloInstruction::CreateGetTupleElement(data_shape_, loop_state, 2)); + + // Create output Tuple. + builder.AddInstruction(HloInstruction::CreateTuple({add0, data1, data2})); + + return builder.Build(); + } + + // Builds a While body computation with read-only tuple element 0. // EX: // Body({in0, in1}) // out0 = in0 @@ -546,6 +455,7 @@ class WhileCopyInsertionTest : public CopyInsertionTest { // Update data GTE(1). auto data = builder.AddInstruction( HloInstruction::CreateGetTupleElement(data_shape_, loop_state, 1)); + // Use 'induction_variable' in computation with no path to output tuple. auto update = builder.AddInstruction( HloInstruction::CreateBroadcast(data_shape_, induction_variable, {8})); @@ -563,31 +473,48 @@ class WhileCopyInsertionTest : public CopyInsertionTest { // out0 = Add(in0, 1) // out1 = Add(in1, {1, 1, 1, 1, 1, 1, 1, 1}) // Tuple(out0, out1) - std::unique_ptr BuildIndependentBodyComputation() { + std::unique_ptr BuildIndependentBodyComputation( + bool nested = false) { auto builder = HloComputation::Builder(TestName() + ".Body"); // Create param instruction to access loop state. + const Shape& loop_state_shape = + nested ? nested_loop_state_shape_ : loop_state_shape_; + auto loop_state = builder.AddInstruction( - HloInstruction::CreateParameter(0, loop_state_shape_, "loop_state")); + HloInstruction::CreateParameter(0, loop_state_shape, "loop_state")); // Update the induction variable GTE(0). auto induction_variable = builder.AddInstruction(HloInstruction::CreateGetTupleElement( induction_variable_shape_, loop_state, 0)); auto inc = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(1))); + HloInstruction::CreateConstant(Literal::CreateR0(1))); // add0 = Add(in0, 1) auto add0 = builder.AddInstruction(HloInstruction::CreateBinary( induction_variable->shape(), HloOpcode::kAdd, induction_variable, inc)); // Update data GTE(1). - auto data = builder.AddInstruction( - HloInstruction::CreateGetTupleElement(data_shape_, loop_state, 1)); - auto update = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR1( - {1.f, 1.f, 1.f, 1.f, 1.f, 1.f, 1.f, 1.f}))); - // add0 = Add(in1, {1, 1, 1, 1, 1, 1, 1, 1}) + HloInstruction* data = nullptr; + if (nested) { + data = builder.AddInstruction(HloInstruction::CreateGetTupleElement( + nested_tuple_shape_, loop_state, 1)); + data = builder.AddInstruction( + HloInstruction::CreateGetTupleElement(data_shape_, data, 0)); + } else { + data = builder.AddInstruction( + HloInstruction::CreateGetTupleElement(data_shape_, loop_state, 1)); + } + auto update = builder.AddInstruction(HloInstruction::CreateConstant( + Literal::CreateR1({1.f, 1.f, 1.f, 1.f, 1.f, 1.f, 1.f, 1.f}))); + // add1 = Add(in1, {1, 1, 1, 1, 1, 1, 1, 1}) auto add1 = builder.AddInstruction(HloInstruction::CreateBinary( data_shape_, HloOpcode::kAdd, data, update)); // Create output Tuple. - builder.AddInstruction(HloInstruction::CreateTuple({add0, add1})); + if (nested) { + auto nested_tuple = + builder.AddInstruction(HloInstruction::CreateTuple({add1, add1})); + builder.AddInstruction(HloInstruction::CreateTuple({add0, nested_tuple})); + } else { + builder.AddInstruction(HloInstruction::CreateTuple({add0, add1})); + } return builder.Build(); } @@ -609,7 +536,7 @@ class WhileCopyInsertionTest : public CopyInsertionTest { auto gte0 = builder.AddInstruction(HloInstruction::CreateGetTupleElement( induction_variable_shape_, loop_state, 0)); auto inc = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(1))); + HloInstruction::CreateConstant(Literal::CreateR0(1))); auto add0 = builder.AddInstruction(HloInstruction::CreateBinary( gte0->shape(), HloOpcode::kAdd, gte0, inc)); @@ -619,9 +546,8 @@ class WhileCopyInsertionTest : public CopyInsertionTest { // GTE(GTE(loop_state, 1), 0) -> Add auto gte10 = builder.AddInstruction( HloInstruction::CreateGetTupleElement(data_shape_, gte1, 0)); - auto update10 = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR1( - {1.f, 1.f, 1.f, 1.f, 1.f, 1.f, 1.f, 1.f}))); + auto update10 = builder.AddInstruction(HloInstruction::CreateConstant( + Literal::CreateR1({1.f, 1.f, 1.f, 1.f, 1.f, 1.f, 1.f, 1.f}))); auto add10 = builder.AddInstruction(HloInstruction::CreateBinary( data_shape_, HloOpcode::kAdd, gte10, update10)); @@ -640,39 +566,39 @@ class WhileCopyInsertionTest : public CopyInsertionTest { // Builds a While instruction using 'condition' and 'body' sub-computations. // Init operand is initialized to zeros of appropriate shape. - void BuildWhileInstruction(HloComputation* condition, HloComputation* body, - bool nested = false) { + HloInstruction* BuildWhileInstruction(HloComputation* condition, + HloComputation* body, + bool nested = false) { auto builder = HloComputation::Builder(TestName() + ".While"); auto induction_var_init = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(0))); + HloInstruction::CreateConstant(Literal::CreateR0(0))); - auto data_init = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR1( - {0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f}))); + auto data_init = builder.AddInstruction(HloInstruction::CreateConstant( + Literal::CreateR1({0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f}))); if (nested) { auto inner_init = builder.AddInstruction( HloInstruction::CreateTuple({data_init, data_init})); auto loop_state_init = builder.AddInstruction( HloInstruction::CreateTuple({induction_var_init, inner_init})); - builder.AddInstruction(HloInstruction::CreateWhile( + auto while_hlo = builder.AddInstruction(HloInstruction::CreateWhile( loop_state_shape_, condition, body, loop_state_init)); - module_.AddEntryComputation(builder.Build()); - return; + module_->AddEntryComputation(builder.Build()); + return while_hlo; } auto loop_state_init = builder.AddInstruction( HloInstruction::CreateTuple({induction_var_init, data_init})); - builder.AddInstruction(HloInstruction::CreateWhile( + auto while_hlo = builder.AddInstruction(HloInstruction::CreateWhile( loop_state_shape_, condition, body, loop_state_init)); - module_.AddEntryComputation(builder.Build()); + module_->AddEntryComputation(builder.Build()); + return while_hlo; } HloInstruction* BuildWhileInstruction_InitPointsToConstant() { auto builder = HloComputation::Builder(TestName() + ".While"); - auto data_init = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR1( - {0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f}))); + auto data_init = builder.AddInstruction(HloInstruction::CreateConstant( + Literal::CreateR1({0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f}))); return BuildWhileInstructionWithCustomInit(loop_state_shape_, data_init, &builder); } @@ -689,11 +615,11 @@ class WhileCopyInsertionTest : public CopyInsertionTest { auto builder = HloComputation::Builder(TestName() + ".While"); auto one = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(1.0))); + HloInstruction::CreateConstant(Literal::CreateR0(1.0))); auto v1 = builder.AddInstruction( HloInstruction::CreateBroadcast(data_shape_, one, {1})); auto zero = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(1.0))); + HloInstruction::CreateConstant(Literal::CreateR0(1.0))); auto v2 = builder.AddInstruction( HloInstruction::CreateBroadcast(data_shape_, zero, {1})); @@ -701,7 +627,7 @@ class WhileCopyInsertionTest : public CopyInsertionTest { auto tuple2 = builder.AddInstruction(HloInstruction::CreateTuple({v2, v1})); auto pred = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(false))); + HloInstruction::CreateConstant(Literal::CreateR0(false))); auto data_init = builder.AddInstruction(HloInstruction::CreateTernary( nested_tuple_shape_, HloOpcode::kSelect, pred, tuple1, tuple2)); @@ -713,7 +639,7 @@ class WhileCopyInsertionTest : public CopyInsertionTest { auto builder = HloComputation::Builder(TestName() + ".While"); auto one = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(1.0))); + HloInstruction::CreateConstant(Literal::CreateR0(1.0))); auto one_vec = builder.AddInstruction( HloInstruction::CreateBroadcast(data_shape_, one, {1})); auto data_init = @@ -726,12 +652,11 @@ class WhileCopyInsertionTest : public CopyInsertionTest { HloInstruction* BuildWhileInstruction_InitPointsToInterfering() { auto builder = HloComputation::Builder(TestName() + ".While"); auto one = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(1.0))); + HloInstruction::CreateConstant(Literal::CreateR0(1.0))); auto data_init = builder.AddInstruction( HloInstruction::CreateBroadcast(data_shape_, one, {1})); - auto one_vec = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR1( - {1.f, 1.f, 1.f, 1.f, 1.f, 1.f, 1.f, 1.f}))); + auto one_vec = builder.AddInstruction(HloInstruction::CreateConstant( + Literal::CreateR1({1.f, 1.f, 1.f, 1.f, 1.f, 1.f, 1.f, 1.f}))); // Take a reference to 'data_init' to make it interfere with while result. builder.AddInstruction(HloInstruction::CreateBinary( data_shape_, HloOpcode::kAdd, data_init, one_vec)); @@ -743,21 +668,23 @@ class WhileCopyInsertionTest : public CopyInsertionTest { HloInstruction* BuildWhileInstructionWithCustomInit( const Shape& loop_state_shape, HloInstruction* data_init, HloComputation::Builder* builder) { + const bool nested = + ShapeUtil::Equal(loop_state_shape, nested_loop_state_shape_); auto induction_var_init = builder->AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(0))); + HloInstruction::CreateConstant(Literal::CreateR0(0))); auto condition = - module_.AddEmbeddedComputation(BuildConditionComputation()); - auto body = - module_.AddEmbeddedComputation(BuildIndependentBodyComputation()); + module_->AddEmbeddedComputation(BuildConditionComputation(nested)); + auto body = module_->AddEmbeddedComputation( + BuildIndependentBodyComputation(nested)); auto loop_state_init = builder->AddInstruction( HloInstruction::CreateTuple({induction_var_init, data_init})); auto while_hlo = builder->AddInstruction(HloInstruction::CreateWhile( loop_state_shape, condition, body, loop_state_init)); - module_.AddEntryComputation(builder->Build()); + module_->AddEntryComputation(builder->Build()); return while_hlo; } - HloModule module_; + std::unique_ptr module_; Shape induction_variable_shape_ = ShapeUtil::MakeShape(S32, {}); Shape data_shape_ = ShapeUtil::MakeShape(F32, {8}); Shape loop_state_shape_ = @@ -779,16 +706,23 @@ class WhileCopyInsertionTest : public CopyInsertionTest { // CopyInsertion pass should not generate any copies. // TEST_F(WhileCopyInsertionTest, IndependentTupleElements) { - auto condition = module_.AddEmbeddedComputation(BuildConditionComputation()); - auto body = module_.AddEmbeddedComputation(BuildIndependentBodyComputation()); - BuildWhileInstruction(condition, body); + auto condition = module_->AddEmbeddedComputation(BuildConditionComputation()); + auto body = + module_->AddEmbeddedComputation(BuildIndependentBodyComputation()); + auto while_hlo = BuildWhileInstruction(condition, body); + const HloInstruction* old_init = while_hlo->operand(0); HloInstruction* old_root = body->root_instruction(); - InsertCopies(&module_); + InsertCopies(module_.get()); HloInstruction* new_root = body->root_instruction(); + const HloInstruction* new_init = while_hlo->operand(0); // No copies should be inserted so root should not be updated. - CHECK_EQ(old_root, new_root); + EXPECT_EQ(old_root, new_root); + + // Both init indices need copies. + EXPECT_THAT(new_init, op::Tuple(op::Copy(old_init->operand(0)), + op::Copy(old_init->operand(1)))); } // Tests while body computation with dependent tuple elements: @@ -798,39 +732,25 @@ TEST_F(WhileCopyInsertionTest, IndependentTupleElements) { // out1 = Add(BCast(in0), in1) // Tuple(out0, out1) // -// CopyInsertion pass should generate: +// CopyInsertion pass should convert the root instruction to: // -// Tuple // old root -// / \ -// GTE(0) GTE(1) -// | | -// Copy | -// \ / -// Tuple // new root +// Tuple(Copy(out0), out1) // TEST_F(WhileCopyInsertionTest, DependentTupleElements) { - auto condition = module_.AddEmbeddedComputation(BuildConditionComputation()); - auto body = module_.AddEmbeddedComputation(BuildDependentBodyComputation()); - BuildWhileInstruction(condition, body); + auto condition = module_->AddEmbeddedComputation(BuildConditionComputation()); + auto body = module_->AddEmbeddedComputation(BuildDependentBodyComputation()); + auto while_hlo = BuildWhileInstruction(condition, body); + const HloInstruction* old_init = while_hlo->operand(0); HloInstruction* old_root = body->root_instruction(); - InsertCopies(&module_); + InsertCopies(module_.get()); HloInstruction* new_root = body->root_instruction(); + const HloInstruction* new_init = while_hlo->operand(0); - // Check all paths from 'new_root' to 'old_root'. - OperandTree op_tree; - op_tree.opcode = HloOpcode::kTuple; - - op_tree.op(0).opcode = HloOpcode::kCopy; - op_tree.op(0).op(0).opcode = HloOpcode::kGetTupleElement; - op_tree.op(0).op(0).op(0).opcode = HloOpcode::kTuple; - op_tree.op(0).op(0).op(0).instruction = old_root; - - op_tree.op(1).opcode = HloOpcode::kGetTupleElement; - op_tree.op(1).op(0).opcode = HloOpcode::kTuple; - op_tree.op(1).op(0).instruction = old_root; - - op_tree.Check(new_root); + EXPECT_THAT(new_root, + op::Tuple(op::Copy(old_root->operand(0)), old_root->operand(1))); + EXPECT_THAT(new_init, op::Tuple(op::Copy(old_init->operand(0)), + op::Copy(old_init->operand(1)))); } // Tests while body computation with read-only tuple element 0: @@ -846,20 +766,113 @@ TEST_F(WhileCopyInsertionTest, DependentTupleElements) { // \ / // TUPLE (root) // -// CopyInsertion pass should not generate any copies. -// +// CopyInsertion pass should not generate any copies for the while body. TEST_F(WhileCopyInsertionTest, DependentTupleElements_OneReadOnly) { - auto condition = module_.AddEmbeddedComputation(BuildConditionComputation()); - auto body = module_.AddEmbeddedComputation( + auto condition = module_->AddEmbeddedComputation(BuildConditionComputation()); + auto body = module_->AddEmbeddedComputation( BuildDependentBodyOneReadOnlyComputation()); - BuildWhileInstruction(condition, body); + auto while_hlo = BuildWhileInstruction(condition, body); + const HloInstruction* old_init = while_hlo->operand(0); HloInstruction* old_root = body->root_instruction(); - InsertCopies(&module_); + InsertCopies(module_.get()); HloInstruction* new_root = body->root_instruction(); + const HloInstruction* new_init = while_hlo->operand(0); - // No copies should be inserted so root should not be updated. - CHECK_EQ(old_root, new_root); + // No copies should be inserted in the body, so root should not be updated. + EXPECT_EQ(old_root, new_root); + + // Both indices need copies, even though Index 0 is read-only, since both are + // constants, which must be copied. + EXPECT_THAT(new_init, op::Tuple(op::Copy(old_init->operand(0)), + op::Copy(old_init->operand(1)))); +} + +// Same as above, but with two while loops, sharing entry parameters. +TEST_F(WhileCopyInsertionTest, + DependentTupleElements_OneReadOnly_TwoLoops_EntryParams) { + auto condition1 = + module_->AddEmbeddedComputation(BuildConditionComputation()); + auto condition2 = + module_->AddEmbeddedComputation(BuildConditionComputation()); + auto body1 = module_->AddEmbeddedComputation( + BuildDependentBodyOneReadOnlyComputation()); + auto body2 = module_->AddEmbeddedComputation( + BuildDependentBodyOneReadOnlyComputation()); + + auto builder = HloComputation::Builder(TestName() + ".While"); + auto iter_param = builder.AddInstruction( + HloInstruction::CreateParameter(0, induction_variable_shape_, "iter")); + auto data_param = builder.AddInstruction( + HloInstruction::CreateParameter(1, data_shape_, "data")); + auto loop_init = builder.AddInstruction( + HloInstruction::CreateTuple({iter_param, data_param})); + + auto while_hlo1 = builder.AddInstruction(HloInstruction::CreateWhile( + loop_state_shape_, condition1, body1, loop_init)); + auto while_hlo2 = builder.AddInstruction(HloInstruction::CreateWhile( + loop_state_shape_, condition2, body2, loop_init)); + module_->AddEntryComputation(builder.Build()); + + InsertCopies(module_.get()); + + // Both while loops alias iter_param, since index 0 is read-only in the body. + EXPECT_EQ(while_hlo1->operand(0)->operand(0), + while_hlo2->operand(0)->operand(0)); + EXPECT_EQ(while_hlo1->operand(0)->operand(0), iter_param); + + // Each while loop gets its own copy of data_param, since index 1 is not + // read-only in the body. + EXPECT_NE(while_hlo1->operand(0)->operand(1), + while_hlo2->operand(0)->operand(1)); + EXPECT_THAT(while_hlo1->operand(0)->operand(1), op::Copy(data_param)); + EXPECT_THAT(while_hlo2->operand(0)->operand(1), op::Copy(data_param)); +} + +// Same as above, but with two while loops, sharing non-parameters. +TEST_F(WhileCopyInsertionTest, + DependentTupleElements_OneReadOnly_TwoLoops_NonParams) { + auto condition1 = + module_->AddEmbeddedComputation(BuildConditionComputation()); + auto condition2 = + module_->AddEmbeddedComputation(BuildConditionComputation()); + auto body1 = module_->AddEmbeddedComputation( + BuildDependentBodyOneReadOnlyComputation()); + auto body2 = module_->AddEmbeddedComputation( + BuildDependentBodyOneReadOnlyComputation()); + + auto builder = HloComputation::Builder(TestName() + ".While"); + auto iter_param = builder.AddInstruction( + HloInstruction::CreateParameter(0, induction_variable_shape_, "iter")); + auto data_param = builder.AddInstruction( + HloInstruction::CreateParameter(1, data_shape_, "data")); + // Add dummy ops to ensure loop_init elements aren't entry parameters. + auto iter_value = builder.AddInstruction(HloInstruction::CreateUnary( + iter_param->shape(), HloOpcode::kExp, iter_param)); + auto data_value = builder.AddInstruction(HloInstruction::CreateUnary( + data_param->shape(), HloOpcode::kExp, data_param)); + auto loop_init = builder.AddInstruction( + HloInstruction::CreateTuple({iter_value, data_value})); + + auto while_hlo1 = builder.AddInstruction(HloInstruction::CreateWhile( + loop_state_shape_, condition1, body1, loop_init)); + auto while_hlo2 = builder.AddInstruction(HloInstruction::CreateWhile( + loop_state_shape_, condition2, body2, loop_init)); + module_->AddEntryComputation(builder.Build()); + + InsertCopies(module_.get()); + + // No copies of iter_value are necessary, since index 0 is read-only in both + // while bodies. + EXPECT_EQ(while_hlo1->operand(0)->operand(0), iter_value); + EXPECT_EQ(while_hlo2->operand(0)->operand(0), iter_value); + + // Each while loop gets its own copy of data_value, since index 1 is not + // read-only in the body. + EXPECT_NE(while_hlo1->operand(0)->operand(1), + while_hlo2->operand(0)->operand(1)); + EXPECT_THAT(while_hlo1->operand(0)->operand(1), op::Copy(data_value)); + EXPECT_THAT(while_hlo2->operand(0)->operand(1), op::Copy(data_value)); } // Tests while body computation with nested tuple elements: @@ -872,7 +885,8 @@ TEST_F(WhileCopyInsertionTest, DependentTupleElements_OneReadOnly) { // Add Reverse // | | // -// CopyInsertion pass should generate: +// CopyInsertion pass will conceptually generate the following, but with the +// actual GTE and Tuple instructions optimized away: // // Tuple // old root // / \ @@ -892,110 +906,47 @@ TEST_F(WhileCopyInsertionTest, DependentTupleElements_OneReadOnly) { // TEST_F(WhileCopyInsertionTest, NestedTupleElements) { auto condition = - module_.AddEmbeddedComputation(BuildConditionComputation(true)); - auto body = module_.AddEmbeddedComputation(BuildNestedBodyComputation()); + module_->AddEmbeddedComputation(BuildConditionComputation(true)); + auto body = module_->AddEmbeddedComputation(BuildNestedBodyComputation()); BuildWhileInstruction(condition, body, true); HloInstruction* old_root = body->root_instruction(); - InsertCopies(&module_); - HloInstruction* new_root = body->root_instruction(); - - // Check all paths from 'new_root' to 'old_root'. - OperandTree op_tree; - op_tree.opcode = HloOpcode::kTuple; - - op_tree.op(0).opcode = HloOpcode::kGetTupleElement; - op_tree.op(0).op(0).opcode = HloOpcode::kTuple; - op_tree.op(0).op(0).instruction = old_root; - - op_tree.op(1).opcode = HloOpcode::kTuple; - - op_tree.op(1).op(0).opcode = HloOpcode::kGetTupleElement; - op_tree.op(1).op(0).op(0).opcode = HloOpcode::kGetTupleElement; - op_tree.op(1).op(0).op(0).op(0).opcode = HloOpcode::kTuple; - op_tree.op(1).op(0).op(0).op(0).instruction = old_root; + InsertCopies(module_.get()); - op_tree.op(1).op(1).opcode = HloOpcode::kCopy; - op_tree.op(1).op(1).op(0).opcode = HloOpcode::kGetTupleElement; - op_tree.op(1).op(1).op(0).op(0).opcode = HloOpcode::kGetTupleElement; - op_tree.op(1).op(1).op(0).op(0).op(0).opcode = HloOpcode::kTuple; - op_tree.op(1).op(1).op(0).op(0).op(0).instruction = old_root; - - op_tree.Check(new_root); + EXPECT_THAT(body->root_instruction(), + op::Tuple(old_root->operand(0), + op::Tuple(old_root->operand(1)->operand(0), + op::Copy(old_root->operand(1)->operand(1))))); } // Tests while init instruction which points-to a constant. // // init = Tuple(Constant(S32, {}), Constant(F32, {8})) // -// CopyInsertion pass should generate: -// -// Tuple // old init -// / \ -// GTE(0) GTE(1) -// | | -// Copy Copy -// \ / -// Tuple // new init +// CopyInsertion pass should add copies for both constants. // TEST_F(WhileCopyInsertionTest, InitPointsToConstant) { auto while_hlo = BuildWhileInstruction_InitPointsToConstant(); auto old_init = while_hlo->operand(0); - InsertCopies(&module_); - auto new_init = while_hlo->operand(0); - - // Check all paths from 'new_init' to 'old_init'. - OperandTree op_tree; - op_tree.opcode = HloOpcode::kTuple; + InsertCopies(module_.get()); - op_tree.op(0).opcode = HloOpcode::kCopy; - op_tree.op(0).op(0).opcode = HloOpcode::kGetTupleElement; - op_tree.op(0).op(0).op(0).opcode = HloOpcode::kTuple; - op_tree.op(0).op(0).op(0).instruction = old_init; - - op_tree.op(1).opcode = HloOpcode::kCopy; - op_tree.op(1).op(0).opcode = HloOpcode::kGetTupleElement; - op_tree.op(1).op(0).op(0).opcode = HloOpcode::kTuple; - op_tree.op(1).op(0).op(0).instruction = old_init; - - op_tree.Check(new_init); + EXPECT_THAT(while_hlo->operand(0), op::Tuple(op::Copy(old_init->operand(0)), + op::Copy(old_init->operand(1)))); } // Tests while init instruction which points-to a parameter. // // init = Tuple(Constant(S32, {}), Parameter(F32, {8})) // -// CopyInsertion pass should generate: -// -// Tuple // old init -// / \ -// GTE(0) GTE(1) -// | | -// Copy Copy -// \ / -// Tuple // new init +// CopyInsertion pass should add copies for both the constant and parameter. // TEST_F(WhileCopyInsertionTest, InitPointsToParameter) { auto while_hlo = BuildWhileInstruction_InitPointsToParameter(); auto old_init = while_hlo->operand(0); - InsertCopies(&module_); - auto new_init = while_hlo->operand(0); + InsertCopies(module_.get()); - // Check all paths from 'new_init' to 'old_init'. - OperandTree op_tree; - op_tree.opcode = HloOpcode::kTuple; - - op_tree.op(0).opcode = HloOpcode::kCopy; - op_tree.op(0).op(0).opcode = HloOpcode::kGetTupleElement; - op_tree.op(0).op(0).op(0).opcode = HloOpcode::kTuple; - op_tree.op(0).op(0).op(0).instruction = old_init; - - op_tree.op(1).opcode = HloOpcode::kCopy; - op_tree.op(1).op(0).opcode = HloOpcode::kGetTupleElement; - op_tree.op(1).op(0).op(0).opcode = HloOpcode::kTuple; - op_tree.op(1).op(0).op(0).instruction = old_init; - - op_tree.Check(new_init); + EXPECT_THAT(while_hlo->operand(0), op::Tuple(op::Copy(old_init->operand(0)), + op::Copy(old_init->operand(1)))); } // Tests while init instruction which has an ambiguous points-to set. @@ -1003,7 +954,8 @@ TEST_F(WhileCopyInsertionTest, InitPointsToParameter) { // select = Select(pred, tuple1, tuple2) // init = Tuple(Constant(S32, {}), Parameter(F32, {8})) // -// CopyInsertion pass should generate: +// CopyInsertion pass will conceptually generate the following, but with some of +// the actual GTE and Tuple instructions optimized away: // // Tuple // old init // / \ @@ -1024,40 +976,22 @@ TEST_F(WhileCopyInsertionTest, InitPointsToParameter) { TEST_F(WhileCopyInsertionTest, InitPointsToAmbiguous) { auto while_hlo = BuildWhileInstruction_InitPointsToAmbiguous(); auto old_init = while_hlo->operand(0); - InsertCopies(&module_); - auto new_init = while_hlo->operand(0); - - // Check all paths from 'new_init' to 'old_init'. - OperandTree op_tree; - op_tree.opcode = HloOpcode::kTuple; - - op_tree.op(0).opcode = HloOpcode::kCopy; - op_tree.op(0).op(0).opcode = HloOpcode::kGetTupleElement; - op_tree.op(0).op(0).op(0).opcode = HloOpcode::kTuple; - op_tree.op(0).op(0).op(0).instruction = old_init; - - op_tree.op(1).opcode = HloOpcode::kTuple; - - op_tree.op(1).op(0).opcode = HloOpcode::kCopy; - op_tree.op(1).op(0).op(0).opcode = HloOpcode::kGetTupleElement; - op_tree.op(1).op(0).op(0).op(0).opcode = HloOpcode::kGetTupleElement; - op_tree.op(1).op(0).op(0).op(0).op(0).opcode = HloOpcode::kTuple; - op_tree.op(1).op(0).op(0).op(0).op(0).instruction = old_init; - - op_tree.op(1).op(1).opcode = HloOpcode::kCopy; - op_tree.op(1).op(1).op(0).opcode = HloOpcode::kGetTupleElement; - op_tree.op(1).op(1).op(0).op(0).opcode = HloOpcode::kGetTupleElement; - op_tree.op(1).op(1).op(0).op(0).op(0).opcode = HloOpcode::kTuple; - op_tree.op(1).op(1).op(0).op(0).op(0).instruction = old_init; - - op_tree.Check(new_init); + InsertCopies(module_.get()); + + EXPECT_THAT( + while_hlo->operand(0), + op::Tuple( + op::Copy(old_init->operand(0)), + op::Tuple(op::Copy(op::GetTupleElement(old_init->operand(1))), + op::Copy(op::GetTupleElement(old_init->operand(1)))))); } // Tests while init instruction which has a non-distinct points-to set. // // init = Tuple(Constant(S32, {}), Tuple({vec_one, vec_one})) // -// CopyInsertion pass should generate: +// CopyInsertion pass will conceptually generate the following, but with some of +// the actual GTE and Tuple instructions optimized away: // // Tuple // old init // / \ @@ -1078,73 +1012,116 @@ TEST_F(WhileCopyInsertionTest, InitPointsToAmbiguous) { TEST_F(WhileCopyInsertionTest, InitPointsToNonDistinct) { auto while_hlo = BuildWhileInstruction_InitPointsToNonDistinct(); auto old_init = while_hlo->operand(0); - InsertCopies(&module_); - auto new_init = while_hlo->operand(0); - - // Check all paths from 'new_init' to 'old_init'. - OperandTree op_tree; - op_tree.opcode = HloOpcode::kTuple; - - op_tree.op(0).opcode = HloOpcode::kCopy; - op_tree.op(0).op(0).opcode = HloOpcode::kGetTupleElement; - op_tree.op(0).op(0).op(0).opcode = HloOpcode::kTuple; - op_tree.op(0).op(0).op(0).instruction = old_init; - - op_tree.op(1).opcode = HloOpcode::kTuple; + InsertCopies(module_.get()); - op_tree.op(1).op(0).opcode = HloOpcode::kCopy; - op_tree.op(1).op(0).op(0).opcode = HloOpcode::kGetTupleElement; - op_tree.op(1).op(0).op(0).op(0).opcode = HloOpcode::kGetTupleElement; - op_tree.op(1).op(0).op(0).op(0).op(0).opcode = HloOpcode::kTuple; - op_tree.op(1).op(0).op(0).op(0).op(0).instruction = old_init; - - op_tree.op(1).op(1).opcode = HloOpcode::kCopy; - op_tree.op(1).op(1).op(0).opcode = HloOpcode::kGetTupleElement; - op_tree.op(1).op(1).op(0).op(0).opcode = HloOpcode::kGetTupleElement; - op_tree.op(1).op(1).op(0).op(0).op(0).opcode = HloOpcode::kTuple; - op_tree.op(1).op(1).op(0).op(0).op(0).instruction = old_init; - - op_tree.Check(new_init); + EXPECT_THAT(while_hlo->operand(0), + op::Tuple(op::Copy(old_init->operand(0)), + op::Tuple(op::Copy(old_init->operand(1)->operand(0)), + op::Copy(old_init->operand(1)->operand(0))))); } -// Tests while init instruction buffer which interfers with while result buffer. +// Tests while init instruction buffer which interferes with while result +// buffer. // // init_data = Broadcast(...) // add_unrelated = Add(init_data) // takes a reference to cause interference // init = Tuple(Constant(S32, {}), init_data)) // -// CopyInsertion pass should generate: -// -// Tuple // old init -// / \ -// GTE(0) GTE(1) -// | | -// Copy Copy -// \ / -// Tuple // new init +// CopyInsertion pass should copy both operands. // TEST_F(WhileCopyInsertionTest, InitPointsToInterfering) { auto while_hlo = BuildWhileInstruction_InitPointsToInterfering(); auto old_init = while_hlo->operand(0); - InsertCopies(&module_); - auto new_init = while_hlo->operand(0); + InsertCopies(module_.get()); - // Check all paths from 'new_init' to 'old_init'. - OperandTree op_tree; - op_tree.opcode = HloOpcode::kTuple; - - op_tree.op(0).opcode = HloOpcode::kCopy; - op_tree.op(0).op(0).opcode = HloOpcode::kGetTupleElement; - op_tree.op(0).op(0).op(0).opcode = HloOpcode::kTuple; - op_tree.op(0).op(0).op(0).instruction = old_init; - - op_tree.op(1).opcode = HloOpcode::kCopy; - op_tree.op(1).op(0).opcode = HloOpcode::kGetTupleElement; - op_tree.op(1).op(0).op(0).opcode = HloOpcode::kTuple; - op_tree.op(1).op(0).op(0).instruction = old_init; + EXPECT_THAT(while_hlo->operand(0), op::Tuple(op::Copy(old_init->operand(0)), + op::Copy(old_init->operand(1)))); +} - op_tree.Check(new_init); +// Tests while init instruction buffer which has a non-distinct points-to set: +// +// init = Tuple(Parameter(S32, {}), Parameter(F32, {8}, +// Parameter(F32, {8}))) +// +// where the second and third parameters are identical *and* the tuple shared +// by another while instruction.. +// +// Verifies that the resulting point-to set is distinct in the resulting Tuple +// (non-identical Copys). In other words, verifies that copy sharing does not +// insert identical copies to the resulting tuple. +TEST_F(WhileCopyInsertionTest, InitPointsToNonDistinctUsedByTwoWhileLoops) { + auto condition1 = + module_->AddEmbeddedComputation(BuildConditionComputation()); + auto condition2 = + module_->AddEmbeddedComputation(BuildConditionComputation()); + // Loop body that outputs tuple comprises two elements dependent on the init + // tuple. + auto body1 = + module_->AddEmbeddedComputation(BuildDependentBodyComputation2()); + auto body2 = + module_->AddEmbeddedComputation(BuildDependentBodyComputation2()); + + auto builder = HloComputation::Builder(TestName() + ".While"); + + auto iter_param = builder.AddInstruction( + HloInstruction::CreateParameter(0, induction_variable_shape_, "iter")); + auto data_param = builder.AddInstruction( + HloInstruction::CreateParameter(1, data_shape_, "data")); + + // Loop init tuple contains two identical parameter buffers. + auto loop_init = builder.AddInstruction( + HloInstruction::CreateTuple({iter_param, data_param, data_param})); + + const Shape& loop_state_shape = ShapeUtil::MakeTupleShape( + {induction_variable_shape_, data_shape_, data_shape_}); + + // Two while loops shares the same loop init tuple. + auto while_hlo1 = builder.AddInstruction(HloInstruction::CreateWhile( + loop_state_shape, condition1, body1, loop_init)); + auto while_hlo2 = builder.AddInstruction(HloInstruction::CreateWhile( + loop_state_shape, condition2, body2, loop_init)); + + module_->AddEntryComputation(builder.Build()); + + auto points_to_analysis = + TuplePointsToAnalysis::Run(module_.get()).ConsumeValueOrDie(); + + // Asserts that the init tuples before copy insertion is non-distinct. + ASSERT_FALSE( + points_to_analysis->GetPointsToSet(while_hlo1->operand(0)).IsDistinct()); + ASSERT_FALSE( + points_to_analysis->GetPointsToSet(while_hlo2->operand(0)).IsDistinct()); + + auto old_init1 = while_hlo1->operand(0); + auto old_init2 = while_hlo2->operand(0); + + InsertCopies(module_.get()); + + EXPECT_THAT(while_hlo1->operand(0), + op::Tuple(op::Copy(old_init1->operand(0)), + op::Copy(old_init1->operand(1)), + op::Copy(old_init1->operand(2)))); + + EXPECT_THAT(while_hlo2->operand(0), + op::Tuple(op::Copy(old_init2->operand(0)), + op::Copy(old_init2->operand(1)), + op::Copy(old_init2->operand(2)))); + + // Verifies the init tuples after copy insertion is distinct. + points_to_analysis = + TuplePointsToAnalysis::Run(module_.get()).ConsumeValueOrDie(); + const auto& points_to1 = + points_to_analysis->GetPointsToSet(while_hlo1->operand(0)); + EXPECT_TRUE(points_to1.IsDistinct()); + + const auto& points_to2 = + points_to_analysis->GetPointsToSet(while_hlo2->operand(0)); + EXPECT_TRUE(points_to2.IsDistinct()); } } // namespace } // namespace xla + +int main(int argc, char** argv) { + return xla::ParseDebugOptionsFlagsAndRunTests(argc, argv); +} diff --git a/tensorflow/compiler/xla/service/cpu/BUILD b/tensorflow/compiler/xla/service/cpu/BUILD index e9963528111994c4918861eaa52ab915fe34fd93..37d523b7573d1023f34a86174add87b83afcc9f0 100644 --- a/tensorflow/compiler/xla/service/cpu/BUILD +++ b/tensorflow/compiler/xla/service/cpu/BUILD @@ -5,10 +5,6 @@ licenses(["notice"]) # Apache 2.0 package( default_visibility = [":friends"], - features = [ - "-layering_check", - "-parse_headers", - ], ) package_group( @@ -38,6 +34,7 @@ cc_library( ":conv_canonicalization", ":cpu_executable", ":cpu_instruction_fusion", + ":cpu_options", ":cpu_parallelization_preparation", ":disassembler", ":ir_emission_utils", @@ -52,25 +49,28 @@ cc_library( "//tensorflow/compiler/xla:types", "//tensorflow/compiler/xla:util", "//tensorflow/compiler/xla:xla_data_proto", - "//tensorflow/compiler/xla/legacy_flags:cpu_compiler_flags", - "//tensorflow/compiler/xla/port:initialize", "//tensorflow/compiler/xla/service:algebraic_simplifier", + "//tensorflow/compiler/xla/service:batchnorm_rewriter", "//tensorflow/compiler/xla/service:buffer_assignment", "//tensorflow/compiler/xla/service:buffer_liveness", - "//tensorflow/compiler/xla/service:compiler", "//tensorflow/compiler/xla/service:copy_insertion", "//tensorflow/compiler/xla/service:executable", + "//tensorflow/compiler/xla/service:flatten_call_graph", "//tensorflow/compiler/xla/service:hlo", "//tensorflow/compiler/xla/service:hlo_constant_folding", "//tensorflow/compiler/xla/service:hlo_cse", "//tensorflow/compiler/xla/service:hlo_dce", - "//tensorflow/compiler/xla/service:hlo_module_config", "//tensorflow/compiler/xla/service:hlo_ordering", "//tensorflow/compiler/xla/service:hlo_pass", "//tensorflow/compiler/xla/service:hlo_pass_pipeline", + "//tensorflow/compiler/xla/service:hlo_proto", + "//tensorflow/compiler/xla/service:hlo_proto_util", + "//tensorflow/compiler/xla/service:hlo_scheduling", "//tensorflow/compiler/xla/service:hlo_subcomputation_unification", "//tensorflow/compiler/xla/service:hlo_verifier", "//tensorflow/compiler/xla/service:inliner", + "//tensorflow/compiler/xla/service:llvm_compiler", + "//tensorflow/compiler/xla/service:reduce_precision_insertion", "//tensorflow/compiler/xla/service:reshape_mover", "//tensorflow/compiler/xla/service:transpose_folding", "//tensorflow/compiler/xla/service/llvm_ir:llvm_util", # fixdeps: keep @@ -97,10 +97,12 @@ cc_library( name = "simple_orc_jit", srcs = ["simple_orc_jit.cc"], hdrs = ["simple_orc_jit.h"], + linkopts = ["-ldl"], deps = [ ":compiler_functor", ":cpu_runtime", ":cpu_runtime_avx", + ":cpu_runtime_neon", ":cpu_runtime_sse4_1", ":disassembler", ":runtime_conv2d", @@ -111,6 +113,7 @@ cc_library( "//tensorflow/compiler/xla:util", "//tensorflow/core:lib", "@llvm//:core", + "@llvm//:execution_engine", "@llvm//:mc", # fixdeps: keep "@llvm//:orc_jit", "@llvm//:support", @@ -137,7 +140,6 @@ cc_library( "//tensorflow/compiler/xla/service:executable", "//tensorflow/compiler/xla/service:hlo", "//tensorflow/compiler/xla/service:hlo_execution_profile", - "//tensorflow/compiler/xla/service:hlo_module_config", "//tensorflow/compiler/xla/service:logical_buffer", "//tensorflow/compiler/xla/service:shaped_buffer", "//tensorflow/compiler/xla/service:tuple_points_to_analysis", @@ -150,9 +152,12 @@ cc_library( cc_library( name = "parallel_cpu_executable", srcs = ["parallel_cpu_executable.cc"], - hdrs = ["parallel_cpu_executable.h"], + hdrs = [ + "parallel_cpu_executable.h", + ], deps = [ ":cpu_runtime", + ":shape_partition", ":simple_orc_jit", "//tensorflow/compiler/xla:shape_util", "//tensorflow/compiler/xla:status_macros", @@ -165,7 +170,6 @@ cc_library( "//tensorflow/compiler/xla/service:executable", "//tensorflow/compiler/xla/service:hlo", "//tensorflow/compiler/xla/service:hlo_execution_profile", - "//tensorflow/compiler/xla/service:hlo_module_config", "//tensorflow/compiler/xla/service:logical_buffer", "//tensorflow/compiler/xla/service:shaped_buffer", "//tensorflow/core:lib", @@ -177,8 +181,11 @@ cc_library( cc_library( name = "ir_emitter", srcs = ["ir_emitter.cc"], - hdrs = ["ir_emitter.h"], + hdrs = [ + "ir_emitter.h", + ], deps = [ + ":cpu_options", ":cpu_runtime", ":dot_op_emitter", ":elemental_ir_emitter", @@ -191,7 +198,6 @@ cc_library( "//tensorflow/compiler/xla:util", "//tensorflow/compiler/xla:window_util", "//tensorflow/compiler/xla:xla_data_proto", - "//tensorflow/compiler/xla/legacy_flags:cpu_runtime_flags", "//tensorflow/compiler/xla/service:buffer_assignment", "//tensorflow/compiler/xla/service:elemental_ir_emitter", "//tensorflow/compiler/xla/service:hlo", @@ -207,6 +213,7 @@ cc_library( "//tensorflow/core:lib", "@llvm//:core", "@llvm//:support", + "@llvm//:target", ], ) @@ -222,8 +229,8 @@ cc_library( "//tensorflow/compiler/xla:types", "//tensorflow/compiler/xla:util", "//tensorflow/compiler/xla:xla_data_proto", - "//tensorflow/compiler/xla/legacy_flags:cpu_runtime_flags", "//tensorflow/compiler/xla/service:hlo", + "//tensorflow/compiler/xla/service:hlo_module_config", "//tensorflow/compiler/xla/service/llvm_ir:ir_array", "//tensorflow/compiler/xla/service/llvm_ir:llvm_loop", "//tensorflow/compiler/xla/service/llvm_ir:llvm_util", @@ -278,17 +285,19 @@ cc_library( deps = [ ":cpu_runtime", ":cpu_runtime_avx", + ":cpu_runtime_neon", ":cpu_runtime_sse4_1", ":disassembler", + ":llvm_ir_runtime", "//tensorflow/compiler/xla:statusor", "//tensorflow/compiler/xla:types", "//tensorflow/compiler/xla:util", - "//tensorflow/compiler/xla/legacy_flags:compiler_functor_flags", - "//tensorflow/compiler/xla/legacy_flags:cpu_runtime_flags", + "//tensorflow/compiler/xla/service:llvm_compiler", "//tensorflow/compiler/xla/service/llvm_ir:llvm_util", "//tensorflow/core:lib", "@llvm//:analysis", "@llvm//:core", + "@llvm//:execution_engine", "@llvm//:ipo", "@llvm//:mc", "@llvm//:object", @@ -302,11 +311,11 @@ cc_library( srcs = ["cpu_runtime_sse4_1.cc"], hdrs = ["cpu_runtime_sse4_1.h"], copts = ["-DEIGEN_AVOID_STL_ARRAY"], + visibility = ["//visibility:public"], deps = [ - "//tensorflow/core:lib", + "//tensorflow/core:framework_lite", "//third_party/eigen3", ], - alwayslink = True, ) cc_library( @@ -314,26 +323,59 @@ cc_library( srcs = ["cpu_runtime_avx.cc"], hdrs = ["cpu_runtime_avx.h"], copts = ["-DEIGEN_AVOID_STL_ARRAY"], + visibility = ["//visibility:public"], deps = [ - "//tensorflow/core:lib", + "//tensorflow/core:framework_lite", + "//third_party/eigen3", + ], +) + +cc_library( + name = "cpu_runtime_neon", + srcs = ["cpu_runtime_neon.cc"], + hdrs = ["cpu_runtime_neon.h"], + # runtime_copts() enables -mfpu=neon + copts = ["-DEIGEN_AVOID_STL_ARRAY"] + runtime_copts(), + visibility = ["//visibility:public"], + deps = [ + "//tensorflow/core:framework_lite", "//third_party/eigen3", ], - alwayslink = True, ) cc_library( name = "cpu_runtime", srcs = [ "cpu_runtime.cc", - "infeed_manager.cc", + "xfeed_manager.cc", ], hdrs = [ "cpu_runtime.h", - "infeed_manager.h", + "xfeed_manager.h", ], copts = runtime_copts(), deps = [ + "//tensorflow/compiler/xla:shape_util", + "//tensorflow/compiler/xla:statusor", "//tensorflow/compiler/xla:types", + "//tensorflow/compiler/xla:xla_data_proto", + "//tensorflow/compiler/xla/service/llvm_ir:llvm_util", + "//tensorflow/core:lib", + ], +) + +cc_library( + name = "llvm_ir_runtime", + srcs = [ + "llvm_ir_runtime.cc", + ], + hdrs = [ + "llvm_ir_runtime.h", + ], + deps = [ + "//tensorflow/core:lib", + "@llvm//:core", + "@llvm//:transform_utils", ], ) @@ -354,6 +396,17 @@ cc_library( ], ) +cc_library( + name = "runtime_matvec", + srcs = ["runtime_matvec.cc"], + hdrs = ["runtime_matvec.h"], + copts = runtime_copts(), + deps = [ + "//tensorflow/core:framework_lite", + "//third_party/eigen3", + ], +) + cc_library( name = "runtime_matmul", srcs = ["runtime_matmul.cc"], @@ -361,6 +414,7 @@ cc_library( copts = runtime_copts(), visibility = ["//visibility:public"], deps = [ + ":runtime_matvec", "//tensorflow/compiler/xla:executable_run_options", "//tensorflow/core:framework_lite", "//third_party/eigen3", @@ -390,6 +444,7 @@ cc_library( copts = runtime_copts(), visibility = ["//visibility:public"], deps = [ + ":runtime_matvec", "//tensorflow/core:framework_lite", "//third_party/eigen3", ], @@ -401,10 +456,12 @@ cc_test( deps = [ ":cpu_runtime", ":runtime_matmul", + ":runtime_single_threaded_matmul", "//tensorflow/compiler/xla:array2d", "//tensorflow/compiler/xla:types", "//tensorflow/compiler/xla:util", "//tensorflow/compiler/xla/client:local_client", + "//tensorflow/core:core_cpu_internal", "//tensorflow/core:lib", "//tensorflow/core:test", "//tensorflow/core:test_main", @@ -413,10 +470,25 @@ cc_test( ) cc_test( - name = "infeed_manager_test", - srcs = ["infeed_manager_test.cc"], + name = "cpu_instruction_fusion_test", + size = "small", + srcs = ["cpu_instruction_fusion_test.cc"], + deps = [ + ":cpu_instruction_fusion", + "//tensorflow/compiler/xla/service:hlo_matchers", + "//tensorflow/compiler/xla/service:transpose_folding", + "//tensorflow/compiler/xla/tests:hlo_test_base", + "//tensorflow/core:lib", + ], +) + +cc_test( + name = "xfeed_manager_test", + size = "small", + srcs = ["xfeed_manager_test.cc"], deps = [ ":cpu_runtime", + "//tensorflow/compiler/xla:shape_util", "//tensorflow/core:lib", "//tensorflow/core:test", "//tensorflow/core:test_main", @@ -436,10 +508,16 @@ cc_library( cc_library( name = "cpu_parallelization_preparation", srcs = ["cpu_parallelization_preparation.cc"], - hdrs = ["cpu_parallelization_preparation.h"], + hdrs = [ + "cpu_parallelization_preparation.h", + ], deps = [ + ":ir_emission_utils", + ":shape_partition", "//tensorflow/compiler/xla:types", + "//tensorflow/compiler/xla:util", "//tensorflow/compiler/xla/service:hlo", + "//tensorflow/compiler/xla/service:hlo_cost_analysis", "//tensorflow/compiler/xla/service:hlo_pass", "//tensorflow/compiler/xla/service:logical_buffer", "//tensorflow/compiler/xla/service:tuple_points_to_analysis", @@ -471,7 +549,6 @@ cc_library( ":cpu_runtime", "//tensorflow/compiler/xla:shape_util", "//tensorflow/compiler/xla:window_util", - "//tensorflow/compiler/xla/legacy_flags:cpu_runtime_flags", "//tensorflow/compiler/xla/service:hlo", ], ) @@ -498,25 +575,56 @@ cc_library( "//tensorflow/compiler/xla:shape_util", "//tensorflow/compiler/xla:util", "//tensorflow/compiler/xla:xla_data_proto", - "//tensorflow/compiler/xla/legacy_flags:cpu_runtime_flags", "//tensorflow/compiler/xla/service:hlo", "//tensorflow/compiler/xla/service:hlo_pass", + "//tensorflow/core:lib", ], ) cc_test( name = "conv_canonicalization_test", + size = "small", srcs = ["conv_canonicalization_test.cc"], deps = [ ":conv_canonicalization", + "//tensorflow/compiler/xla:test", "//tensorflow/compiler/xla:test_helpers", "//tensorflow/compiler/xla:util", "//tensorflow/compiler/xla/service:hlo", "//tensorflow/compiler/xla/tests:hlo_test_base", + ], +) + +cc_library( + name = "shape_partition", + srcs = ["shape_partition.cc"], + hdrs = ["shape_partition.h"], + deps = [ + "//tensorflow/compiler/xla:shape_util", + ], +) + +cc_test( + name = "shape_partition_test", + srcs = ["shape_partition_test.cc"], + deps = [ + ":shape_partition", + "//tensorflow/compiler/xla:test_helpers", + "//tensorflow/compiler/xla:util", + "//tensorflow/compiler/xla/tests:hlo_test_base", "//tensorflow/core:test_main", ], ) +cc_library( + name = "cpu_options", + srcs = ["cpu_options.cc"], + hdrs = ["cpu_options.h"], + deps = [ + "//tensorflow/compiler/xla/service:hlo_module_config", + ], +) + # ----------------------------------------------------------------------------- filegroup( diff --git a/tensorflow/compiler/xla/service/cpu/build_defs.bzl b/tensorflow/compiler/xla/service/cpu/build_defs.bzl index b4b52197516dc083809014c9882bf7845f3723ac..e78330b21689fdd818cd97128bbcaaa9e0118602 100644 --- a/tensorflow/compiler/xla/service/cpu/build_defs.bzl +++ b/tensorflow/compiler/xla/service/cpu/build_defs.bzl @@ -1,11 +1,12 @@ """build_defs for service/cpu.""" + def runtime_copts(): """Returns copts used for CPU runtime libraries.""" - return (["-DEIGEN_AVOID_STL_ARRAY"] + - select({ - "//tensorflow:android_arm": ["-mfpu=neon"], - "//conditions:default": []}) + - select({ - "//tensorflow:android": ["-O2"], - "//conditions:default": []})) + return (["-DEIGEN_AVOID_STL_ARRAY"] + select({ + "//tensorflow:android_arm": ["-mfpu=neon"], + "//conditions:default": [] + }) + select({ + "//tensorflow:android": ["-O2"], + "//conditions:default": [] + })) diff --git a/tensorflow/compiler/xla/service/cpu/compiler_functor.cc b/tensorflow/compiler/xla/service/cpu/compiler_functor.cc index 89b3302bca052586ddf6cdd63c4cf2483e51b8f8..141582c0690474d27cd6917dd6031d33004be5d0 100644 --- a/tensorflow/compiler/xla/service/cpu/compiler_functor.cc +++ b/tensorflow/compiler/xla/service/cpu/compiler_functor.cc @@ -22,30 +22,29 @@ limitations under the License. #include #include -#include "external/llvm/include/llvm/ADT/StringRef.h" -#include "external/llvm/include/llvm/Analysis/TargetLibraryInfo.h" -#include "external/llvm/include/llvm/Analysis/TargetTransformInfo.h" -#include "external/llvm/include/llvm/ExecutionEngine/ObjectMemoryBuffer.h" -#include "external/llvm/include/llvm/IR/LegacyPassManager.h" -#include "external/llvm/include/llvm/IR/Verifier.h" -#include "external/llvm/include/llvm/MC/MCContext.h" -#include "external/llvm/include/llvm/Object/ObjectFile.h" -#include "external/llvm/include/llvm/Support/raw_ostream.h" -#include "external/llvm/include/llvm/Target/TargetMachine.h" -#include "external/llvm/include/llvm/Transforms/IPO.h" -#include "external/llvm/include/llvm/Transforms/IPO/AlwaysInliner.h" -#include "external/llvm/include/llvm/Transforms/IPO/PassManagerBuilder.h" -#include "tensorflow/compiler/xla/legacy_flags/compiler_functor_flags.h" -#include "tensorflow/compiler/xla/legacy_flags/cpu_runtime_flags.h" +#include "llvm/ADT/StringRef.h" +#include "llvm/Analysis/TargetLibraryInfo.h" +#include "llvm/Analysis/TargetTransformInfo.h" +#include "llvm/ExecutionEngine/ObjectMemoryBuffer.h" +#include "llvm/IR/LegacyPassManager.h" +#include "llvm/IR/Verifier.h" +#include "llvm/MC/MCContext.h" +#include "llvm/Object/ObjectFile.h" +#include "llvm/Support/raw_ostream.h" +#include "llvm/Target/TargetMachine.h" +#include "llvm/Transforms/IPO.h" +#include "llvm/Transforms/IPO/AlwaysInliner.h" +#include "llvm/Transforms/IPO/PassManagerBuilder.h" #include "tensorflow/compiler/xla/ptr_util.h" #include "tensorflow/compiler/xla/service/cpu/cpu_runtime.h" #include "tensorflow/compiler/xla/service/cpu/cpu_runtime_avx.h" +#include "tensorflow/compiler/xla/service/cpu/cpu_runtime_neon.h" #include "tensorflow/compiler/xla/service/cpu/cpu_runtime_sse4_1.h" +#include "tensorflow/compiler/xla/service/cpu/llvm_ir_runtime.h" #include "tensorflow/compiler/xla/service/llvm_ir/llvm_util.h" #include "tensorflow/compiler/xla/statusor.h" #include "tensorflow/compiler/xla/types.h" #include "tensorflow/compiler/xla/util.h" -#include "tensorflow/core/platform/env.h" #include "tensorflow/core/platform/logging.h" namespace xla { @@ -56,37 +55,108 @@ CompilerFunctor::AllIntrinsics() { VectorIntrinsics intrinsics; intrinsics.sse_intrinsics = true; intrinsics.avx_intrinsics = true; + intrinsics.neon_intrinsics = true; return intrinsics; } +/* Create filtered versions of the LLVM Pass Managers to filter out some +of the expensive passes. +Profiling: + learning/brain/google/xla/benchmarks:inception_cpu_benchmark + learning/brain/google/xla/benchmarks:cifarnet +pointed to LICM and IndVarSimplify as the hottest passes. +LICM is known to exhibit O(n^2) time in the number of instructions. +IndVarSimplify is slow due to SCEV. If loops are emitted in canonical form, +this pass is not necessary. +Disabling these as a starting point. +*/ +// TODO(b/64227304) Creating a custom pass pipeline will replace this. + +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 { + if (disable_expensive_passes_) { + llvm::StringRef PassName = p->getPassName(); + if (PassName.contains("LICM") || PassName.contains("IndVarSimplify") || + PassName.contains("LoopUnroll")) { + return; + } + } + llvm::legacy::FunctionPassManager::add(p); + } + + private: + bool disable_expensive_passes_; +}; + +class FilteredPassManager : public llvm::legacy::PassManager { + public: + explicit FilteredPassManager(bool disable_expensive_passes) + : disable_expensive_passes_(disable_expensive_passes) {} + void add(llvm::Pass* p) override { + if (disable_expensive_passes_) { + llvm::StringRef PassName = p->getPassName(); + if (PassName.contains("LICM") || PassName.contains("IndVarSimplify") || + PassName.contains("LoopUnroll")) { + return; + } + } + llvm::legacy::PassManager::add(p); + } + + private: + bool disable_expensive_passes_; +}; + llvm::object::OwningBinary CompilerFunctor:: operator()(llvm::Module& module) const { - llvm::legacy::PassManager module_passes; - llvm::legacy::FunctionPassManager function_passes(&module); + FilteredPassManager module_passes(disable_expensive_passes_); + FilteredFunctionPassManager function_passes(&module, + disable_expensive_passes_); VLOG(2) << "IR before optimizations"; XLA_VLOG_LINES(2, llvm_ir::DumpModuleToString(module)); - legacy_flags::CompilerFunctorFlags* flags = - legacy_flags::GetCompilerFunctorFlags(); - string dump_path = flags->xla_debug_cpu_dump_ir; - if (!dump_path.empty()) { - std::unique_ptr f; - TF_CHECK_OK(tensorflow::Env::Default()->NewAppendableFile(dump_path, &f)); - TF_CHECK_OK(f->Append(llvm_ir::DumpModuleToString(module))); - TF_CHECK_OK(f->Close()); + + if (pre_optimization_hook_) { + TF_CHECK_OK(pre_optimization_hook_(module)); } + // Add the appropriate TargetLibraryInfo and TargetTransformInfo. + AddTargetInfoPasses(&module_passes); + // Build up optimization pipeline. - AddOptimizationPasses(&module_passes, &function_passes); + if (optimize_for_size_) { + // Optimizing for size turns on -O2 level optimizations. + // + // TODO(b/64153864): Although the code generator supports size_level = 2 to + // turn on more aggressive code size optimizations than size_level = 1, we + // pass size_level = 1 because in many cases a size_level of 2 does + // worse. Investigate why. + AddOptimizationPasses(&module_passes, &function_passes, /*opt_level=*/2, + /*size_level=*/1); + } else { + AddOptimizationPasses(&module_passes, &function_passes, + /*opt_level=*/opt_level_, /*size_level=*/0); + } // Run optimization passes on module. function_passes.doInitialization(); + + CHECK(!llvm::verifyModule(module, &llvm::dbgs())); + for (auto func = module.begin(); func != module.end(); ++func) { function_passes.run(*func); } function_passes.doFinalization(); module_passes.run(module); + CHECK(!llvm::verifyModule(module, &llvm::dbgs())); + + runtime::RewriteIRRuntimeFunctions(&module, enable_fast_math_); + // Buffer for holding machine code prior to constructing the ObjectFile. llvm::SmallVector stream_buffer; llvm::raw_svector_ostream ostream(stream_buffer); @@ -94,6 +164,10 @@ operator()(llvm::Module& module) const { VLOG(2) << "IR after optimizations"; XLA_VLOG_LINES(2, llvm_ir::DumpModuleToString(module)); + if (post_optimization_hook_) { + TF_CHECK_OK(post_optimization_hook_(module)); + } + // Generate code. llvm::MCContext* mc_context; llvm::legacy::PassManager codegen_passes; @@ -111,10 +185,12 @@ operator()(llvm::Module& module) const { std::unique_ptr object_file = std::move(object_file_or_error.get()); if (VLOG_IS_ON(2)) { - StatusOr disassembly_status = + StatusOr disassembly_status = disassembler_->DisassembleObjectFile(*object_file); if (disassembly_status.ok()) { - XLA_VLOG_LINES(2, disassembly_status.ValueOrDie()); + auto result = disassembly_status.ValueOrDie(); + XLA_VLOG_LINES(2, result.text); + VLOG(2) << "compiled code size: " << result.code_size_bytes << " bytes"; } } @@ -129,57 +205,85 @@ std::vector VectorFunctionsForTargetLibraryInfoImpl( CompilerFunctor::VectorIntrinsics const& available_intrinsics) { std::vector vector_functions; - const llvm::VecDesc four_wide_vector_functions[] = { - {"expf", runtime::kExpV4F32, 4}, - {"llvm.exp.f32", runtime::kExpV4F32, 4}, + const llvm::VecDesc four_wide_vector_functions_neon[] = { + {"expf", runtime::kExpV4F32NEONSymbolName, 4}, + {"llvm.exp.f32", runtime::kExpV4F32NEONSymbolName, 4}, + + {"logf", runtime::kLogV4F32NEONSymbolName, 4}, + {"llvm.log.f32", runtime::kLogV4F32NEONSymbolName, 4}, + }; + + const llvm::VecDesc four_wide_vector_functions_sse[] = { + {"expf", runtime::kExpV4F32SSESymbolName, 4}, + {"llvm.exp.f32", runtime::kExpV4F32SSESymbolName, 4}, - {"logf", runtime::kLogV4F32, 4}, - {"llvm.log.f32", runtime::kLogV4F32, 4}, + {"logf", runtime::kLogV4F32SSESymbolName, 4}, + {"llvm.log.f32", runtime::kLogV4F32SSESymbolName, 4}, + }; + + const llvm::VecDesc eight_wide_vector_functions_avx[] = { + {"expf", runtime::kExpV8F32AVXSymbolName, 8}, + {"llvm.exp.f32", runtime::kExpV8F32AVXSymbolName, 8}, - {"tanhf", runtime::kTanhV4F32, 4}, - {"llvm.tanh.f32", runtime::kTanhV4F32, 4}, + {"logf", runtime::kLogV8F32AVXSymbolName, 8}, + {"llvm.log.f32", runtime::kLogV8F32AVXSymbolName, 8}, }; - const llvm::VecDesc eight_wide_vector_functions[] = { - {"expf", runtime::kExpV8F32, 8}, - {"llvm.exp.f32", runtime::kExpV8F32, 8}, + // These functions are generated by XLA as LLVM IR, so they're always + // available. + const llvm::VecDesc ir_vector_functions[] = { + {"tanhf", runtime::kTanhV4F32SymbolName, 4}, + {"llvm.tanh.f32", runtime::kTanhV4F32SymbolName, 4}, - {"logf", runtime::kLogV8F32, 8}, - {"llvm.log.f32", runtime::kLogV8F32, 8}, + {"tanhf", runtime::kTanhV8F32SymbolName, 8}, + {"llvm.tanh.f32", runtime::kTanhV8F32SymbolName, 8}, + }; - {"tanhf", runtime::kTanhV8F32, 8}, - {"llvm.tanh.f32", runtime::kTanhV8F32, 8}, + llvm::SmallVector features; + feature_string.split(features, ',', -1, /*KeepEmpty=*/false); + auto has_feature = [&features](const llvm::StringRef feature) { + return std::find(features.begin(), features.end(), feature) != + features.end(); }; - // Our vectorized library calls are currently implement by calling into Eigen. - // As such, only emit calls to these routines if --xla_cpu_use_eigen is - // enabled. - legacy_flags::CpuRuntimeFlags* flags = legacy_flags::GetCpuRuntimeFlags(); - if (flags->xla_cpu_use_eigen && - (arch == llvm::Triple::x86 || llvm::Triple::x86_64)) { - llvm::SmallVector features; - feature_string.split(features, ',', -1, /*KeepEmpty=*/false); - if (std::find(features.begin(), features.end(), "+sse4.1") != - features.end() && - available_intrinsics.sse_intrinsics) { - vector_functions.insert(vector_functions.end(), - std::begin(four_wide_vector_functions), - std::end(four_wide_vector_functions)); + switch (arch) { + case llvm::Triple::x86: + case llvm::Triple::x86_64: { + if (has_feature("+sse4.1") && available_intrinsics.sse_intrinsics) { + vector_functions.insert(vector_functions.end(), + std::begin(four_wide_vector_functions_sse), + std::end(four_wide_vector_functions_sse)); + } + if (has_feature("+avx") && available_intrinsics.avx_intrinsics) { + vector_functions.insert(vector_functions.end(), + std::begin(eight_wide_vector_functions_avx), + std::end(eight_wide_vector_functions_avx)); + } + break; } - if (std::find(features.begin(), features.end(), "+avx") != features.end() && - available_intrinsics.avx_intrinsics) { - vector_functions.insert(vector_functions.end(), - std::begin(eight_wide_vector_functions), - std::end(eight_wide_vector_functions)); + case llvm::Triple::arm: + case llvm::Triple::aarch64: { + if (has_feature("+neon") && available_intrinsics.neon_intrinsics) { + vector_functions.insert(vector_functions.end(), + std::begin(four_wide_vector_functions_neon), + std::end(four_wide_vector_functions_neon)); + } + break; } + default: + break; } + + vector_functions.insert(vector_functions.end(), + std::begin(ir_vector_functions), + std::end(ir_vector_functions)); + return vector_functions; } } // namespace -void CompilerFunctor::AddOptimizationPasses( - llvm::legacy::PassManagerBase* module_passes, - llvm::legacy::FunctionPassManager* function_passes) const { +void CompilerFunctor::AddTargetInfoPasses( + llvm::legacy::PassManagerBase* passes) const { llvm::Triple target_triple(target_machine_->getTargetTriple()); auto target_library_info_impl = MakeUnique(target_triple); @@ -187,18 +291,21 @@ void CompilerFunctor::AddOptimizationPasses( VectorFunctionsForTargetLibraryInfoImpl( target_triple.getArch(), target_machine_->getTargetFeatureString(), available_intrinsics_)); - module_passes->add( + passes->add( new llvm::TargetLibraryInfoWrapperPass(*target_library_info_impl)); - module_passes->add(createTargetTransformInfoWrapperPass( + passes->add(createTargetTransformInfoWrapperPass( target_machine_->getTargetIRAnalysis())); +} - module_passes->add(llvm::createVerifierPass()); - +void CompilerFunctor::AddOptimizationPasses( + llvm::legacy::PassManagerBase* module_passes, + llvm::legacy::FunctionPassManager* function_passes, unsigned opt_level, + unsigned size_level) const { llvm::PassManagerBuilder builder; - builder.OptLevel = opt_level_; - builder.SizeLevel = 0; + builder.OptLevel = opt_level; + builder.SizeLevel = size_level; - if (opt_level_ > 1) { + if (opt_level > 1) { builder.Inliner = llvm::createFunctionInliningPass(); } else { // Only inline functions marked with "alwaysinline". @@ -206,14 +313,12 @@ void CompilerFunctor::AddOptimizationPasses( } builder.DisableUnitAtATime = false; - builder.DisableUnrollLoops = opt_level_ == 0; - builder.LoopVectorize = opt_level_ > 0; - builder.SLPVectorize = opt_level_ > 1; + builder.DisableUnrollLoops = opt_level == 0; + builder.LoopVectorize = opt_level > 0 && size_level == 0; + builder.SLPVectorize = opt_level > 1 && size_level == 0; builder.populateFunctionPassManager(*function_passes); builder.populateModulePassManager(*module_passes); - - module_passes->add(llvm::createVerifierPass()); } } // namespace cpu diff --git a/tensorflow/compiler/xla/service/cpu/compiler_functor.h b/tensorflow/compiler/xla/service/cpu/compiler_functor.h index 17dadebe975b936b7d5d7a78ac69b890d9c8e7ac..8cdd049e7b773bdc455db627ff1749997d621ee4 100644 --- a/tensorflow/compiler/xla/service/cpu/compiler_functor.h +++ b/tensorflow/compiler/xla/service/cpu/compiler_functor.h @@ -16,11 +16,12 @@ limitations under the License. #ifndef TENSORFLOW_COMPILER_XLA_SERVICE_CPU_COMPILER_FUNCTOR_H_ #define TENSORFLOW_COMPILER_XLA_SERVICE_CPU_COMPILER_FUNCTOR_H_ -#include "external/llvm/include/llvm/IR/LegacyPassManager.h" -#include "external/llvm/include/llvm/IR/Module.h" -#include "external/llvm/include/llvm/Object/ObjectFile.h" -#include "external/llvm/include/llvm/Target/TargetMachine.h" +#include "llvm/IR/LegacyPassManager.h" +#include "llvm/IR/Module.h" +#include "llvm/Object/ObjectFile.h" +#include "llvm/Target/TargetMachine.h" #include "tensorflow/compiler/xla/service/cpu/disassembler.h" +#include "tensorflow/compiler/xla/service/llvm_compiler.h" #include "tensorflow/core/platform/logging.h" namespace xla { @@ -34,33 +35,52 @@ class CompilerFunctor { struct VectorIntrinsics { bool sse_intrinsics; bool avx_intrinsics; + bool neon_intrinsics; }; // Returns a VectorIntrinsics where all intrinsics are available. static VectorIntrinsics AllIntrinsics(); - explicit CompilerFunctor(llvm::TargetMachine* target_machine, - const Disassembler* disassembler, int opt_level, - const VectorIntrinsics& available_intrinsics) + explicit CompilerFunctor( + llvm::TargetMachine* target_machine, const Disassembler* disassembler, + int opt_level, bool optimize_for_size, bool enable_fast_math, + bool disable_expensive_passes, + const VectorIntrinsics& available_intrinsics, + LLVMCompiler::ModuleHook pre_optimization_hook = nullptr, + LLVMCompiler::ModuleHook post_optimization_hook = nullptr) : target_machine_(target_machine), disassembler_(CHECK_NOTNULL(disassembler)), opt_level_(opt_level), - available_intrinsics_(available_intrinsics) {} + optimize_for_size_(optimize_for_size), + enable_fast_math_(enable_fast_math), + disable_expensive_passes_(disable_expensive_passes), + available_intrinsics_(available_intrinsics), + pre_optimization_hook_(pre_optimization_hook), + post_optimization_hook_(post_optimization_hook) {} // Compile a Module to an ObjectFile. llvm::object::OwningBinary operator()( llvm::Module& module) const; // NOLINT private: + // Populates the given pass manager with TargetLibraryInfo and + // TargetTransformInfo passes. + void AddTargetInfoPasses(llvm::legacy::PassManagerBase* passes) const; + // Populates the given pass managers based on the optimization level. - void AddOptimizationPasses( - llvm::legacy::PassManagerBase* module_passes, - llvm::legacy::FunctionPassManager* function_passes) const; + void AddOptimizationPasses(llvm::legacy::PassManagerBase* module_passes, + llvm::legacy::FunctionPassManager* function_passes, + unsigned opt_level, unsigned size_level) const; llvm::TargetMachine* target_machine_; const Disassembler* disassembler_; const unsigned opt_level_; + const bool optimize_for_size_; + const bool enable_fast_math_; + const bool disable_expensive_passes_; const VectorIntrinsics available_intrinsics_; + LLVMCompiler::ModuleHook pre_optimization_hook_; + LLVMCompiler::ModuleHook post_optimization_hook_; }; } // namespace cpu diff --git a/tensorflow/compiler/xla/service/cpu/conv_canonicalization.cc b/tensorflow/compiler/xla/service/cpu/conv_canonicalization.cc index cdf43587b683e4e22d14d4fc08fa3705bc636de8..069979c6611e90ed2d95cbbe341198577cdf56cf 100644 --- a/tensorflow/compiler/xla/service/cpu/conv_canonicalization.cc +++ b/tensorflow/compiler/xla/service/cpu/conv_canonicalization.cc @@ -15,7 +15,6 @@ limitations under the License. #include "tensorflow/compiler/xla/service/cpu/conv_canonicalization.h" -#include "tensorflow/compiler/xla/legacy_flags/cpu_runtime_flags.h" #include "tensorflow/compiler/xla/service/cpu/cpu_runtime.h" #include "tensorflow/compiler/xla/service/cpu/ir_emission_utils.h" #include "tensorflow/compiler/xla/service/hlo_computation.h" @@ -30,11 +29,6 @@ namespace xla { namespace cpu { StatusOr ConvCanonicalization::Run(HloModule* module) { - legacy_flags::CpuRuntimeFlags* flags = legacy_flags::GetCpuRuntimeFlags(); - if (!flags->xla_cpu_use_eigen) { - return false; - } - bool changed = false; for (HloInstruction* hlo : module->entry_computation()->MakeInstructionPostOrder()) { diff --git a/tensorflow/compiler/xla/service/cpu/conv_canonicalization_test.cc b/tensorflow/compiler/xla/service/cpu/conv_canonicalization_test.cc index d18141af83e4653e18d3b0118d0892f41db5b69b..ec992f15e63b29ee67d16b6d841fedffd9c90f5b 100644 --- a/tensorflow/compiler/xla/service/cpu/conv_canonicalization_test.cc +++ b/tensorflow/compiler/xla/service/cpu/conv_canonicalization_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.h" +#include "tensorflow/compiler/xla/test.h" #include "tensorflow/compiler/xla/tests/hlo_test_base.h" #include "tensorflow/compiler/xla/util.h" @@ -28,6 +29,8 @@ limitations under the License. namespace xla { namespace cpu { +using ::testing::ElementsAre; + class ConvCanonicalizationTest : public HloTestBase { public: ConvCanonicalizationTest() { @@ -56,11 +59,11 @@ TEST_F(ConvCanonicalizationTest, NonCanonicalToCanonical) { auto builder = HloComputation::Builder(TestName()); // The input dimensions are in CNHW order. auto input = builder.AddInstruction(HloInstruction::CreateConstant( - LiteralUtil::CreateR4FromArray4D(Array4D( + Literal::CreateR4FromArray4D(Array4D( kInputFeatureCount, kBatchSize, kInputSize, kInputSize)))); // The kernel dimensions are in OIHW order. auto kernel = builder.AddInstruction(HloInstruction::CreateConstant( - LiteralUtil::CreateR4FromArray4D(Array4D( + Literal::CreateR4FromArray4D(Array4D( kOutputFeatureCount, kInputFeatureCount, kWindowSize, kWindowSize)))); ConvolutionDimensionNumbers dnums; @@ -78,7 +81,7 @@ TEST_F(ConvCanonicalizationTest, NonCanonicalToCanonical) { F32, {kOutputFeatureCount, kBatchSize, output_size, output_size}), input, kernel, conv_window_, dnums)); - auto module = MakeUnique(TestName()); + auto module = CreateNewModule(); HloComputation* entry_computation = module->AddEntryComputation(builder.Build()); @@ -96,28 +99,25 @@ TEST_F(ConvCanonicalizationTest, NonCanonicalToCanonical) { // The input is in CNHW order. input_reshape should produce // NHWC for the convolution to hit the Eigen fast path. - EXPECT_TRUE(ContainersEqual(input_reshape->dimensions(), - std::vector({1, 2, 3, 0}))); + EXPECT_THAT(input_reshape->dimensions(), ElementsAre(1, 2, 3, 0)); // The kernel is in OIHW order. kernel_reshape should produce // HWIO for the convolution to hit the Eigen fast path. - EXPECT_TRUE(ContainersEqual(kernel_reshape->dimensions(), - std::vector({2, 3, 1, 0}))); + EXPECT_THAT(kernel_reshape->dimensions(), ElementsAre(2, 3, 1, 0)); // The output of the canonical convolution is in NHWC order (the same as // input_reshape's order). output_reshape should restore that order to the // order of the computation root (CNHW). - EXPECT_TRUE(ContainersEqual(output_reshape->dimensions(), - std::vector({3, 0, 1, 2}))); + EXPECT_THAT(output_reshape->dimensions(), ElementsAre(3, 0, 1, 2)); } TEST_F(ConvCanonicalizationTest, CanonicalStaysTheSame) { auto builder = HloComputation::Builder(TestName()); // The input dimensions are in NHWC order. auto input = builder.AddInstruction(HloInstruction::CreateConstant( - LiteralUtil::CreateR4FromArray4D(Array4D( + Literal::CreateR4FromArray4D(Array4D( kBatchSize, kInputSize, kInputSize, kInputFeatureCount)))); // The kernel dimensions are in HWIO order. auto kernel = builder.AddInstruction(HloInstruction::CreateConstant( - LiteralUtil::CreateR4FromArray4D(Array4D( + Literal::CreateR4FromArray4D(Array4D( kWindowSize, kWindowSize, kInputFeatureCount, kOutputFeatureCount)))); ConvolutionDimensionNumbers dnums; @@ -135,7 +135,7 @@ TEST_F(ConvCanonicalizationTest, CanonicalStaysTheSame) { F32, {kBatchSize, output_size, output_size, kOutputFeatureCount}), input, kernel, conv_window_, dnums)); - auto module = MakeUnique(TestName()); + auto module = CreateNewModule(); module->AddEntryComputation(builder.Build()); ConvCanonicalization conv_canonicalization; @@ -144,3 +144,7 @@ TEST_F(ConvCanonicalizationTest, CanonicalStaysTheSame) { } // namespace cpu } // namespace xla + +int main(int argc, char** argv) { + return xla::ParseDebugOptionsFlagsAndRunTests(argc, argv); +} diff --git a/tensorflow/compiler/xla/service/cpu/cpu_compiler.cc b/tensorflow/compiler/xla/service/cpu/cpu_compiler.cc index c5433d4b89d7ccab0f04e9ab2787ce150417b669..839fe4848882e506a13122f90eae306cd509d964 100644 --- a/tensorflow/compiler/xla/service/cpu/cpu_compiler.cc +++ b/tensorflow/compiler/xla/service/cpu/cpu_compiler.cc @@ -18,6 +18,7 @@ limitations under the License. #include #include #include +#include // NOLINT(build/c++11): only using std::call_once, not mutex. #include #include #include @@ -25,24 +26,23 @@ limitations under the License. // IWYU pragma: no_include "llvm/Config/Disassemblers.def.inc" // IWYU pragma: no_include "llvm/Config/Targets.def.inc" -#include "external/llvm/include/llvm/ADT/StringRef.h" -#include "external/llvm/include/llvm/ADT/Triple.h" -#include "external/llvm/include/llvm/IR/Function.h" -#include "external/llvm/include/llvm/IR/LLVMContext.h" -#include "external/llvm/include/llvm/IR/Module.h" -#include "external/llvm/include/llvm/Object/ObjectFile.h" -#include "external/llvm/include/llvm/Support/CommandLine.h" -#include "external/llvm/include/llvm/Support/TargetRegistry.h" -#include "external/llvm/include/llvm/Support/TargetSelect.h" -#include "external/llvm/include/llvm/Target/TargetMachine.h" -#include "external/llvm/include/llvm/Target/TargetOptions.h" -#include "tensorflow/compiler/xla/legacy_flags/cpu_compiler_flags.h" +#include "llvm/ADT/StringRef.h" +#include "llvm/ADT/Triple.h" +#include "llvm/IR/Function.h" +#include "llvm/IR/LLVMContext.h" +#include "llvm/IR/Module.h" +#include "llvm/Object/ObjectFile.h" +#include "llvm/Support/CommandLine.h" +#include "llvm/Support/TargetRegistry.h" +#include "llvm/Support/TargetSelect.h" +#include "llvm/Target/TargetMachine.h" +#include "llvm/Target/TargetOptions.h" #include "tensorflow/compiler/xla/literal_util.h" #include "tensorflow/compiler/xla/map_util.h" -#include "tensorflow/compiler/xla/port/initialize.h" #include "tensorflow/compiler/xla/protobuf_util.h" #include "tensorflow/compiler/xla/ptr_util.h" #include "tensorflow/compiler/xla/service/algebraic_simplifier.h" +#include "tensorflow/compiler/xla/service/batchnorm_rewriter.h" #include "tensorflow/compiler/xla/service/buffer_assignment.h" #include "tensorflow/compiler/xla/service/buffer_liveness.h" #include "tensorflow/compiler/xla/service/copy_insertion.h" @@ -50,6 +50,7 @@ limitations under the License. #include "tensorflow/compiler/xla/service/cpu/conv_canonicalization.h" #include "tensorflow/compiler/xla/service/cpu/cpu_executable.h" #include "tensorflow/compiler/xla/service/cpu/cpu_instruction_fusion.h" +#include "tensorflow/compiler/xla/service/cpu/cpu_options.h" #include "tensorflow/compiler/xla/service/cpu/cpu_parallelization_preparation.h" #include "tensorflow/compiler/xla/service/cpu/disassembler.h" #include "tensorflow/compiler/xla/service/cpu/ir_emission_utils.h" @@ -58,6 +59,8 @@ limitations under the License. #include "tensorflow/compiler/xla/service/cpu/parallel_cpu_executable.h" #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/flatten_call_graph.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" @@ -67,10 +70,13 @@ limitations under the License. #include "tensorflow/compiler/xla/service/hlo_ordering.h" #include "tensorflow/compiler/xla/service/hlo_pass_fix.h" #include "tensorflow/compiler/xla/service/hlo_pass_pipeline.h" +#include "tensorflow/compiler/xla/service/hlo_proto_util.h" +#include "tensorflow/compiler/xla/service/hlo_scheduling.h" #include "tensorflow/compiler/xla/service/hlo_subcomputation_unification.h" #include "tensorflow/compiler/xla/service/hlo_verifier.h" #include "tensorflow/compiler/xla/service/inliner.h" #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/transpose_folding.h" #include "tensorflow/compiler/xla/status_macros.h" @@ -78,7 +84,10 @@ 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/io/path.h" #include "tensorflow/core/lib/strings/str_util.h" +#include "tensorflow/core/lib/strings/strcat.h" +#include "tensorflow/core/platform/env.h" namespace se = ::perftools::gputools; @@ -142,22 +151,44 @@ CpuCompiler::CpuCompiler() { LLVMInitializePowerPCTargetMC(); LLVMInitializePowerPCAsmPrinter(); LLVMInitializePowerPCDisassembler(); +} + +namespace { - // LLVM command-line flags are global, so set them during initialization. - legacy_flags::CpuCompilerFlags* flags = legacy_flags::GetCpuCompilerFlags(); - if (!flags->xla_cpu_llvm_cl_opts.empty()) { - std::vector opts = - tensorflow::str_util::Split(flags->xla_cpu_llvm_cl_opts, ','); +// LLVM makes certain options configurable only through its command-line +// options; it provide the ParseCommandLineOptions function that lets us set +// flags at runtime. However, since these flags are global we want to avoid +// multiple invocations of the LLVM compilation pipeline with a different set of +// flags. Therefore, we only pass command-line flags to LLVM once, before the +// first module is compiled. +std::once_flag llvm_command_line_options_initialized; + +void InitializeLLVMCommandLineOptions(const HloModuleConfig& config) { + auto options = config.debug_options().xla_backend_extra_options(); + if (!options.empty()) { + std::vector fake_argv_storage; + fake_argv_storage.push_back(""); + for (const auto& it : options) { + // Skip options the XLA backend itself consumes. + if (!tensorflow::StringPiece(it.first).starts_with("xla_")) { + if (it.second.empty()) { + fake_argv_storage.push_back(it.first); + } else { + fake_argv_storage.push_back(it.first + "=" + it.second); + } + } + } + + VLOG(2) << "Passing argv to LLVM:"; std::vector fake_argv; - fake_argv.push_back(""); - for (const string& opt : opts) { - fake_argv.push_back(opt.c_str()); + for (const auto& s : fake_argv_storage) { + fake_argv.push_back(s.c_str()); + VLOG(2) << s; } llvm::cl::ParseCommandLineOptions(fake_argv.size(), &fake_argv[0]); } } -namespace { // This visitor records which HLO instructions should have profiling information // recorded. class CollectProfileCandidates : public DfsHloVisitorWithDefault { @@ -181,6 +212,15 @@ class CollectProfileCandidates : public DfsHloVisitorWithDefault { hlo_to_profile_idx_->insert({hlo_instruction, hlo_to_profile_idx_->size()}); return Status::OK(); } + + Status HandleCall(HloInstruction* call) override { + TF_RETURN_IF_ERROR(DefaultAction(call)); + CollectProfileCandidates candidates_for_call(hlo_to_profile_idx_); + TF_RETURN_IF_ERROR( + call->to_apply()->root_instruction()->Accept(&candidates_for_call)); + return Status::OK(); + } + // Skip constants, there is nothing to profile. Status HandleConstant(HloInstruction* /*constant*/, const Literal& /*literal*/) override { @@ -192,16 +232,16 @@ class CollectProfileCandidates : public DfsHloVisitorWithDefault { } // It is important to recurse for "while" or else we risk overly coarse // profiling information. - Status HandleWhile(HloInstruction* xla_while, HloInstruction* /*init*/, - HloComputation* condition, HloComputation* body) override { + Status HandleWhile(HloInstruction* xla_while) override { TF_RETURN_IF_ERROR(DefaultAction(xla_while)); CollectProfileCandidates candidates_for_condition(hlo_to_profile_idx_); - TF_RETURN_IF_ERROR( - condition->root_instruction()->Accept(&candidates_for_condition)); + TF_RETURN_IF_ERROR(xla_while->while_condition()->root_instruction()->Accept( + &candidates_for_condition)); CollectProfileCandidates candidates_for_body(hlo_to_profile_idx_); - TF_RETURN_IF_ERROR(body->root_instruction()->Accept(&candidates_for_body)); + TF_RETURN_IF_ERROR(xla_while->while_body()->root_instruction()->Accept( + &candidates_for_body)); return Status::OK(); } @@ -210,12 +250,14 @@ class CollectProfileCandidates : public DfsHloVisitorWithDefault { }; } // namespace -Status CpuCompiler::RunHloPasses(HloModule* hlo_module, - HloModuleConfig* module_config, - HloDumper dump_hlo) { +Status CpuCompiler::RunHloPasses(HloModule* module) { // Optimization pipeline. - HloPassPipeline pipeline("CPU", dump_hlo); - pipeline.AddInvariantChecker(); + HloPassPipeline pipeline("CPU"); + pipeline.AddInvariantChecker(ShapeSizeBytesFunction()); + + ReducePrecisionInsertion::AddPasses( + &pipeline, module->config().debug_options(), + ReducePrecisionInsertion::PassTiming::BEFORE_OPTIMIZATION); // TODO(b/35786417): Re-enable inliner pass after fixing the bug and deciding // where we will take this pass in future. @@ -223,8 +265,13 @@ Status CpuCompiler::RunHloPasses(HloModule* hlo_module, pipeline.AddPass(); { - auto& pass = pipeline.AddPass>("simplification", - dump_hlo); + auto& pass = + pipeline.AddPass>("simplification"); + pass.AddPass( + /*rewrite_training_op=*/true, + /*rewrite_inference_op=*/true, + /*rewrite_grad_op=*/true, + /*use_fusion=*/false); pass.AddPass( /*is_layout_sensitive=*/false, [](const Shape&, const Shape&) { return false; }, @@ -232,12 +279,23 @@ Status CpuCompiler::RunHloPasses(HloModule* hlo_module, pass.AddPass(); pass.AddPass(); } - pipeline.AddPass(PotentiallyImplementedAsEigenDot); - pipeline.AddPass(); + pipeline.AddPass( + [](const HloInstruction& dot, + const TransposeFolding::OperandIndices& candidate_operands) { + return PotentiallyImplementedAsEigenDot(dot) + ? candidate_operands + : TransposeFolding::OperandIndices{}; + }, + TransposeFolding::NeverFoldTranspose); pipeline.AddPass(/*is_layout_sensitive=*/false); pipeline.AddPass(); + + ReducePrecisionInsertion::AddPasses( + &pipeline, module->config().debug_options(), + ReducePrecisionInsertion::PassTiming::AFTER_FUSION); + pipeline.AddPass( - module_config->mutable_entry_computation_layout()); + module->mutable_entry_computation_layout()); // The LayoutAssignment pass may leave behind kCopy instructions which are // duplicate or NOPs, so remove them with algebraic simplification and CSE. pipeline.AddPass>( @@ -246,83 +304,183 @@ Status CpuCompiler::RunHloPasses(HloModule* hlo_module, /*enable_dot_simplification=*/false); pipeline.AddPass(/*is_layout_sensitive=*/true); // Outline ops in the entry computation into calls to subcomputations. - legacy_flags::CpuCompilerFlags* flags = legacy_flags::GetCpuCompilerFlags(); - if (flags->xla_cpu_parallel) { - pipeline.AddPass(); + const int max_parallelism = + module->config().intra_op_parallelism_threads() > 0 + ? module->config().intra_op_parallelism_threads() + : tensorflow::port::NumSchedulableCPUs(); + if (options::CpuParallelBackendRequested(module->config())) { + pipeline.AddPass(max_parallelism, + ShapeSizeBytesFunction()); } - // Copy insertion should be performed immediately before IR emission to - // avoid inserting unnecessary copies (later pass adds an instruction which - // materializes the value) or missing a necessary copy (later pass removes - // an instruction which materializes a value). + // Copy insertion should be performed immediately before IR emission to avoid + // inserting unnecessary copies (later pass adds an instruction which + // materializes the value) or missing a necessary copy (later pass removes an + // instruction which materializes a value). DCE must be run immediately before + // (and sometime after) copy insertion, to avoid dead code from interfering + // with the rewrites. + pipeline.AddPass(); pipeline.AddPass(); - if (flags->xla_cpu_parallel) { + if (options::CpuParallelBackendRequested(module->config())) { // Re-run the outlining, in case any copies were inserted into the entry // computation. - pipeline.AddPass(); + pipeline.AddPass(max_parallelism, + ShapeSizeBytesFunction()); } pipeline.AddPass(); - return pipeline.Run(hlo_module).status(); + pipeline.AddPass(); + return pipeline.Run(module).status(); } namespace { // Align buffers to 16-byte boundaries. constexpr int64 kMemoryAlignment = 16; +auto memory_alignment = [](LogicalBuffer::Color) { return kMemoryAlignment; }; llvm::TargetOptions CompilerTargetOptions( - const HloModuleConfig& execution_options) { + const HloModuleConfig& module_config) { llvm::TargetOptions target_options; - llvm_ir::SetTargetOptions(execution_options, &target_options); + llvm_ir::SetTargetOptions( + /*fast_math_enabled=*/module_config.debug_options() + .xla_enable_fast_math(), + &target_options); return target_options; } -llvm::CodeGenOpt::Level CodeGenOptLevel() { - legacy_flags::CpuCompilerFlags* flags = legacy_flags::GetCpuCompilerFlags(); - switch (flags->xla_cpu_llvm_opt_level) { +llvm::CodeGenOpt::Level CodeGenOptLevel(const HloModuleConfig& module_config) { + VLOG(2) << "backend_optimization_level: " + << module_config.debug_options().xla_backend_optimization_level(); + switch (module_config.debug_options().xla_backend_optimization_level()) { case 1: return llvm::CodeGenOpt::Less; case 2: return llvm::CodeGenOpt::Default; - break; case 3: return llvm::CodeGenOpt::Aggressive; - break; default: return llvm::CodeGenOpt::None; } } +Status AppendIRToFile(const string& file_name, const string& ir_module_string) { + std::unique_ptr f; + TF_RETURN_IF_ERROR( + tensorflow::Env::Default()->NewAppendableFile(file_name, &f)); + TF_RETURN_IF_ERROR(f->Append(ir_module_string)); + TF_RETURN_IF_ERROR(f->Close()); + return Status::OK(); +} + +Status InitializeModuleHooks( + const HloModule& module, + const LLVMCompiler::ModuleHook& user_pre_optimization_hook, + const LLVMCompiler::ModuleHook& user_post_optimization_hook, + LLVMCompiler::ModuleHook* pre_optimization_ir_hook, + LLVMCompiler::ModuleHook* post_optimization_ir_hook) { + const string& dump_ir_to = module.config().debug_options().xla_dump_ir_to(); + if (dump_ir_to.empty()) { + *pre_optimization_ir_hook = user_pre_optimization_hook; + *post_optimization_ir_hook = user_post_optimization_hook; + return Status::OK(); + } + + // Initialize the output directory and create the output file names. + TF_RETURN_IF_ERROR( + tensorflow::Env::Default()->RecursivelyCreateDir(dump_ir_to)); + string safe_file_name_base = module.name(); + std::replace_if(safe_file_name_base.begin(), safe_file_name_base.end(), + [](char c) { return c == '/' || c == '\\'; }, '_'); + + string unoptimized_ir_file_name = tensorflow::io::JoinPath( + dump_ir_to, + tensorflow::strings::StrCat("ir-", safe_file_name_base, "-no-opt.ll")); + string optimized_ir_file_name = tensorflow::io::JoinPath( + dump_ir_to, + tensorflow::strings::StrCat("ir-", safe_file_name_base, "-opt.ll")); + + // Create the IR hooks. If applicable, each IR hook does the following: + // * Call the user supplied module hook. + // * Write to the output directory. Files will be appended to. We still want + // to append to avoid overwriting possibly important information due to + // operator error. + + *pre_optimization_ir_hook = + [user_pre_optimization_hook, + unoptimized_ir_file_name](const llvm::Module& module) { + if (user_pre_optimization_hook) { + TF_RETURN_IF_ERROR(user_pre_optimization_hook(module)); + } + TF_RETURN_IF_ERROR(AppendIRToFile(unoptimized_ir_file_name, + llvm_ir::DumpModuleToString(module))); + return Status::OK(); + }; + + *post_optimization_ir_hook = + [user_post_optimization_hook, + optimized_ir_file_name](const llvm::Module& module) { + if (user_post_optimization_hook) { + TF_RETURN_IF_ERROR(user_post_optimization_hook(module)); + } + TF_RETURN_IF_ERROR(AppendIRToFile(optimized_ir_file_name, + llvm_ir::DumpModuleToString(module))); + return Status::OK(); + }; + + return Status::OK(); +} + } // namespace StatusOr> CpuCompiler::Compile( - std::unique_ptr hlo_module, - std::unique_ptr module_config, HloDumper dump_hlo, - se::StreamExecutor* stream_exec) { + std::unique_ptr module, se::StreamExecutor* stream_exec) { + VLOG(1) << "Compiling: " << module->name(); TF_RET_CHECK(stream_exec != nullptr); + std::call_once(llvm_command_line_options_initialized, + &InitializeLLVMCommandLineOptions, module->config()); + + ModuleHook pre_optimization_ir_hook; + ModuleHook post_optimization_ir_hook; + TF_RETURN_IF_ERROR(InitializeModuleHooks( + *module, user_pre_optimization_hook_, user_post_optimization_hook_, + &pre_optimization_ir_hook, &post_optimization_ir_hook)); // Compile must be thread-safe so create a new LLVM context for the module. auto llvm_context = MakeUnique(); auto llvm_module = MakeUnique("__compute_module", *llvm_context); - auto jit = MakeUnique(CompilerTargetOptions(*module_config), - CodeGenOptLevel()); + + auto jit = MakeUnique( + CompilerTargetOptions(module->config()), + CodeGenOptLevel(module->config()), + options::OptimizeForSizeRequested(module->config()), + module->config().debug_options().xla_enable_fast_math(), + module->config().debug_options().xla_llvm_disable_expensive_passes(), + pre_optimization_ir_hook, post_optimization_ir_hook); llvm_module->setDataLayout(jit->data_layout()); llvm_module->setTargetTriple(jit->target_triple().getTriple()); - TF_RETURN_IF_ERROR( - RunHloPasses(hlo_module.get(), module_config.get(), dump_hlo)); + TF_RETURN_IF_ERROR(RunHloPasses(module.get())); - HloComputation* computation = hlo_module->entry_computation(); + HloComputation* computation = module->entry_computation(); std::unordered_map hlo_to_profile_idx; - if (module_config->hlo_profiling_enabled()) { + if (module->config().hlo_profiling_enabled()) { TF_ASSIGN_OR_RETURN( hlo_to_profile_idx, CollectProfileCandidates::GetCandidatesForComputation(computation)); } std::unique_ptr cpu_executable; - legacy_flags::CpuCompilerFlags* flags = legacy_flags::GetCpuCompilerFlags(); - if (flags->xla_cpu_parallel) { + + // Cache these flags here since we'll want to access them after the module's + // ownership is std::moved. + const bool embed_ir_in_executable = + module->config().debug_options().xla_embed_ir_in_executable(); + const string dump_debug_json_to = + module->config().debug_options().xla_dump_debug_json_to(); + + if (options::CpuParallelBackendRequested(module->config())) { + VLOG(1) << "Using parallel cpu backend"; + // Run buffer analysis on the HLO graph. This analysis figures out which // temporary buffers are required to run the computation. // DependencyHloOrdering is used for the parallel emitter because the order @@ -331,12 +489,15 @@ StatusOr> CpuCompiler::Compile( // uses data dependencies for determining order. TF_ASSIGN_OR_RETURN( std::unique_ptr assignment, - BufferAssigner::Run(hlo_module.get(), - MakeUnique(hlo_module.get()), - [this](const LogicalBuffer& buffer) { - return ShapeSizeBytes(buffer.shape()); - }, - kMemoryAlignment)); + BufferAssigner::Run(module.get(), + MakeUnique(module.get()), + BufferSizeBytesFunction(), memory_alignment)); + + if (!dump_debug_json_to.empty()) { + HloProto proto = MakeHloProto(*module, *assignment); + TF_RETURN_IF_ERROR(protobuf_util::DumpJsonToDirectory( + proto, dump_debug_json_to, module->name())); + } // If we are using the parallel CPU backend, we need to create map from // HloInstruction to the corresponding generated function name. @@ -351,8 +512,8 @@ StatusOr> CpuCompiler::Compile( if (instruction->opcode() == HloOpcode::kConstant) { // Copy the constant out of the ProtocolBuffer so that we can give it a // higher alignment. - const void* data = LiteralUtil::InternalData(instruction->literal()); - int64 size = ShapeSizeBytes(instruction->shape()); + const void* data = instruction->literal().InternalData(); + int64 size = CpuExecutable::ShapeSizeBytes(instruction->shape()); auto iter = aligned_constants.emplace( instruction, MakeUnique(size)); CHECK_EQ(iter.second, true); @@ -363,17 +524,21 @@ StatusOr> CpuCompiler::Compile( // The parallel preparation should have ensured that the top-level // computation consists solely of Call instructions. TF_RET_CHECK(instruction->opcode() == HloOpcode::kCall) - << hlo_module->ToString(); + << module->ToString(); HloComputation* to_apply = instruction->to_apply(); parallel_computations.emplace(to_apply, instruction); } - IrEmitter ir_emitter(*hlo_module, *module_config, *assignment, - llvm_module.get(), &hlo_to_profile_idx); + IrEmitter ir_emitter(*module, *assignment, llvm_module.get(), + &hlo_to_profile_idx, jit->target_machine()); + std::unique_ptr> function_names( new std::map()); for (auto embedded_computation : computation->MakeEmbeddedComputationsList()) { + if (embedded_computation->IsFusionComputation()) { + continue; + } auto parallel_computation_iter = parallel_computations.find(embedded_computation); // All parallel computations are considered to be an entry computation for @@ -384,7 +549,8 @@ StatusOr> CpuCompiler::Compile( llvm::Function * ir_function, ir_emitter.EmitComputation( embedded_computation, embedded_computation->name(), - /*is_entry_computation=*/computation_is_parallel)); + /*is_entry_computation=*/computation_is_parallel, + /*instruction_order=*/nullptr)); // If this computation is parallel, remember it in the function name map. // This way we know what function to execute when we try to run code for // the Call instruction. @@ -396,52 +562,58 @@ StatusOr> CpuCompiler::Compile( } string ir_module_string; - if (flags->xla_cpu_embed_ir) { + if (embed_ir_in_executable) { ir_module_string = llvm_ir::DumpModuleToString(*llvm_module); } // JIT compile the LLVM IR module to in-memory machine code. jit->AddModule(std::move(llvm_module)); cpu_executable.reset(new ParallelCpuExecutable( - std::move(jit), std::move(assignment), std::move(hlo_module), - std::move(module_config), std::move(function_names), - std::move(hlo_to_profile_idx), std::move(aligned_constants))); + std::move(jit), std::move(assignment), std::move(module), + std::move(function_names), std::move(hlo_to_profile_idx), + std::move(aligned_constants))); - if (flags->xla_cpu_embed_ir) { + if (embed_ir_in_executable) { static_cast(*cpu_executable) .set_ir_module_string(ir_module_string); } } else { + VLOG(1) << "Using sequential cpu backend"; + // 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( SequentialHloOrdering::HloModuleSequence module_sequence, - CreateMemoryMinimizingSequence(*hlo_module, - [this](const LogicalBuffer& buffer) { - return ShapeSizeBytes(buffer.shape()); - })); + CreateMemoryMinimizingSequence(*module, BufferSizeBytesFunction())); // Run buffer analysis on the HLO graph. This analysis figures out which // temporary buffers are required to run the computation. TF_ASSIGN_OR_RETURN( std::unique_ptr assignment, - BufferAssigner::Run(hlo_module.get(), - MakeUnique(hlo_module.get(), - module_sequence), - [this](const LogicalBuffer& buffer) { - return ShapeSizeBytes(buffer.shape()); - }, - kMemoryAlignment)); + BufferAssigner::Run( + module.get(), + MakeUnique(module.get(), module_sequence), + BufferSizeBytesFunction(), memory_alignment)); + + if (!dump_debug_json_to.empty()) { + HloProto proto = MakeHloProto(*module, *assignment); + TF_RETURN_IF_ERROR(protobuf_util::DumpJsonToDirectory( + proto, dump_debug_json_to, module->name())); + } // Each computation is a single function. Emit all embedded computations // before the entry computation. The order of computations returned from // GetEmbeddedComputations guarantees that a called computation occurs // before a caller computation. - IrEmitter ir_emitter(*hlo_module, *module_config, *assignment, - llvm_module.get(), &hlo_to_profile_idx); + IrEmitter ir_emitter(*module, *assignment, llvm_module.get(), + &hlo_to_profile_idx, jit->target_machine()); + for (auto embedded_computation : computation->MakeEmbeddedComputationsList()) { + if (embedded_computation->IsFusionComputation()) { + continue; + } TF_RETURN_IF_ERROR( ir_emitter .EmitComputation(embedded_computation, @@ -460,51 +632,51 @@ StatusOr> CpuCompiler::Compile( string function_name = llvm_ir::AsString(entry_function->getName()); string ir_module_string; - if (flags->xla_cpu_embed_ir) { + if (embed_ir_in_executable) { ir_module_string = llvm_ir::DumpModuleToString(*llvm_module); } // JIT compile the LLVM IR module to in-memory machine code. jit->AddModule(std::move(llvm_module)); - cpu_executable.reset( - new CpuExecutable(std::move(jit), std::move(assignment), - std::move(hlo_module), std::move(module_config), - function_name, std::move(hlo_to_profile_idx))); + cpu_executable.reset(new CpuExecutable( + std::move(jit), std::move(assignment), std::move(module), function_name, + std::move(hlo_to_profile_idx))); - if (flags->xla_cpu_embed_ir) { + if (embed_ir_in_executable) { static_cast(*cpu_executable) .set_ir_module_string(ir_module_string); } } + VLOG(1) << "Compilation finished"; return std::move(cpu_executable); } StatusOr>> CpuCompiler::Compile( - std::vector> hlo_modules, - std::vector> module_configs, - HloDumper dump_hlos, std::vector stream_execs) { + std::vector> modules, + std::vector stream_execs) { return Unimplemented( "Compilation of multiple HLO modules is not yet supported on CPU."); } StatusOr>> -CpuCompiler::CompileAheadOfTime( - std::vector> hlo_modules, - std::vector> module_configs, - HloDumper dump_hlo, const AotCompilationOptions& aot_options) { - TF_RET_CHECK(hlo_modules.size() == module_configs.size()); - TF_RET_CHECK(!hlo_modules.empty()); +CpuCompiler::CompileAheadOfTime(std::vector> modules, + const AotCompilationOptions& aot_options) { + TF_RET_CHECK(!modules.empty()); + std::call_once(llvm_command_line_options_initialized, + &InitializeLLVMCommandLineOptions, modules[0]->config()); // We can pass just one llvm::TargetOptions when we compile the LLVM module, // so we bail if the configs have conflicting flags. At the moment, the only // flag that needs to be consistent is fast-math. - bool fast_math_disabled = module_configs[0]->fast_math_disabled(); - for (const auto& module_config : module_configs) { - if (module_config->fast_math_disabled() != fast_math_disabled) { + const bool fast_math_enabled = + modules[0]->config().debug_options().xla_enable_fast_math(); + for (const auto& module : modules) { + if (module->config().debug_options().xla_enable_fast_math() != + fast_math_enabled) { return InvalidArgument( "All HLO module configs must have the same value for " - "fast_math_disabled."); + "xla_enable_fast_math."); } } @@ -555,12 +727,11 @@ CpuCompiler::CompileAheadOfTime( } llvm::StringRef cpu_name = llvm_ir::AsStringRef(options.cpu_name()); llvm::StringRef features = llvm_ir::AsStringRef(options.features()); - llvm::CodeGenOpt::Level opt_level = CodeGenOptLevel(); - std::unique_ptr target_machine = - WrapUnique(target->createTargetMachine( - triple.getTriple(), cpu_name, features, - CompilerTargetOptions(*module_configs[0]), reloc_model, - llvm::CodeModel::Default, opt_level)); + llvm::CodeGenOpt::Level opt_level = CodeGenOptLevel(modules[0]->config()); + std::unique_ptr target_machine = WrapUnique( + target->createTargetMachine(triple.getTriple(), cpu_name, features, + CompilerTargetOptions(modules[0]->config()), + reloc_model, llvm::None, opt_level)); // Compile must be thread-safe so create a new LLVM context for the module. llvm::LLVMContext llvm_context; @@ -575,36 +746,40 @@ CpuCompiler::CompileAheadOfTime( } std::vector> results; - for (size_t i = 0; i < hlo_modules.size(); ++i) { - HloModule* hlo_module = hlo_modules[i].get(); - HloModuleConfig* module_config = module_configs[i].get(); + for (size_t i = 0; i < modules.size(); ++i) { + HloModule* module = modules[i].get(); + VLOG(1) << "Compiling ahead-of-time: " << module->name(); - TF_RETURN_IF_ERROR(RunHloPasses(hlo_module, module_config, dump_hlo)); + TF_RETURN_IF_ERROR(RunHloPasses(module)); TF_ASSIGN_OR_RETURN( SequentialHloOrdering::HloModuleSequence module_sequence, - CreateMemoryMinimizingSequence(*hlo_module, - [this](const LogicalBuffer& buffer) { - return ShapeSizeBytes(buffer.shape()); - })); + CreateMemoryMinimizingSequence(*module, BufferSizeBytesFunction())); // Run buffer analysis on the HLO graph. This analysis figures out which // temporary buffers are required to run the computation. TF_ASSIGN_OR_RETURN( std::unique_ptr assignment, BufferAssigner::Run( - hlo_module, - MakeUnique(hlo_module, module_sequence), - [this](const LogicalBuffer& buffer) { - return ShapeSizeBytes(buffer.shape()); - }, - kMemoryAlignment)); - - IrEmitter ir_emitter(*hlo_module, *module_config, *assignment, &llvm_module, - /*hlo_to_profile_idx=*/nullptr); - HloComputation* computation = hlo_module->entry_computation(); + module, MakeUnique(module, module_sequence), + BufferSizeBytesFunction(), memory_alignment)); + + const string dump_debug_json_to = + module->config().debug_options().xla_dump_debug_json_to(); + if (!dump_debug_json_to.empty()) { + HloProto proto = MakeHloProto(*module, *assignment); + TF_RETURN_IF_ERROR(protobuf_util::DumpJsonToDirectory( + proto, dump_debug_json_to, module->name())); + } + + IrEmitter ir_emitter(*module, *assignment, &llvm_module, + /*hlo_to_profile_idx=*/nullptr, target_machine.get()); + HloComputation* computation = module->entry_computation(); for (auto embedded_computation : computation->MakeEmbeddedComputationsList()) { + if (embedded_computation->IsFusionComputation()) { + continue; + } TF_RETURN_IF_ERROR( ir_emitter .EmitComputation(embedded_computation, @@ -617,14 +792,25 @@ CpuCompiler::CompileAheadOfTime( TF_ASSIGN_OR_RETURN( llvm::Function * entry_function, ir_emitter.EmitComputation(computation, entry_point_name, - /*is_entry_computation=*/true)); + /*is_entry_computation=*/true, + &module_sequence.at(computation))); entry_function->setName(llvm_ir::AsStringRef(entry_point_name)); + ModuleHook pre_optimization_ir_dump_hook; + ModuleHook post_optimization_ir_dump_hook; + TF_RETURN_IF_ERROR(InitializeModuleHooks( + *module, user_pre_optimization_hook_, user_post_optimization_hook_, + &pre_optimization_ir_dump_hook, &post_optimization_ir_dump_hook)); + Disassembler disassembler(*target_machine); - CompilerFunctor compiler_functor(target_machine.get(), &disassembler, - opt_level, - CompilerFunctor::AllIntrinsics()); + CompilerFunctor compiler_functor( + target_machine.get(), &disassembler, opt_level, + options::OptimizeForSizeRequested(module->config()), + module->config().debug_options().xla_enable_fast_math(), + module->config().debug_options().xla_llvm_disable_expensive_passes(), + CompilerFunctor::AllIntrinsics(), pre_optimization_ir_dump_hook, + post_optimization_ir_dump_hook); llvm::object::OwningBinary object_file = compiler_functor(llvm_module); llvm::StringRef object_file_data_ref = object_file.getBinary()->getData(); @@ -654,6 +840,8 @@ CpuCompiler::CompileAheadOfTime( std::move(object_file_data), std::move(buffer_sizes), result_slice.index())); } + + VLOG(1) << "Compilation finished"; return std::move(results); } @@ -661,19 +849,17 @@ se::Platform::Id CpuCompiler::PlatformId() const { return se::host::kHostPlatformId; } -int64 CpuCompiler::ShapeSizeBytes(const Shape& shape) const { - // On the cpu, opaques are pointers. - if (ShapeUtil::IsOpaque(shape)) { - return sizeof(void*); - } - return ShapeUtil::ByteSizeOf(shape, sizeof(void*)); +HloCostAnalysis::ShapeSizeFunction CpuCompiler::ShapeSizeBytesFunction() const { + return CpuExecutable::ShapeSizeBytes; } } // namespace cpu } // namespace xla -REGISTER_MODULE_INITIALIZER(cpu_compiler, { +static bool InitModule() { xla::Compiler::RegisterCompilerFactory(se::host::kHostPlatformId, []() { return xla::MakeUnique(); }); -}); + return true; +} +static bool module_initialized = InitModule(); diff --git a/tensorflow/compiler/xla/service/cpu/cpu_compiler.h b/tensorflow/compiler/xla/service/cpu/cpu_compiler.h index a32aa84ea51123f76551ad617cc914a53d4ca4d1..bd3541500dae9d9d59c56bfb062912a1b85c2219 100644 --- a/tensorflow/compiler/xla/service/cpu/cpu_compiler.h +++ b/tensorflow/compiler/xla/service/cpu/cpu_compiler.h @@ -18,10 +18,9 @@ limitations under the License. #include -#include "tensorflow/compiler/xla/service/compiler.h" #include "tensorflow/compiler/xla/service/executable.h" #include "tensorflow/compiler/xla/service/hlo_module.h" -#include "tensorflow/compiler/xla/service/hlo_module_config.h" +#include "tensorflow/compiler/xla/service/llvm_compiler.h" #include "tensorflow/compiler/xla/statusor.h" #include "tensorflow/core/lib/gtl/array_slice.h" #include "tensorflow/core/platform/macros.h" @@ -33,8 +32,6 @@ namespace cpu { // This class wraps the configurability options that LLVM exposes including: the // target triple, the target cpu and the target features. It also includes the // desired linkage name for the computation entry point. -// Note that the optimization level can be controlled by the -// --xla_cpu_llvm_opt_level flag. class CpuAotCompilationOptions : public AotCompilationOptions { public: // Relocation models available for compilation. @@ -107,31 +104,26 @@ class CpuAotCompilationResult : public AotCompilationResult { // The compiler translates XLA HLO code into LLVM IR and uses LLVM's JIT // infrastructure to create an executable "blob" that can then be returned // wrapped in CpuExecutable and actually invoked. -class CpuCompiler : public Compiler { +class CpuCompiler : public LLVMCompiler { public: CpuCompiler(); ~CpuCompiler() override {} StatusOr> Compile( - std::unique_ptr hlo_module, - std::unique_ptr module_config, HloDumper dump_hlo, + std::unique_ptr module, perftools::gputools::StreamExecutor* stream_exec) override; StatusOr>> Compile( - std::vector> hlo_module, - std::vector> module_config, - HloDumper dump_hlo, + std::vector> modules, std::vector stream_exec) override; StatusOr>> - CompileAheadOfTime( - std::vector> module, - std::vector> module_config, - HloDumper dump_hlo, const AotCompilationOptions& options) override; + CompileAheadOfTime(std::vector> modules, + const AotCompilationOptions& options) override; perftools::gputools::Platform::Id PlatformId() const override; - int64 ShapeSizeBytes(const Shape& shape) const override; + HloCostAnalysis::ShapeSizeFunction ShapeSizeBytesFunction() const override; private: // Initialize the LLVM target. @@ -139,8 +131,7 @@ class CpuCompiler : public Compiler { // Runs the HLO passes which are necessary for both optimizations and // correctness. - Status RunHloPasses(HloModule* hlo_module, HloModuleConfig* module_config, - HloDumper dump_hlo); + Status RunHloPasses(HloModule* module); TF_DISALLOW_COPY_AND_ASSIGN(CpuCompiler); }; diff --git a/tensorflow/compiler/xla/service/cpu/cpu_executable.cc b/tensorflow/compiler/xla/service/cpu/cpu_executable.cc index 88283e6010ea784e2a977a80adbe6315782f7fdc..6cc1d65c7afe50cbe7ee84a3d8c5cbfee1993f9a 100644 --- a/tensorflow/compiler/xla/service/cpu/cpu_executable.cc +++ b/tensorflow/compiler/xla/service/cpu/cpu_executable.cc @@ -22,12 +22,11 @@ limitations under the License. #include #include -#include "external/llvm/include/llvm/ExecutionEngine/Orc/IRCompileLayer.h" +#include "llvm/ExecutionEngine/Orc/IRCompileLayer.h" #include "tensorflow/compiler/xla/service/buffer_assignment.h" #include "tensorflow/compiler/xla/service/computation_layout.h" #include "tensorflow/compiler/xla/service/hlo_computation.h" #include "tensorflow/compiler/xla/service/hlo_module.h" -#include "tensorflow/compiler/xla/service/hlo_module_config.h" #include "tensorflow/compiler/xla/service/logical_buffer.h" #include "tensorflow/compiler/xla/service/shaped_buffer.h" #include "tensorflow/compiler/xla/shape_tree.h" @@ -53,11 +52,9 @@ namespace cpu { CpuExecutable::CpuExecutable( std::unique_ptr jit, std::unique_ptr assignment, - std::unique_ptr hlo_module, - std::unique_ptr module_config, - const string& entry_function_name, + std::unique_ptr hlo_module, const string& entry_function_name, std::unordered_map hlo_to_profile_idx) - : Executable(std::move(hlo_module), std::move(module_config)), + : Executable(std::move(hlo_module)), jit_(std::move(jit)), assignment_(std::move(assignment)), hlo_to_profile_idx_(std::move(hlo_to_profile_idx)) { @@ -69,7 +66,8 @@ CpuExecutable::CpuExecutable( CHECK(sym) << "Symbol " << entry_function_name << " not found."; // getAddress can do work under the hood in the jit, so it needs to be // guarded by the mutex. - compute_function_ = reinterpret_cast(sym.getAddress()); + compute_function_ = + reinterpret_cast(cantFail(sym.getAddress())); } // Given a pointer to an output buffer (following the CPU JIT calling @@ -314,11 +312,11 @@ StatusOr> CpuExecutable::ExecuteOnStream( std::vector buffers_in_result(assignment_->Allocations().size(), false); TF_RETURN_IF_ERROR( result_buffer->mutable_shape_index_to_buffer_entry() - ->ForEachMutableElement( + ->ForEachMutableElementWithStatus( [&buffers, &buffers_in_result, &result_buffer, this]( - const ShapeIndex& index, bool is_leaf, size_t* buffer_entry) { - if (is_leaf) { - const std::vector& sources = + const ShapeIndex& index, size_t* buffer_entry) { + if (ShapeUtil::IsLeafIndex(result_buffer->shape(), index)) { + const auto& sources = this->GetRootPointsToSet().element(index); // The points to set is unambiguous so the set should be a // singleton. @@ -369,10 +367,22 @@ CpuExecutable::ExecuteAsyncOnStream( "Asynchronous execution on stream is not yet supported on CPU."); } +/*static*/ int64 CpuExecutable::ShapeSizeBytes(const Shape& shape) { + // On the cpu, opaques are pointers. + if (ShapeUtil::IsOpaque(shape)) { + return sizeof(void*); + } + return ShapeUtil::ByteSizeOf(shape, sizeof(void*)); +} + const PointsToSet& CpuExecutable::GetRootPointsToSet() const { return assignment_->points_to_analysis().GetPointsToSet( module().entry_computation()->root_instruction()); } +std::unique_ptr CpuExecutable::CreateCostAnalysis() const { + return MakeUnique(ShapeSizeBytes); +} + } // namespace cpu } // namespace xla diff --git a/tensorflow/compiler/xla/service/cpu/cpu_executable.h b/tensorflow/compiler/xla/service/cpu/cpu_executable.h index b04b4e8dd1fd23839a4684f72622e32eca9c3730..a64537eaa3e3baefaefcc618ac971b7559badd94 100644 --- a/tensorflow/compiler/xla/service/cpu/cpu_executable.h +++ b/tensorflow/compiler/xla/service/cpu/cpu_executable.h @@ -29,7 +29,6 @@ limitations under the License. #include "tensorflow/compiler/xla/service/hlo_execution_profile.h" #include "tensorflow/compiler/xla/service/hlo_instruction.h" #include "tensorflow/compiler/xla/service/hlo_module.h" -#include "tensorflow/compiler/xla/service/hlo_module_config.h" #include "tensorflow/compiler/xla/service/shaped_buffer.h" #include "tensorflow/compiler/xla/service/tuple_points_to_analysis.h" #include "tensorflow/compiler/xla/statusor.h" @@ -52,7 +51,6 @@ class CpuExecutable : public Executable { std::unique_ptr jit, std::unique_ptr assignment, std::unique_ptr hlo_module, - std::unique_ptr module_config, const string& entry_function_name, std::unordered_map hlo_to_profile_idx); ~CpuExecutable() override {} @@ -80,6 +78,15 @@ class CpuExecutable : public Executable { ir_module_string_ = ir_module_string; } + const Status EqualOrFail(const Executable& executable) { + // TODO(b/62952745) Implement equality test on CPU executable. + return Unimplemented("Equality test on CPU executable is not implemented."); + } + + static int64 ShapeSizeBytes(const Shape& shape); + + std::unique_ptr CreateCostAnalysis() const override; + private: // Allocate buffers required for execution and assign them to the elements of // "buffers". "buffers" should be sized to the number of buffers in buffer diff --git a/tensorflow/compiler/xla/service/cpu/cpu_instruction_fusion.cc b/tensorflow/compiler/xla/service/cpu/cpu_instruction_fusion.cc index 240da35ef190eb7080947ab7d1da91d8d2dd8973..eb08bbe08e40ebbf48795166d17d8cd6a602472e 100644 --- a/tensorflow/compiler/xla/service/cpu/cpu_instruction_fusion.cc +++ b/tensorflow/compiler/xla/service/cpu/cpu_instruction_fusion.cc @@ -14,30 +14,93 @@ limitations under the License. ==============================================================================*/ #include "tensorflow/compiler/xla/service/cpu/cpu_instruction_fusion.h" - #include "tensorflow/compiler/xla/service/hlo_opcode.h" namespace xla { namespace cpu { +namespace { + +int64 BytesInDimension(const Shape& shape, int64 dimension) { + return ShapeUtil::ByteSizeOfPrimitiveType(shape.element_type()) * + shape.dimensions(dimension); +} + +bool IsFusile(const HloInstruction& hlo) { + // These are the only ones we fuse since we rely on effective elemental IR + // generation. + return (hlo.opcode() == HloOpcode::kBroadcast || + hlo.opcode() == HloOpcode::kReshape || + hlo.opcode() == HloOpcode::kBitcast || + hlo.opcode() == HloOpcode::kReverse || + hlo.opcode() == HloOpcode::kSlice || + hlo.opcode() == HloOpcode::kDynamicSlice || + hlo.opcode() == HloOpcode::kTranspose || hlo.IsElementwise()); +} + +} // namespace + bool CpuInstructionFusion::ShouldFuse(HloInstruction* consumer, int64 operand_index) { HloInstruction* producer = consumer->mutable_operand(operand_index); - // Condition for consumer: must be elementwise or a fusion op - // (which necessarily only contains elementwise operations) - if (!(consumer->opcode() == HloOpcode::kFusion || - consumer->IsElementwise())) { + constexpr int kFusionThresholdBytes = 16 * 1024; + + if (!IsFusile(*producer)) { return false; } // Producer or consumer cannot be Map. Maps are technically elementwise but // of a slightly different form (call instead of a computation). These are not // yet supported in the CPU backend. - return producer->IsElementwise() && producer->operand_count() > 0 && - producer->opcode() != HloOpcode::kMap && - consumer->opcode() != HloOpcode::kMap && - InstructionFusion::ShouldFuse(consumer, operand_index); + if (producer->opcode() == HloOpcode::kMap || + consumer->opcode() == HloOpcode::kMap) { + return false; + } + + // TODO(b/28644064): see if the "producer->operand_count() == 0" check is + // necessary. + if (producer->operand_count() == 0 || + !InstructionFusion::ShouldFuse(consumer, operand_index)) { + return false; + } + + // Output fusion is not currently supported on CPUs. + if (producer->opcode() == HloOpcode::kFusion) { + return false; + } + + if (consumer->opcode() == HloOpcode::kDot) { + // In the general case we call out to optimized "black box" GEMM routines + // for Dot, which precludes fusion. However, in very specific cases, we try + // to fuse Dot operations by generating an elemental dot implementation. + // + // We need to be careful and conservative here since any benefit we get from + // fusion can easily be overshadowed by the overhead of a naive GEMM + // algorithm in the IR. + const Shape& output_shape = consumer->shape(); + if (output_shape.dimensions_size() == 2) { + // We fuse in cases where we have dot([A,B],[B,1]) or dot([1,A],[A,B]) and + // fusion can get rid of the larger tensor. We assume that a naive + // traversal of a small enough (to fit in L1) column or row tensor is + // "good enough" from the perspective of cache management; and calling out + // to an optimized GEMM kernel is not a huge win. + if (output_shape.dimensions(0) == 1 && operand_index == 1 && + BytesInDimension(output_shape, 1) < kFusionThresholdBytes) { + return true; + } else if (output_shape.dimensions(1) == 1 && operand_index == 0 && + BytesInDimension(output_shape, 0) < kFusionThresholdBytes) { + return true; + } + } + } + + // InstructionFusion::ShouldFuse above only allows kLoop and kInput fusions. + // The CPU backend does not create kInput fusions, so we only expect to see + // kLoop here. + CHECK(consumer->opcode() != HloOpcode::kFusion || + consumer->fusion_kind() == HloInstruction::FusionKind::kLoop); + return consumer->opcode() == HloOpcode::kFusion || consumer->IsElementwise(); } } // namespace cpu diff --git a/tensorflow/compiler/xla/service/cpu/cpu_instruction_fusion_test.cc b/tensorflow/compiler/xla/service/cpu/cpu_instruction_fusion_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..75a665c2228c03bef0972e69cf7e4c066ea67200 --- /dev/null +++ b/tensorflow/compiler/xla/service/cpu/cpu_instruction_fusion_test.cc @@ -0,0 +1,411 @@ +/* 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/cpu/cpu_instruction_fusion.h" + +#include +#include + +#include "tensorflow/compiler/xla/service/hlo_matchers.h" +#include "tensorflow/compiler/xla/service/transpose_folding.h" +#include "tensorflow/compiler/xla/tests/hlo_test_base.h" +#include "tensorflow/core/lib/gtl/array_slice.h" + +namespace op = xla::testing::opcode_matchers; + +namespace xla { +namespace cpu { +namespace { + +using InstructionFusionTest = HloTestBase; + +TEST_F(InstructionFusionTest, DotOperationFusion_Basic_0) { + HloComputation::Builder builder(TestName()); + HloInstruction* arg0 = builder.AddInstruction(HloInstruction::CreateParameter( + 0, ShapeUtil::MakeShape(F32, {1024, 256}), "arg0")); + HloInstruction* arg1 = builder.AddInstruction(HloInstruction::CreateParameter( + 1, ShapeUtil::MakeShape(F32, {256, 1}), "arg1")); + + HloInstruction* exp0 = builder.AddInstruction(HloInstruction::CreateUnary( + ShapeUtil::MakeShape(S32, {1024, 256}), HloOpcode::kExp, arg0)); + HloInstruction* dot = builder.AddInstruction(HloInstruction::CreateBinary( + ShapeUtil::MakeShape(F32, {1024, 1}), HloOpcode::kDot, exp0, arg1)); + + auto module = CreateNewModule(); + auto computation = module->AddEntryComputation(builder.Build()); + EXPECT_EQ(dot, computation->root_instruction()); + EXPECT_TRUE(CpuInstructionFusion().Run(module.get()).ValueOrDie()); + EXPECT_THAT(computation->root_instruction(), op::Fusion()); +} + +TEST_F(InstructionFusionTest, DotOperationFusion_Basic_1) { + HloComputation::Builder builder(TestName()); + HloInstruction* arg0 = builder.AddInstruction(HloInstruction::CreateParameter( + 0, ShapeUtil::MakeShape(F32, {1, 256}), "arg0")); + HloInstruction* arg1 = builder.AddInstruction(HloInstruction::CreateParameter( + 1, ShapeUtil::MakeShape(F32, {256, 1024}), "arg1")); + + HloInstruction* exp1 = builder.AddInstruction(HloInstruction::CreateUnary( + ShapeUtil::MakeShape(S32, {256, 1024}), HloOpcode::kExp, arg1)); + HloInstruction* dot = builder.AddInstruction(HloInstruction::CreateBinary( + ShapeUtil::MakeShape(F32, {1, 1024}), HloOpcode::kDot, arg0, exp1)); + + auto module = CreateNewModule(); + auto computation = module->AddEntryComputation(builder.Build()); + EXPECT_EQ(dot, computation->root_instruction()); + EXPECT_TRUE(CpuInstructionFusion().Run(module.get()).ValueOrDie()); + EXPECT_THAT(computation->root_instruction(), op::Fusion()); +} + +TEST_F(InstructionFusionTest, DotOperationFusion_Bitcast) { + HloComputation::Builder builder(TestName()); + HloInstruction* arg0 = builder.AddInstruction(HloInstruction::CreateParameter( + 0, ShapeUtil::MakeShape(F32, {2, 512, 2, 128}), "arg0")); + HloInstruction* arg1 = builder.AddInstruction(HloInstruction::CreateParameter( + 1, ShapeUtil::MakeShape(F32, {256, 1}), "arg1")); + + HloInstruction* exp0 = builder.AddInstruction(HloInstruction::CreateUnary( + ShapeUtil::MakeShape(S32, {2, 512, 2, 128}), HloOpcode::kExp, arg0)); + HloInstruction* bitcast0 = builder.AddInstruction(HloInstruction::CreateUnary( + ShapeUtil::MakeShape(S32, {1024, 256}), HloOpcode::kBitcast, exp0)); + HloInstruction* dot = builder.AddInstruction(HloInstruction::CreateBinary( + ShapeUtil::MakeShape(F32, {1024, 1}), HloOpcode::kDot, bitcast0, arg1)); + + auto module = CreateNewModule(); + auto computation = module->AddEntryComputation(builder.Build()); + EXPECT_EQ(dot, computation->root_instruction()); + EXPECT_TRUE(CpuInstructionFusion().Run(module.get()).ValueOrDie()); + EXPECT_THAT(computation->root_instruction(), op::Fusion()); +} + +TEST_F(InstructionFusionTest, DotOperationFusion_Reshape) { + HloComputation::Builder builder(TestName()); + HloInstruction* arg0 = builder.AddInstruction(HloInstruction::CreateParameter( + 0, ShapeUtil::MakeShape(F32, {2, 512, 2, 128}), "arg0")); + HloInstruction* arg1 = builder.AddInstruction(HloInstruction::CreateParameter( + 1, ShapeUtil::MakeShape(F32, {256, 1}), "arg1")); + + HloInstruction* exp0 = builder.AddInstruction(HloInstruction::CreateUnary( + ShapeUtil::MakeShape(S32, {2, 512, 2, 128}), HloOpcode::kExp, arg0)); + HloInstruction* reshape0 = + builder.AddInstruction(HloInstruction::CreateReshape( + ShapeUtil::MakeShape(S32, {1024, 256}), exp0)); + HloInstruction* dot = builder.AddInstruction(HloInstruction::CreateBinary( + ShapeUtil::MakeShape(F32, {1024, 1}), HloOpcode::kDot, reshape0, arg1)); + + auto module = CreateNewModule(); + auto computation = module->AddEntryComputation(builder.Build()); + EXPECT_EQ(dot, computation->root_instruction()); + EXPECT_TRUE(CpuInstructionFusion().Run(module.get()).ValueOrDie()); + EXPECT_THAT(computation->root_instruction(), op::Fusion()); +} + +TEST_F(InstructionFusionTest, DotOperationFusion_TooLarge) { + HloComputation::Builder builder(TestName()); + HloInstruction* arg0 = builder.AddInstruction(HloInstruction::CreateParameter( + 0, ShapeUtil::MakeShape(F32, {1, 32 * 1024}), "arg0")); + HloInstruction* arg1 = builder.AddInstruction(HloInstruction::CreateParameter( + 1, ShapeUtil::MakeShape(F32, {256, 32 * 1024}), "arg1")); + + HloInstruction* exp1 = builder.AddInstruction(HloInstruction::CreateUnary( + ShapeUtil::MakeShape(S32, {256, 32 * 1024}), HloOpcode::kExp, arg1)); + HloInstruction* dot = builder.AddInstruction(HloInstruction::CreateBinary( + ShapeUtil::MakeShape(F32, {1, 32 * 1024}), HloOpcode::kDot, arg0, exp1)); + + auto module = CreateNewModule(); + auto computation = module->AddEntryComputation(builder.Build()); + EXPECT_EQ(dot, computation->root_instruction()); + EXPECT_FALSE(CpuInstructionFusion().Run(module.get()).ValueOrDie()); + EXPECT_EQ(dot, computation->root_instruction()); +} + +TEST_F(InstructionFusionTest, DotOperationFusion_ElementReuse) { + HloComputation::Builder builder(TestName()); + HloInstruction* arg0 = builder.AddInstruction(HloInstruction::CreateParameter( + 0, ShapeUtil::MakeShape(F32, {2, 256}), "arg0")); + HloInstruction* arg1 = builder.AddInstruction(HloInstruction::CreateParameter( + 1, ShapeUtil::MakeShape(F32, {256, 1024}), "arg1")); + + HloInstruction* exp1 = builder.AddInstruction(HloInstruction::CreateUnary( + ShapeUtil::MakeShape(S32, {256, 1024}), HloOpcode::kExp, arg1)); + HloInstruction* dot = builder.AddInstruction(HloInstruction::CreateBinary( + ShapeUtil::MakeShape(F32, {2, 1024}), HloOpcode::kDot, arg0, exp1)); + + auto module = CreateNewModule(); + auto computation = module->AddEntryComputation(builder.Build()); + EXPECT_EQ(dot, computation->root_instruction()); + EXPECT_FALSE(CpuInstructionFusion().Run(module.get()).ValueOrDie()); + EXPECT_EQ(dot, computation->root_instruction()); +} + +TEST_F(InstructionFusionTest, DotOperationFusion_TransposeFusion) { + HloComputation::Builder builder(TestName()); + HloInstruction* arg0 = builder.AddInstruction(HloInstruction::CreateParameter( + 0, ShapeUtil::MakeShape(F32, {1, 256}), "arg0")); + HloInstruction* arg1 = builder.AddInstruction(HloInstruction::CreateParameter( + 1, ShapeUtil::MakeShape(F32, {1024, 256}), "arg1")); + + HloInstruction* exp1 = builder.AddInstruction(HloInstruction::CreateUnary( + ShapeUtil::MakeShape(S32, {1024, 256}), HloOpcode::kExp, arg1)); + HloInstruction* transpose1 = + builder.AddInstruction(HloInstruction::CreateTranspose( + ShapeUtil::MakeShape(S32, {256, 1024}), exp1, {1, 0})); + builder.AddInstruction(HloInstruction::CreateBinary( + ShapeUtil::MakeShape(F32, {1, 1024}), HloOpcode::kDot, arg0, transpose1)); + + auto module = CreateNewModule(); + auto computation = module->AddEntryComputation(builder.Build()); + TransposeFolding transpose_folding( + [](const HloInstruction& dot, + const TransposeFolding::OperandIndices& candidate_operands) { + return candidate_operands; + }, + TransposeFolding::NeverFoldTranspose); + EXPECT_TRUE(transpose_folding.Run(module.get()).ValueOrDie()); + EXPECT_EQ(computation->root_instruction()->opcode(), HloOpcode::kFusion); + EXPECT_EQ(computation->root_instruction()->fusion_kind(), + HloInstruction::FusionKind::kTransposeDot); + EXPECT_FALSE(CpuInstructionFusion().Run(module.get()).ValueOrDie()); + EXPECT_EQ(computation->root_instruction()->opcode(), HloOpcode::kFusion); + EXPECT_EQ(computation->root_instruction()->fusion_kind(), + HloInstruction::FusionKind::kTransposeDot); +} + +class OpcodeFusionTest : public InstructionFusionTest { + protected: + // Runs CPU instruction fusion on the given module, and tests that the result + // contains a fused op at the root with exactly the given multiset of opcodes. + void RunFusionAndCheckOpcodesWereFused( + HloModule* module, const std::multiset& expected_opcodes) { + auto computation = module->entry_computation(); + auto did_fusion = CpuInstructionFusion().Run(module); + ASSERT_TRUE(did_fusion.ok()); + EXPECT_TRUE(did_fusion.ValueOrDie()); + + HloInstruction* root = computation->root_instruction(); + ASSERT_THAT(root, op::Fusion()); + EXPECT_EQ(root->fusion_kind(), HloInstruction::FusionKind::kLoop); + + std::vector fused_opcodes(root->fused_instructions().size()); + std::transform(root->fused_instructions().begin(), + root->fused_instructions().end(), fused_opcodes.begin(), + [](const std::unique_ptr& hlo) { + return hlo->opcode(); + }); + + EXPECT_EQ( + std::multiset(fused_opcodes.begin(), fused_opcodes.end()), + expected_opcodes); + } +}; + +TEST_F(OpcodeFusionTest, Exponential_Bitcast_Negate) { + HloComputation::Builder builder(TestName()); + Shape param_shape = ShapeUtil::MakeShape(F32, {1, 4}); + Shape result_shape = ShapeUtil::MakeShape(F32, {4}); + HloInstruction* param0 = builder.AddInstruction( + HloInstruction::CreateParameter(0, param_shape, "param")); + // InstructionFusion::ShouldFuse() precludes fusing a bitcast whose operand + // is a parameter, so create an operand between the parameter and bitcast. + HloInstruction* exp1 = builder.AddInstruction( + HloInstruction::CreateUnary(param_shape, HloOpcode::kExp, param0)); + HloInstruction* bitcast2 = builder.AddInstruction( + HloInstruction::CreateUnary(result_shape, HloOpcode::kBitcast, exp1)); + builder.AddInstruction( + HloInstruction::CreateUnary(result_shape, HloOpcode::kNegate, bitcast2)); + + auto module = CreateNewModule(); + module->AddEntryComputation(builder.Build()); + + RunFusionAndCheckOpcodesWereFused( + module.get(), {HloOpcode::kNegate, HloOpcode::kBitcast, HloOpcode::kExp, + HloOpcode::kParameter}); +} + +TEST_F(OpcodeFusionTest, Broadcast_Bitcast_DynamicSlice_Tanh) { + HloComputation::Builder builder(TestName()); + Shape param_shape = ShapeUtil::MakeShape(F32, {8}); + Shape starts_shape = ShapeUtil::MakeShape(F32, {2}); + Shape broadcast_shape = ShapeUtil::MakeShape(F32, {1, 8, 8}); + Shape bitcast_shape = ShapeUtil::MakeShape(F32, {8, 8}); + Shape dynamic_slice_shape = ShapeUtil::MakeShape(F32, {4, 4}); + HloInstruction* param0 = builder.AddInstruction( + HloInstruction::CreateParameter(0, param_shape, "param")); + HloInstruction* param1 = builder.AddInstruction( + HloInstruction::CreateParameter(1, starts_shape, "starts")); + HloInstruction* broadcast2 = builder.AddInstruction( + HloInstruction::CreateBroadcast(broadcast_shape, param0, {1})); + HloInstruction* bitcast3 = builder.AddInstruction(HloInstruction::CreateUnary( + bitcast_shape, HloOpcode::kBitcast, broadcast2)); + HloInstruction* dynamic_slice4 = + builder.AddInstruction(HloInstruction::CreateDynamicSlice( + dynamic_slice_shape, bitcast3, param1, {4, 4})); + builder.AddInstruction(HloInstruction::CreateUnary( + dynamic_slice_shape, HloOpcode::kTanh, dynamic_slice4)); + + auto module = CreateNewModule(); + module->AddEntryComputation(builder.Build()); + + RunFusionAndCheckOpcodesWereFused( + module.get(), + {HloOpcode::kTanh, HloOpcode::kDynamicSlice, HloOpcode::kBitcast, + HloOpcode::kBroadcast, HloOpcode::kParameter, HloOpcode::kParameter}); +} + +TEST_F(OpcodeFusionTest, Broadcast_Negate) { + HloComputation::Builder builder(TestName()); + Shape param_shape = ShapeUtil::MakeShape(F32, {8}); + Shape result_shape = ShapeUtil::MakeShape(F32, {8, 8}); + HloInstruction* param0 = builder.AddInstruction( + HloInstruction::CreateParameter(0, param_shape, "param")); + HloInstruction* broadcast1 = builder.AddInstruction( + HloInstruction::CreateBroadcast(result_shape, param0, {1})); + builder.AddInstruction(HloInstruction::CreateUnary( + result_shape, HloOpcode::kNegate, broadcast1)); + + auto module = CreateNewModule(); + module->AddEntryComputation(builder.Build()); + + RunFusionAndCheckOpcodesWereFused( + module.get(), + {HloOpcode::kNegate, HloOpcode::kBroadcast, HloOpcode::kParameter}); +} + +TEST_F(OpcodeFusionTest, DynamicSlice_Negate) { + HloComputation::Builder builder(TestName()); + Shape param_shape = ShapeUtil::MakeShape(F32, {4}); + Shape slice_shape = ShapeUtil::MakeShape(F32, {1}); + 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})); + builder.AddInstruction(HloInstruction::CreateUnary( + result_shape, HloOpcode::kNegate, dynamic_slice2)); + + auto module = CreateNewModule(); + module->AddEntryComputation(builder.Build()); + + RunFusionAndCheckOpcodesWereFused( + module.get(), {HloOpcode::kNegate, HloOpcode::kDynamicSlice, + HloOpcode::kParameter, HloOpcode::kParameter}); +} + +TEST_F(OpcodeFusionTest, Exponential_Negate) { + HloComputation::Builder builder(TestName()); + Shape param_shape = ShapeUtil::MakeShape(F32, {4}); + HloInstruction* param0 = builder.AddInstruction( + HloInstruction::CreateParameter(0, param_shape, "param")); + HloInstruction* exp1 = builder.AddInstruction( + HloInstruction::CreateUnary(param_shape, HloOpcode::kExp, param0)); + builder.AddInstruction( + HloInstruction::CreateUnary(param_shape, HloOpcode::kNegate, exp1)); + + auto module = CreateNewModule(); + module->AddEntryComputation(builder.Build()); + + RunFusionAndCheckOpcodesWereFused( + module.get(), + {HloOpcode::kNegate, HloOpcode::kExp, HloOpcode::kParameter}); +} + +TEST_F(OpcodeFusionTest, Reshape_Negate) { + HloComputation::Builder builder(TestName()); + Shape param_shape = ShapeUtil::MakeShape(F32, {4, 4}); + Shape result_shape = ShapeUtil::MakeShape(F32, {16}); + HloInstruction* param0 = builder.AddInstruction( + HloInstruction::CreateParameter(0, param_shape, "param")); + HloInstruction* reshape1 = builder.AddInstruction( + HloInstruction::CreateReshape(result_shape, param0)); + builder.AddInstruction( + HloInstruction::CreateUnary(result_shape, HloOpcode::kNegate, reshape1)); + + auto module = CreateNewModule(); + module->AddEntryComputation(builder.Build()); + + RunFusionAndCheckOpcodesWereFused( + module.get(), + {HloOpcode::kNegate, HloOpcode::kReshape, HloOpcode::kParameter}); +} + +TEST_F(OpcodeFusionTest, Reverse_Negate) { + HloComputation::Builder builder(TestName()); + Shape param_shape = ShapeUtil::MakeShape(F32, {8}); + HloInstruction* param0 = builder.AddInstruction( + HloInstruction::CreateParameter(0, param_shape, "param")); + HloInstruction* reverse1 = builder.AddInstruction( + HloInstruction::CreateReverse(param_shape, param0, {0})); + builder.AddInstruction( + HloInstruction::CreateUnary(param_shape, HloOpcode::kNegate, reverse1)); + + auto module = CreateNewModule(); + module->AddEntryComputation(builder.Build()); + + RunFusionAndCheckOpcodesWereFused( + module.get(), + {HloOpcode::kNegate, HloOpcode::kReverse, HloOpcode::kParameter}); +} + +TEST_F(OpcodeFusionTest, Slice_Negate) { + HloComputation::Builder builder(TestName()); + Shape param_shape = ShapeUtil::MakeShape(F32, {4}); + Shape slice_shape = ShapeUtil::MakeShape(F32, {2}); + HloInstruction* param0 = builder.AddInstruction( + HloInstruction::CreateParameter(0, param_shape, "param")); + HloInstruction* slice1 = builder.AddInstruction( + HloInstruction::CreateSlice(slice_shape, param0, {0}, {4}, {2})); + builder.AddInstruction(HloInstruction::CreateUnary( + ShapeUtil::MakeShape(S32, {2}), HloOpcode::kNegate, slice1)); + + auto module = CreateNewModule(); + module->AddEntryComputation(builder.Build()); + + RunFusionAndCheckOpcodesWereFused( + module.get(), + {HloOpcode::kNegate, HloOpcode::kSlice, HloOpcode::kParameter}); +} + +TEST_F(OpcodeFusionTest, Exponential_Transpose_Negate) { + HloComputation::Builder builder(TestName()); + Shape param_shape = ShapeUtil::MakeShape(F32, {3, 4}); + Shape result_shape = ShapeUtil::MakeShape(F32, {4, 3}); + HloInstruction* param0 = builder.AddInstruction( + HloInstruction::CreateParameter(0, param_shape, "param")); + // InstructionFusion::ShouldFuse() precludes fusing a transpose whose operand + // is a parameter, so create an operand between the parameter and transpose. + HloInstruction* exp1 = builder.AddInstruction( + HloInstruction::CreateUnary(param_shape, HloOpcode::kExp, param0)); + HloInstruction* transpose2 = builder.AddInstruction( + HloInstruction::CreateTranspose(result_shape, exp1, {1, 0})); + builder.AddInstruction(HloInstruction::CreateUnary( + result_shape, HloOpcode::kNegate, transpose2)); + + auto module = CreateNewModule(); + module->AddEntryComputation(builder.Build()); + + RunFusionAndCheckOpcodesWereFused( + module.get(), {HloOpcode::kNegate, HloOpcode::kTranspose, HloOpcode::kExp, + HloOpcode::kParameter}); +} + +} // namespace +} // namespace cpu +} // namespace xla + +int main(int argc, char** argv) { + return xla::ParseDebugOptionsFlagsAndRunTests(argc, argv); +} diff --git a/tensorflow/compiler/xla/service/cpu/cpu_options.cc b/tensorflow/compiler/xla/service/cpu/cpu_options.cc new file mode 100644 index 0000000000000000000000000000000000000000..dba140d1120bc5502d2039e1663b9bf035d8d66a --- /dev/null +++ b/tensorflow/compiler/xla/service/cpu/cpu_options.cc @@ -0,0 +1,50 @@ +/* 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/cpu/cpu_options.h" + +namespace { + +const char* const kXlaParallelCpuOption = "xla_cpu_parallel"; +const char* const kXlaOptimizeForSizeCpuOption = "xla_cpu_optimize_for_size"; +const char* const kXlaDisableVectorizedReduce = "xla_disable_vectorized_reduce"; + +} // namespace + +namespace xla { +namespace cpu { +namespace options { + +bool CpuParallelBackendRequested(const HloModuleConfig& config) { + const auto& extra_options_map = + config.debug_options().xla_backend_extra_options(); + return extra_options_map.count(kXlaParallelCpuOption) > 0; +} + +bool OptimizeForSizeRequested(const HloModuleConfig& config) { + const auto& extra_options_map = + config.debug_options().xla_backend_extra_options(); + return extra_options_map.count(kXlaOptimizeForSizeCpuOption) > 0; +} + +bool VectorizedReduceDisabled(const HloModuleConfig& config) { + const auto& extra_options_map = + config.debug_options().xla_backend_extra_options(); + return extra_options_map.count(kXlaOptimizeForSizeCpuOption) > 0; +} + +} // namespace options +} // namespace cpu +} // namespace xla diff --git a/tensorflow/compiler/xla/service/cpu/cpu_options.h b/tensorflow/compiler/xla/service/cpu/cpu_options.h new file mode 100644 index 0000000000000000000000000000000000000000..5dc24ebc7b8661092e3bc27c4f30fda1e497e41b --- /dev/null +++ b/tensorflow/compiler/xla/service/cpu/cpu_options.h @@ -0,0 +1,35 @@ +/* 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_CPU_CPU_OPTIONS_H_ +#define TENSORFLOW_COMPILER_XLA_SERVICE_CPU_CPU_OPTIONS_H_ + +#include "tensorflow/compiler/xla/service/hlo_module_config.h" + +// Helper functions for querying options that are specific to the CPU backend. + +namespace xla { +namespace cpu { +namespace options { + +bool CpuParallelBackendRequested(const HloModuleConfig& config); +bool OptimizeForSizeRequested(const HloModuleConfig& config); +bool VectorizedReduceDisabled(const HloModuleConfig& config); + +} // namespace options +} // namespace cpu +} // namespace xla + +#endif // TENSORFLOW_COMPILER_XLA_SERVICE_CPU_CPU_OPTIONS_H_ diff --git a/tensorflow/compiler/xla/service/cpu/cpu_parallelization_preparation.cc b/tensorflow/compiler/xla/service/cpu/cpu_parallelization_preparation.cc index f6b1dcae75a773811f8c652dea36b7f3ca36e901..20ee4f12e53a16b76d39d0151bc5b8ca4475f7ab 100644 --- a/tensorflow/compiler/xla/service/cpu/cpu_parallelization_preparation.cc +++ b/tensorflow/compiler/xla/service/cpu/cpu_parallelization_preparation.cc @@ -15,19 +15,28 @@ limitations under the License. #include "tensorflow/compiler/xla/service/cpu/cpu_parallelization_preparation.h" +#include "tensorflow/compiler/xla/map_util.h" +#include "tensorflow/compiler/xla/service/cpu/ir_emission_utils.h" +#include "tensorflow/compiler/xla/service/cpu/shape_partition.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/logical_buffer.h" #include "tensorflow/compiler/xla/service/tuple_points_to_analysis.h" #include "tensorflow/compiler/xla/types.h" +#include "tensorflow/compiler/xla/util.h" #include "tensorflow/core/lib/strings/strcat.h" namespace xla { namespace cpu { StatusOr ParallelizationPreparation::Run(HloModule* module) { + XLA_VLOG_LINES(2, "ParallelizationPreparation ENTRY"); + XLA_VLOG_LINES(2, module->ToString()); + bool changed = false; + TF_ASSIGN_OR_RETURN(changed, RunParallelTaskAssignment(module)); + HloComputation* entry_computation = module->entry_computation(); std::unordered_set outlined; std::vector instructions_to_outline; @@ -44,16 +53,24 @@ StatusOr ParallelizationPreparation::Run(HloModule* module) { instruction->opcode() == HloOpcode::kConstant) { continue; } + + // Outline 'instruction' in isolation if it was assigned parallel tasks. + if (OutlineParallelizableInstruction(instruction)) { + outlined.insert(instruction); + changed = true; + continue; + } + instructions_to_outline.clear(); HloInstruction* outline_candidate = instruction; instructions_to_outline.push_back(outline_candidate); bool all_bitcasts = outline_candidate->opcode() == HloOpcode::kBitcast; // Outline sole users with the current instruction. - while (outline_candidate->users().size() == 1) { + while (CanOutlineWithUser(outline_candidate)) { HloInstruction* prior_candidate = outline_candidate; outline_candidate = *outline_candidate->users().begin(); - all_bitcasts |= outline_candidate->opcode() == HloOpcode::kBitcast; + all_bitcasts &= outline_candidate->opcode() == HloOpcode::kBitcast; if (std::any_of(outline_candidate->operands().begin(), outline_candidate->operands().end(), [&](const HloInstruction* operand) { @@ -108,6 +125,9 @@ StatusOr ParallelizationPreparation::Run(HloModule* module) { TF_ASSIGN_OR_RETURN(auto points_to_analysis, TuplePointsToAnalysis::Run(module)); for (auto& computation : module->computations()) { + if (computation->IsFusionComputation()) { + continue; + } HloInstruction* root = computation->root_instruction(); // Copy root instruction if it does not define its own top-level buffer. // TODO(b/32885001) Remove these copies (at least for the unambiguous case). @@ -120,8 +140,136 @@ StatusOr ParallelizationPreparation::Run(HloModule* module) { changed = true; } } + + XLA_VLOG_LINES(2, "ParallelizationPreparation EXIT"); + XLA_VLOG_LINES(2, module->ToString()); return changed; } +StatusOr ParallelizationPreparation::RunParallelTaskAssignment( + HloModule* module) { + VLOG(1) << "RunParallelTaskAssignment max_parallelism_: " << max_parallelism_; + bool changed = false; + // Run cost analysis on entry computation. + HloCostAnalysis cost_analysis(shape_size_); + HloComputation* computation = module->entry_computation(); + Status cost_status = computation->root_instruction()->Accept(&cost_analysis); + for (auto& instruction : computation->instructions()) { + // Currently, we do not assign parallel tasks to instructions with at least + // one of the following properties: + // *) Internal threading (library calls to kConv, kDot, and kCustomCall). + // *) Emit custom loops (kSelectAndScatter, FusionKind::kTransposeDot). + // *) Tuple-shaped. + // TODO(b/27458679) Parallelize instructions which are skipped here. + if (instruction->opcode() == HloOpcode::kParameter || + instruction->opcode() == HloOpcode::kConstant || + instruction->opcode() == HloOpcode::kCall || + instruction->opcode() == HloOpcode::kCustomCall || + instruction->opcode() == HloOpcode::kSelectAndScatter || + (instruction->opcode() == HloOpcode::kConvolution && + PotentiallyImplementedAsEigenConvolution(*instruction)) || + PotentiallyImplementedAsEigenDot(*instruction) || + (instruction->opcode() == HloOpcode::kFusion && + instruction->fusion_kind() != HloInstruction::FusionKind::kLoop) || + ShapeUtil::IsTuple(instruction->shape())) { + continue; + } + + // Calculate target parallel task count in [1, max_parallelism_]. + const int64 target_parallel_task_count = GetTargetParallelTaskCount( + cost_status.ok() ? &cost_analysis : nullptr, instruction.get()); + if (target_parallel_task_count == 1) { + continue; + } + + // Assign feasible dimension partitions (based on actual dimension sizes). + auto dim_partition_counts = ShapePartitionAssigner(instruction->shape()) + .Run(target_parallel_task_count); + const int64 total_partition_count = + ShapePartitionAssigner::GetTotalPartitionCount(dim_partition_counts); + if (total_partition_count <= 1) { + // Feasible partition calculation resulting in no partitioning, so skip. + continue; + } + VLOG(2) << "Assigning parallel task count: " << total_partition_count + << " to instruction: " << instruction->name(); + // Map 'instruction' to assigned dimension partitioning. + instruction->set_outer_dimension_partitions(dim_partition_counts); + } + + return changed; +} + +int64 ParallelizationPreparation::GetTargetParallelTaskCount( + const HloCostAnalysis* cost_analysis, HloInstruction* instruction) { + // Default to a simple cost model based on hlo size and typical L2 cache size. + // Note that 'cost_analysis' can be 'nullptr' if HloCostAnalysis returns an + // error status (likely because HLOs like CustomCall are not yet implemented + // in the HloCostAnalysis). + int64 instruction_cost = shape_size_(instruction->shape()); + int64 min_cost_per_thread = 256LL << 10; // 256KB L2 Cache size. + if (cost_analysis != nullptr) { + // Calculate the instruction cost in cycles. + // TODO(29630486) Improve on this linear cost model. + // Consider making 'min_cost_per_thread' be a function of the target + // bandwidth limit for instructions with low arithmetic complexity. + instruction_cost = 1 * cost_analysis->flop_count(*instruction) + + 2 * cost_analysis->transcendental_count(*instruction) + + 10 * cost_analysis->bytes_accessed(*instruction); + // Minimum per-thread cost is 100us of work on a 2GHz core. + min_cost_per_thread = 100000; + } + // Return target parallel task count in [1, max_parallelism_]. + return std::min(max_parallelism_, + std::max(1LL, instruction_cost / min_cost_per_thread)); +} + +bool ParallelizationPreparation::OutlineParallelizableInstruction( + HloInstruction* instruction) { + if (instruction->outer_dimension_partitions().empty()) { + return false; + } + // Store dimension partition counts before outlining (which clones + // 'instruction'). + std::vector dim_partition_counts = + instruction->outer_dimension_partitions(); + // Outline 'instruction' in its own sub-computation. + HloModule* module = instruction->parent()->parent(); + auto* call = module->OutlineExpressionFromComputation( + {instruction}, tensorflow::strings::StrCat("pp_", instruction->name()), + module->entry_computation()); + // Map previously assigned 'dim_partition_counts' to cloned root instruction. + VLOG(1) << "Outlining parallelizable" + << " caller: " << call->name() + << " callee: " << call->to_apply()->root_instruction()->name(); + call->to_apply()->root_instruction()->set_outer_dimension_partitions( + dim_partition_counts); + return true; +} + +bool ParallelizationPreparation::CanOutlineWithUser( + HloInstruction* instruction) { + if (instruction->users().size() != 1) { + // Do not outline 'instruction' with multiple users. + return false; + } + if (AssignedParallelTasks(instruction) || + AssignedParallelTasks(*instruction->users().begin())) { + // Do not outline if 'instruction' (or user) were assigned parallel tasks. + return false; + } + return true; +} + +bool ParallelizationPreparation::AssignedParallelTasks( + HloInstruction* instruction) { + return !instruction->outer_dimension_partitions().empty() || + (instruction->opcode() == HloOpcode::kCall && + !instruction->to_apply() + ->root_instruction() + ->outer_dimension_partitions() + .empty()); +} + } // namespace cpu } // namespace xla diff --git a/tensorflow/compiler/xla/service/cpu/cpu_parallelization_preparation.h b/tensorflow/compiler/xla/service/cpu/cpu_parallelization_preparation.h index 62999f5686db2e4db3ace0c5580bd156edbfa994..d53fc461509cad51778dba37922212731236952f 100644 --- a/tensorflow/compiler/xla/service/cpu/cpu_parallelization_preparation.h +++ b/tensorflow/compiler/xla/service/cpu/cpu_parallelization_preparation.h @@ -16,6 +16,7 @@ limitations under the License. #ifndef TENSORFLOW_COMPILER_XLA_SERVICE_CPU_CPU_PARALLELIZATION_PREPARATION_H_ #define TENSORFLOW_COMPILER_XLA_SERVICE_CPU_CPU_PARALLELIZATION_PREPARATION_H_ +#include "tensorflow/compiler/xla/service/hlo_cost_analysis.h" #include "tensorflow/compiler/xla/service/hlo_module.h" #include "tensorflow/compiler/xla/service/hlo_pass_interface.h" @@ -32,14 +33,51 @@ namespace cpu { // handle While constructs. class ParallelizationPreparation : public HloPassInterface { public: + // 'max_parallelism': the maximum parallel task count per instruction. + // 'shape_size': shape size function used by HloCostAnalysis during parallel + // task assignment. + ParallelizationPreparation( + const int64 max_parallelism, + const HloCostAnalysis::ShapeSizeFunction& shape_size) + : max_parallelism_(max_parallelism), shape_size_(shape_size) {} ~ParallelizationPreparation() override {} + tensorflow::StringPiece name() const override { return "cpu-parallel-prepare"; } - // Run instruction fusion on the given computation. Returns whether the + // Run parallel preparation on the given computation. Returns whether the // computation was changed. StatusOr Run(HloModule* module) override; + + private: + // Assigns parallel task partitions to conformant instructions in 'module'. + // Returns true on success or error status otherwise. + StatusOr RunParallelTaskAssignment(HloModule* module); + + // Returns the target parallel task count for 'instruction'. + // Utilizes 'cost_analysis' if non-null. + // Otherwise defaults to a simple HLO output size-based cost model. + int64 GetTargetParallelTaskCount(const HloCostAnalysis* cost_analysis, + HloInstruction* instruction); + + // Outlines 'instruction' from entry computation, if it had + // been assigned parallel tasks in an earlier pass through the computation. + // Returns true if 'instruction' was successfully outlined, false otherwise. + bool OutlineParallelizableInstruction(HloInstruction* instruction); + + // Returns true if 'instruction' can be outlined into the same sub-computation + // with its single user (parallelizable instructions are not outlined with + // each other). Returns false otherwise. + bool CanOutlineWithUser(HloInstruction* instruction); + + // Returns true if 'instruction' (or the root of the sub-computation that + // 'instruction' calls) has had parallel tasks assigned in earlier pass. + // Returns false otherwise. + bool AssignedParallelTasks(HloInstruction* instruction); + + const int64 max_parallelism_; + const HloCostAnalysis::ShapeSizeFunction shape_size_; }; } // namespace cpu diff --git a/tensorflow/compiler/xla/service/cpu/cpu_runtime.cc b/tensorflow/compiler/xla/service/cpu/cpu_runtime.cc index 8e06f0520edfb05c7ec606dcb8e85c5ef997c2c0..c7155b858bda5e5640e9a6719fb394ca1360d128 100644 --- a/tensorflow/compiler/xla/service/cpu/cpu_runtime.cc +++ b/tensorflow/compiler/xla/service/cpu/cpu_runtime.cc @@ -15,38 +15,121 @@ limitations under the License. #include "tensorflow/compiler/xla/service/cpu/cpu_runtime.h" -#include #include +#include "tensorflow/compiler/xla/service/llvm_ir/llvm_util.h" #include "tensorflow/core/platform/logging.h" #include "tensorflow/core/platform/macros.h" +#include "tensorflow/core/platform/types.h" namespace xla { namespace cpu { namespace runtime { -InfeedManager* GetInfeedManager() { - static InfeedManager* manager = new InfeedManager; +XfeedManager* GetXfeedManager() { + static XfeedManager* manager = new XfeedManager; return manager; } +extern const char* const kEigenMatMulF32SymbolName = + "__xla_cpu_runtime_EigenMatMulF32"; +extern const char* const kEigenMatMulF64SymbolName = + "__xla_cpu_runtime_EigenMatMulF64"; +extern const char* const kEigenConvF32SymbolName = + "__xla_cpu_runtime_EigenConvF32"; +extern const char* const kEigenSingleThreadedMatMulF32SymbolName = + "__xla_cpu_runtime_EigenSingleThreadedMatMulF32"; +extern const char* const kEigenSingleThreadedMatMulF64SymbolName = + "__xla_cpu_runtime_EigenSingleThreadedMatMulF64"; +extern const char* const kEigenSingleThreadedConvF32SymbolName = + "__xla_cpu_runtime_EigenSingleThreadedConvF32"; +extern const char* const kAcquireInfeedBufferForDequeueSymbolName = + "__xla_cpu_runtime_AcquireInfeedBufferForDequeue"; +extern const char* const kReleaseInfeedBufferAfterDequeueSymbolName = + "__xla_cpu_runtime_ReleaseInfeedBufferAfterDequeue"; +extern const char* const kAcquireOutfeedBufferForPopulationSymbolName = + "__xla_cpu_runtime_AcquireOutfeedBufferForPopulation"; +extern const char* const kReleaseOutfeedBufferAfterPopulationSymbolName = + "__xla_cpu_runtime_ReleaseOutfeedBufferAfterPopulation"; +extern const char* const kXlaCpuRuntimeSymbolNamePrefix = "__xla_cpu_runtime_"; } // namespace runtime } // namespace cpu } // namespace xla -void* __xla_cpu_runtime_AcquireInfeedBufferForDequeue( - xla::int32 buffer_length) { - xla::cpu::runtime::InfeedManager* infeed = - xla::cpu::runtime::GetInfeedManager(); +namespace { + +tensorflow::string ShapeString(const void* shape_ptr, xla::int32 shape_length) { + xla::StatusOr shape = + xla::llvm_ir::DecodeSelfDescribingShapeConstant(shape_ptr, shape_length); + if (shape.ok()) { + return xla::ShapeUtil::HumanStringWithLayout(shape.ValueOrDie()); + } + return ""; +} + +} // namespace + +void* __xla_cpu_runtime_AcquireInfeedBufferForDequeue(xla::int32 buffer_length, + const void* shape, + xla::int32 shape_length) { + if (VLOG_IS_ON(2)) { + LOG(INFO) << "AcquireInfeedBufferForDequeue: " + << ShapeString(shape, shape_length); + } + xla::cpu::runtime::XfeedManager* xfeed = xla::cpu::runtime::GetXfeedManager(); + // Wait until there's a buffer to dequeue. + xla::cpu::runtime::XfeedBuffer* buffer = + xfeed->infeed()->BlockingDequeueBuffer(); + CHECK_EQ(buffer->length(), buffer_length) + << "XLA program infeed request buffer size " << buffer_length + << " did not match the runtime's infed buffer length " << buffer->length() + << "; program reports desired shape: " + << ShapeString(shape, shape_length); + return buffer->data(); +} + +void __xla_cpu_runtime_ReleaseInfeedBufferAfterDequeue( + xla::int32 buffer_length, void* buffer_ptr, const void* shape_ptr, + xla::int32 shape_length) { + if (VLOG_IS_ON(2)) { + LOG(INFO) << "ReleaseInfeedBufferAfterDeque: " + << ShapeString(shape_ptr, shape_length); + } + xla::cpu::runtime::XfeedManager* xfeed = xla::cpu::runtime::GetXfeedManager(); + xla::StatusOr shape = + xla::llvm_ir::DecodeSelfDescribingShapeConstant(shape_ptr, shape_length); + xfeed->infeed()->ReleaseCurrentBuffer(buffer_length, buffer_ptr, + std::move(shape)); +} + +void* __xla_cpu_runtime_AcquireOutfeedBufferForPopulation( + xla::int32 buffer_length, const void* shape_ptr, xla::int32 shape_length) { + if (VLOG_IS_ON(2)) { + LOG(INFO) << "AcquireOutfeedBufferForPopulation: " + << ShapeString(shape_ptr, shape_length); + } + xla::cpu::runtime::XfeedManager* xfeed = xla::cpu::runtime::GetXfeedManager(); // Wait until there's a buffer to dequeue. - xla::cpu::runtime::InfeedBuffer* buffer = infeed->BlockingDequeueBuffer(); - CHECK_EQ(buffer->length(), buffer_length); + xla::cpu::runtime::XfeedBuffer* buffer = + xfeed->outfeed()->BlockingDequeueBuffer(); + CHECK_EQ(buffer->length(), buffer_length) + << "XLA program outfeed request buffer size " << buffer_length + << " did not match the runtime's outfeed buffer length " + << buffer->length() << "; program reports outfed shape: " + << ShapeString(shape_ptr, shape_length); return buffer->data(); } -void __xla_cpu_runtime_ReleaseInfeedBufferAfterDequeue(xla::int32 buffer_length, - void* buffer_ptr) { - xla::cpu::runtime::InfeedManager* infeed = - xla::cpu::runtime::GetInfeedManager(); - infeed->ReleaseCurrentBuffer(buffer_length, buffer_ptr); +void __xla_cpu_runtime_ReleaseOutfeedBufferAfterPopulation( + xla::int32 buffer_length, void* buffer_ptr, const void* shape_ptr, + xla::int32 shape_length) { + if (VLOG_IS_ON(2)) { + LOG(INFO) << "ReleaseOutfeedBufferAfterPopulation: " + << ShapeString(shape_ptr, shape_length); + } + xla::cpu::runtime::XfeedManager* xfeed = xla::cpu::runtime::GetXfeedManager(); + xla::StatusOr shape = + xla::llvm_ir::DecodeSelfDescribingShapeConstant(shape_ptr, shape_length); + xfeed->outfeed()->ReleaseCurrentBuffer(buffer_length, buffer_ptr, + std::move(shape)); } diff --git a/tensorflow/compiler/xla/service/cpu/cpu_runtime.h b/tensorflow/compiler/xla/service/cpu/cpu_runtime.h index 8eae2102305a3898c244a356d383184139e9208e..29feb7267fe97f6876827b6cbfa6217a0cecf238 100644 --- a/tensorflow/compiler/xla/service/cpu/cpu_runtime.h +++ b/tensorflow/compiler/xla/service/cpu/cpu_runtime.h @@ -26,7 +26,7 @@ limitations under the License. #ifndef TENSORFLOW_COMPILER_XLA_SERVICE_CPU_CPU_RUNTIME_H_ #define TENSORFLOW_COMPILER_XLA_SERVICE_CPU_CPU_RUNTIME_H_ -#include "tensorflow/compiler/xla/service/cpu/infeed_manager.h" +#include "tensorflow/compiler/xla/service/cpu/xfeed_manager.h" #include "tensorflow/compiler/xla/types.h" namespace xla { @@ -41,22 +41,23 @@ namespace runtime { // the actual symbol. // 2. When using ahead-of-time compilation, the linker can resolve the name // because it is a symbol in the cpu_runtime library. -constexpr char kEigenMatmulF32SymbolName[] = "__xla_cpu_runtime_EigenMatMulF32"; -constexpr char kEigenMatmulF64SymbolName[] = "__xla_cpu_runtime_EigenMatMulF64"; -constexpr char kEigenConvF32SymbolName[] = "__xla_cpu_runtime_EigenConvF32"; -constexpr char kEigenSingleThreadedMatmulF32SymbolName[] = - "__xla_cpu_runtime_EigenSingleThreadedMatMulF32"; -constexpr char kEigenSingleThreadedMatmulF64SymbolName[] = - "__xla_cpu_runtime_EigenSingleThreadedMatMulF64"; -constexpr char kEigenSingleThreadedConvF32SymbolName[] = - "__xla_cpu_runtime_EigenSingleThreadedConvF32"; -constexpr char kAcquireInfeedBufferForDequeueSymbolName[] = - "__xla_cpu_runtime_AcquireInfeedBufferForDequeue"; -constexpr char kReleaseInfeedBufferAfterDequeueSymbolName[] = - "__xla_cpu_runtime_ReleaseInfeedBufferAfterDequeue"; +extern const char* const kEigenMatMulF32SymbolName; +extern const char* const kEigenMatMulF64SymbolName; +extern const char* const kEigenConvF32SymbolName; +extern const char* const kEigenSingleThreadedMatMulF32SymbolName; +extern const char* const kEigenSingleThreadedMatMulF64SymbolName; +extern const char* const kEigenSingleThreadedConvF32SymbolName; +extern const char* const kAcquireInfeedBufferForDequeueSymbolName; +extern const char* const kReleaseInfeedBufferAfterDequeueSymbolName; +extern const char* const kAcquireOutfeedBufferForPopulationSymbolName; +extern const char* const kReleaseOutfeedBufferAfterPopulationSymbolName; + +// All symbol names for XLA CPU runtime functions need to start with this +// prefix. +extern const char* const kXlaCpuRuntimeSymbolNamePrefix; // Returns the infeed manager used by the CPU runtime. -InfeedManager* GetInfeedManager(); +XfeedManager* GetXfeedManager(); } // namespace runtime } // namespace cpu @@ -64,13 +65,19 @@ InfeedManager* GetInfeedManager(); extern "C" { +// Note: in the runtime entry points below, the shape pointer and shape_length +// reflect values that can be deserialized via +// llvm_ir::DecodeSelfDescribingShapeConstant. This is the way we pass reified +// type information from the generated program to the runtime, which helps check +// the type safety and contract for the emitted-code/runtime communication. + // Blocks until the next infeed buffer is ready to be dequeued, then // returns it. Fails catastrophically if the next enqueued buffer is // not of the correct length in bytes. Checking the shape rather than // the length would be more exact, but the length check is chosen as a // tradeoff between error checking and speed/simplicity. extern void* __xla_cpu_runtime_AcquireInfeedBufferForDequeue( - xla::int32 buffer_length); + xla::int32 buffer_length, const void* shape, xla::int32 shape_length); // Relinquishes the next infeed buffer that was returned by // __xla_cpu_runtime_AcquireInfeedBufferForDequeue. Once this call @@ -85,7 +92,27 @@ extern void* __xla_cpu_runtime_AcquireInfeedBufferForDequeue( // implemented we will add support for multiple outstanding buffers // that can be returned out of order. extern void __xla_cpu_runtime_ReleaseInfeedBufferAfterDequeue( - xla::int32 buffer_length, void* buffer_ptr); -} + xla::int32 buffer_length, void* buffer_ptr, const void* shape_ptr, + xla::int32 shape_length); + +// Blocks until the next outfeed buffer is available to be populated, then +// returns it. +extern void* __xla_cpu_runtime_AcquireOutfeedBufferForPopulation( + xla::int32 buffer_length, const void* shape_ptr, xla::int32 shape_length); + +// Relinquishes the outfeed buffer after it has been populated. +// buffer_ptr must have been previously returned by +// __xla_cpu_runtime_AcquireOutfeedBufferForPopulation. +// Once this call completes, buffer_ptr may no longer be accessed. +// buffer_length must match the length passed to the call to +// __xla_cpu_runtime_AcquireInfeedBufferForDequeue that returned +// buffer_ptr. This function must be called before the next buffer is +// acquired, i.e., there may only be one outstanding outfeed buffer in +// use by the runtime. +extern void __xla_cpu_runtime_ReleaseOutfeedBufferAfterPopulation( + xla::int32 buffer_length, void* buffer_ptr, const void* shape_ptr, + xla::int32 shape_length); + +} // extern "C" #endif // TENSORFLOW_COMPILER_XLA_SERVICE_CPU_CPU_RUNTIME_H_ diff --git a/tensorflow/compiler/xla/service/cpu/cpu_runtime_avx.cc b/tensorflow/compiler/xla/service/cpu/cpu_runtime_avx.cc index 646254887c83fcaff8fd5def9fafc8ff17d03d32..181deedde71bab3cb9ef1820a88de557131b9311 100644 --- a/tensorflow/compiler/xla/service/cpu/cpu_runtime_avx.cc +++ b/tensorflow/compiler/xla/service/cpu/cpu_runtime_avx.cc @@ -19,17 +19,24 @@ limitations under the License. #include "third_party/eigen3/Eigen/Core" +#ifdef __AVX__ +xla::cpu::runtime::V8F32AVX __xla_cpu_runtime_ExpV8F32AVX( + xla::cpu::runtime::V8F32AVX x) { + return Eigen::internal::pexp(x); +} + +xla::cpu::runtime::V8F32AVX __xla_cpu_runtime_LogV8F32AVX( + xla::cpu::runtime::V8F32AVX x) { + return Eigen::internal::plog(x); +} +#endif // __AVX__ + namespace xla { namespace cpu { namespace runtime { -#ifdef __AVX__ -V8F32 ExpV8F32(V8F32 x) { return Eigen::internal::pexp(x); } - -V8F32 LogV8F32(V8F32 x) { return Eigen::internal::plog(x); } - -V8F32 TanhV8F32(V8F32 x) { return Eigen::internal::ptanh(x); } -#endif // __AVX__ +const char *const kExpV8F32AVXSymbolName = "__xla_cpu_runtime_ExpV8F32AVX"; +const char *const kLogV8F32AVXSymbolName = "__xla_cpu_runtime_LogV8F32AVX"; } // namespace runtime } // namespace cpu diff --git a/tensorflow/compiler/xla/service/cpu/cpu_runtime_avx.h b/tensorflow/compiler/xla/service/cpu/cpu_runtime_avx.h index 89721aaf835eec5e4a8be0fbabb310b084065825..acfada8540d89bb098bb0b04e109441e2123e678 100644 --- a/tensorflow/compiler/xla/service/cpu/cpu_runtime_avx.h +++ b/tensorflow/compiler/xla/service/cpu/cpu_runtime_avx.h @@ -28,23 +28,24 @@ namespace xla { namespace cpu { namespace runtime { -constexpr char kExpV8F32[] = "__xla_cpu_runtime_ExpV8F32"; -constexpr char kLogV8F32[] = "__xla_cpu_runtime_LogV8F32"; -constexpr char kTanhV8F32[] = "__xla_cpu_runtime_TanhV8F32"; +extern const char *const kExpV8F32AVXSymbolName; +extern const char *const kLogV8F32AVXSymbolName; -typedef float V8F32 __attribute__((__vector_size__(32))); +typedef float V8F32AVX __attribute__((__vector_size__(32))); +} // namespace runtime +} // namespace cpu +} // namespace xla + +extern "C" { // The following functions are vectorized versions of a selection of libm // library functions. // References to these functions are created by the LLVM vectorizer. -V8F32 ExpV8F32(V8F32 x) TF_ATTRIBUTE_WEAK; - -V8F32 LogV8F32(V8F32 x) TF_ATTRIBUTE_WEAK; +xla::cpu::runtime::V8F32AVX __xla_cpu_runtime_ExpV8F32AVX( + xla::cpu::runtime::V8F32AVX x) TF_ATTRIBUTE_WEAK; -V8F32 TanhV8F32(V8F32 x) TF_ATTRIBUTE_WEAK; - -} // namespace runtime -} // namespace cpu -} // namespace xla +xla::cpu::runtime::V8F32AVX __xla_cpu_runtime_LogV8F32AVX( + xla::cpu::runtime::V8F32AVX x) TF_ATTRIBUTE_WEAK; +} #endif // TENSORFLOW_COMPILER_XLA_SERVICE_CPU_CPU_RUNTIME_AVX_H_ diff --git a/tensorflow/compiler/xla/service/cpu/cpu_runtime_neon.cc b/tensorflow/compiler/xla/service/cpu/cpu_runtime_neon.cc new file mode 100644 index 0000000000000000000000000000000000000000..abe792b2787ce8baf56ee62585a0ab886d922a23 --- /dev/null +++ b/tensorflow/compiler/xla/service/cpu/cpu_runtime_neon.cc @@ -0,0 +1,46 @@ +/* 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/cpu/cpu_runtime_neon.h" + +#define EIGEN_USE_THREADS + +#include "third_party/eigen3/Eigen/Core" + +#ifdef __ARM_NEON__ + +xla::cpu::runtime::V4F32NEON __xla_cpu_runtime_ExpV4F32NEON( + xla::cpu::runtime::V4F32NEON x) { + return Eigen::internal::pexp(x); +} + +xla::cpu::runtime::V4F32NEON __xla_cpu_runtime_LogV4F32NEON( + xla::cpu::runtime::V4F32NEON x) { + Eigen::internal::Packet4f p = x; + return Eigen::internal::plog(p); +} + +#endif // __ARM_NEON__ + +namespace xla { +namespace cpu { +namespace runtime { + +const char *const kExpV4F32NEONSymbolName = "__xla_cpu_runtime_ExpV4F32NEON"; +const char *const kLogV4F32NEONSymbolName = "__xla_cpu_runtime_LogV4F32NEON"; + +} // namespace runtime +} // namespace cpu +} // namespace xla diff --git a/tensorflow/compiler/xla/service/cpu/cpu_runtime_neon.h b/tensorflow/compiler/xla/service/cpu/cpu_runtime_neon.h new file mode 100644 index 0000000000000000000000000000000000000000..75cb16b273973d2bf665d378084343fd612a2941 --- /dev/null +++ b/tensorflow/compiler/xla/service/cpu/cpu_runtime_neon.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_CPU_CPU_RUNTIME_NEON_H_ +#define TENSORFLOW_COMPILER_XLA_SERVICE_CPU_CPU_RUNTIME_NEON_H_ + +// This header declares functions which may be called by the generated code on +// the CPU. Calls to these functions must be resolved explicitly in the JIT in +// xla::cpu::SimpleResolver. + +#include "tensorflow/core/platform/macros.h" + +#ifdef __ARM_NEON__ +// For the other runtimes (AVX, SSE4.1) we define the vector type directly using +// __attribute__((__vector_size__(*))). Unfortunately, the typedef for the ARM +// NEON SIMD types is not portable, so the type has to come from +#include +#endif // __ARM_NEON__ + +namespace xla { +namespace cpu { +namespace runtime { + +extern const char *const kExpV4F32NEONSymbolName; +extern const char *const kLogV4F32NEONSymbolName; + +#ifdef __ARM_NEON__ +typedef float32x4_t V4F32NEON; +#else +// On non-ARM platforms ensure the declaration is present +struct V4F32NEON; +#endif // __ARM_NEON__ + +} // namespace runtime +} // namespace cpu +} // namespace xla + +extern "C" { + +// The following functions are vectorized versions of a selection of libm +// library functions. +// References to these functions are created by the LLVM vectorizer. +xla::cpu::runtime::V4F32NEON __xla_cpu_runtime_ExpV4F32NEON( + xla::cpu::runtime::V4F32NEON x) TF_ATTRIBUTE_WEAK; + +xla::cpu::runtime::V4F32NEON __xla_cpu_runtime_LogV4F32NEON( + xla::cpu::runtime::V4F32NEON x) TF_ATTRIBUTE_WEAK; +} + +#endif // TENSORFLOW_COMPILER_XLA_SERVICE_CPU_CPU_RUNTIME_NEON_H_ diff --git a/tensorflow/compiler/xla/service/cpu/cpu_runtime_sse4_1.cc b/tensorflow/compiler/xla/service/cpu/cpu_runtime_sse4_1.cc index 69d04427c60b0d8db8a8859b4abff9bfa7e93260..a9a45db5a424d2faecbd437542c41fbd7fdf0bb8 100644 --- a/tensorflow/compiler/xla/service/cpu/cpu_runtime_sse4_1.cc +++ b/tensorflow/compiler/xla/service/cpu/cpu_runtime_sse4_1.cc @@ -19,29 +19,29 @@ limitations under the License. #include "third_party/eigen3/Eigen/Core" -namespace xla { -namespace cpu { -namespace runtime { - #ifdef __SSE4_1__ -V4F32 ExpV4F32(V4F32 x) { +xla::cpu::runtime::V4F32SSE __xla_cpu_runtime_ExpV4F32SSE( + xla::cpu::runtime::V4F32SSE x) { Eigen::internal::Packet4f p = x; return Eigen::internal::pexp(p); } -V4F32 LogV4F32(V4F32 x) { +xla::cpu::runtime::V4F32SSE __xla_cpu_runtime_LogV4F32SSE( + xla::cpu::runtime::V4F32SSE x) { Eigen::internal::Packet4f p = x; return Eigen::internal::plog(p); } -V4F32 TanhV4F32(V4F32 x) { - Eigen::internal::Packet4f p = x; - return Eigen::internal::ptanh(p); -} - #endif // __SSE4_1__ +namespace xla { +namespace cpu { +namespace runtime { + +const char *const kExpV4F32SSESymbolName = "__xla_cpu_runtime_ExpV4F32SSE"; +const char *const kLogV4F32SSESymbolName = "__xla_cpu_runtime_LogV4F32SSE"; + } // namespace runtime } // namespace cpu } // namespace xla diff --git a/tensorflow/compiler/xla/service/cpu/cpu_runtime_sse4_1.h b/tensorflow/compiler/xla/service/cpu/cpu_runtime_sse4_1.h index ded206f90a076ba81643799c07e3f3a7d481eaf2..96587d10d2b86e14ff6a7400fdf14ca0d994ddc5 100644 --- a/tensorflow/compiler/xla/service/cpu/cpu_runtime_sse4_1.h +++ b/tensorflow/compiler/xla/service/cpu/cpu_runtime_sse4_1.h @@ -28,23 +28,25 @@ namespace xla { namespace cpu { namespace runtime { -constexpr char kExpV4F32[] = "__xla_cpu_runtime_ExpV4F32"; -constexpr char kLogV4F32[] = "__xla_cpu_runtime_LogV4F32"; -constexpr char kTanhV4F32[] = "__xla_cpu_runtime_TanhV4F32"; +extern const char *const kExpV4F32SSESymbolName; +extern const char *const kLogV4F32SSESymbolName; -typedef float V4F32 __attribute__((__vector_size__(16))); +typedef float V4F32SSE __attribute__((__vector_size__(16))); + +} // namespace runtime +} // namespace cpu +} // namespace xla + +extern "C" { // The following functions are vectorized versions of a selection of libm // library functions. // References to these functions are created by the LLVM vectorizer. -V4F32 ExpV4F32(V4F32 x) TF_ATTRIBUTE_WEAK; +xla::cpu::runtime::V4F32SSE __xla_cpu_runtime_ExpV4F32SSE( + xla::cpu::runtime::V4F32SSE x) TF_ATTRIBUTE_WEAK; -V4F32 LogV4F32(V4F32 x) TF_ATTRIBUTE_WEAK; - -V4F32 TanhV4F32(V4F32 x) TF_ATTRIBUTE_WEAK; - -} // namespace runtime -} // namespace cpu -} // namespace xla +xla::cpu::runtime::V4F32SSE __xla_cpu_runtime_LogV4F32SSE( + xla::cpu::runtime::V4F32SSE x) TF_ATTRIBUTE_WEAK; +} #endif // TENSORFLOW_COMPILER_XLA_SERVICE_CPU_CPU_RUNTIME_SSE4_1_H_ diff --git a/tensorflow/compiler/xla/service/cpu/cpu_runtime_test.cc b/tensorflow/compiler/xla/service/cpu/cpu_runtime_test.cc index 52eed7dbad27e01f93de67e0f6e838de7dc7e1c5..f8e260dd90149405fff7beefba3f7fe83b75d4b6 100644 --- a/tensorflow/compiler/xla/service/cpu/cpu_runtime_test.cc +++ b/tensorflow/compiler/xla/service/cpu/cpu_runtime_test.cc @@ -17,6 +17,7 @@ limitations under the License. #include #include +#include #define EIGEN_USE_THREADS @@ -25,8 +26,10 @@ limitations under the License. #include "tensorflow/compiler/xla/client/local_client.h" #include "tensorflow/compiler/xla/ptr_util.h" #include "tensorflow/compiler/xla/service/cpu/runtime_matmul.h" +#include "tensorflow/compiler/xla/service/cpu/runtime_single_threaded_matmul.h" #include "tensorflow/compiler/xla/types.h" #include "tensorflow/core/common_runtime/eigen_thread_pool.h" +#include "tensorflow/core/lib/strings/stringprintf.h" #include "tensorflow/core/platform/env.h" #include "tensorflow/core/platform/logging.h" #include "tensorflow/core/platform/test.h" @@ -75,14 +78,8 @@ void CheckMatrixMultiply(const Array2D& a, const Array2D& b, std::unique_ptr> EigenMatrixMultiply(const Array2D& a, const Array2D& b, bool transpose_lhs, - bool transpose_rhs) { - tensorflow::thread::ThreadPool pool(tensorflow::Env::Default(), "XLAEigen", - 2); - tensorflow::EigenThreadPoolWrapper tp(&pool); - Eigen::ThreadPoolDevice device(&tp, tp.NumThreads()); - ExecutableRunOptions run_options; - run_options.set_intra_op_thread_pool(&device); - + bool transpose_rhs, + bool single_threaded) { CHECK_EQ(a.width(), b.height()); int64 m = a.height(); int64 n = b.width(); @@ -98,41 +95,81 @@ std::unique_ptr> EigenMatrixMultiply(const Array2D& a, // Since we're going to transpose c before returning it. Swap the order of the // dimension sizes to ensure the returned array is properly dimensioned. auto c_transpose = MakeUnique>(n, m); - __xla_cpu_runtime_EigenMatMulF32(&run_options, c_transpose->data(), - a_transpose->data(), b_transpose->data(), m, - n, k, transpose_lhs, transpose_rhs); + if (single_threaded) { + __xla_cpu_runtime_EigenSingleThreadedMatMulF32( + nullptr, c_transpose->data(), a_transpose->data(), b_transpose->data(), + m, n, k, transpose_lhs, transpose_rhs); + } else { + tensorflow::thread::ThreadPool pool(tensorflow::Env::Default(), "XLAEigen", + 2); + tensorflow::EigenThreadPoolWrapper tp(&pool); + Eigen::ThreadPoolDevice device(&tp, tp.NumThreads()); + ExecutableRunOptions run_options; + run_options.set_intra_op_thread_pool(&device); + + __xla_cpu_runtime_EigenMatMulF32(&run_options, c_transpose->data(), + a_transpose->data(), b_transpose->data(), + m, n, k, transpose_lhs, transpose_rhs); + } return MaybeTransposeArray2D(*c_transpose, true); } -TEST_F(CpuRuntimeTest, SmallEigenMatmul) { - Array2D a({{1.0f, 2.0f}, {3.0f, 4.0f}}); - Array2D b({{5.0f, -1.0f, 3.0f}, {2.0f, 6.0f, 4.0f}}); - - for (bool transpose_lhs : {false, true}) { - for (bool transpose_rhs : {false, true}) { - auto c = EigenMatrixMultiply(a, b, transpose_lhs, transpose_rhs); - - LOG(INFO) << "a = " << a.ToString(); - LOG(INFO) << "b = " << b.ToString(); - LOG(INFO) << "c = " << c->ToString(); - - CheckMatrixMultiply(a, b, *c); - } +struct MatMulShape { + int64 m; + int64 k; + int64 n; +}; + +MatMulShape MatMulShapes[] = { + MatMulShape{2, 2, 3}, MatMulShape{256, 512, 1024}, + MatMulShape{128, 128, 1}, MatMulShape{1, 128, 128}, + MatMulShape{1, 32, 128}, MatMulShape{1, 32, 16}, + MatMulShape{32, 16, 1}, MatMulShape{32, 128, 1}, +}; + +// This takes 4 parameters: +// * shape of the matmul +// * transpose_lhs +// * transpose_rhs +// * single_threaded +using EigenMatMulTestParam = std::tuple; + +class EigenMatMulTest + : public CpuRuntimeTest, + public ::testing::WithParamInterface { + public: + static string Name( + const ::testing::TestParamInfo& info) { + MatMulShape shape = std::get<0>(info.param); + bool transpose_lhs = std::get<1>(info.param); + bool transpose_rhs = std::get<2>(info.param); + bool single_threaded = std::get<3>(info.param); + + return tensorflow::strings::Printf( + "MatMul_%lld_%lld_%lld_%s%s%s_threaded", shape.m, shape.k, shape.n, + transpose_lhs ? "Tlhs_" : "", transpose_rhs ? "Trhs_" : "", + single_threaded ? "single" : "multi"); } +}; // namespace xla + +TEST_P(EigenMatMulTest, DoIt) { + MatMulShape shape = std::get<0>(GetParam()); + bool transpose_lhs = std::get<1>(GetParam()); + bool transpose_rhs = std::get<2>(GetParam()); + bool single_threaded = std::get<3>(GetParam()); + + auto a = MakeLinspaceArray2D(0.0, 1.0, shape.m, shape.k); + auto b = MakeLinspaceArray2D(-2.0, 2.0, shape.k, shape.n); + auto c = EigenMatrixMultiply(*a, *b, transpose_lhs, transpose_rhs, + single_threaded); + CheckMatrixMultiply(*a, *b, *c); } -TEST_F(CpuRuntimeTest, LargeEigenMatmul) { - auto a = MakeLinspaceArray2D(0.0, 1.0, 256, 512); - auto b = MakeLinspaceArray2D(-2.0, 2.0, 512, 1024); - - for (bool transpose_lhs : {false, true}) { - for (bool transpose_rhs : {false, true}) { - auto c = EigenMatrixMultiply(*a, *b, transpose_lhs, transpose_rhs); - - CheckMatrixMultiply(*a, *b, *c); - } - } -} +INSTANTIATE_TEST_CASE_P(EigenMatMulTestInstantiaion, EigenMatMulTest, + ::testing::Combine(::testing::ValuesIn(MatMulShapes), + ::testing::Bool(), ::testing::Bool(), + ::testing::Bool()), + EigenMatMulTest::Name); } // namespace } // namespace xla diff --git a/tensorflow/compiler/xla/service/cpu/disassembler.cc b/tensorflow/compiler/xla/service/cpu/disassembler.cc index f0dcce56b40d5b4021d1c1396ce96cfa3aa6e813..dbd2860cc8aec79b3804b561e0d48713e07a7aa2 100644 --- a/tensorflow/compiler/xla/service/cpu/disassembler.cc +++ b/tensorflow/compiler/xla/service/cpu/disassembler.cc @@ -21,9 +21,9 @@ limitations under the License. #include #include -#include "external/llvm/include/llvm/MC/MCInst.h" -#include "external/llvm/include/llvm/Support/TargetRegistry.h" -#include "external/llvm/include/llvm/Support/raw_ostream.h" +#include "llvm/MC/MCInst.h" +#include "llvm/Support/TargetRegistry.h" +#include "llvm/Support/raw_ostream.h" #include "tensorflow/compiler/xla/status_macros.h" #include "tensorflow/compiler/xla/types.h" #include "tensorflow/compiler/xla/util.h" @@ -52,7 +52,7 @@ Disassembler::Disassembler(const llvm::TargetMachine& target_machine) } // This code is based on llvm-objdump in llvm/tools. -StatusOr Disassembler::DisassembleObjectFile( +StatusOr Disassembler::DisassembleObjectFile( const llvm::object::ObjectFile& object_file) const { if (disassembler_ == nullptr) { return NotFound("could not find a disassembler for this platform"); @@ -60,6 +60,7 @@ StatusOr Disassembler::DisassembleObjectFile( std::string buffer_string; llvm::raw_string_ostream ostream(buffer_string); + uint64_t code_size_bytes = 0; // Iterate through sections. Disassemble symbols of the text section(s). for (auto& section : object_file.sections()) { @@ -131,6 +132,9 @@ StatusOr Disassembler::DisassembleObjectFile( TF_RET_CHECK(name_or_error); ostream << name_or_error.get().str() << ":\n"; + // Update the code size statistic. + code_size_bytes += end_index - start_index; + // Disassemble symbol instruction-by-instruction. uint64_t index = start_index; while (index < end_index) { @@ -175,7 +179,8 @@ StatusOr Disassembler::DisassembleObjectFile( } ostream.flush(); - return string(buffer_string.data(), buffer_string.length()); + return DisassemblerResult{ + string(buffer_string.data(), buffer_string.length()), code_size_bytes}; } } // namespace cpu diff --git a/tensorflow/compiler/xla/service/cpu/disassembler.h b/tensorflow/compiler/xla/service/cpu/disassembler.h index e90f26fc827b6a96599a8d0c486dd4f23a4851d0..b6feaa7e45cee26eb7f850081bd1fad2cb63b15c 100644 --- a/tensorflow/compiler/xla/service/cpu/disassembler.h +++ b/tensorflow/compiler/xla/service/cpu/disassembler.h @@ -19,33 +19,43 @@ limitations under the License. #include #include -#include "external/llvm/include/llvm/MC/MCContext.h" -#include "external/llvm/include/llvm/MC/MCDisassembler/MCDisassembler.h" -#include "external/llvm/include/llvm/MC/MCInstPrinter.h" -#include "external/llvm/include/llvm/MC/MCInstrAnalysis.h" -#include "external/llvm/include/llvm/MC/MCObjectFileInfo.h" -#include "external/llvm/include/llvm/MC/MCSubtargetInfo.h" -#include "external/llvm/include/llvm/Object/ObjectFile.h" -#include "external/llvm/include/llvm/Target/TargetMachine.h" +#include "llvm/MC/MCContext.h" +#include "llvm/MC/MCDisassembler/MCDisassembler.h" +#include "llvm/MC/MCInstPrinter.h" +#include "llvm/MC/MCInstrAnalysis.h" +#include "llvm/MC/MCObjectFileInfo.h" +#include "llvm/MC/MCSubtargetInfo.h" +#include "llvm/Object/ObjectFile.h" +#include "llvm/Target/TargetMachine.h" #include "tensorflow/compiler/xla/statusor.h" #include "tensorflow/compiler/xla/types.h" namespace xla { namespace cpu { +struct DisassemblerResult { + DisassemblerResult(const string& text, size_t code_size_bytes) + : text(text), code_size_bytes(code_size_bytes) {} + + // The dissassembled text sections of the object file. + string text; + // The total number of bytes of executable code in the object file. + uint64_t code_size_bytes; +}; + // Class for disassembling object files (and potentially other constructs) into -// X86 assembly. Builds all the LLVM disassembly and instruction printing +// x86/ARM assembly. Builds all the LLVM disassembly and instruction printing // constructs from a given TargetMachine. class Disassembler { public: explicit Disassembler(const llvm::TargetMachine& target_machine); - // Returns a string containing the disassembled text sections of the given - // object file. + // Returns a DisassemblerResult for the given object file, containing the + // disassembled code. // // If we couldnt' retrieve a disassembler for this platform, an error status // is returned. - StatusOr DisassembleObjectFile( + StatusOr DisassembleObjectFile( const llvm::object::ObjectFile& object_file) const; private: diff --git a/tensorflow/compiler/xla/service/cpu/dot_op_emitter.cc b/tensorflow/compiler/xla/service/cpu/dot_op_emitter.cc index 420f9cebc5b1ded365c20079589ebc79a03b3164..c9f0f115c09e92be014af567c99d446dbd265a4b 100644 --- a/tensorflow/compiler/xla/service/cpu/dot_op_emitter.cc +++ b/tensorflow/compiler/xla/service/cpu/dot_op_emitter.cc @@ -18,13 +18,13 @@ limitations under the License. #include #include -#include "external/llvm/include/llvm/IR/BasicBlock.h" -#include "external/llvm/include/llvm/IR/Instructions.h" -#include "external/llvm/include/llvm/IR/Module.h" -#include "external/llvm/include/llvm/IR/Value.h" -#include "tensorflow/compiler/xla/legacy_flags/cpu_runtime_flags.h" +#include "llvm/IR/BasicBlock.h" +#include "llvm/IR/Instructions.h" +#include "llvm/IR/Module.h" +#include "llvm/IR/Value.h" #include "tensorflow/compiler/xla/service/cpu/cpu_runtime.h" #include "tensorflow/compiler/xla/service/cpu/ir_emission_utils.h" +#include "tensorflow/compiler/xla/service/hlo_module.h" #include "tensorflow/compiler/xla/service/llvm_ir/llvm_util.h" #include "tensorflow/compiler/xla/shape_util.h" #include "tensorflow/compiler/xla/status_macros.h" @@ -44,7 +44,8 @@ DotOpEmitter::DotOpEmitter(const HloInstruction& dot, bool transpose_lhs, const llvm_ir::IrArray& lhs_array, const llvm_ir::IrArray& rhs_array, llvm::Value* executable_run_options_value, - llvm::IRBuilder<>* ir_builder) + llvm::IRBuilder<>* ir_builder, + const HloModuleConfig& hlo_module_config) : dot_(dot), transpose_lhs_(transpose_lhs), transpose_rhs_(transpose_rhs), @@ -52,18 +53,20 @@ DotOpEmitter::DotOpEmitter(const HloInstruction& dot, bool transpose_lhs, lhs_array_(lhs_array), rhs_array_(rhs_array), executable_run_options_value_(executable_run_options_value), - ir_builder_(ir_builder) {} + ir_builder_(ir_builder), + hlo_module_config_(hlo_module_config) {} /* static */ tensorflow::Status DotOpEmitter::EmitDotOperation( const HloInstruction& dot, bool transpose_lhs, bool transpose_rhs, 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<>* ir_builder) { + llvm::Value* executable_run_options_value, llvm::IRBuilder<>* ir_builder, + const HloModuleConfig& hlo_module_config) { PrimitiveType type = target_array.GetShape().element_type(); TF_RET_CHECK(F32 == type || F64 == type); DotOpEmitter dot_emitter(dot, transpose_lhs, transpose_rhs, target_array, lhs_array, rhs_array, executable_run_options_value, - ir_builder); + ir_builder, hlo_module_config); return dot_emitter.Emit(); } @@ -233,22 +236,22 @@ tensorflow::Status DotOpEmitter::EmitCallToRuntime() { // The two transpose_... parameters are actually booleans, but we use int32 // to avoid target-dependent calling convention details. - legacy_flags::CpuRuntimeFlags* flags = legacy_flags::GetCpuRuntimeFlags(); - bool multi_threaded = flags->xla_cpu_multi_thread_eigen; + bool multi_threaded_eigen = + hlo_module_config_.debug_options().xla_cpu_multi_thread_eigen(); PrimitiveType type = target_array_.GetShape().element_type(); llvm::Type* float_type; const char* fn_name; switch (type) { case F32: - fn_name = multi_threaded - ? runtime::kEigenMatmulF32SymbolName - : runtime::kEigenSingleThreadedMatmulF32SymbolName; + fn_name = multi_threaded_eigen + ? runtime::kEigenMatMulF32SymbolName + : runtime::kEigenSingleThreadedMatMulF32SymbolName; float_type = ir_builder_->getFloatTy(); break; case F64: - fn_name = multi_threaded - ? runtime::kEigenMatmulF64SymbolName - : runtime::kEigenSingleThreadedMatmulF64SymbolName; + fn_name = multi_threaded_eigen + ? runtime::kEigenMatMulF64SymbolName + : runtime::kEigenSingleThreadedMatMulF64SymbolName; float_type = ir_builder_->getDoubleTy(); break; default: @@ -280,6 +283,10 @@ tensorflow::Status DotOpEmitter::EmitCallToRuntime() { // // (A x B)^T = B^T x A^T // + // The connection between this identity and memory layout is that the + // transpose operation can also be considered as an operation that changes the + // memory layout of a matrix from row-major to column-major or vice versa. + // // Effectively this involves swapping the 'lhs' with 'rhs' and 'm' with 'n'. const Shape& lhs_shape = lhs_array_.GetShape(); diff --git a/tensorflow/compiler/xla/service/cpu/dot_op_emitter.h b/tensorflow/compiler/xla/service/cpu/dot_op_emitter.h index 44dfe5f2a91222d99907e31062fb1d8f74aed3ff..cfc10660453c822635d68270c053977fca779ee1 100644 --- a/tensorflow/compiler/xla/service/cpu/dot_op_emitter.h +++ b/tensorflow/compiler/xla/service/cpu/dot_op_emitter.h @@ -16,8 +16,9 @@ limitations under the License. #ifndef TENSORFLOW_COMPILER_XLA_SERVICE_CPU_DOT_OP_EMITTER_H_ #define TENSORFLOW_COMPILER_XLA_SERVICE_CPU_DOT_OP_EMITTER_H_ -#include "external/llvm/include/llvm/IR/IRBuilder.h" +#include "llvm/IR/IRBuilder.h" #include "tensorflow/compiler/xla/service/hlo_instruction.h" +#include "tensorflow/compiler/xla/service/hlo_module_config.h" #include "tensorflow/compiler/xla/service/llvm_ir/ir_array.h" #include "tensorflow/compiler/xla/service/llvm_ir/llvm_loop.h" #include "tensorflow/compiler/xla/types.h" @@ -39,7 +40,8 @@ class DotOpEmitter { const HloInstruction& dot, bool transpose_lhs, bool transpose_rhs, 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<>* ir_builder); + llvm::Value* executable_run_options_value, llvm::IRBuilder<>* ir_builder, + const HloModuleConfig& hlo_module_config); private: DotOpEmitter(const HloInstruction& dot, bool transpose_lhs, @@ -47,7 +49,8 @@ class DotOpEmitter { const llvm_ir::IrArray& lhs_array, const llvm_ir::IrArray& rhs_array, llvm::Value* executable_run_options_value, - llvm::IRBuilder<>* ir_builder); + llvm::IRBuilder<>* ir_builder, + const HloModuleConfig& hlo_module_config); // Emits the IR to perform the dot operation. tensorflow::Status Emit(); @@ -82,6 +85,7 @@ class DotOpEmitter { const llvm_ir::IrArray& rhs_array_; llvm::Value* executable_run_options_value_; llvm::IRBuilder<>* ir_builder_; + const HloModuleConfig& hlo_module_config_; }; } // namespace cpu diff --git a/tensorflow/compiler/xla/service/cpu/elemental_ir_emitter.cc b/tensorflow/compiler/xla/service/cpu/elemental_ir_emitter.cc index 9b46c35b4161e4e97fb5c68c0d0bba1b004aef84..fe447adf893d8c254828dba73d74856392042892 100644 --- a/tensorflow/compiler/xla/service/cpu/elemental_ir_emitter.cc +++ b/tensorflow/compiler/xla/service/cpu/elemental_ir_emitter.cc @@ -17,8 +17,8 @@ limitations under the License. #include -#include "external/llvm/include/llvm/IR/Instructions.h" -#include "external/llvm/include/llvm/IR/Module.h" +#include "llvm/IR/Instructions.h" +#include "llvm/IR/Module.h" #include "tensorflow/compiler/xla/service/hlo_opcode.h" #include "tensorflow/compiler/xla/service/llvm_ir/llvm_util.h" #include "tensorflow/compiler/xla/types.h" diff --git a/tensorflow/compiler/xla/service/cpu/elemental_ir_emitter.h b/tensorflow/compiler/xla/service/cpu/elemental_ir_emitter.h index 5160217674fc2c69e92220c3730fed7af03739d5..6f9d6a24b4c73a03560da6db373be19d5737fa91 100644 --- a/tensorflow/compiler/xla/service/cpu/elemental_ir_emitter.h +++ b/tensorflow/compiler/xla/service/cpu/elemental_ir_emitter.h @@ -16,9 +16,9 @@ limitations under the License. #ifndef TENSORFLOW_COMPILER_XLA_SERVICE_CPU_ELEMENTAL_IR_EMITTER_H_ #define TENSORFLOW_COMPILER_XLA_SERVICE_CPU_ELEMENTAL_IR_EMITTER_H_ -#include "external/llvm/include/llvm/IR/IRBuilder.h" -#include "external/llvm/include/llvm/IR/Module.h" -#include "external/llvm/include/llvm/IR/Value.h" +#include "llvm/IR/IRBuilder.h" +#include "llvm/IR/Module.h" +#include "llvm/IR/Value.h" #include "tensorflow/compiler/xla/service/elemental_ir_emitter.h" #include "tensorflow/compiler/xla/service/hlo_instruction.h" #include "tensorflow/compiler/xla/statusor.h" diff --git a/tensorflow/compiler/xla/service/cpu/infeed_manager.cc b/tensorflow/compiler/xla/service/cpu/infeed_manager.cc deleted file mode 100644 index 14c882a06ee9fdfc66f3d6db55146431634dd85e..0000000000000000000000000000000000000000 --- a/tensorflow/compiler/xla/service/cpu/infeed_manager.cc +++ /dev/null @@ -1,72 +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/service/cpu/infeed_manager.h" - -#include "tensorflow/core/platform/logging.h" - -namespace xla { -namespace cpu { -namespace runtime { - -InfeedBuffer::~InfeedBuffer() = default; - -InfeedManager::InfeedManager() : current_buffer_(nullptr) {} - -void InfeedManager::Reset() { - tensorflow::mutex_lock l(mu_); - CHECK(!current_buffer_); - for (auto buffer : enqueued_buffer_) { - buffer->Done(); - } - enqueued_buffer_.clear(); -} - -void InfeedManager::EnqueueBuffer(InfeedBuffer* buffer) { - tensorflow::mutex_lock l(mu_); - bool was_empty = enqueued_buffer_.empty(); - enqueued_buffer_.push_back(buffer); - if (was_empty) { - // This has the potential to suffer from the notified thread - // immediately trying and failing to acquire mu_, but seems - // preferable to the alternative of notifying outside the lock - // on every enqueue. - cv_.notify_one(); - } -} - -InfeedBuffer* InfeedManager::BlockingDequeueBuffer() { - tensorflow::mutex_lock l(mu_); - while (enqueued_buffer_.empty()) { - cv_.wait(l); - } - CHECK(!current_buffer_); - current_buffer_ = enqueued_buffer_.front(); - enqueued_buffer_.pop_front(); - return current_buffer_; -} - -void InfeedManager::ReleaseCurrentBuffer(int32 length, void* data) { - tensorflow::mutex_lock l(mu_); - CHECK(current_buffer_); - CHECK_EQ(length, current_buffer_->length()); - CHECK_EQ(data, current_buffer_->data()); - current_buffer_->Done(); - current_buffer_ = nullptr; -} - -} // namespace runtime -} // namespace cpu -} // namespace xla diff --git a/tensorflow/compiler/xla/service/cpu/infeed_manager.h b/tensorflow/compiler/xla/service/cpu/infeed_manager.h deleted file mode 100644 index 77472746e659b2ddbd9b54a036775ebdd0084fdd..0000000000000000000000000000000000000000 --- a/tensorflow/compiler/xla/service/cpu/infeed_manager.h +++ /dev/null @@ -1,92 +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 header declares the abstract class for the infeed manager that -// is used by the CPU runtime to transfer buffers into an executing -// CPU computation, e.g., to feed data into a while loop. - -#ifndef TENSORFLOW_COMPILER_XLA_SERVICE_CPU_INFEED_MANAGER_H_ -#define TENSORFLOW_COMPILER_XLA_SERVICE_CPU_INFEED_MANAGER_H_ - -#include - -#include "tensorflow/compiler/xla/types.h" -#include "tensorflow/core/platform/mutex.h" - -namespace xla { -namespace cpu { -namespace runtime { - -// Abstract class defining an infeed buffer that is passed to the -// runtime by a client. The client manages the storage of the buffer. -class InfeedBuffer { - public: - virtual ~InfeedBuffer(); - - virtual int32 length() = 0; - virtual void* data() = 0; - virtual void Done() = 0; -}; - -// Client-side class used to enqueue infeed buffers. -class InfeedManager { - public: - InfeedManager(); - - // Calls the completion callback for any enqueued buffers that have - // not been dequeued by the runtime, and empties the infeed - // queue. Reset may not be called while a runtime computation is - // processing a dequeued buffer. The only safe way to ensure this - // condition is to call Reset when no computation is taking place. - void Reset(); - - // Adds buffer to the infeed queue. buffer->Done will be called when - // the buffer will no longer be accessed by the InfeedManager, - // either as a result of a call to Reset or because the runtime has - // dequeued and used the buffer. - void EnqueueBuffer(InfeedBuffer* buffer); - - // Blocks until the infeed queue is non-empty, then returns the - // buffer at the head of the queue. Sets the current buffer to be - // the returned buffer. It is an error to call BlockingDequeueBuffer - // if there is an unreleased current buffer, i.e., - // ReleaseCurrentBuffer must be called between calls to - // BlockingDequeueBuffer. - InfeedBuffer* BlockingDequeueBuffer(); - - // Releases the current buffer, which is the last buffer returned by - // BlockingDequeuBuffer and not yet released. length and data must - // match the buffer->length() and buffer->data() for the current - // buffer. - void ReleaseCurrentBuffer(int32 length, void* data); - - private: - tensorflow::mutex mu_; - // Condition variable that is signaled every time a buffer is - // enqueued to an empty queue. - tensorflow::condition_variable cv_; - // InfeedBuffer* queue contents are not owned, but buffer->Done must - // be called when the buffer is no longer needed by the runtime. - std::deque enqueued_buffer_; - // If non-NULL, the buffer that is currently being processed by the - // runtime. Not owned. - InfeedBuffer* current_buffer_; -}; - -} // namespace runtime -} // namespace cpu -} // namespace xla - -#endif // TENSORFLOW_COMPILER_XLA_SERVICE_CPU_INFEED_MANAGER_H_ diff --git a/tensorflow/compiler/xla/service/cpu/infeed_manager_test.cc b/tensorflow/compiler/xla/service/cpu/infeed_manager_test.cc deleted file mode 100644 index c65d8216606a1caa561adea5a83c8f1aa2c82906..0000000000000000000000000000000000000000 --- a/tensorflow/compiler/xla/service/cpu/infeed_manager_test.cc +++ /dev/null @@ -1,102 +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/service/cpu/infeed_manager.h" - -#include - -#include "tensorflow/compiler/xla/service/cpu/cpu_runtime.h" -#include "tensorflow/core/lib/core/threadpool.h" -#include "tensorflow/core/platform/env.h" -#include "tensorflow/core/platform/logging.h" -#include "tensorflow/core/platform/test.h" - -namespace xla { -namespace { - -class InfeedManagerTest : public ::testing::Test {}; - -class TestInfeedBuffer : public cpu::runtime::InfeedBuffer { - public: - explicit TestInfeedBuffer(int32 length) - : done_called_(false), length_(length) {} - ~TestInfeedBuffer() override { EXPECT_TRUE(done_called_); } - - int32 length() override { return length_; } - void* data() override { return nullptr; } - void Done() override { - CHECK(!done_called_); - done_called_ = true; - } - - private: - bool done_called_; - int32 length_; -}; - -void ProcessNextBuffer(int32 length) { - void* buffer = __xla_cpu_runtime_AcquireInfeedBufferForDequeue(length); - __xla_cpu_runtime_ReleaseInfeedBufferAfterDequeue(length, buffer); -} - -TEST_F(InfeedManagerTest, SingleThreadedSequential) { - TestInfeedBuffer* a = new TestInfeedBuffer(64); - TestInfeedBuffer* b = new TestInfeedBuffer(32); - - cpu::runtime::InfeedManager* infeed = cpu::runtime::GetInfeedManager(); - - infeed->EnqueueBuffer(a); - infeed->EnqueueBuffer(b); - ProcessNextBuffer(a->length()); - ProcessNextBuffer(b->length()); -} - -TEST_F(InfeedManagerTest, SingleThreadedInterleaved) { - TestInfeedBuffer* a = new TestInfeedBuffer(64); - TestInfeedBuffer* b = new TestInfeedBuffer(32); - - cpu::runtime::InfeedManager* infeed = cpu::runtime::GetInfeedManager(); - - infeed->EnqueueBuffer(a); - ProcessNextBuffer(a->length()); - infeed->EnqueueBuffer(b); - ProcessNextBuffer(b->length()); -} - -TEST_F(InfeedManagerTest, MultiThreaded) { - tensorflow::thread::ThreadPool pool(tensorflow::Env::Default(), "test", 2); - - cpu::runtime::InfeedManager* infeed = cpu::runtime::GetInfeedManager(); - - const int32 length = 64; - - pool.Schedule([infeed]() { - // Spin for 100 milliseconds - int64 start_micros = tensorflow::Env::Default()->NowMicros(); - while (true) { - int64 end_micros = tensorflow::Env::Default()->NowMicros(); - if ((end_micros - start_micros) >= 100000) { // 100 ms - break; - } - } - TestInfeedBuffer* a = new TestInfeedBuffer(length); - infeed->EnqueueBuffer(a); - }); - - ProcessNextBuffer(length); -} - -} // namespace -} // namespace xla diff --git a/tensorflow/compiler/xla/service/cpu/ir_emission_utils.cc b/tensorflow/compiler/xla/service/cpu/ir_emission_utils.cc index 2d855d0eb1e9448707b3916d20803cebf2ebabe4..859329e2c1ddca9dbea14c16b67f63d4803b6acd 100644 --- a/tensorflow/compiler/xla/service/cpu/ir_emission_utils.cc +++ b/tensorflow/compiler/xla/service/cpu/ir_emission_utils.cc @@ -16,7 +16,6 @@ limitations under the License. #include "tensorflow/compiler/xla/service/cpu/ir_emission_utils.h" #include "tensorflow/compiler/xla/layout_util.h" -#include "tensorflow/compiler/xla/legacy_flags/cpu_runtime_flags.h" #include "tensorflow/compiler/xla/service/cpu/cpu_runtime.h" #include "tensorflow/compiler/xla/shape_util.h" #include "tensorflow/compiler/xla/window_util.h" @@ -26,11 +25,6 @@ namespace cpu { bool PotentiallyImplementedAsEigenConvolution( const HloInstruction& convolution) { - legacy_flags::CpuRuntimeFlags* flags = legacy_flags::GetCpuRuntimeFlags(); - if (!flags->xla_cpu_use_eigen) { - return false; - } - // The following conditions are necessary (but not sufficient) for // implementing `convolution` with Eigen convolution: // - the input and kernel have a non-zero number of elements. @@ -82,11 +76,6 @@ bool AreValidGemmShapes(const Shape& lhs_shape, const Shape& rhs_shape, } // namespace bool PotentiallyImplementedAsEigenDot(const HloInstruction& hlo) { - legacy_flags::CpuRuntimeFlags* flags = legacy_flags::GetCpuRuntimeFlags(); - if (!flags->xla_cpu_use_eigen) { - return false; - } - // For certain types of Dot, we can call Eigen if (hlo.opcode() == HloOpcode::kDot) { const Shape& lhs_shape = hlo.operand(0)->shape(); diff --git a/tensorflow/compiler/xla/service/cpu/ir_emitter.cc b/tensorflow/compiler/xla/service/cpu/ir_emitter.cc index 51c6dc4426f8c40d60ba933ce0a31f8fb9d927c1..c5275ede651bb5e4a35a4e14a9baf966cc036040 100644 --- a/tensorflow/compiler/xla/service/cpu/ir_emitter.cc +++ b/tensorflow/compiler/xla/service/cpu/ir_emitter.cc @@ -26,16 +26,18 @@ limitations under the License. #include "tensorflow/core/platform/logging.h" // IWYU pragma: no_include "llvm/IR/Intrinsics.gen.inc" -#include "external/llvm/include/llvm/IR/BasicBlock.h" -#include "external/llvm/include/llvm/IR/Constants.h" -#include "external/llvm/include/llvm/IR/GlobalVariable.h" -#include "external/llvm/include/llvm/IR/Instructions.h" -#include "external/llvm/include/llvm/IR/Intrinsics.h" -#include "external/llvm/include/llvm/IR/LLVMContext.h" +#include "llvm/IR/BasicBlock.h" +#include "llvm/IR/Constants.h" +#include "llvm/IR/GlobalVariable.h" +#include "llvm/IR/Instructions.h" +#include "llvm/IR/Intrinsics.h" +#include "llvm/IR/LLVMContext.h" +#include "llvm/Target/TargetRegisterInfo.h" +#include "llvm/Target/TargetSubtargetInfo.h" #include "tensorflow/compiler/xla/layout_util.h" -#include "tensorflow/compiler/xla/legacy_flags/cpu_runtime_flags.h" #include "tensorflow/compiler/xla/map_util.h" #include "tensorflow/compiler/xla/service/buffer_assignment.h" +#include "tensorflow/compiler/xla/service/cpu/cpu_options.h" #include "tensorflow/compiler/xla/service/cpu/cpu_runtime.h" #include "tensorflow/compiler/xla/service/cpu/dot_op_emitter.h" #include "tensorflow/compiler/xla/service/cpu/elemental_ir_emitter.h" @@ -52,9 +54,12 @@ limitations under the License. #include "tensorflow/compiler/xla/types.h" #include "tensorflow/compiler/xla/util.h" #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/gtl/array_slice.h" +#include "tensorflow/core/lib/gtl/flatset.h" #include "tensorflow/core/lib/strings/strcat.h" +#include "tensorflow/core/lib/strings/stringprintf.h" namespace xla { @@ -63,37 +68,51 @@ using llvm_ir::SetToFirstInsertPoint; namespace cpu { IrEmitter::IrEmitter( - const HloModule& hlo_module, const HloModuleConfig& hlo_module_config, - const BufferAssignment& assignment, llvm::Module* llvm_module, - const std::unordered_map* hlo_to_profile_idx) + const HloModule& hlo_module, const BufferAssignment& assignment, + llvm::Module* llvm_module, + const std::unordered_map* hlo_to_profile_idx, + llvm::TargetMachine* target_machine) : assignment_(assignment), module_(llvm_module), arch_type_(llvm::Triple(llvm_module->getTargetTriple()).getArch()), ir_builder_(llvm_module->getContext()), hlo_to_profile_idx_(hlo_to_profile_idx), alias_analysis_(hlo_module, assignment, &llvm_module->getContext()), - hlo_module_config_(hlo_module_config) { - ir_builder_.setFastMathFlags(llvm_ir::GetFastMathFlags(hlo_module_config)); + hlo_module_config_(hlo_module.config()), + parallel_cpu_backend_( + options::CpuParallelBackendRequested(hlo_module_config_)), + is_top_level_computation_(false), + target_machine_features_(target_machine) { + ir_builder_.setFastMathFlags(llvm_ir::GetFastMathFlags( + /*fast_math_enabled=*/hlo_module_config_.debug_options() + .xla_enable_fast_math())); } StatusOr IrEmitter::EmitComputation( HloComputation* computation, const string& function_name_prefix, - bool is_entry_computation, + bool is_top_level_computation, std::vector* instruction_order) { string function_name = name_uniquer_.GetUniqueName(function_name_prefix); - VLOG(2) << "Emitting IR for CPU function [" << function_name_prefix << "]"; - InitializeIrFunction(function_name, is_entry_computation); + VLOG(2) << "Emitting IR for CPU function [" << function_name_prefix + << "]; ordered? " << (instruction_order != nullptr); + is_top_level_computation_ = is_top_level_computation; + num_dynamic_loop_bounds_ = 0; + if (!computation->root_instruction()->outer_dimension_partitions().empty()) { + num_dynamic_loop_bounds_ = + computation->root_instruction()->outer_dimension_partitions().size(); + } + + InitializeIrFunction(function_name); // The rdtscp instruction is x86 specific. We will fallback to LLVM's generic // readcyclecounter if it is unavailable. bool use_rdtscp = arch_type_ == llvm::Triple::ArchType::x86 || arch_type_ == llvm::Triple::ArchType::x86_64; - profiling_state_ = ProfilingState(is_entry_computation, use_rdtscp, + profiling_state_ = ProfilingState(is_top_level_computation_, use_rdtscp, GetProfileCountersArgument()); - if (instruction_order != nullptr) { - TF_RETURN_IF_ERROR(computation->root_instruction()->AcceptOrdered( - this, *instruction_order)); + if (instruction_order == nullptr) { + TF_RETURN_IF_ERROR(computation->Accept(this)); } else { - TF_RETURN_IF_ERROR(computation->root_instruction()->Accept(this)); + TF_RETURN_IF_ERROR(computation->AcceptOrdered(this, *instruction_order)); } InsertOrDie(&emitted_functions_, computation, compute_function_); @@ -106,11 +125,10 @@ static llvm::Argument* GetArg(llvm::Function* f, int idx) { return &*arg_iter; } -void IrEmitter::InitializeIrFunction(const string& function_name, - bool is_entry_computation) { +void IrEmitter::InitializeIrFunction(const string& function_name) { // The function signature is: // void function(i8* retval, i8* run_options, i8** params, i8** temps, - // i64* prof_counters) + // i64* dynamic_loop_bounds, i64* prof_counters) // // retval: points to the returned value. // params: address of an array with pointers to parameters. @@ -150,6 +168,10 @@ void IrEmitter::InitializeIrFunction(const string& function_name, // | temp 0 | | temp 1 | | temp N-1 | // \---------/ \---------/ \-----------/ // + // /--------------------------------------------\ + // dynamic loop bounds -> | outer_dim0_start | outer_dim0_limit | .....| + // (elided for aot) \--------------------------------------------/ + // // /---------------------------------------------\ // prof counters -> | counter 0 | counter 1 | ..... | counter N-1 | // (elided for aot) \---------------------------------------------/ @@ -162,6 +184,9 @@ void IrEmitter::InitializeIrFunction(const string& function_name, llvm::Type* i64_ptr_type = llvm::Type::getInt64PtrTy(module_->getContext()); std::vector compute_function_params( {i8_ptr_type, i8_ptr_type, i8_ptr_ptr_type, i8_ptr_ptr_type}); + if (IsParallelContext()) { + compute_function_params.push_back(i64_ptr_type); + } if (hlo_to_profile_idx_) { compute_function_params.push_back(i64_ptr_type); } @@ -174,8 +199,8 @@ void IrEmitter::InitializeIrFunction(const string& function_name, // a-priori that embedded functions (non-entry functions) will not have its // name resolved, give it local linkage. llvm::Function::LinkageTypes linkage = - is_entry_computation ? llvm::GlobalValue::ExternalLinkage - : llvm::GlobalValue::InternalLinkage; + is_top_level_computation_ ? llvm::GlobalValue::ExternalLinkage + : llvm::GlobalValue::InternalLinkage; compute_function_ = llvm::Function::Create(/*Ty=*/compute_function_type, /*Linkage=*/linkage, /*Name=*/function_name.c_str(), @@ -188,6 +213,9 @@ void IrEmitter::InitializeIrFunction(const string& function_name, (++arg_iter)->setName("run_options"); (++arg_iter)->setName("params"); (++arg_iter)->setName("temps"); + if (IsParallelContext()) { + (++arg_iter)->setName("dynamic_loop_bounds"); + } if (hlo_to_profile_idx_) { (++arg_iter)->setName("prof_counters"); } @@ -201,7 +229,15 @@ void IrEmitter::InitializeIrFunction(const string& function_name, if (&argument == retval) { continue; } - compute_function_->setDoesNotAlias(argument.getArgNo() + 1); + compute_function_->addAttribute(argument.getArgNo() + 1, + llvm::Attribute::NoAlias); + } + + // Add the optize attribute to the function if optimizing for size. This + // controls internal behavior of some optimization passes (e.g. loop + // unrolling). + if (options::OptimizeForSizeRequested(hlo_module_config_)) { + compute_function_->addFnAttr(llvm::Attribute::OptimizeForSize); } ir_builder_.SetInsertPoint(llvm::BasicBlock::Create( @@ -239,12 +275,12 @@ Status IrEmitter::HandleConstant(HloInstruction* constant, return Status::OK(); } -Status IrEmitter::HandleCopy(HloInstruction* copy, HloInstruction* operand) { +Status IrEmitter::HandleCopy(HloInstruction* copy) { if (ShapeUtil::IsTuple(copy->shape())) { // kCopy shallow copies a tuple so just memcpy the top-level buffer. TF_ASSIGN_OR_RETURN(llvm::Value * copy_value, EmitTargetAddressForOp(copy)); emitted_value_[copy] = copy_value; - return EmitMemcpy(*operand, *copy); + return EmitMemcpy(*(copy->operand(0)), *copy); } else { // Use the elemental emitter for non-tuple shapes. return DefaultAction(copy); @@ -355,63 +391,158 @@ Status IrEmitter::HandleSelect(HloInstruction* select, HloInstruction* pred, Status IrEmitter::HandleInfeed(HloInstruction* infeed) { VLOG(2) << "HandleInfeed: " << infeed->ToString(); + const Shape& shape = infeed->shape(); + + // The infeed operation produces data (dequeued from the infeed queue) at this + // address, which has been provided by buffer assignment. + TF_ASSIGN_OR_RETURN(llvm::Value * target_address, + EmitTargetAddressForOp(infeed)); + + if (ShapeUtil::IsTuple(shape)) { + TF_RET_CHECK(!ShapeUtil::IsNestedTuple(shape)); + + // For a tuple, we first copy each of the internal elements to + // their corresponding target locations. We then construct the + // tuple outer buffer containing pointers to the internal + // elements. + std::vector tuple_element_addresses; + for (int64 i = 0; i < shape.tuple_shapes_size(); ++i) { + TF_ASSIGN_OR_RETURN(BufferAllocation::Slice buffer, + assignment_.GetUniqueSlice(infeed, {i})); + + const Shape& tuple_element_shape = + ShapeUtil::GetTupleElementShape(shape, i); + + // Only the outer tuple buffer's target address is obtained from + // EmitTargetAddressForOp to handle the case when Infeed is the + // root instruction. Target addresses for internal elements can + // be obtained from EmitTempBufferPointer. + llvm::Value* tuple_element_address = + EmitTempBufferPointer(buffer, tuple_element_shape); + + TF_RETURN_IF_ERROR(EmitXfeedTransfer( + XfeedKind::kInfeed, tuple_element_shape, tuple_element_address)); + + tuple_element_addresses.push_back(tuple_element_address); + } + + llvm_ir::EmitTuple(llvm_ir::IrArray(target_address, shape), + tuple_element_addresses, &ir_builder_); + } else { + TF_RETURN_IF_ERROR( + EmitXfeedTransfer(XfeedKind::kInfeed, shape, target_address)); + } + + emitted_value_[infeed] = target_address; + + return Status::OK(); +} + +Status IrEmitter::EmitXfeedTransfer(XfeedKind kind, const Shape& shape, + llvm::Value* program_buffer_address) { + int64 length = ByteSizeOf(shape); + if (length <= 0 || length > std::numeric_limits::max()) { + return InvalidArgument( + "xfeed (infeed or outfeed) buffer length %lld is outside the valid " + "size range", + length); + } + int32 length_32 = static_cast(length); + + int32 shape_length; + TF_ASSIGN_OR_RETURN(llvm::Value * shape_ptr, + llvm_ir::EncodeSelfDescribingShapeConstant( + shape, &shape_length, &ir_builder_)); + // The signature of the acquire infeed buffer function is: // // (void*)(int32 length); - llvm::Type* i8_ptr_type = llvm::Type::getInt8PtrTy(module_->getContext()); llvm::Type* int32_type = ir_builder_.getInt32Ty(); - llvm::FunctionType* acquire_type = - llvm::FunctionType::get(i8_ptr_type, {int32_type}, - /*isVarArg=*/false); + llvm::Type* i8_ptr_type = llvm::Type::getInt8PtrTy(module_->getContext()); + llvm::FunctionType* acquire_type = llvm::FunctionType::get( + i8_ptr_type, {int32_type, i8_ptr_type, int32_type}, + /*isVarArg=*/false); - llvm::Function* acquire_func = - llvm::cast(module_->getOrInsertFunction( - runtime::kAcquireInfeedBufferForDequeueSymbolName, acquire_type)); + llvm::Function* acquire_func; + if (kind == XfeedKind::kInfeed) { + acquire_func = llvm::cast(module_->getOrInsertFunction( + runtime::kAcquireInfeedBufferForDequeueSymbolName, acquire_type)); + } else { + acquire_func = llvm::cast(module_->getOrInsertFunction( + runtime::kAcquireOutfeedBufferForPopulationSymbolName, acquire_type)); + } acquire_func->setCallingConv(llvm::CallingConv::C); // The signature of the release infeed buffer function is: // // (void)(int32 length, void* buffer); llvm::FunctionType* release_type = llvm::FunctionType::get( - ir_builder_.getVoidTy(), {int32_type, i8_ptr_type}, + ir_builder_.getVoidTy(), + {int32_type, i8_ptr_type, i8_ptr_type, int32_type}, /*isVarArg=*/false); - llvm::Function* release_func = - llvm::cast(module_->getOrInsertFunction( - runtime::kReleaseInfeedBufferAfterDequeueSymbolName, release_type)); + llvm::Function* release_func; + if (kind == XfeedKind::kInfeed) { + release_func = llvm::cast(module_->getOrInsertFunction( + runtime::kReleaseInfeedBufferAfterDequeueSymbolName, release_type)); + } else { + release_func = llvm::cast(module_->getOrInsertFunction( + runtime::kReleaseOutfeedBufferAfterPopulationSymbolName, release_type)); + } release_func->setCallingConv(llvm::CallingConv::C); - const Shape& shape = infeed->shape(); - int64 length = ByteSizeOf(shape); - if (length > std::numeric_limits::max()) { - return InvalidArgument("infeed buffer length %lld is too large", length); + // Implementation note: this call informs the runtime that it wants a buffer + // of size exactly 'length_32', and the runtime is responsible for + // check-failing the process if there is a mismatch, versus passing us back a + // buffer that we might overrun. + llvm::Value* acquired_pointer = ir_builder_.CreateCall( + acquire_func, {ir_builder_.getInt32(length_32), shape_ptr, + ir_builder_.getInt32(shape_length)}); + + if (kind == XfeedKind::kInfeed) { + // Copy to the program buffer address from the acquired buffer. + ir_builder_.CreateMemCpy(program_buffer_address, acquired_pointer, + length_32, 1); + } else { + // Outfeed -- copy from the in-program address to the acquired buffer. + ir_builder_.CreateMemCpy(acquired_pointer, program_buffer_address, + length_32, 1); } - int32 length_32 = static_cast(length); - - llvm::Value* acquired_pointer = - ir_builder_.CreateCall(acquire_func, {ir_builder_.getInt32(length_32)}); - - TF_ASSIGN_OR_RETURN(llvm::Value * target_address, - EmitTargetAddressForOp(infeed)); - - ir_builder_.CreateMemCpy(target_address, acquired_pointer, length_32, 1); ir_builder_.CreateCall(release_func, - {ir_builder_.getInt32(length_32), acquired_pointer}); - - emitted_value_[infeed] = target_address; + {ir_builder_.getInt32(length_32), acquired_pointer, + shape_ptr, ir_builder_.getInt32(shape_length)}); return Status::OK(); } Status IrEmitter::HandleOutfeed(HloInstruction* outfeed) { - // TODO(b/34359662): Implement outfeed on CPU. - return Unimplemented("Outfeed is not supported on CPU (b/34359662)."); + HloInstruction* operand = outfeed->operands()[0]; + const Shape& operand_shape = operand->shape(); + + llvm::Value* value = GetEmittedValueFor(operand); + if (!ShapeUtil::IsTuple(operand_shape)) { + return EmitXfeedTransfer(XfeedKind::kOutfeed, operand_shape, value); + } + + TF_RET_CHECK(!ShapeUtil::IsNestedTuple(operand_shape)); + + for (int64 i = 0; i < operand_shape.tuple_shapes_size(); ++i) { + const Shape& tuple_element_shape = + ShapeUtil::GetTupleElementShape(operand_shape, i); + llvm::Value* tuple_element = llvm_ir::EmitGetTupleElement( + tuple_element_shape, i, MinimumAlignmentForShape(tuple_element_shape), + value, &ir_builder_); + TF_RETURN_IF_ERROR(EmitXfeedTransfer(XfeedKind::kOutfeed, + tuple_element_shape, tuple_element)); + } + + return Status::OK(); } Status IrEmitter::HandleSort(HloInstruction* sort, HloInstruction* operand) { // TODO(b/26783907): Implement sort on CPU. - return Unimplemented("Sort is not supported on GPU (b/26783907)."); + return Unimplemented("Sort is not supported on CPU (b/26783907)."); } Status IrEmitter::HandleTuple( @@ -757,7 +888,8 @@ Status IrEmitter::HandleDot(HloInstruction* dot, HloInstruction* lhs, // Dot operation is complicated so we delegate to a helper class. TF_RETURN_IF_ERROR(DotOpEmitter::EmitDotOperation( *dot, /*transpose_lhs=*/false, /*transpose_rhs=*/false, target_array, - lhs_array, rhs_array, GetExecutableRunOptionsArgument(), &ir_builder_)); + lhs_array, rhs_array, GetExecutableRunOptionsArgument(), &ir_builder_, + hlo_module_config_)); emitted_value_[dot] = target_address; return Status::OK(); @@ -842,9 +974,10 @@ Status IrEmitter::HandleConvolution(HloInstruction* convolution, int64_type, int64_type, int64_type, int64_type, int64_type, int64_type, int64_type, int64_type}, /*isVarArg=*/false); - legacy_flags::CpuRuntimeFlags* flags = legacy_flags::GetCpuRuntimeFlags(); + bool multi_threaded_eigen = + hlo_module_config_.debug_options().xla_cpu_multi_thread_eigen(); const char* fn_name = - (flags->xla_cpu_multi_thread_eigen + (multi_threaded_eigen ? runtime::kEigenConvF32SymbolName : runtime::kEigenSingleThreadedConvF32SymbolName); llvm::Function* conv_func = llvm::cast( @@ -1036,6 +1169,237 @@ Status IrEmitter::HandleCrossReplicaSum(HloInstruction* crs) { "Cross replica sum not implemented on CPU. See b/33011107."); } +// 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::HandleBatchNormTraining(HloInstruction* batch_norm_training) { + // The output of BatchNormTraining is a tuple of three element: + // - An N-dimensional array containing normalized values. + // - A 1 dimensional array containing the mean value for each feature. + // - A 1 dimensional array containing the variance value for each feature. + HloInstruction* operand = batch_norm_training->operands()[0]; + HloInstruction* scale = batch_norm_training->operands()[1]; + HloInstruction* offset = batch_norm_training->operands()[2]; + float epsilon = batch_norm_training->epsilon(); + int64 feature_index = batch_norm_training->feature_index(); + TF_RET_CHECK(ShapeUtil::IsTuple(batch_norm_training->shape()) && + ShapeUtil::TupleElementCount(batch_norm_training->shape()) == 3); + + const Shape& output_shape = + ShapeUtil::GetTupleElementShape(batch_norm_training->shape(), 0); + const Shape& feature_shape = + ShapeUtil::GetTupleElementShape(batch_norm_training->shape(), 1); + + // Reduce vector of the non-feature dimensions. + std::vector dimensions_to_reduce; + + for (int64 i = 0; i < operand->shape().dimensions_size(); ++i) { + if (i != feature_index) { + dimensions_to_reduce.push_back(i); + } + } + + // Get the second and third allocations in the output tuple, which should be + // used to store the result of mean and variance value calculation. + TF_ASSIGN_OR_RETURN( + const BufferAllocation::Slice slice_mean, + assignment_.GetUniqueSlice(batch_norm_training, /*index=*/{1})); + TF_ASSIGN_OR_RETURN( + const BufferAllocation::Slice slice_var, + assignment_.GetUniqueSlice(batch_norm_training, /*index=*/{2})); + const int feature_count = output_shape.dimensions(feature_index); + const int size_in_elements = ShapeUtil::ElementsIn(output_shape); + TF_RET_CHECK(ShapeUtil::ElementsIn(operand->shape()) == size_in_elements); + const int elements_per_feature = size_in_elements / feature_count; + + llvm::Value* mean = EmitTempBufferPointer(slice_mean, feature_shape); + llvm_ir::IrArray mean_array(mean, feature_shape); + + llvm::Value* var = EmitTempBufferPointer(slice_var, feature_shape); + llvm_ir::IrArray var_array(var, feature_shape); + + // This loop calculates mean and variance for each feature. + // + // In theory this could be swapped by multi-output fusion. We will evaluate + // this when it's ready. + // + // For variance calculation, we use a simplified formula so we can fuse the + // computation into the same loop to calculate mean: Var=E(X^2) - E(X)^2. + TF_RETURN_IF_ERROR( + llvm_ir::LoopEmitter( + [this, operand, dimensions_to_reduce, feature_shape, var_array, + elements_per_feature](const llvm_ir::IrArray::Index& index) { + PrimitiveType element_type = operand->shape().element_type(); + // Used to calculate E(X). + llvm::Value* sum_address = llvm_ir::EmitAllocaAtFunctionEntry( + llvm_ir::PrimitiveTypeToIrType(element_type, &ir_builder_), + "sum_address", &ir_builder_, + MinimumAlignmentForPrimitiveType(element_type)); + + // Used to calculate E(X^2). + llvm::Value* sum_square_address = + llvm_ir::EmitAllocaAtFunctionEntry( + llvm_ir::PrimitiveTypeToIrType(element_type, &ir_builder_), + "sum_square_address", &ir_builder_, + MinimumAlignmentForPrimitiveType(element_type)); + + ir_builder_.CreateStore( + llvm::ConstantFP::get(ir_builder_.getFloatTy(), 0.0), + sum_address); + + ir_builder_.CreateStore( + llvm::ConstantFP::get(ir_builder_.getFloatTy(), 0.0), + sum_square_address); + + llvm_ir::ForLoopNest loops(&ir_builder_); + + const llvm_ir::IrArray::Index reduced_dims_index = + loops.AddLoopsForShapeOnDimensions( + operand->shape(), dimensions_to_reduce, "reduction_dim"); + + SetToFirstInsertPoint(loops.GetInnerLoopBodyBasicBlock(), + &ir_builder_); + + llvm_ir::IrArray operand_array(GetIrArrayForOp(operand)); + llvm_ir::IrArray::Index input_index = + FillReducedDimensionIndex(reduced_dims_index, index); + llvm::Value* new_value = + operand_array.EmitReadArrayElement(input_index, &ir_builder_); + + llvm::Value* new_value_square = + ir_builder_.CreateFMul(new_value, new_value); + + llvm::Value* current_sum = ir_builder_.CreateLoad(sum_address); + llvm::Value* current_sum_square = + ir_builder_.CreateLoad(sum_square_address); + // Update sum. + ir_builder_.CreateStore( + ir_builder_.CreateFAdd(current_sum, new_value), sum_address); + + // Update sum square. + ir_builder_.CreateStore( + ir_builder_.CreateFAdd(current_sum_square, new_value_square), + sum_square_address); + + SetToFirstInsertPoint(loops.GetOuterLoopExitBasicBlock(), + &ir_builder_); + + llvm::Value* sum = ir_builder_.CreateLoad(sum_address); + llvm::Value* elements_per_feature_value = llvm::ConstantFP::get( + ir_builder_.getFloatTy(), elements_per_feature); + llvm::Value* mean = + ir_builder_.CreateFDiv(sum, elements_per_feature_value); + llvm::Value* mean_square = ir_builder_.CreateFMul(mean, mean); + llvm::Value* sum_square = + ir_builder_.CreateLoad(sum_square_address); + + // Var=E(X^2) - E(X)^2. + llvm::Value* var = ir_builder_.CreateFSub( + ir_builder_.CreateFDiv(sum_square, elements_per_feature_value), + mean_square); + + var_array.EmitWriteArrayElement(index, var, &ir_builder_); + return mean; + }, + mean_array, &ir_builder_) + .EmitLoop()); + + TF_ASSIGN_OR_RETURN(llvm::Value * target_address, + EmitTargetAddressForOp(batch_norm_training)); + + TF_ASSIGN_OR_RETURN( + const BufferAllocation::Slice slice, + assignment_.GetUniqueSlice(batch_norm_training, /*index=*/{0})); + + llvm::Value* normalized = EmitTempBufferPointer(slice, output_shape); + + llvm_ir::IrArray target_array(normalized, output_shape); + + AddAliasingInformationToIrArray(*batch_norm_training, &target_array); + + TF_RETURN_IF_ERROR( + llvm_ir::LoopEmitter( + [this, mean_array, var_array, epsilon, operand, dimensions_to_reduce, + feature_index, offset, scale](const llvm_ir::IrArray::Index& index) { + // The following logic normalizes the input value, scales and shifts + // it: + // + // normalized = (input - mean) / sqrt(variance + epsilon) + // result = normalized * scale + offset + + // Current index in the feature dimension. + llvm_ir::IrArray::Index feature_index_value(1, + index[feature_index]); + + llvm::Value* mean = mean_array.EmitReadArrayElement( + feature_index_value, &ir_builder_); + llvm::Value* var = var_array.EmitReadArrayElement( + feature_index_value, &ir_builder_); + + llvm_ir::IrArray operand_array(GetIrArrayForOp(operand)); + llvm::Value* input = + operand_array.EmitReadArrayElement(index, &ir_builder_); + + llvm::Value* variance_with_epsilon = ir_builder_.CreateFAdd( + var, llvm::ConstantFP::get(ir_builder_.getFloatTy(), epsilon)); + llvm::Function* func_llvm_sqrt = llvm::Intrinsic::getDeclaration( + module_, llvm::Intrinsic::sqrt, {ir_builder_.getFloatTy()}); + llvm::Value* variance_sqrt = + ir_builder_.CreateCall(func_llvm_sqrt, {variance_with_epsilon}); + llvm::Value* normalized = ir_builder_.CreateFDiv( + ir_builder_.CreateFSub(input, mean), variance_sqrt); + llvm_ir::IrArray offset_array(GetIrArrayForOp(offset)); + llvm::Value* offset = offset_array.EmitReadArrayElement( + feature_index_value, &ir_builder_); + llvm_ir::IrArray scale_array(GetIrArrayForOp(scale)); + llvm::Value* scale = scale_array.EmitReadArrayElement( + feature_index_value, &ir_builder_); + llvm::Value* result = ir_builder_.CreateFAdd( + ir_builder_.CreateFMul(normalized, scale), offset); + + return result; + }, + target_array, &ir_builder_) + .EmitLoop()); + + llvm_ir::EmitTuple( + llvm_ir::IrArray(target_address, batch_norm_training->shape()), + {normalized, mean, var}, &ir_builder_); + emitted_value_[batch_norm_training] = target_address; + + return Status::OK(); +} + +Status IrEmitter::HandleBatchNormGrad(HloInstruction* batch_norm_grad) { + // TODO(b/62843645) Implement BatchNormGrad on CPU backend. + return Unimplemented( + "BatchNormGrad is not implemented on CPU. See b/62843645."); +} + Status IrEmitter::HandleParameter(HloInstruction* parameter) { VLOG(2) << "HandleParameter: " << parameter->ToString(); auto param_number = parameter->parameter_number(); @@ -1070,10 +1434,450 @@ Status IrEmitter::HandleParameter(HloInstruction* parameter) { return Status::OK(); } +IrEmitter::ReductionGenerator IrEmitter::MatchReductionGenerator( + HloComputation* function, string* failure_reason) const { + CHECK_EQ(function->num_parameters(), 2); + + auto root_instruction = function->root_instruction(); + CHECK(ShapeUtil::IsScalar(root_instruction->shape())); + + if (root_instruction->operand_count() != 2) { + *failure_reason = "root instruction is not a binary operation"; + return nullptr; + } + + const Shape& root_shape = root_instruction->shape(); + bool root_is_floating_point = ShapeUtil::ElementIsFloating(root_shape); + bool root_is_integral = ShapeUtil::ElementIsIntegral(root_shape); + bool root_is_signed = ShapeUtil::ElementIsSigned(root_shape); + + auto lhs = root_instruction->operand(0); + auto rhs = root_instruction->operand(1); + + auto param_0 = function->parameter_instruction(0); + auto param_1 = function->parameter_instruction(1); + if (!(lhs == param_0 && rhs == param_1) && + !(rhs == param_0 && lhs == param_1)) { + *failure_reason = + "root instruction is not a binary operation on the incoming arguments"; + return nullptr; + } + + CHECK(ShapeUtil::IsScalar(lhs->shape()) && ShapeUtil::IsScalar(rhs->shape())); + + // This is visually similar to ElementalIrEmitter, though conceptually we're + // doing something different here. ElementalIrEmitter emits scalar operations + // while these emit scalar or vector operations depending on the type of the + // operands. + switch (root_instruction->opcode()) { + default: + *failure_reason = "did not recognize root instruction opcode"; + return nullptr; + + case HloOpcode::kAdd: + return [root_is_integral](llvm::IRBuilder<>* ir_builder, llvm::Value* lhs, + llvm::Value* rhs) { + return root_is_integral ? ir_builder->CreateAdd(lhs, rhs) + : ir_builder->CreateFAdd(lhs, rhs); + }; + + case HloOpcode::kMultiply: + return [root_is_integral](llvm::IRBuilder<>* ir_builder, llvm::Value* lhs, + llvm::Value* rhs) { + return root_is_integral ? ir_builder->CreateMul(lhs, rhs) + : ir_builder->CreateFMul(lhs, rhs); + }; + + case HloOpcode::kLogicalAnd: + return [](llvm::IRBuilder<>* ir_builder, llvm::Value* lhs, + llvm::Value* rhs) { return ir_builder->CreateAnd(lhs, rhs); }; + + case HloOpcode::kLogicalOr: + return [](llvm::IRBuilder<>* ir_builder, llvm::Value* lhs, + llvm::Value* rhs) { return ir_builder->CreateOr(lhs, rhs); }; + + case HloOpcode::kMaximum: + return [root_is_floating_point, root_is_signed]( + llvm::IRBuilder<>* ir_builder, llvm::Value* lhs, + llvm::Value* rhs) { + if (root_is_floating_point) { + return llvm_ir::EmitCallToIntrinsic(llvm::Intrinsic::maxnum, + {lhs, rhs}, {lhs->getType()}, + ir_builder); + } + + return ir_builder->CreateSelect( + ir_builder->CreateICmp(root_is_signed ? llvm::ICmpInst::ICMP_SGE + : llvm::ICmpInst::ICMP_UGE, + lhs, rhs), + lhs, rhs); + }; + + case HloOpcode::kMinimum: + return [root_is_floating_point, root_is_signed]( + llvm::IRBuilder<>* ir_builder, llvm::Value* lhs, + llvm::Value* rhs) { + if (root_is_floating_point) { + return llvm_ir::EmitCallToIntrinsic(llvm::Intrinsic::minnum, + {lhs, rhs}, {lhs->getType()}, + ir_builder); + } + + return ir_builder->CreateSelect( + ir_builder->CreateICmp(root_is_signed ? llvm::ICmpInst::ICMP_SLE + : llvm::ICmpInst::ICMP_ULE, + lhs, rhs), + lhs, rhs); + }; + } +} + +IrEmitter::ShardedVectorType IrEmitter::CreateShardedVectorType( + PrimitiveType element_type, unsigned element_count) { + // Here we assume that the largest register is a vector register. + int max_vector_register_size_in_bytes = + target_machine_features_.largest_register_size_in_bytes( + compute_function_); + + int vector_register_size_in_elements = + max_vector_register_size_in_bytes / + ShapeUtil::ByteSizeOfPrimitiveType(element_type); + + ShardedVectorType sharded_vector_type; + llvm::Type* element_ir_type = + llvm_ir::PrimitiveTypeToIrType(element_type, &ir_builder_); + + for (int i = 0, e = 1 + tensorflow::Log2Ceiling(element_count); i < e; i++) { + // For every power of two present in element_count, we generate one or more + // vector or scalar types. + const unsigned current_size_fragment = 1u << i; + if (!(element_count & current_size_fragment)) { + // Power of two not present in element_count. + continue; + } + + if (current_size_fragment == 1) { + // Single element, use a scalar type. + sharded_vector_type.push_back(element_ir_type); + continue; + } + + // Lower "current_size_fragment" number of elements using (as few as + // possible) vector registers. + + if (current_size_fragment >= vector_register_size_in_elements) { + auto vector_type = llvm::VectorType::get( + element_ir_type, vector_register_size_in_elements); + sharded_vector_type.insert( + sharded_vector_type.end(), + current_size_fragment / vector_register_size_in_elements, + vector_type); + + // Both current_size_fragment and vector_register_size_in_elements are + // powers of two. + CHECK_EQ(current_size_fragment % vector_register_size_in_elements, 0); + continue; + } + + // For now we assume that vector_register_size_in_elements and lower powers + // of two are all legal vector sizes (or at least can be lowered easily by + // LLVM). + sharded_vector_type.push_back( + llvm::VectorType::get(element_ir_type, current_size_fragment)); + } + return sharded_vector_type; +} + +StatusOr +IrEmitter::EmitInnerLoopForVectorizedReduction( + const ReductionGenerator& reduction_generator, + const llvm_ir::IrArray::Index& output_index, + const ShardedVectorType& accumulator_type, HloInstruction* init_value, + HloInstruction* arg, tensorflow::gtl::ArraySlice dimensions, + unsigned element_alignment) { + ShardedVector accumulator; + accumulator.reserve(accumulator_type.size()); + for (auto accumulator_shard_type : accumulator_type) { + accumulator.push_back(llvm_ir::EmitAllocaAtFunctionEntry( + accumulator_shard_type, "accumulator", &ir_builder_, 0)); + } + + llvm::Value* init_value_ssa = + ir_builder_.CreateLoad(GetEmittedValueFor(init_value)); + + for (llvm::Value* accumulator_shard : accumulator) { + llvm::Value* initial_value; + auto shard_type = accumulator_shard->getType()->getPointerElementType(); + if (auto vector_type = llvm::dyn_cast(shard_type)) { + initial_value = ir_builder_.CreateVectorSplat( + vector_type->getNumElements(), init_value_ssa); + } else { + initial_value = init_value_ssa; + } + + ir_builder_.CreateAlignedStore(initial_value, accumulator_shard, + element_alignment); + } + + llvm_ir::ForLoopNest reduction_loop_nest(&ir_builder_); + llvm_ir::IrArray::Index reduced_dims_index = + reduction_loop_nest.AddLoopsForShapeOnDimensions(arg->shape(), dimensions, + "reduction_dim"); + + SetToFirstInsertPoint(reduction_loop_nest.GetInnerLoopBodyBasicBlock(), + &ir_builder_); + + llvm_ir::IrArray arg_array(GetIrArrayForOp(arg)); + llvm_ir::IrArray::Index input_index = reduced_dims_index; + llvm_ir::IrArray::Index::const_iterator it = output_index.begin(); + + for (size_t i = 0; i < input_index.size(); ++i) { + if (input_index[i] == nullptr) { + input_index[i] = *it++; + } + } + CHECK(output_index.end() == it); + + llvm::Value* input_address = ir_builder_.CreateBitCast( + arg_array.EmitArrayElementAddress(input_index, &ir_builder_), + ir_builder_.getInt8PtrTy()); + + for (int i = 0; i < accumulator.size(); i++) { + auto input_address_typed = + ir_builder_.CreateBitCast(input_address, accumulator[i]->getType()); + auto current_accumulator_value = + ir_builder_.CreateAlignedLoad(accumulator[i], element_alignment); + auto addend = + ir_builder_.CreateAlignedLoad(input_address_typed, element_alignment); + arg_array.AnnotateLoadStoreInstructionWithMetadata(addend); + + auto reduced_result = + reduction_generator(&ir_builder_, current_accumulator_value, addend); + ir_builder_.CreateAlignedStore(reduced_result, accumulator[i], + element_alignment); + + if (i != (accumulator.size() - 1)) { + input_address = ir_builder_.CreateConstInBoundsGEP1_32( + reduced_result->getType(), input_address_typed, 1); + } + } + + SetToFirstInsertPoint(reduction_loop_nest.GetOuterLoopExitBasicBlock(), + &ir_builder_); + + ShardedVector result_ssa; + result_ssa.reserve(accumulator.size()); + for (auto accumulator_shard : accumulator) { + result_ssa.push_back( + ir_builder_.CreateAlignedLoad(accumulator_shard, element_alignment)); + } + return result_ssa; +} + +void IrEmitter::EmitShardedVectorStore( + llvm::Value* store_address, const std::vector& value_to_store, + const int alignment, const llvm_ir::IrArray& containing_array) { + for (int i = 0; i < value_to_store.size(); i++) { + auto store_address_typed = ir_builder_.CreateBitCast( + store_address, + llvm::PointerType::getUnqual(value_to_store[i]->getType())); + + auto store_instruction = ir_builder_.CreateAlignedStore( + value_to_store[i], store_address_typed, alignment); + containing_array.AnnotateLoadStoreInstructionWithMetadata( + store_instruction); + + if (i != (value_to_store.size() - 1)) { + store_address = ir_builder_.CreateConstInBoundsGEP1_32( + value_to_store[i]->getType(), store_address_typed, 1); + } + } +} + +namespace { +// TODO(sanjoy): This is duplicated in tensorflow/core/lib/core/arena.cc. +// Extract out a common implementation to tensorflow/core/lib/math/math_util.h +uint32 GCD(uint32 x, uint32 y) { + while (y != 0) { + uint32 r = x % y; + x = y; + y = r; + } + return x; +} +} // namespace + +StatusOr IrEmitter::EmitVectorizedReduce( + HloInstruction* reduce, HloInstruction* arg, HloInstruction* init_value, + tensorflow::gtl::ArraySlice dimensions, HloComputation* function, + string* failure_reason) { + ReductionGenerator reduction_generator = + MatchReductionGenerator(function, failure_reason); + if (!reduction_generator) { + return false; + } + + int vectorization_factor_in_bytes = + target_machine_features_.vectorization_factor_in_bytes(); + + // We try to process vectorization_factor elements at the same time. + const int vectorization_factor = + vectorization_factor_in_bytes / + ShapeUtil::ByteSizeOfPrimitiveType(reduce->shape().element_type()); + + bool is_reduction_over_minor_dimension = + std::find(dimensions.begin(), dimensions.end(), + arg->shape().layout().minor_to_major(0)) != dimensions.end(); + + unsigned element_alignment = + GCD(ShapeUtil::ByteSizeOfPrimitiveType(reduce->shape().element_type()), + MinimumAlignmentForPrimitiveType(reduce->shape().element_type())); + + if (is_reduction_over_minor_dimension) { + // TODO(sanjoy): Implement vectorized reduction over the minor dimension. + *failure_reason = "reduction over minor dimension not implemented"; + return false; + } + + CHECK(!ShapeUtil::IsTuple(reduce->shape())); + + // We know we're not reducing over the most minor dimension, which means we + // can lower the reduction loop as: + // + // 1. We're reducing over dimensions R0, R1. + // 2. D0 is the most minor dimension. + // 3. VS is the vectorization stride (we want to reduce this many elements at + // once) + // + // for (d1 in D1) { + // for (d0 in D0 with stride VS) { + // vector_acc = init + // for (r1 in R1) { + // for (r0 in R0) { + // vector_acc = elementwise_reduce(vector_acc, input[d1, d0, r1, r0] + // } + // } + // output[d1, d0] = vector_acc + // } + // } + + llvm_ir::ForLoopNest loop_nest(&ir_builder_); + llvm_ir::IrArray::Index array_index(reduce->shape().dimensions_size()); + for (int i = reduce->shape().layout().minor_to_major_size() - 1; i > 0; --i) { + int64 dimension = reduce->shape().layout().minor_to_major(i); + int64 start_index = 0; + int64 end_index = reduce->shape().dimensions(dimension); + std::unique_ptr loop = + loop_nest.AddLoop(start_index, end_index, + tensorflow::strings::Printf("dim.%lld", dimension)); + array_index[dimension] = loop->GetIndVarValue(); + } + + int64 innermost_dimension = reduce->shape().layout().minor_to_major(0); + int64 innermost_dimension_size = + reduce->shape().dimensions(innermost_dimension); + + if (llvm::BasicBlock* innermost_body_bb = + loop_nest.GetInnerLoopBodyBasicBlock()) { + SetToFirstInsertPoint(innermost_body_bb, &ir_builder_); + } + + auto outermost_loop_exit_block = loop_nest.GetOuterLoopExitBasicBlock(); + + if (innermost_dimension_size >= vectorization_factor) { + int64 start_index = 0; + int64 end_index = (innermost_dimension_size / vectorization_factor) * + vectorization_factor; + std::unique_ptr loop = loop_nest.AddLoop( + start_index, end_index, vectorization_factor, + tensorflow::strings::Printf("dim.%lld", innermost_dimension)); + array_index[innermost_dimension] = loop->GetIndVarValue(); + + SetToFirstInsertPoint(loop->GetBodyBasicBlock(), &ir_builder_); + + ShardedVectorType vector_type = CreateShardedVectorType( + reduce->shape().element_type(), vectorization_factor); + TF_ASSIGN_OR_RETURN(std::vector accumulator, + EmitInnerLoopForVectorizedReduction( + reduction_generator, array_index, vector_type, + init_value, arg, dimensions, element_alignment)); + + TF_ASSIGN_OR_RETURN(llvm::Value * target_address, + EmitTargetAddressForOp(reduce)); + llvm_ir::IrArray target_array(target_address, reduce->shape()); + AddAliasingInformationToIrArray(*reduce, &target_array); + llvm::Value* output_address = + target_array.EmitArrayElementAddress(array_index, &ir_builder_); + EmitShardedVectorStore(output_address, accumulator, element_alignment, + target_array); + + if (auto exit_terminator = loop->GetExitBasicBlock()->getTerminator()) { + CHECK_GT(reduce->shape().layout().minor_to_major_size(), 1); + ir_builder_.SetInsertPoint(exit_terminator); + } else { + CHECK_EQ(reduce->shape().layout().minor_to_major_size(), 1); + ir_builder_.SetInsertPoint(loop->GetExitBasicBlock()); + } + } + + // Since we increment the stride for the inner dimension by more than 1, we + // may need to peel out an "epilogue" iteration to get the remaining elements + // in the following case: + if (innermost_dimension_size % vectorization_factor) { + // TODO(b/63775531): Consider using a scalar loop here to save on code size. + array_index[innermost_dimension] = + ir_builder_.getInt64(innermost_dimension_size - + (innermost_dimension_size % vectorization_factor)); + + ShardedVectorType vector_type = CreateShardedVectorType( + reduce->shape().element_type(), + innermost_dimension_size % vectorization_factor); + TF_ASSIGN_OR_RETURN(std::vector accumulator, + EmitInnerLoopForVectorizedReduction( + reduction_generator, array_index, vector_type, + init_value, arg, dimensions, element_alignment)); + + TF_ASSIGN_OR_RETURN(llvm::Value * target_address, + EmitTargetAddressForOp(reduce)); + llvm_ir::IrArray target_array(target_address, reduce->shape()); + AddAliasingInformationToIrArray(*reduce, &target_array); + llvm::Value* output_address = + target_array.EmitArrayElementAddress(array_index, &ir_builder_); + EmitShardedVectorStore(output_address, accumulator, element_alignment, + target_array); + } + + if (outermost_loop_exit_block) { + ir_builder_.SetInsertPoint(outermost_loop_exit_block); + } + + TF_ASSIGN_OR_RETURN(llvm::Value * target_address, + EmitTargetAddressForOp(reduce)); + + emitted_value_[reduce] = target_address; + return true; +} + Status IrEmitter::HandleReduce(HloInstruction* reduce, HloInstruction* arg, HloInstruction* init_value, tensorflow::gtl::ArraySlice dimensions, HloComputation* function) { + if (!options::VectorizedReduceDisabled(hlo_module_config_)) { + string vectorization_failure_reason; + TF_ASSIGN_OR_RETURN( + bool vectorization_successful, + EmitVectorizedReduce(reduce, arg, init_value, dimensions, function, + &vectorization_failure_reason)); + if (vectorization_successful) { + VLOG(1) << "Successfully vectorized reduction " << reduce->ToString() + << "\n"; + return Status::OK(); + } else { + VLOG(1) << "Could not vectorize reduction " << reduce->ToString() << ": " + << vectorization_failure_reason; + } + } + // The called computation should have been emitted previously. llvm::Function* reducer_function = FindOrDie(emitted_functions_, function); return EmitTargetElementLoop( @@ -1136,6 +1940,271 @@ Status IrEmitter::HandleSend(HloInstruction* send) { return Unimplemented("Send is not implemented on CPU. See b/33942983."); } +Status IrEmitter::HandleSlice(HloInstruction* slice, HloInstruction* operand) { + VLOG(2) << "HandleSlice: " << slice->ToString(); + + // The code below emits a sequential loop nest. For the parallel backend, use + // EmitParallelTargetElementLoop() which respects dynamic loop bounds. + if (ShouldEmitParallelLoopFor(*slice)) { + return DefaultAction(slice); + } + + // The code below assumes the layouts are equal. + if (!LayoutUtil::Equal(operand->shape().layout(), slice->shape().layout())) { + return DefaultAction(slice); + } + + TF_ASSIGN_OR_RETURN(llvm::Value * target_address, + EmitTargetAddressForOp(slice)); + emitted_value_[slice] = target_address; + + if (ShapeUtil::HasZeroElements(slice->shape())) { + return Status::OK(); + } + + const Layout& layout = operand->shape().layout(); + const int64 num_dims = operand->shape().dimensions_size(); + + // The slice lowering finds maximal contiguous blocks of memory that can be + // copied from the source to the target. This is done by looking at the + // source/target layout in minor to major order and do the following: + // + // * Find an initial segment of dimensions along which the slice uses the + // whole dimension. These are the "inner" dimensions and can be folded into + // the memcpy. + // + // * Of the remaining dimensions decide which ones require loops. + // + // * Implement the memcpy within the innermost loop. + + tensorflow::gtl::FlatSet inner_dims; + for (int64 dim : layout.minor_to_major()) { + if (operand->shape().dimensions(dim) != slice->shape().dimensions(dim)) { + break; + } + inner_dims.insert(dim); + } + + const bool is_trivial_copy = (inner_dims.size() == num_dims); + if (is_trivial_copy) { + if (ShapeUtil::IsEffectiveScalar(slice->shape())) { + return DefaultAction(slice); + } else { + return EmitMemcpy(*slice, *operand); + } + } + + // 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); }, + operand->shape()); + + const int64 primitive_elements_per_logical_element = + ShapeUtil::ElementsIn(logical_element_shape); + + // memcpy_dim is the innermost (in terms of layout) dimension for which the + // slice does *not* just copy all the elements along the dimension. + const int64 memcpy_dim = layout.minor_to_major(inner_dims.size()); + + const bool memcpy_is_contiguous = slice->slice_strides(memcpy_dim) == 1; + // The number of logical elements that can be copied in a single call + // to memcpy. We can only copy 1 element at a time if there is a non-trivial + // stride. + const int64 memcpy_logical_elements = + memcpy_is_contiguous + ? slice->slice_limits(memcpy_dim) - slice->slice_starts(memcpy_dim) + : 1; + + // Determine the dimensions that get lowered as loops. + std::vector outer_dims; + for (int64 i = 0; i < num_dims - inner_dims.size() - 1; ++i) { + outer_dims.push_back(LayoutUtil::Major(layout, i)); + } + + // Is the slice along the memcpy dimension contiguous? If not, then memcpy_dim + // needs to be wrapped around a loop as well. + if (!memcpy_is_contiguous) { + outer_dims.push_back(memcpy_dim); + } + + llvm_ir::IrArray target_array(target_address, slice->shape()); + AddAliasingInformationToIrArray(*slice, &target_array); + + const int64 num_outer_loops = outer_dims.size(); + llvm_ir::ForLoopNest loops(&ir_builder_); + llvm_ir::IrArray::Index target_index = + loops.AddLoopsForShapeOnDimensions(slice->shape(), outer_dims, "slice"); + + // Only the indices for the outer dimensions have been initialized in + // target_index. The rest of the indices should get initialized to 0, since + // for the rest of the dimensions the copy writes to the full dimension. + std::replace(target_index.begin(), target_index.end(), + static_cast(nullptr), + static_cast(ir_builder_.getInt64(0))); + + if (num_outer_loops > 0) { + SetToFirstInsertPoint(loops.GetInnerLoopBodyBasicBlock(), &ir_builder_); + } + + llvm_ir::IrArray source_array(GetEmittedValueFor(operand), operand->shape()); + + const llvm_ir::IrArray::Index source_index = target_index.SourceIndexOfSlice( + /*shape=*/slice->shape(), /*starts=*/slice->slice_starts(), + /*strides=*/slice->slice_strides(), /*builder=*/&ir_builder_); + + llvm::Value* memcpy_dest = target_array.EmitArrayElementAddress( + target_index, &ir_builder_, "slice.dest"); + llvm::Value* memcpy_source = source_array.EmitArrayElementAddress( + source_index, &ir_builder_, "slice.source"); + + const int64 memcpy_elements = + primitive_elements_per_logical_element * memcpy_logical_elements; + + EmitTransferElements(memcpy_dest, memcpy_source, memcpy_elements, + slice->shape().element_type(), target_array, + source_array); + + if (VLOG_IS_ON(2)) { + const int64 memcpy_bytes = + ShapeUtil::ByteSizeOf(logical_element_shape) * memcpy_elements; + VLOG(2) << " emitted copy of " << memcpy_bytes << " bytes inside " + << num_outer_loops << " loops"; + } + + if (num_outer_loops > 0) { + SetToFirstInsertPoint(loops.GetOuterLoopExitBasicBlock(), &ir_builder_); + } + + return Status::OK(); +} + +Status IrEmitter::HandleDynamicSlice(HloInstruction* dynamic_slice, + HloInstruction* operand, + HloInstruction* /*start_indices*/) { + if (ShapeUtil::IsScalar(dynamic_slice->shape())) { + TF_ASSIGN_OR_RETURN(llvm::Value * target_address, + EmitTargetAddressForOp(dynamic_slice)); + emitted_value_[dynamic_slice] = target_address; + return EmitMemcpy(*operand, *dynamic_slice); + } + return DefaultAction(dynamic_slice); +} + +namespace { + +// Returns the first non-GetTupleElement ancestor instruction of 'hlo'. +// If the first non-GTE ancestor is tuple-shaped, populates 'index' with the +// (possibly nested) tuple indices used on the path from ancestor to 'hlo'. +const HloInstruction* LatestNonGteAncestorAndIndex(const HloInstruction* hlo, + ShapeIndex* index) { + if (hlo->opcode() == HloOpcode::kGetTupleElement) { + const auto* operand = LatestNonGteAncestorAndIndex(hlo->operand(0), index); + index->push_back(hlo->tuple_index()); + return operand; + } + return hlo; +} + +// Checks if we can emit code for DynamicUpdateSlice to update data in-place. +// Returns true if operand 0 of DynamicUpdateSlice and its output buffer +// share the same buffer allocation. +// Returns false otherwise. +// TODO(b/64142684) Share code with GPU implementation. +bool CanUpdateDynamicSliceInPlace(const BufferAssignment& assignment, + HloInstruction* dynamic_update_slice) { + CHECK_EQ(HloOpcode::kDynamicUpdateSlice, dynamic_update_slice->opcode()); + + // Walk DynamicUpdateSlice operand(0) to parameter and get its + // associated operand. See if it shares an allocation with this operand. + ShapeIndex index; + auto* operand = + LatestNonGteAncestorAndIndex(dynamic_update_slice->operand(0), &index); + if (operand->opcode() != HloOpcode::kParameter) { + return false; + } + + BufferAllocation::Slice operand_slice = + assignment.GetUniqueSlice(operand, index).ConsumeValueOrDie(); + + BufferAllocation::Slice dynamic_update_slice_slice = + assignment.GetUniqueTopLevelSlice(dynamic_update_slice) + .ConsumeValueOrDie(); + + return operand_slice == dynamic_update_slice_slice; +} + +} // namespace + +Status IrEmitter::HandleDynamicUpdateSlice(HloInstruction* dynamic_update_slice, + HloInstruction* operand, + HloInstruction* update, + HloInstruction* start_indices) { + if (ShapeUtil::IsScalar(dynamic_update_slice->shape())) { + TF_ASSIGN_OR_RETURN(llvm::Value * target_address, + EmitTargetAddressForOp(dynamic_update_slice)); + emitted_value_[dynamic_update_slice] = target_address; + return EmitMemcpy(*update, *dynamic_update_slice); + } else if (CanUpdateDynamicSliceInPlace(assignment_, dynamic_update_slice)) { + VLOG(2) << "Emitting HandleDynamicUpdateSlice in-place."; + // DynamicUpdateSlice's operand(0) and 'fusion' output share the same + // BufferAllocation::Slice, so it is safe to emit code to update the slice + // 'in-place'. This avoids copying data outside of the slice update region. + // TODO(b/64142684) Implement in-place update for fused DynamicUpdateSlice. + + // Emit IR to read dynamic start indices from 'start_indices'. + const int64 rank = ShapeUtil::Rank(operand->shape()); + llvm_ir::IrArray::Index start_index(rank); + for (int64 i = 0; i < rank; ++i) { + llvm_ir::IrArray::Index dim_index({ir_builder_.getInt64(i)}); + llvm_ir::IrArray start_indices_array(GetIrArrayForOp(start_indices)); + start_index[i] = + start_indices_array.EmitReadArrayElement(dim_index, &ir_builder_); + } + + // Create loop body emitter which emits code to do the following: + // *) Map requested 'index' and slice 'start_index' to input/output shape + // as 'output_index'. + // *) Reads value from 'update'. + // *) Writes value to input/output array at 'output_index'. + auto loop_body_emitter = + [&](const llvm_ir::IrArray::Index& index) -> Status { + // Calculate 'output_index' at which to write value from update. + llvm_ir::IrArray::Index output_index(rank); + for (int64 i = 0; i < rank; ++i) { + // Emit IR which computes: + // output_index = (start_index + index) % dim_size + llvm::Value* dim_size = llvm::ConstantInt::get( + index[i]->getType(), operand->shape().dimensions(i)); + llvm::Value* start_index0 = ir_builder_.CreateZExtOrBitCast( + start_index[i], index[i]->getType()); + output_index[i] = ir_builder_.CreateURem( + ir_builder_.CreateAdd(start_index0, index[i]), dim_size); + } + + // Read value from 'update'. + llvm_ir::IrArray update_array(GetIrArrayForOp(update)); + llvm::Value* update_data = + update_array.EmitReadArrayElement(index, &ir_builder_); + + // Write value to output array. + llvm_ir::IrArray(GetEmittedValueFor(operand), operand->shape()) + .EmitWriteArrayElement(output_index, update_data, &ir_builder_); + return Status::OK(); + }; + + TF_RETURN_IF_ERROR( + llvm_ir::LoopEmitter(loop_body_emitter, update->shape(), &ir_builder_) + .EmitLoop()); + + TF_ASSIGN_OR_RETURN(llvm::Value * dynamic_update_slice_address, + EmitTargetAddressForOp(dynamic_update_slice)); + emitted_value_[dynamic_update_slice] = dynamic_update_slice_address; + return Status::OK(); + } + return DefaultAction(dynamic_update_slice); +} + Status IrEmitter::HandleRecv(HloInstruction* recv) { // TODO(b/33942983): Support Send/Recv on CPU. return Unimplemented("Recv is not implemented on CPU. See b/33942983."); @@ -1245,7 +2314,7 @@ Status IrEmitter::HandleFusion(HloInstruction* fusion) { TF_RETURN_IF_ERROR(DotOpEmitter::EmitDotOperation( *dot, dot->operand(0)->IsRank2Transpose(), dot->operand(1)->IsRank2Transpose(), target_array, lhs_array, rhs_array, - GetExecutableRunOptionsArgument(), &ir_builder_)); + GetExecutableRunOptionsArgument(), &ir_builder_, hlo_module_config_)); emitted_value_[fusion] = target_address; return Status::OK(); @@ -1265,13 +2334,12 @@ Status IrEmitter::HandleFusion(HloInstruction* fusion) { } } -Status IrEmitter::HandleCall( - HloInstruction* call, tensorflow::gtl::ArraySlice operands, - HloComputation* computation) { +Status IrEmitter::HandleCall(HloInstruction* call) { + HloComputation* computation = call->to_apply(); llvm::Function* call_ir_function = FindOrDie(emitted_functions_, computation); std::vector parameter_addresses; - for (HloInstruction* operand : operands) { + for (const HloInstruction* operand : call->operands()) { parameter_addresses.push_back(GetEmittedValueFor(operand)); } @@ -1322,14 +2390,14 @@ Status IrEmitter::HandleCustomCall( return Status::OK(); } -Status IrEmitter::HandleWhile(HloInstruction* xla_while, HloInstruction* init, - HloComputation* condition, HloComputation* body) { +Status IrEmitter::HandleWhile(HloInstruction* xla_while) { // Precondition: Condition computation must return a scalar bool. + HloComputation* condition = xla_while->while_condition(); TF_RET_CHECK(ShapeUtil::IsScalar(condition->root_instruction()->shape()) && condition->root_instruction()->shape().element_type() == PRED) << "While condition computation must return bool"; // Check that all while-related buffers share an allocation slice. - TF_RETURN_IF_ERROR(ShapeUtil::ForEachSubshape( + TF_RETURN_IF_ERROR(ShapeUtil::ForEachSubshapeWithStatus( xla_while->shape(), [this, &xla_while](const Shape& /*subshape*/, const ShapeIndex& index) -> Status { @@ -1361,12 +2429,14 @@ Status IrEmitter::HandleWhile(HloInstruction* xla_while, HloInstruction* init, })); // Set emitted value to that of 'init' with which it shares an allocation. + const HloInstruction* init = xla_while->operand(0); emitted_value_[xla_while] = GetEmittedValueFor(init); // The called computation should have been emitted previously. llvm::Function* condition_ir_function = FindOrDie(emitted_functions_, condition); - llvm::Function* body_ir_function = FindOrDie(emitted_functions_, body); + llvm::Function* body_ir_function = + FindOrDie(emitted_functions_, xla_while->while_body()); // Generating: // while (Condition(while_result)) { @@ -1415,6 +2485,173 @@ Status IrEmitter::HandleWhile(HloInstruction* xla_while, HloInstruction* init, return Status::OK(); } +StatusOr IrEmitter::EmitFastConcatenate( + HloInstruction* concatenate, + tensorflow::gtl::ArraySlice operands, + string* failure_reason) { + if (ShouldEmitParallelLoopFor(*concatenate)) { + *failure_reason = + "cannot generate memcpy-based concat for the parallel CPU backend"; + return false; + } + + const Shape& output_shape = concatenate->shape(); + for (auto* op : operands) { + if (!LayoutUtil::Equal(op->shape().layout(), output_shape.layout())) { + *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). + + int64 concat_dim = concatenate->dimensions(0); + const Layout& output_layout = output_shape.layout(); + auto concat_dim_layout_itr = + std::find(output_layout.minor_to_major().begin(), + output_layout.minor_to_major().end(), concat_dim); + + std::vector inner_dims(output_layout.minor_to_major().begin(), + concat_dim_layout_itr); + std::vector outer_dims(std::next(concat_dim_layout_itr), + output_layout.minor_to_major().end()); + + llvm::Type* i8_ptr_type = ir_builder_.getInt8PtrTy(); + llvm::Type* i8_type = ir_builder_.getInt8Ty(); + + TF_ASSIGN_OR_RETURN(llvm::Value * target_address, + EmitTargetAddressForOp(concatenate)); + + llvm_ir::IrArray target_array(target_address, output_shape); + + llvm_ir::ForLoopNest loops(&ir_builder_); + llvm_ir::IrArray::Index outer_dims_index = + loops.AddLoopsForShapeOnDimensions(output_shape, outer_dims, "concat"); + std::replace(outer_dims_index.begin(), outer_dims_index.end(), + static_cast(nullptr), + static_cast(ir_builder_.getInt64(0))); + + if (!outer_dims.empty()) { + SetToFirstInsertPoint(loops.GetInnerLoopBodyBasicBlock(), &ir_builder_); + } + + PrimitiveType primitive_type = output_shape.element_type(); + unsigned primitive_type_size = + ShapeUtil::ByteSizeOfPrimitiveType(primitive_type); + + AddAliasingInformationToIrArray(*concatenate, &target_array); + + // Contiguous subregions from each operand to the concatenate contribute to a + // contiguous subregion in the target buffer starting at target_region_begin. + llvm::Value* target_region_begin = ir_builder_.CreateBitCast( + target_array.EmitArrayElementAddress(outer_dims_index, &ir_builder_, + "target_region"), + i8_ptr_type); + int64 byte_offset_into_target_region = 0; + + int64 inner_dims_product = + std::accumulate(inner_dims.begin(), inner_dims.end(), 1l, + [&](int64 product, int64 inner_dim) { + return product * output_shape.dimensions(inner_dim); + }); + + // For each operand, emit a memcpy from the operand to the target of size + // equal to the product of inner dimensions. + for (HloInstruction* operand : operands) { + const Shape& input_shape = operand->shape(); + llvm_ir::IrArray source_array(GetEmittedValueFor(operand), input_shape); + AddAliasingInformationToIrArray(*operand, &source_array); + + llvm::Value* copy_source_address = ir_builder_.CreateBitCast( + source_array.EmitArrayElementAddress(outer_dims_index, &ir_builder_, + "src_addr"), + i8_ptr_type); + + llvm::Value* copy_target_address = ir_builder_.CreateGEP( + i8_type, target_region_begin, + ir_builder_.getInt64(byte_offset_into_target_region)); + + EmitTransferElements( + copy_target_address, copy_source_address, + inner_dims_product * input_shape.dimensions(concat_dim), primitive_type, + target_array, source_array); + + byte_offset_into_target_region += inner_dims_product * + input_shape.dimensions(concat_dim) * + primitive_type_size; + } + + if (!outer_dims.empty()) { + SetToFirstInsertPoint(loops.GetOuterLoopExitBasicBlock(), &ir_builder_); + } + + emitted_value_[concatenate] = target_address; + + return true; +} + +void IrEmitter::EmitTransferElements(llvm::Value* target, llvm::Value* source, + int64 element_count, + PrimitiveType primitive_type, + const llvm_ir::IrArray& target_array, + const llvm_ir::IrArray& source_array) { + unsigned primitive_type_size = + ShapeUtil::ByteSizeOfPrimitiveType(primitive_type); + unsigned element_alignment = GCD( + primitive_type_size, MinimumAlignmentForPrimitiveType(primitive_type)); + llvm::Type* primitive_ptr_type = llvm::PointerType::getUnqual( + llvm_ir::PrimitiveTypeToIrType(primitive_type, &ir_builder_)); + + if (element_count == 1) { + auto* load_instruction = ir_builder_.CreateAlignedLoad( + ir_builder_.CreateBitCast(source, primitive_ptr_type), + element_alignment); + source_array.AnnotateLoadStoreInstructionWithMetadata(load_instruction); + auto* store_instruction = ir_builder_.CreateAlignedStore( + load_instruction, ir_builder_.CreateBitCast(target, primitive_ptr_type), + element_alignment); + target_array.AnnotateLoadStoreInstructionWithMetadata(store_instruction); + } else { + auto* memcpy_instruction = ir_builder_.CreateMemCpy( + target, source, element_count * primitive_type_size, element_alignment); + + // The memcpy does the load and the store internally. The aliasing related + // metadata has to reflect that. + std::map merged_metadata = + llvm_ir::MergeMetadata(&module_->getContext(), source_array.metadata(), + target_array.metadata()); + for (const auto& kind_md_pair : merged_metadata) { + memcpy_instruction->setMetadata(kind_md_pair.first, kind_md_pair.second); + } + } +} + +Status IrEmitter::HandleConcatenate( + HloInstruction* concatenate, + tensorflow::gtl::ArraySlice operands) { + string failure_reason; + TF_ASSIGN_OR_RETURN( + bool successful, + EmitFastConcatenate(concatenate, operands, &failure_reason)); + if (successful) { + VLOG(1) << "Emitted fast concatenate for " << concatenate->ToString(); + return Status::OK(); + } + + VLOG(1) << "Could not emit fast concatenate for " << concatenate->ToString() + << ": " << failure_reason; + + return DefaultAction(concatenate); +} + Status IrEmitter::FinishVisit(HloInstruction* root) { // When this method is called, we should have already emitted an IR value for // the root (return) op. The IR value holds the address of the buffer holding @@ -1426,7 +2663,23 @@ Status IrEmitter::FinishVisit(HloInstruction* root) { llvm::Value* root_value = GetEmittedValueFor(root); VLOG(2) << " value: " << llvm_ir::DumpToString(*root_value); - if (auto* prof_counter = GetProfileCounterFor(/*hlo=*/nullptr)) { + // For the parallel cpu backend, we record the total for each embedded + // computation callee with its caller kCall HLO. + HloInstruction* hlo_to_lookup = nullptr; + if (IsParallelContext()) { + auto* computation = root->parent(); + auto* entry_computation = computation->parent()->entry_computation(); + if (computation != entry_computation) { + for (auto& instruction : entry_computation->instructions()) { + if (instruction->opcode() == HloOpcode::kCall && + instruction->to_apply()->root_instruction() == root) { + hlo_to_lookup = instruction.get(); + break; + } + } + } + } + if (auto* prof_counter = GetProfileCounterFor(hlo_to_lookup)) { profiling_state_.RecordCompleteComputation(&ir_builder_, prof_counter); } @@ -1521,7 +2774,7 @@ void IrEmitter::ProfilingState::RecordCycleDelta(llvm::IRBuilder<>* ir_builder, void IrEmitter::ProfilingState::RecordCompleteComputation( llvm::IRBuilder<>* ir_builder, llvm::Value* prof_counter) { - if (is_entry_computation_ && last_read_cycle_end_ && + if (is_top_level_computation_ && last_read_cycle_end_ && first_read_cycle_start_) { UpdateProfileCounter(ir_builder, prof_counter, last_read_cycle_end_, first_read_cycle_start_); @@ -1529,6 +2782,7 @@ void IrEmitter::ProfilingState::RecordCompleteComputation( } Status IrEmitter::Preprocess(HloInstruction* hlo) { + VLOG(3) << "Visiting: " << hlo->ToString(); if (hlo_to_profile_idx_ && hlo_to_profile_idx_->count(hlo)) { profiling_state_.RecordCycleStart(&ir_builder_, hlo); } @@ -1567,13 +2821,24 @@ llvm::Argument* IrEmitter::GetResultArgument() { } llvm::Argument* IrEmitter::GetProfileCountersArgument() { - return hlo_to_profile_idx_ ? GetArg(compute_function_, 4) : nullptr; + const int64 arg_index = IsParallelContext() ? 5 : 4; + return hlo_to_profile_idx_ ? GetArg(compute_function_, arg_index) : nullptr; } llvm::Value* IrEmitter::GetTempBuffersArgument() { return GetArg(compute_function_, 3); } +llvm::Value* IrEmitter::GetDynamicLoopBound(const int64 offset) { + CHECK_GT(num_dynamic_loop_bounds_, 0); + CHECK_LT(offset, num_dynamic_loop_bounds_ * 2); + llvm::Argument* loop_bounds_arg = GetArg(compute_function_, 4); + string name = tensorflow::strings::StrCat("dynamic_loop_bound_", offset); + return ir_builder_.CreateLoad( + ir_builder_.CreateGEP(loop_bounds_arg, ir_builder_.getInt64(offset), + llvm_ir::AsStringRef(name))); +} + llvm::Value* IrEmitter::GetExecutableRunOptionsArgument() { return GetArg(compute_function_, 1); } @@ -1606,11 +2871,14 @@ llvm::Value* IrEmitter::EmitTempBufferPointer( GetTempBuffersArgument(), slice.index(), &ir_builder_); llvm::LoadInst* tempbuf_address_base = ir_builder_.CreateLoad(tempbuf_address_ptr); - // Loading the address of a buffer is invariant of the point at which the - // load is executed in the program because we never reassign buffers. - tempbuf_address_base->setMetadata( - llvm::LLVMContext::MD_invariant_load, - llvm::MDNode::get(tempbuf_address_base->getContext(), /*MDs=*/{})); + if (hlo_module_config_.debug_options() + .xla_llvm_enable_invariant_load_metadata()) { + // Loading the address of a buffer is invariant of the point at which the + // load is executed in the program because we never reassign buffers. + tempbuf_address_base->setMetadata( + llvm::LLVMContext::MD_invariant_load, + llvm::MDNode::get(tempbuf_address_base->getContext(), /*MDs=*/{})); + } llvm_ir::SetTbaaForInstruction(tempbuf_address_base, target_shape, /*is_pointer_to=*/true); AttachAlignmentMetadataForLoad(tempbuf_address_base, allocation.size()); @@ -1700,18 +2968,17 @@ llvm::Value* IrEmitter::EmitArrayFunctionCall( } StatusOr IrEmitter::EmitTargetAddressForOp( - const HloInstruction* op) { - const Shape& target_shape = op->shape(); - if (op == op->parent()->root_instruction()) { + const HloInstruction* op, const ShapeIndex& shape_index) { + const Shape& target_shape = ShapeUtil::GetSubshape(op->shape(), shape_index); + if (op == op->parent()->root_instruction() && shape_index.empty()) { // For the root node, we write directly to the output buffer of the // function. llvm::Argument* retval = GetResultArgument(); - if (!ShapeUtil::HasZeroElements(target_shape)) { + if (!ShapeUtil::IsNil(target_shape)) { llvm::AttrBuilder attr_builder; attr_builder.addAlignmentAttr(MinimumAlignmentForShape(target_shape)); attr_builder.addDereferenceableAttr(ByteSizeOf(target_shape)); - retval->addAttr(llvm::AttributeList::get( - retval->getContext(), retval->getArgNo() + 1, attr_builder)); + retval->addAttrs(attr_builder); } return ir_builder_.CreateBitCast(retval, IrShapeType(target_shape)->getPointerTo()); @@ -1735,21 +3002,109 @@ Status IrEmitter::EmitTargetElementLoop( TF_ASSIGN_OR_RETURN(llvm::Value * target_address, EmitTargetAddressForOp(target_op)); VLOG(2) << " target address: " << llvm_ir::DumpToString(*target_address); - llvm_ir::IrArray target_array(target_address, target_shape); - AddAliasingInformationToIrArray(*target_op, &target_array); - TF_RETURN_IF_ERROR( - llvm_ir::LoopEmitter(element_generator, target_array, &ir_builder_) - .EmitLoop()); + if (target_op->IsMultiOutputFusion()) { + // For multiple outputs fusion, we need to emit each operand and the root. + TF_RET_CHECK(num_dynamic_loop_bounds_ == 0); + std::vector output_arrays; + for (int64 i = 0; i < ShapeUtil::TupleElementCount(target_shape); ++i) { + TF_ASSIGN_OR_RETURN(BufferAllocation::Slice slice, + assignment_.GetUniqueSlice(target_op, {i})); + const Shape& element_shape = ShapeUtil::GetSubshape(target_shape, {i}); + llvm::Value* op_target_address = + EmitTempBufferPointer(slice, element_shape); + output_arrays.push_back( + llvm_ir::IrArray(op_target_address, element_shape)); + } + TF_RETURN_IF_ERROR( + llvm_ir::LoopEmitter(element_generator, output_arrays, &ir_builder_) + .EmitLoop()); + + std::vector tuple_operand_ptrs; + for (int64 i = 0; i < output_arrays.size(); ++i) { + tuple_operand_ptrs.push_back(output_arrays[i].GetBasePointer()); + } + llvm_ir::EmitTuple(llvm_ir::IrArray(target_address, target_shape), + tuple_operand_ptrs, &ir_builder_); + + } else { + llvm_ir::IrArray target_array(target_address, target_shape); + AddAliasingInformationToIrArray(*target_op, &target_array); + + if (ShouldEmitParallelLoopFor(*target_op)) { + TF_RETURN_IF_ERROR(EmitParallelTargetElementLoop( + target_shape, element_generator, &target_array)); + } else { + TF_RETURN_IF_ERROR( + llvm_ir::LoopEmitter(element_generator, target_array, &ir_builder_) + .EmitLoop()); + } + } + emitted_value_[target_op] = target_address; return Status::OK(); } +Status IrEmitter::EmitParallelTargetElementLoop( + const Shape& target_shape, + const llvm_ir::ElementGenerator& element_generator, + llvm_ir::IrArray* target_array) { + CHECK(!ShapeUtil::IsTuple(target_shape)); + CHECK(!ShapeUtil::IsScalar(target_shape)); + + // Emit code to read dynamic loop bounds from function argument 4. + std::vector dynamic_loop_bounds(2 * num_dynamic_loop_bounds_); + for (int i = 0; i < 2 * num_dynamic_loop_bounds_; ++i) { + dynamic_loop_bounds[i] = GetDynamicLoopBound(i); + } + + llvm_ir::ForLoopNest loop_nest(&ir_builder_); + const int64 num_dims = target_shape.dimensions_size(); + llvm_ir::IrArray::Index array_index(num_dims); + + // Add loops from outer-most to inner-most dimensions. + for (int i = target_shape.layout().minor_to_major_size() - 1; i >= 0; --i) { + const int64 dimension = target_shape.layout().minor_to_major(i); + const int bounds_index = num_dims - 1 - i; + if (bounds_index < num_dynamic_loop_bounds_) { + // Emit dynamic loop bounds for this dimension. Dynamic loop bounds + // are read from ir function dynamic loop bounds argument. + llvm::Value* start_index = dynamic_loop_bounds[bounds_index * 2 + 0]; + llvm::Value* end_index = dynamic_loop_bounds[bounds_index * 2 + 1]; + + std::unique_ptr loop = loop_nest.AddLoop( + /*suffix=*/tensorflow::strings::Printf("dim.%lld", dimension), + start_index, end_index); + array_index[dimension] = loop->GetIndVarValue(); + } else { + // Emit static loop bounds for this dimension. + std::unique_ptr loop = loop_nest.AddLoop( + /*start_index=*/0, + /*end_index=*/target_shape.dimensions(dimension), + /*suffix=*/tensorflow::strings::Printf("dim.%lld", dimension)); + array_index[dimension] = loop->GetIndVarValue(); + } + } + // Point IR builder at inner loop BB. + SetToFirstInsertPoint(loop_nest.GetInnerLoopBodyBasicBlock(), &ir_builder_); + + // Emit loop body. + TF_ASSIGN_OR_RETURN(llvm::Value * target_element, + element_generator(array_index)); + target_array->EmitWriteArrayElement(array_index, target_element, + &ir_builder_); + // Point IR builder at outer loop exit BB. + SetToFirstInsertPoint(loop_nest.GetOuterLoopExitBasicBlock(), &ir_builder_); + + return Status::OK(); +} + Status IrEmitter::EmitMemcpy(const HloInstruction& source, const HloInstruction& destination) { llvm::Value* source_value = GetEmittedValueFor(&source); llvm::Value* destination_value = GetEmittedValueFor(&destination); int64 source_size = ByteSizeOf(source.shape()); + // TODO(b/63762267): Be more aggressive about specifying alignment. ir_builder_.CreateMemCpy(destination_value, source_value, source_size, 1); return Status::OK(); } @@ -1787,5 +3142,36 @@ Status IrEmitter::DefaultAction(HloInstruction* hlo) { hlo, elemental_emitter.MakeElementGenerator(hlo, operand_to_generator)); } +unsigned TargetMachineFeatures::largest_register_size_in_bytes( + llvm::Function* function) { + auto itr = largest_register_size_in_bytes_.find(function); + if (itr != largest_register_size_in_bytes_.end()) { + return itr->second; + } + + int result = largest_register_size_in_bytes_impl(function); + + InsertOrDie(&largest_register_size_in_bytes_, function, result); + DCHECK_EQ(result, largest_register_size_in_bytes_.begin()->second); + return result; +} + +unsigned TargetMachineFeatures::largest_register_size_in_bytes_impl( + llvm::Function* function) const { + auto register_info = + target_machine_->getSubtargetImpl(*function)->getRegisterInfo(); + + unsigned largest_register_size = 0; + for (const llvm::TargetRegisterClass* register_class : + register_info->regclasses()) { + if (register_class->isAllocatable()) { + largest_register_size = + std::max(largest_register_size, + register_info->getRegSizeInBits(*register_class)); + } + } + + return largest_register_size / 8; +} } // namespace cpu } // namespace xla diff --git a/tensorflow/compiler/xla/service/cpu/ir_emitter.h b/tensorflow/compiler/xla/service/cpu/ir_emitter.h index 66bae457e3741332f23abc7d54b8d775aa193ca9..9eb777c731cacce0a39b62fa374d1537d843b542 100644 --- a/tensorflow/compiler/xla/service/cpu/ir_emitter.h +++ b/tensorflow/compiler/xla/service/cpu/ir_emitter.h @@ -22,11 +22,11 @@ limitations under the License. #include #include -#include "external/llvm/include/llvm/ADT/Triple.h" -#include "external/llvm/include/llvm/IR/Function.h" -#include "external/llvm/include/llvm/IR/IRBuilder.h" -#include "external/llvm/include/llvm/IR/Module.h" -#include "external/llvm/include/llvm/IR/Value.h" +#include "llvm/ADT/Triple.h" +#include "llvm/IR/Function.h" +#include "llvm/IR/IRBuilder.h" +#include "llvm/IR/Module.h" +#include "llvm/IR/Value.h" #include "tensorflow/compiler/xla/service/buffer_assignment.h" #include "tensorflow/compiler/xla/service/dfs_hlo_visitor_with_default.h" #include "tensorflow/compiler/xla/service/hlo_computation.h" @@ -41,12 +41,55 @@ limitations under the License. #include "tensorflow/compiler/xla/xla_data.pb.h" #include "tensorflow/core/lib/core/stringpiece.h" #include "tensorflow/core/lib/gtl/array_slice.h" +#include "tensorflow/core/lib/gtl/flatmap.h" #include "tensorflow/core/platform/macros.h" #include "tensorflow/core/platform/types.h" namespace xla { namespace cpu { +// Wraps an llvm::TargetMachine and parses out some information that feeds into +// code LLVM IR generation decisions. +// +// Ideally we'd be able to use llvm::TargetTransformInfo here (since its +// interface is pretty much a perfect fit for our use case), but obtaining an +// instance of llvm::TargetTransformInfo outside an LLVM pass pipeline without +// super-ugly hacks is difficult. +// +// TODO(b/27457097): See if the LLVM community will be receptive to exposing an +// API that lets us directly create and use llvm::TargetTransformInfo instances +// outside of a pass manager. +class TargetMachineFeatures { + public: + TargetMachineFeatures(llvm::TargetMachine* target_machine) + : target_machine_(target_machine) {} + + // Return the vectorization factor, which is the number of bytes of data + // explicitly vectorized routines will try to process at once. + int vectorization_factor_in_bytes() const { + // Ideally this should be a function of the cache line size (which we can + // get from llvm::TargetTransformInfo::getCacheLineSize) of the target + // machine. Guess a value of 128 bytes for now. + return 128; + } + + // Return the size of the largest register size in bytes. We need to pass in + // "function" since llvm functions can contain annotations for specializing + // them to specific micro-architectures (though currently XLA does not use + // this functionality). + // + // Ideally we should have been able to use + // llvm::TargetTransformInfo::getRegisterBitWidth(true) here. + unsigned largest_register_size_in_bytes(llvm::Function* function); + + private: + unsigned largest_register_size_in_bytes_impl(llvm::Function* function) const; + + tensorflow::gtl::FlatMap + largest_register_size_in_bytes_; + llvm::TargetMachine* target_machine_; +}; + // This class is the top-level API for the XLA HLO --> LLVM IR compiler. It // implements the DfsHloVisitor interface and emits HLO computations as LLVM IR // functions. @@ -60,24 +103,34 @@ class IrEmitter : public DfsHloVisitorWithDefault { // llvm_module: the LLVM module to emit IR into. // hlo_to_profile_idx: the mapping from HLO to its index in the profiling // array. - IrEmitter(const HloModule& hlo_module, const HloModuleConfig& module_config, - const BufferAssignment& assignment, llvm::Module* llvm_module, + IrEmitter(const HloModule& hlo_module, const BufferAssignment& assignment, + llvm::Module* llvm_module, const std::unordered_map* - hlo_to_profile_idx); + hlo_to_profile_idx, + llvm::TargetMachine* target_machine); ~IrEmitter() override; // Emit and return the given HLO computation as an LLVM IR - // function. function_name_prefix is the desired name of the function. If the - // name is not unique among already emitted functions then a suffix is - // appended to make the name unique. is_entry_computation indicates that this - // is the entry computation of the HLO module. If 'instruction_order' is given - // then the HLO instructions are emitted in the given order. In this case, - // 'instruction_order' must be a topological sort of the set of nodes - // accessible from the root of the computation. + // function. + // + // function_name_prefix is the desired name of the function. If the name is + // not unique among already emitted functions then a suffix is appended to + // make the name unique. + // + // 'is_top_level_computation' has the following meanings for each CPU backend: + // *) sequential: indicates that this is the entry computation of the HLO + // module. + // *) parallel: indices that this is the callee of a kCall HLO in the entry + // computation of the HLO module. + // + // If 'instruction_order' is not NULL, then the HLO instructions are emitted + // in the given order. In this case, 'instruction_order' must be a + // topological sort of the set of nodes accessible from the root of the + // computation. StatusOr EmitComputation( HloComputation* computation, const string& function_name_prefix, - bool is_entry_computation, - std::vector* instruction_order = nullptr); + bool is_top_level_computation, + std::vector* instruction_order); protected: // @@ -90,7 +143,7 @@ class IrEmitter : public DfsHloVisitorWithDefault { Status HandleBitcast(HloInstruction* bitcast) override; Status HandleConstant(HloInstruction* constant, const Literal& literal) override; - Status HandleCopy(HloInstruction* copy, HloInstruction* operand) override; + Status HandleCopy(HloInstruction* copy) override; Status HandleGetTupleElement(HloInstruction* get_tuple_element, HloInstruction* operand) override; Status HandleSelect(HloInstruction* select, HloInstruction* pred, @@ -100,9 +153,11 @@ class IrEmitter : public DfsHloVisitorWithDefault { HloInstruction* rhs) override; Status HandleConvolution(HloInstruction* convolution, HloInstruction* lhs, HloInstruction* rhs, const Window& window) override; + Status HandleBatchNormTraining(HloInstruction* batch_norm_training) override; + Status HandleBatchNormGrad(HloInstruction* batch_norm_grad) override; Status HandleCrossReplicaSum(HloInstruction* crs) override; Status HandleInfeed(HloInstruction* infeed) override; - Status HandleOutfeed(HloInstruction* infeed) override; + Status HandleOutfeed(HloInstruction* outfeed) override; Status HandleSort(HloInstruction* sort, HloInstruction* operand) override; Status HandleParameter(HloInstruction* parameter) override; Status HandleReduce(HloInstruction* reduce, HloInstruction* arg, @@ -114,6 +169,15 @@ class IrEmitter : public DfsHloVisitorWithDefault { HloComputation* function) override; Status HandleSelectAndScatter(HloInstruction* instruction) override; Status HandleSend(HloInstruction* send) override; + Status HandleSlice(HloInstruction* slice, + HloInstruction* /*operand*/) override; + Status HandleDynamicSlice(HloInstruction* dynamic_slice, + HloInstruction* /*operand*/, + HloInstruction* /*start_indices*/) override; + Status HandleDynamicUpdateSlice(HloInstruction* dynamic_update_slice, + HloInstruction* /*operand*/, + HloInstruction* /*update*/, + HloInstruction* /*start_indices*/) override; Status HandleRecv(HloInstruction* recv) override; Status HandlePad(HloInstruction* pad) override; Status HandleTuple( @@ -125,14 +189,14 @@ class IrEmitter : public DfsHloVisitorWithDefault { HloComputation* function, tensorflow::gtl::ArraySlice static_operands) override; Status HandleFusion(HloInstruction* fusion) override; - Status HandleCall(HloInstruction* call, - tensorflow::gtl::ArraySlice operands, - HloComputation* computation) override; + Status HandleCall(HloInstruction* call) override; Status HandleCustomCall(HloInstruction* custom_call, tensorflow::gtl::ArraySlice operands, tensorflow::StringPiece custom_call_target) override; - Status HandleWhile(HloInstruction* xla_while, HloInstruction* init, - HloComputation* condition, HloComputation* body) override; + Status HandleWhile(HloInstruction* xla_while) override; + Status HandleConcatenate( + HloInstruction* concatenate, + tensorflow::gtl::ArraySlice operands) override; Status FinishVisit(HloInstruction* root) override; Status Preprocess(HloInstruction* hlo) override; @@ -140,8 +204,7 @@ class IrEmitter : public DfsHloVisitorWithDefault { private: // Private helper to initialize an IR function for the computation. - void InitializeIrFunction(const string& function_name, - bool is_entry_computation); + void InitializeIrFunction(const string& function_name); // Convenience function to generate a GEP into the profile counter parameter // which would correspond to the index for a given HLO. @@ -180,6 +243,11 @@ class IrEmitter : public DfsHloVisitorWithDefault { // computation function being emitted by this emitter. llvm::Value* GetTempBuffersArgument(); + // Emit ir to read and return the ir value for the dynamic loop bound at + // 'offset' from the "dynamic_loop_bounds" argument of the computation + // function being emitted by this emitter. + llvm::Value* GetDynamicLoopBound(const int64 offset); + // Emits code that computes the address of the given temporary buffer to the // function. target_shape is the shape of this temporary buffer. // The returned Value's type is a pointer to element_type. @@ -250,6 +318,15 @@ class IrEmitter : public DfsHloVisitorWithDefault { HloInstruction* target_op, const llvm_ir::ElementGenerator& element_generator); + // Emit IR to perform a computation for every element in a partition/slice of + // 'target_shape'. The loop bounds for the outer-dimension partitions are + // passed into the compute function as a runtime argument (accessible from + // GetDynamicLoopBound). + Status EmitParallelTargetElementLoop( + const Shape& target_shape, + const llvm_ir::ElementGenerator& element_generator, + llvm_ir::IrArray* target_array); + // Emits a memcpy from the source instruction's result value to the // destination's. Both source and destination must have an entry in the // emitted_value_ table. @@ -259,7 +336,8 @@ class IrEmitter : public DfsHloVisitorWithDefault { // Emit IR to compute the target address of the buffer for the given op. // The returned Value is a pointer to a IR type that represents the op's // element type. - StatusOr EmitTargetAddressForOp(const HloInstruction* op); + StatusOr EmitTargetAddressForOp( + const HloInstruction* op, const ShapeIndex& shape_index = {}); // Structurizes "array_elements" into an MD array that represents "shape". // This is a recursive function, and "dimension_index" indicates the index of @@ -269,6 +347,86 @@ class IrEmitter : public DfsHloVisitorWithDefault { const std::vector& array_elements, const Shape& shape, int64 dimension_index); + // Tries to codegen a reduction operation using vectorized instructions. + // Returns true if successful, and false on failure. On failure, sets + // "failure_reason" to a string describing why it could not vectorize the + // reduction. + // + // TODO(sanjoy): Some of the things we do here can be abstracted out into + // concepts that generalize over other vectorizable operations. We should + // consider pulling out these abstractions into a VectorizingIrEmitter or + // something similar. + StatusOr EmitVectorizedReduce( + HloInstruction* reduce, HloInstruction* arg, HloInstruction* init_value, + tensorflow::gtl::ArraySlice dimensions, HloComputation* function, + string* failure_reason); + + // We'd like to keep one or two one cache-line's worth of data in registers + // without generating IR with illegal (e.g. excessively large or + // non-power-of-two) vector types. We do this by introducing a layer of + // abstraction: we introduce a high level vector-like concept called a + // "sharded vector" that models data paralleism, and is mapped to a sequence + // scalar and vector llvm::Value s. + // + // For example, we can represent 29 f32 elements by a sharded vector mapped to + // a sequence of LLVM values of types [<16 x f32>, <8 x f32>, <4 x f32>, f32]. + // Note that the last element is scalar. + // + // There is no requirement on the ordering or the uniqueness of the elements + // mapped to sharded vectors -- we allow repeated elements, and we allow + // elements to appear in any order. + using ShardedVector = std::vector; + + // A sharded vector type is the element-wise llvm::Type's of some + // ShardedVector. + using ShardedVectorType = std::vector; + + // Create a sharded vector type corresponding to a "element_count" long + // sequence of "element_type" values. + ShardedVectorType CreateShardedVectorType(PrimitiveType element_type, + unsigned element_count); + + // Emit LLVM IR to store the sharded vector "value_to_store" to + // "store_address". + void EmitShardedVectorStore(llvm::Value* store_address, + const ShardedVector& value_to_store, + const int alignment, + const llvm_ir::IrArray& containing_array); + + using ReductionGenerator = std ::function*, llvm::Value*, llvm::Value*)>; + + // Tries to match the reduction function "function" to a known reduction + // pattern. Returns a non-null ReductionGenerator on a successful match, + // which can be used to generate the LLVM IR corresponding to said reduction. + // On failure, this stores a reason string into "failure_reason". + ReductionGenerator MatchReductionGenerator(HloComputation* function, + string* failure_reason) const; + + // Emits the inner loop nest that runs the reduction. Helper function for + // EmitVectorizedReduce. + StatusOr EmitInnerLoopForVectorizedReduction( + const ReductionGenerator& reduction_generator, + const llvm_ir::IrArray::Index& output_index, + const ShardedVectorType& accumulator_type, HloInstruction* init_value, + HloInstruction* arg, tensorflow::gtl::ArraySlice dimensions, + unsigned element_alignment); + + // Tries to emit a fast concatenate operation using memcpy. Returns true if + // successful, and false on failure. On failure, sets "failure_reason" to a + // string describing why it could not emit a fast concatenate. + StatusOr EmitFastConcatenate( + HloInstruction* concatenate, + tensorflow::gtl::ArraySlice operands, + string* failure_reason); + + // Emits LLVM IR to transfer "element_count" elements of type "primitive_type" + // from the address "source" to the address "target". + void EmitTransferElements(llvm::Value* target, llvm::Value* source, + int64 element_count, PrimitiveType primitive_type, + const llvm_ir::IrArray& target_array, + const llvm_ir::IrArray& source_array); + // Name of the computation entry function. This function serves as the // top-level "main" of the computation and will be invoked by the JIT. string entry_function_name_; @@ -307,17 +465,29 @@ class IrEmitter : public DfsHloVisitorWithDefault { llvm_ir::AliasAnalysis alias_analysis_; + // The number of root instruction outer dimensions used in parallel loop + // emission (EmitParallelTargetElementLoop). + int64 num_dynamic_loop_bounds_ = 0; + + // Returns whether the given instruction should be emitted as a parallel loop. + bool ShouldEmitParallelLoopFor(const HloInstruction& op) const { + // Emit parallel loop for root instruction if dynamic outer-dimension loop + // bounds were specified. + return num_dynamic_loop_bounds_ > 0 && + op.parent()->root_instruction() == &op; + } + // This struct contains all the state needed to emit instructions for // profiling a computation. class ProfilingState { public: ProfilingState() - : is_entry_computation_(false), + : is_top_level_computation_(false), use_rdtscp_(false), prof_counters_(nullptr) {} - ProfilingState(bool is_entry_computation, bool use_rdtscp, + ProfilingState(bool is_top_level_computation, bool use_rdtscp, llvm::Argument* prof_counters) - : is_entry_computation_(is_entry_computation), + : is_top_level_computation_(is_top_level_computation), use_rdtscp_(use_rdtscp), prof_counters_(prof_counters) {} @@ -342,7 +512,7 @@ class IrEmitter : public DfsHloVisitorWithDefault { private: // Is this IrEmitter for a top-level computation? - bool is_entry_computation_; + bool is_top_level_computation_; // Should we use the x86-specific rdtscp or the generic readcyclecounter // intrinsic? @@ -392,8 +562,30 @@ class IrEmitter : public DfsHloVisitorWithDefault { // Returns the number of bytes within the shape. int64 ByteSizeOf(const Shape& shape) const; + enum class XfeedKind { + kInfeed, + kOutfeed, + }; + + // Emit IR to transfer between a {infeed,outfeed} buffer and an in-program + // address. + Status EmitXfeedTransfer(XfeedKind kind, const Shape& shape, + llvm::Value* program_buffer_address); + + // Returns true if the current function being emitted is called in a + // parallel context (returns false otherwise). + bool IsParallelContext() { + return parallel_cpu_backend_ && is_top_level_computation_; + } + const HloModuleConfig& hlo_module_config_; + const bool parallel_cpu_backend_; + + bool is_top_level_computation_; + + TargetMachineFeatures target_machine_features_; + TF_DISALLOW_COPY_AND_ASSIGN(IrEmitter); }; diff --git a/tensorflow/compiler/xla/service/cpu/llvm_ir_runtime.cc b/tensorflow/compiler/xla/service/cpu/llvm_ir_runtime.cc new file mode 100644 index 0000000000000000000000000000000000000000..424306a194b583c66777f3af71ca13932a3676a8 --- /dev/null +++ b/tensorflow/compiler/xla/service/cpu/llvm_ir_runtime.cc @@ -0,0 +1,170 @@ +/* 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/cpu/llvm_ir_runtime.h" + +#include "llvm/IR/Function.h" +#include "llvm/IR/IRBuilder.h" +#include "llvm/IR/Intrinsics.h" +#include "llvm/IR/Verifier.h" +#include "llvm/Transforms/Utils/Cloning.h" +#include "tensorflow/core/platform/logging.h" + +namespace xla { +namespace cpu { +namespace runtime { + +const char* const kTanhV4F32SymbolName = "__xla_cpu_runtime_TanhV4F32"; +const char* const kTanhV8F32SymbolName = "__xla_cpu_runtime_TanhV8F32"; + +namespace { +llvm::Value* EmitFMinOrMax(llvm::IRBuilder<>* ir_builder, llvm::Module* module, + llvm::Type* vector_type, llvm::Value* lhs, + llvm::Value* rhs, bool is_min, + bool enable_fast_math) { + if (enable_fast_math) { + // Using an unordered comparison lets LLVM generate a vminps / vmaxps + // instruction on x86. vminps/vmaxps choose the second operand if either + // operand is a NaN and thus don't accurately implement the semantics of the + // minnum and maxnum intrinsics, necessitating different IR emission. + // + // We can _probably_ do this even when fast math is disabled, but we can + // certainly do this if fast math is enabled (and nnan applies). + auto* compare = ir_builder->CreateFCmp( + is_min ? llvm::FCmpInst::FCMP_ULE : llvm::FCmpInst::FCMP_UGE, lhs, rhs); + return ir_builder->CreateSelect(compare, lhs, rhs); + } else { + llvm::Function* intrinsic = llvm::Intrinsic::getDeclaration( + module, is_min ? llvm::Intrinsic::minnum : llvm::Intrinsic::maxnum, + vector_type); + return ir_builder->CreateCall(intrinsic, {lhs, rhs}); + } +} + +llvm::Function* EmitVectorF32TanhIfNeeded(llvm::Module* module, + llvm::StringRef function_name, + int vector_width, + bool enable_fast_math) { + llvm::Function* vector_tanh_function = module->getFunction(function_name); + if (vector_tanh_function == nullptr) { + // If the function declaration is not present in the module, there can't be + // any calls to resolve. Don't emit the function in this case. + return nullptr; + } + + llvm::LLVMContext* context = &module->getContext(); + llvm::Type* float_type = llvm::Type::getFloatTy(*context); + llvm::VectorType* vector_type = + llvm::VectorType::get(float_type, vector_width); + + llvm::BasicBlock* vector_tanh_body = + llvm::BasicBlock::Create(*context, "body", vector_tanh_function); + + llvm::IRBuilder<> ir_builder(vector_tanh_body); + + llvm::FastMathFlags fast_math_flags; + fast_math_flags.setUnsafeAlgebra(); + ir_builder.setFastMathFlags(fast_math_flags); + + auto emit_fmin = [&](llvm::Value* lhs, llvm::Value* rhs) { + return EmitFMinOrMax(&ir_builder, module, vector_type, lhs, rhs, + /*is_min=*/true, + /*enable_fast_math=*/enable_fast_math); + }; + auto emit_fmax = [&](llvm::Value* lhs, llvm::Value* rhs) { + return EmitFMinOrMax(&ir_builder, module, vector_type, lhs, rhs, + /*is_min=*/false, + /*enable_fast_math=*/enable_fast_math); + }; + + llvm::Value* input = &*vector_tanh_function->arg_begin(); + CHECK_EQ(input->getType(), vector_type); + + // This implements the same rational interpolant as implemented in Eigen3. + llvm::Value* input_clamped = + emit_fmin(emit_fmax(input, llvm::ConstantFP::get(vector_type, -9.0)), + llvm::ConstantFP::get(vector_type, 9.0)); + + std::array numerator_coeffs( + {{-2.76076847742355e-16f, 2.00018790482477e-13f, -8.60467152213735e-11f, + 5.12229709037114e-08f, 1.48572235717979e-05f, 6.37261928875436e-04f, + 4.89352455891786e-03f}}); + + std::array denominator_coeffs( + {{1.19825839466702e-06f, 1.18534705686654e-04f, 2.26843463243900e-03f, + 4.89352518554385e-03f}}); + + llvm::Value* input_squared = + ir_builder.CreateFMul(input_clamped, input_clamped); + llvm::Value* numerator = + llvm::ConstantFP::get(vector_type, numerator_coeffs[0]); + for (int i = 1; i < numerator_coeffs.size(); i++) { + numerator = ir_builder.CreateFAdd( + ir_builder.CreateFMul(input_squared, numerator), + llvm::ConstantFP::get(vector_type, numerator_coeffs[i])); + } + numerator = ir_builder.CreateFMul(input_clamped, numerator); + + llvm::Value* denominator = + llvm::ConstantFP::get(vector_type, denominator_coeffs[0]); + for (int i = 1; i < denominator_coeffs.size(); i++) { + denominator = ir_builder.CreateFAdd( + ir_builder.CreateFMul(input_squared, denominator), + llvm::ConstantFP::get(vector_type, denominator_coeffs[i])); + } + + llvm::Value* result = ir_builder.CreateFDiv(numerator, denominator); + ir_builder.CreateRet(result); + + DCHECK(!llvm::verifyFunction(*vector_tanh_function)); + return vector_tanh_function; +} +} // namespace + +void RewriteIRRuntimeFunctions(llvm::Module* module, bool enable_fast_math) { + auto* tanh_v4f32 = + EmitVectorF32TanhIfNeeded(module, kTanhV4F32SymbolName, + /*vector_width=*/4, enable_fast_math); + auto* tanh_v8f32 = + EmitVectorF32TanhIfNeeded(module, kTanhV8F32SymbolName, + /*vector_width=*/8, enable_fast_math); + + // Gather all the call sites, force inline them and then delete the vector + // function bodies. + + std::vector calls_to_inline; + for (auto* function : {tanh_v4f32, tanh_v8f32}) { + if (function != nullptr) { + for (auto* user : function->users()) { + calls_to_inline.push_back(llvm::cast(user)); + } + } + } + + for (auto* call_to_inline : calls_to_inline) { + llvm::InlineFunctionInfo inline_function_info; + CHECK(llvm::InlineFunction(call_to_inline, inline_function_info)); + } + + for (auto* function : {tanh_v4f32, tanh_v8f32}) { + if (function != nullptr) { + function->eraseFromParent(); + } + } +} + +} // namespace runtime +} // namespace cpu +} // namespace xla diff --git a/tensorflow/compiler/xla/service/cpu/llvm_ir_runtime.h b/tensorflow/compiler/xla/service/cpu/llvm_ir_runtime.h new file mode 100644 index 0000000000000000000000000000000000000000..3082b39b634e3fb533e3b7b8f13c98c0140c6b03 --- /dev/null +++ b/tensorflow/compiler/xla/service/cpu/llvm_ir_runtime.h @@ -0,0 +1,42 @@ +/* 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_CPU_LLVM_IR_RUNTINE_H_ +#define TENSORFLOW_COMPILER_XLA_SERVICE_CPU_LLVM_IR_RUNTINE_H_ + +#include "llvm/IR/Module.h" + +namespace xla { +namespace cpu { +namespace runtime { + +extern const char* const kTanhV4F32SymbolName; +extern const char* const kTanhV8F32SymbolName; + +// The following CPU runtime functions have LLVM-IR only implementations: +// +// - __xla_cpu_runtime_TanhV4F32 +// - __xla_cpu_runtime_TanhV8F32 +// +// |LinkIRRuntimeFunctions| rewrites calls to these functions into generic LLVM +// IR. + +void RewriteIRRuntimeFunctions(llvm::Module* module, bool enable_fast_math); + +} // namespace runtime +} // namespace cpu +} // namespace xla + +#endif // TENSORFLOW_COMPILER_XLA_SERVICE_CPU_LLVM_IR_RUNTINE_H_ diff --git a/tensorflow/compiler/xla/service/cpu/parallel_cpu_executable.cc b/tensorflow/compiler/xla/service/cpu/parallel_cpu_executable.cc index 7a4723e8d75588d8ccb711892b4082024695e444..40fa3a67bdec3953003ba8f98f2a19a9082a82c5 100644 --- a/tensorflow/compiler/xla/service/cpu/parallel_cpu_executable.cc +++ b/tensorflow/compiler/xla/service/cpu/parallel_cpu_executable.cc @@ -24,13 +24,13 @@ limitations under the License. #include #include -#include "external/llvm/include/llvm/ExecutionEngine/Orc/IRCompileLayer.h" +#include "llvm/ExecutionEngine/Orc/IRCompileLayer.h" #include "tensorflow/compiler/xla/map_util.h" #include "tensorflow/compiler/xla/service/buffer_assignment.h" #include "tensorflow/compiler/xla/service/cpu/cpu_runtime.h" +#include "tensorflow/compiler/xla/service/cpu/shape_partition.h" #include "tensorflow/compiler/xla/service/hlo_computation.h" #include "tensorflow/compiler/xla/service/hlo_module.h" -#include "tensorflow/compiler/xla/service/hlo_module_config.h" #include "tensorflow/compiler/xla/service/hlo_opcode.h" #include "tensorflow/compiler/xla/service/logical_buffer.h" #include "tensorflow/compiler/xla/service/shaped_buffer.h" @@ -58,12 +58,11 @@ ParallelCpuExecutable::ParallelCpuExecutable( std::unique_ptr jit, std::unique_ptr assignment, std::unique_ptr hlo_module, - std::unique_ptr module_config, std::unique_ptr> function_names, std::unordered_map hlo_to_profile_idx, std::unordered_map> aligned_constants) - : Executable(std::move(hlo_module), std::move(module_config)), + : Executable(std::move(hlo_module)), jit_(std::move(jit)), assignment_(std::move(assignment)), functions_names_(std::move(function_names)), @@ -72,7 +71,7 @@ ParallelCpuExecutable::ParallelCpuExecutable( // Type of the computation function we expect in the JIT. using ComputeFunctionType = void (*)(void*, const void*, const void**, void**, - uint64*); + int64*, uint64*); // Given a pointer to an output buffer (following the CPU JIT calling // conventions), mark addresses that are "live". The initial pointer itself is @@ -97,6 +96,232 @@ static void MarkLiveAddressesInOutput( } } +namespace { + +// Executor manages the concurrent execution of 'functions' for instructions +// in 'pending' on 'thread_pool' (storing resulting data in 'results'). +class Executor { + public: + Executor(const std::map& functions, + const ServiceExecutableRunOptions* run_options, + std::list* pending, + std::map* results, void** temps_array, + uint64* profile_counters_array, BufferAssignment* assignment) + : functions_(functions), + run_options_(run_options), + pending_(pending), + results_(results), + temps_array_(temps_array), + profile_counters_array_(profile_counters_array), + thread_pool_(CHECK_NOTNULL(run_options_->xla_intra_op_thread_pool())), + assignment_(assignment) {} + + // Executes pending list of instructions on thread pool. + // Returns OK status on success, error status otherwise. + Status Run(); + + private: + // Schedules a parallel invocation of compute function for 'instruction' on + // 'thread_pool_', storing result in 'result_buffer'. + // If 'partition_buffers' is non-null, parallel task will be invoked on + // per-dimension partition [start, limit) values stored in + // 'partition_buffers'. + void Schedule(HloInstruction* instruction, int64* partition_buffers, + void* result_buffer); + + // Returns true if 'instruction' has been assigned parallel tasks (returns + // false otherwise). + bool HasParallelTasks(HloInstruction* instruction); + + // Returns in 'partition_buffers' the partition [size, limit) for each + // dimension. + int64* GetPartitionBuffers( + const std::vector>& partition); + + // Returns array of result buffers for all operands in 'instruction'. + const void** GetOperandBuffers(HloInstruction* instruction); + + // Arguments passed into Executor. + const std::map& functions_; + const ServiceExecutableRunOptions* run_options_; + std::list* pending_; + std::map* results_; + void** temps_array_; + uint64* profile_counters_array_; + tensorflow::thread::ThreadPool* thread_pool_; + BufferAssignment* assignment_; + + // Members used to manage instruction execution. + tensorflow::mutex completion_queue_lock_; + tensorflow::condition_variable completion_queue_cv_; + std::deque completion_queue_; + int64 instructions_in_flight_ = 0; + std::unordered_map tasks_in_flight_; +}; + +Status Executor::Run() { + while (!pending_->empty() || instructions_in_flight_ > 0) { + auto pending_it = pending_->begin(); + while (pending_it != pending_->end()) { + HloInstruction* instruction = *pending_it; + // Skip pending instructions whose operands aren't ready. + if (std::any_of(instruction->operands().begin(), + instruction->operands().end(), + [&](HloInstruction* operand) { + return !ContainsKey(*results_, operand); + })) { + ++pending_it; + continue; + } + + // Get 'result_buffer' reference to result buffer for 'instruction'. + TF_ASSIGN_OR_RETURN(const BufferAllocation::Slice result_slice, + assignment_->GetUniqueTopLevelSlice(instruction)); + void* result_buffer = + static_cast(temps_array_[result_slice.index()]) + + result_slice.offset(); + + if (HasParallelTasks(instruction)) { + // 'instruction' has been assigned parallel task partitions. + CHECK_EQ(HloOpcode::kCall, instruction->opcode()); + HloInstruction* root = instruction->to_apply()->root_instruction(); + + // Create ShapePartitionIterator to iterate through all outer dimension + // partitions of 'instruction'. + ShapePartitionIterator partition_iterator( + root->shape(), root->outer_dimension_partitions()); + + const int64 partition_count = + partition_iterator.GetTotalPartitionCount(); + + // Record total parallel task count for 'instruction' before dispatch. + { + tensorflow::mutex_lock l(completion_queue_lock_); + tasks_in_flight_.insert(std::make_pair(instruction, partition_count)); + VLOG(2) << "Schedule PARALLEL" + << " instruction: " << instruction->name() + << " instruction.callee: " + << instruction->to_apply()->root_instruction()->name() + << " partition_count: " << partition_count; + } + + for (int64 i = 0; i < partition_count; ++i) { + // Get partition [start, limit) for each dimension. + auto partition_buffers = + GetPartitionBuffers(partition_iterator.GetPartition(i)); + Schedule(instruction, partition_buffers, result_buffer); + } + + } else { + // Set tasks in-flight to '1' for sequential instruction execution. + { + tensorflow::mutex_lock l(completion_queue_lock_); + tasks_in_flight_.insert(std::make_pair(instruction, 1)); + VLOG(2) << "Schedule SEQUENTIAL" + << " instruction: " << instruction->name() + << " instruction.callee: " + << instruction->to_apply()->root_instruction()->name(); + } + Schedule(instruction, nullptr, result_buffer); + } + + ++instructions_in_flight_; + pending_it = pending_->erase(pending_it); + } + // Wait for a completed HLO instruction to be present in the queue. We will + // pop it out of the queue and make the result available to its users. + HloInstruction* instruction; + do { + tensorflow::mutex_lock l(completion_queue_lock_); + if (completion_queue_.empty()) { + completion_queue_cv_.wait(l); + } + if (!completion_queue_.empty()) { + instruction = completion_queue_.front(); + completion_queue_.pop_front(); + break; + } + } while (true); + TF_ASSIGN_OR_RETURN(const BufferAllocation::Slice result_slice, + assignment_->GetUniqueTopLevelSlice(instruction)); + void* result_buffer = + static_cast(temps_array_[result_slice.index()]) + + result_slice.offset(); + InsertOrDie(results_, instruction, result_buffer); + --instructions_in_flight_; + } + return Status::OK(); +} + +void Executor::Schedule(HloInstruction* instruction, int64* partition_buffers, + void* result_buffer) { + // The thread pool entry takes ownership of |operand_buffers|. + auto operand_buffers = GetOperandBuffers(instruction); + + auto function = FindOrDie(functions_, instruction); + const auto* exec_run_options = &run_options_->run_options(); + thread_pool_->Schedule([this, instruction, result_buffer, operand_buffers, + partition_buffers, exec_run_options, function]() { + function(result_buffer, exec_run_options, operand_buffers, temps_array_, + partition_buffers, profile_counters_array_); + + delete[] operand_buffers; + delete[] partition_buffers; + // Push the completed HLO instruction on the queue, the main + // thread will pop it off and potentially launch more work which + // uses the result. + // TODO(b/27458679) Consider alternative task scheduling and synchronization + // schemes. For example, we could avoid the overhead associate with the + // condvar here if the thread just dequed the next instruction to execute + // on completion. + { + tensorflow::mutex_lock l(completion_queue_lock_); + // Decrement in-flight task count for this completion. + if (--FindOrDie(tasks_in_flight_, instruction) == 0) { + completion_queue_.push_back(instruction); + completion_queue_cv_.notify_all(); + tasks_in_flight_.erase(instruction); + } + } + }); +} + +int64* Executor::GetPartitionBuffers( + const std::vector>& partition) { + // Return in 'partition_buffers' partition [size, limit) for each dimension. + auto partition_buffers = new int64[partition.size() * 2]; + for (int i = 0; i < partition.size(); ++i) { + partition_buffers[2 * i + 0] = partition[i].first; + partition_buffers[2 * i + 1] = partition[i].first + partition[i].second; + } + return partition_buffers; +} + +bool Executor::HasParallelTasks(HloInstruction* instruction) { + return instruction->opcode() == HloOpcode::kCall && + !instruction->to_apply() + ->root_instruction() + ->outer_dimension_partitions() + .empty(); +} + +const void** Executor::GetOperandBuffers(HloInstruction* instruction) { + // We cannot use a move-only RAII type like std::unique_ptr because the + // list of operands is allocated on the main thread and transferred to the + // worker via the lambda passed to enqueue_function. In order for the + // lambda to take ownership, we would need to use generalized lambda + // capture which is a feature new to C++14. + // TODO(b/27458679) Avoid dynamic allocations in Executor. + auto operand_buffers = new const void*[instruction->operand_count()]; + std::transform(instruction->operands().begin(), instruction->operands().end(), + operand_buffers, [this](HloInstruction* operand) { + return FindOrDie(*results_, operand); + }); + return operand_buffers; +} + +} // namespace + Status ParallelCpuExecutable::AllocateBuffers( DeviceMemoryAllocator* memory_allocator, int device_ordinal, std::vector* buffers) { @@ -146,7 +371,7 @@ Status ParallelCpuExecutable::AllocateBuffers( } Status ParallelCpuExecutable::ExecuteComputeFunctions( - const ExecutableRunOptions* run_options, + const ServiceExecutableRunOptions* run_options, tensorflow::gtl::ArraySlice arguments, tensorflow::gtl::ArraySlice buffers, HloExecutionProfile* hlo_execution_profile) { @@ -160,7 +385,7 @@ Status ParallelCpuExecutable::ExecuteComputeFunctions( } Status ParallelCpuExecutable::ExecuteComputeFunctions( - const ExecutableRunOptions* run_options, + const ServiceExecutableRunOptions* run_options, tensorflow::gtl::ArraySlice arguments, tensorflow::gtl::ArraySlice buffers, HloExecutionProfile* hlo_execution_profile) { @@ -182,8 +407,9 @@ Status ParallelCpuExecutable::ExecuteComputeFunctions( HloInstruction* instruction = entry.first; llvm::JITSymbol sym = jit_->FindSymbol(entry.second); TF_RET_CHECK(sym); - InsertOrDie(&functions, instruction, - reinterpret_cast(sym.getAddress())); + InsertOrDie( + &functions, instruction, + reinterpret_cast(cantFail(sym.getAddress()))); } // Map containing pointers to result buffers for each instruction. @@ -212,87 +438,16 @@ Status ParallelCpuExecutable::ExecuteComputeFunctions( } } - void** temps_array = buffer_pointers.data(); - uint64* profile_counters_array = profile_counters.data(); - auto* thread_pool = CHECK_NOTNULL(run_options->inter_op_thread_pool()); - tensorflow::mutex completion_queue_lock; - tensorflow::condition_variable completion_queue_cv; - std::deque completion_queue; - int64 instructions_in_flight = 0; - while (!pending.empty() || instructions_in_flight > 0) { - auto pending_it = pending.begin(); - while (pending_it != pending.end()) { - HloInstruction* instruction = *pending_it; - // Skip pending instructions whose operands aren't ready. - if (std::any_of(instruction->operands().begin(), - instruction->operands().end(), - [&](HloInstruction* operand) { - return !ContainsKey(results, operand); - })) { - ++pending_it; - continue; - } + // TODO(b/27458679) Manage scheduling based on in-flight concurrency limits. + // For example, if we expect a library conv/matmul call to run at max + // concurrency, we should not dispatch runnable instructions until the + // library call is finished (to avoid expensive cache invalidation). + Executor executor(functions, run_options, &pending, &results, + buffer_pointers.data(), profile_counters.data(), + assignment_.get()); - TF_ASSIGN_OR_RETURN(const BufferAllocation::Slice result_slice, - assignment_->GetUniqueTopLevelSlice(instruction)); - void* result_buffer = - static_cast(temps_array[result_slice.index()]) + - result_slice.offset(); - // We cannot use a move-only RAII type like std::unique_ptr because the - // list of operands is allocated on the main thread and transferred to the - // worker via the lambda passed to enqueue_function. In order for the - // lambda to take ownership, we would need to use generalized lambda - // capture which is a feature new to C++14. - auto operand_buffers = new const void*[instruction->operand_count()]; - std::transform(instruction->operands().begin(), - instruction->operands().end(), operand_buffers, - [&results](HloInstruction* operand) { - return FindOrDie(results, operand); - }); - auto function = FindOrDie(functions, instruction); - // The thread pool entry takes ownership of |operand_buffers|. - thread_pool->Schedule([instruction, &completion_queue, - &completion_queue_lock, &completion_queue_cv, - result_buffer, run_options, operand_buffers, - temps_array, profile_counters_array, function] { - function(result_buffer, run_options, operand_buffers, temps_array, - profile_counters_array); - delete[] operand_buffers; - // Push the completed HLO instruction on the queue, the main thread - // will pop it off and potentially launch more work which uses the - // result. - { - tensorflow::mutex_lock l(completion_queue_lock); - completion_queue.push_back(instruction); - completion_queue_cv.notify_all(); - } - }); + TF_RETURN_IF_ERROR(executor.Run()); - ++instructions_in_flight; - pending_it = pending.erase(pending_it); - } - // Wait for a completed HLO instruction to be present in the queue. We will - // pop it out of the queue and make the result available to its users. - HloInstruction* instruction; - do { - tensorflow::mutex_lock l(completion_queue_lock); - if (completion_queue.empty()) { - completion_queue_cv.wait(l); - } - if (!completion_queue.empty()) { - instruction = completion_queue.front(); - completion_queue.pop_front(); - break; - } - } while (1); - TF_ASSIGN_OR_RETURN(const BufferAllocation::Slice result_slice, - assignment_->GetUniqueTopLevelSlice(instruction)); - void* result_buffer = - static_cast(temps_array[result_slice.index()]) + - result_slice.offset(); - InsertOrDie(&results, instruction, result_buffer); - --instructions_in_flight; - } uint64 end_micros = tensorflow::Env::Default()->NowMicros(); { @@ -345,9 +500,8 @@ ParallelCpuExecutable::ExecuteOnStream( const BufferAllocation::Index result_index = result_slice.index(); VLOG(3) << "result index: " << result_index; - TF_RETURN_IF_ERROR(ExecuteComputeFunctions(&run_options->run_options(), - arguments, device_allocations, - hlo_execution_profile)); + TF_RETURN_IF_ERROR(ExecuteComputeFunctions( + run_options, arguments, device_allocations, hlo_execution_profile)); // Mark the buffers that are actually live (used in the output) when the // computation finishes executing. @@ -400,19 +554,19 @@ StatusOr> ParallelCpuExecutable::ExecuteOnStream( TF_RETURN_IF_ERROR(AllocateBuffers( memory_allocator, stream->parent()->device_ordinal(), &buffers)); - TF_RETURN_IF_ERROR(ExecuteComputeFunctions( - &run_options->run_options(), arguments, buffers, hlo_execution_profile)); + TF_RETURN_IF_ERROR(ExecuteComputeFunctions(run_options, arguments, buffers, + hlo_execution_profile)); // Copy DeviceMemoryBase values which contain the array(s) of the result into // the respective location in ShapedBuffer which is returned to the caller. std::vector buffers_in_result(assignment_->Allocations().size(), false); TF_RETURN_IF_ERROR( result_buffer->mutable_shape_index_to_buffer_entry() - ->ForEachMutableElement( + ->ForEachMutableElementWithStatus( [&buffers, &buffers_in_result, &result_buffer, this]( - const ShapeIndex& index, bool is_leaf, size_t* buffer_entry) { - if (is_leaf) { - const std::vector& sources = + const ShapeIndex& index, size_t* buffer_entry) { + if (ShapeUtil::IsLeafIndex(result_buffer->shape(), index)) { + const auto& sources = this->GetRootPointsToSet().element(index); // The points to set is unambiguous so the set should be a // singleton. @@ -468,5 +622,10 @@ const PointsToSet& ParallelCpuExecutable::GetRootPointsToSet() const { module().entry_computation()->root_instruction()); } +std::unique_ptr ParallelCpuExecutable::CreateCostAnalysis() + const { + return MakeUnique(ShapeSizeBytes); +} + } // namespace cpu } // namespace xla diff --git a/tensorflow/compiler/xla/service/cpu/parallel_cpu_executable.h b/tensorflow/compiler/xla/service/cpu/parallel_cpu_executable.h index 7223de9f0798365138cdb26ca9dce07cd0e474e3..d9200e13ed2ae8ed8afc4e4c7475e72aed4ae3c7 100644 --- a/tensorflow/compiler/xla/service/cpu/parallel_cpu_executable.h +++ b/tensorflow/compiler/xla/service/cpu/parallel_cpu_executable.h @@ -29,7 +29,6 @@ limitations under the License. #include "tensorflow/compiler/xla/service/hlo_execution_profile.h" #include "tensorflow/compiler/xla/service/hlo_instruction.h" #include "tensorflow/compiler/xla/service/hlo_module.h" -#include "tensorflow/compiler/xla/service/hlo_module_config.h" #include "tensorflow/compiler/xla/service/shaped_buffer.h" #include "tensorflow/compiler/xla/statusor.h" #include "tensorflow/compiler/xla/types.h" @@ -52,7 +51,6 @@ class ParallelCpuExecutable : public Executable { std::unique_ptr jit, std::unique_ptr assignment, std::unique_ptr hlo_module, - std::unique_ptr module_config, std::unique_ptr> instruction_functions, std::unordered_map hlo_to_profile_idx, std::unordered_map CreateCostAnalysis() const override; + private: // Allocate buffers required for execution and assign them to the elements of // "buffers". "buffers" should be sized to the number of buffers in buffer @@ -96,14 +110,14 @@ class ParallelCpuExecutable : public Executable { // Calls the generated functions in 'function_names_', performing the // computation with the given arguments using the supplied buffers. Status ExecuteComputeFunctions( - const ExecutableRunOptions* run_options, + const ServiceExecutableRunOptions* run_options, tensorflow::gtl::ArraySlice arguments, tensorflow::gtl::ArraySlice buffers, HloExecutionProfile* hlo_execution_profile); Status ExecuteComputeFunctions( - const ExecutableRunOptions* run_options, + const ServiceExecutableRunOptions* run_options, tensorflow::gtl::ArraySlice arguments, tensorflow::gtl::ArraySlice buffers, diff --git a/tensorflow/compiler/xla/service/cpu/runtime_matmul.cc b/tensorflow/compiler/xla/service/cpu/runtime_matmul.cc index 677080a8623224cdd65e35b3116ae57b7b3b3ca2..bff57d33ae23fbba8c664cbd18df77e4c35eb592 100644 --- a/tensorflow/compiler/xla/service/cpu/runtime_matmul.cc +++ b/tensorflow/compiler/xla/service/cpu/runtime_matmul.cc @@ -19,6 +19,7 @@ limitations under the License. #include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor" #include "tensorflow/compiler/xla/executable_run_options.h" +#include "tensorflow/compiler/xla/service/cpu/runtime_matvec.h" #include "tensorflow/core/platform/types.h" using tensorflow::int32; @@ -54,7 +55,7 @@ void MatMul(const void* run_options_ptr, T* out, T* lhs, T* rhs, int64 m, int lhs_contract_dim = transpose_lhs ? 0 : 1; int rhs_contract_dim = transpose_rhs ? 1 : 0; const Eigen::array dims( - DimPair(lhs_contract_dim, rhs_contract_dim)); + {DimPair(lhs_contract_dim, rhs_contract_dim)}); // Matrix multiply is a special case of the "contract" operation where // the contraction is performed along dimension 1 of the lhs and dimension @@ -68,14 +69,24 @@ 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) { - MatMul(run_options_ptr, out, lhs, rhs, m, n, k, transpose_lhs, - transpose_rhs); + if (m == 1 || n == 1) { + // Despite being single threaded, this version of matrix * vector is faster. + xla::EigenMatVecF32(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); + } } 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) { - MatMul(run_options_ptr, out, lhs, rhs, m, n, k, transpose_lhs, - transpose_rhs); + if (m == 1 || n == 1) { + // Despite being single threaded, this version of matrix * vector is faster. + xla::EigenMatVecF64(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); + } } diff --git a/tensorflow/compiler/xla/service/cpu/runtime_matvec.cc b/tensorflow/compiler/xla/service/cpu/runtime_matvec.cc new file mode 100644 index 0000000000000000000000000000000000000000..435820cdd36e2a906d9dfbe2555f4c0df623c729 --- /dev/null +++ b/tensorflow/compiler/xla/service/cpu/runtime_matvec.cc @@ -0,0 +1,110 @@ +/* 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 "third_party/eigen3/Eigen/Core" +#include "tensorflow/compiler/xla/service/cpu/runtime_matvec.h" + +using tensorflow::int32; +using tensorflow::int64; + +namespace { + +// Does mat * x or mat^T * x. +template +void MatVec(T* out_buf, T* mat_buf, T* x_buf, int64 rows, int64 cols, + int32 transpose) { + // Use an Eigen Matrix instead of a Tensor, as the GEMV from Matrix seems to + // be faster (b/30223679). See also: the matmul op kernel in TensorFlow, + // which implements the same optimization. + using Matrix = Eigen::Matrix; + using MatrixMap = Eigen::Map; + + using Vector = Eigen::Matrix; + using VectorMap = Eigen::Map; + + auto x = VectorMap(x_buf, cols); + auto out = VectorMap(out_buf, rows); + + int64 mat_rows = rows; + int64 mat_cols = cols; + + if (transpose) { + std::swap(mat_rows, mat_cols); + } + + auto mat = MatrixMap(mat_buf, mat_rows, mat_cols); + + if (transpose) { + out = mat.transpose() * x; + } else { + out = mat * x; + } +} + +// Converts matmul-style args to matvec. +template +void DispatchMatVec(T* out, T* lhs, T* rhs, int64 m, int64 n, int64 k, + int32 transpose_lhs, int32 transpose_rhs) { + // If the input is in the form x * A, where x is the vector, then bring A back + // over to the left hand side. We make use of the identity + // + // (x * A)^T = A^T * x^T + // + // We do not need to take the transpose of x or of the result since taking + // the transpose of a vector does not change the memory layout. + const int64 cols = k; + + T* mat; + T* vec; + int64 rows; + bool transpose_mat; + + bool is_mat_vec = (n == 1); + + if (is_mat_vec) { + mat = lhs; + vec = rhs; + rows = m; + transpose_mat = transpose_lhs; + } else { + mat = rhs; + vec = lhs; + rows = n; + transpose_mat = !transpose_rhs; + } + + MatVec(out, mat, vec, rows, cols, transpose_mat); +} + +} // namespace + +namespace xla { + +void EigenMatVecF32(float* out, float* lhs, float* rhs, int64 m, int64 n, + int64 k, int32 transpose_lhs, int32 transpose_rhs) { + assert((m == 1 || n == 1) && "not a matrix-vector multiply"); + DispatchMatVec(out, lhs, rhs, m, n, k, transpose_lhs, transpose_rhs); +} + +void EigenMatVecF64(double* out, double* lhs, double* rhs, int64 m, int64 n, + int64 k, int32 transpose_lhs, int32 transpose_rhs) { + assert((m == 1 || n == 1) && "not a matrix-vector multiply"); + DispatchMatVec(out, lhs, rhs, m, n, k, transpose_lhs, transpose_rhs); +} + +} // namespace xla diff --git a/tensorflow/compiler/xla/service/cpu/runtime_matvec.h b/tensorflow/compiler/xla/service/cpu/runtime_matvec.h new file mode 100644 index 0000000000000000000000000000000000000000..cb7e0a81f09e2702de565012e1fcac8b7cd841ab --- /dev/null +++ b/tensorflow/compiler/xla/service/cpu/runtime_matvec.h @@ -0,0 +1,45 @@ +/* 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 THIRD_PARTY_TENSORFLOW_COMPILER_XLA_SERVICE_CPU_RUNTIME_MATVEC_H_ +#define THIRD_PARTY_TENSORFLOW_COMPILER_XLA_SERVICE_CPU_RUNTIME_MATVEC_H_ + +#include "tensorflow/core/platform/types.h" + +namespace xla { + +// Performs a matrix-vector multiplication using Eigen. 'lhs' and 'rhs' are +// pointers to buffers containing input matrices in column-major order. 'out' is +// a pointer to a buffer sufficiently large to hold the result of the +// operation. Following standard nomenclature: lhs is m x k, rhs is k x n, and +// out is m x n. +// +// This requires that m = 1 or n = 1. +// +// TODO(b/64684907): Compare runtime performance of these functions with dot +// simplification. +void EigenMatVecF32(float* out, float* lhs, float* rhs, tensorflow::int64 m, + tensorflow::int64 n, tensorflow::int64 k, + tensorflow::int32 transpose_lhs, + tensorflow::int32 transpose_rhs); + +void EigenMatVecF64(double* out, double* lhs, double* rhs, tensorflow::int64 m, + tensorflow::int64 n, tensorflow::int64 k, + tensorflow::int32 transpose_lhs, + tensorflow::int32 transpose_rhs); + +} // namespace xla + +#endif // THIRD_PARTY_TENSORFLOW_COMPILER_XLA_SERVICE_CPU_RUNTIME_MATVEC_H_ diff --git a/tensorflow/compiler/xla/service/cpu/runtime_single_threaded_matmul.cc b/tensorflow/compiler/xla/service/cpu/runtime_single_threaded_matmul.cc index 384a978873de89526f43556296aaa51c46ac1d3f..ee8eb081556d60fcf6537b1036a4a5825c4c7bf6 100644 --- a/tensorflow/compiler/xla/service/cpu/runtime_single_threaded_matmul.cc +++ b/tensorflow/compiler/xla/service/cpu/runtime_single_threaded_matmul.cc @@ -16,6 +16,7 @@ limitations under the License. #include "tensorflow/compiler/xla/service/cpu/runtime_single_threaded_matmul.h" #include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor" +#include "tensorflow/compiler/xla/service/cpu/runtime_matvec.h" #include "tensorflow/core/platform/types.h" using tensorflow::int32; @@ -48,7 +49,7 @@ void MatMul(const void* run_options_ptr, T* out, T* lhs, T* rhs, int64 m, int lhs_contract_dim = transpose_lhs ? 0 : 1; int rhs_contract_dim = transpose_rhs ? 1 : 0; const Eigen::array dims( - DimPair(lhs_contract_dim, rhs_contract_dim)); + {DimPair(lhs_contract_dim, rhs_contract_dim)}); // Matrix multiply is a special case of the "contract" operation where // the contraction is performed along dimension 1 of the lhs and dimension @@ -61,13 +62,21 @@ void MatMul(const void* run_options_ptr, T* out, T* lhs, T* rhs, int64 m, void __xla_cpu_runtime_EigenSingleThreadedMatMulF32( const void* run_options_ptr, float* out, float* lhs, float* rhs, int64 m, int64 n, int64 k, int32 transpose_lhs, int32 transpose_rhs) { - MatMul(run_options_ptr, out, lhs, rhs, m, n, k, transpose_lhs, - transpose_rhs); + if (m == 1 || n == 1) { + xla::EigenMatVecF32(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); + } } void __xla_cpu_runtime_EigenSingleThreadedMatMulF64( const void* run_options_ptr, double* out, double* lhs, double* rhs, int64 m, int64 n, int64 k, int32 transpose_lhs, int32 transpose_rhs) { - MatMul(run_options_ptr, out, lhs, rhs, m, n, k, transpose_lhs, - transpose_rhs); + if (m == 1 || n == 1) { + xla::EigenMatVecF64(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); + } } diff --git a/tensorflow/compiler/xla/service/cpu/sample_harness.cc b/tensorflow/compiler/xla/service/cpu/sample_harness.cc index 8f1ce82d49a1c7cabfb62bf30e69faedc0318138..b3f4609d465efb4df8921abb684bafd263fe040f 100644 --- a/tensorflow/compiler/xla/service/cpu/sample_harness.cc +++ b/tensorflow/compiler/xla/service/cpu/sample_harness.cc @@ -38,13 +38,12 @@ int main(int argc, char** argv) { // Transfer parameters. std::unique_ptr param0_literal = - xla::LiteralUtil::CreateR1({1.1f, 2.2f, 3.3f, 5.5f}); + xla::Literal::CreateR1({1.1f, 2.2f, 3.3f, 5.5f}); std::unique_ptr param0_data = client->TransferToServer(*param0_literal).ConsumeValueOrDie(); - std::unique_ptr param1_literal = - xla::LiteralUtil::CreateR2( - {{3.1f, 4.2f, 7.3f, 9.5f}, {1.1f, 2.2f, 3.3f, 4.4f}}); + std::unique_ptr param1_literal = xla::Literal::CreateR2( + {{3.1f, 4.2f, 7.3f, 9.5f}, {1.1f, 2.2f, 3.3f, 4.4f}}); std::unique_ptr param1_data = client->TransferToServer(*param1_literal).ConsumeValueOrDie(); @@ -69,7 +68,7 @@ int main(int argc, char** argv) { LOG(INFO) << tensorflow::strings::Printf("computation took %lldns", profile.compute_time_ns()); - LOG(INFO) << xla::LiteralUtil::ToString(*actual); + LOG(INFO) << actual->ToString(); return 0; } diff --git a/tensorflow/compiler/xla/service/cpu/shape_partition.cc b/tensorflow/compiler/xla/service/cpu/shape_partition.cc new file mode 100644 index 0000000000000000000000000000000000000000..61b408b8c24dded134218110d4e219c31f1685a8 --- /dev/null +++ b/tensorflow/compiler/xla/service/cpu/shape_partition.cc @@ -0,0 +1,160 @@ +/* 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/cpu/shape_partition.h" + +namespace xla { +namespace cpu { + +std::vector ShapePartitionAssigner::Run(int64 target_partition_count) { + // Gather outer-most dims where dim_size >= 'target_partition_count'. + // Note: always leave inner-dim static for vectorization/optimizations. + std::vector outer_dims; + int64 outer_dim_size = 1; + // TODO(b/27458679) Consider reserving enough minor dimensions (based on + // target vector register width) to enable vector instructions. + for (int i = shape_.layout().minor_to_major_size() - 1; i >= 1; --i) { + const int64 dimension = shape_.layout().minor_to_major(i); + outer_dims.push_back(dimension); + outer_dim_size *= shape_.dimensions(dimension); + if (outer_dim_size >= target_partition_count) { + break; + } + } + + // Clip target partition count if outer dim size is insufficient to cover. + target_partition_count = std::min(outer_dim_size, target_partition_count); + + // Calculate the target number of partitions per-dimension, by factoring + // 'target_partition_count' into 'num_outer_dims' equal terms. + // EX: + // *) target_partition_count = 16 + // *) out_dim_count = 2 + // *) target_dim_partition_count = 16 ^ (1.0 / 2) == 4 + const int64 target_dim_partition_count = std::pow( + static_cast(target_partition_count), 1.0 / outer_dims.size()); + + // Assign feasible dimension partitions based on 'target_dim_partition_count' + // and actual dimension sizes from 'shape_'. + std::vector dimension_partition_counts(outer_dims.size()); + for (int64 i = 0; i < outer_dims.size(); ++i) { + dimension_partition_counts[i] = + std::min(static_cast(shape_.dimensions(outer_dims[i])), + target_dim_partition_count); + } + + // Check if total partition count is below 'target_partition_count'. + // This can occur if some dimensions in 'shape_' are below the + // 'target_dim_partition_count' threshold. + if (GetTotalPartitionCount(dimension_partition_counts) < + target_partition_count) { + // Assign additional partitions (greedily to outer dimensions), if doing + // so would keep the total number of partitions <= 'target_partition_count', + // using one pass over 'dimension_partition_counts'. + for (int64 i = 0; i < dimension_partition_counts.size(); ++i) { + const int64 current_dim_partition_count = dimension_partition_counts[i]; + const int64 other_dims_partition_count = + GetTotalPartitionCount(dimension_partition_counts) / + current_dim_partition_count; + // Constraint: (current + additional) * other <= target + // Calculate: additional = target / other - current + int64 additional_partition_count = + target_partition_count / other_dims_partition_count - + current_dim_partition_count; + // Clip 'additional_partition_count' by current dimension size. + additional_partition_count = std::min( + shape_.dimensions(outer_dims[i]) - dimension_partition_counts[i], + additional_partition_count); + if (additional_partition_count > 0) { + dimension_partition_counts[i] += additional_partition_count; + } + } + } + + return dimension_partition_counts; +} + +int64 ShapePartitionAssigner::GetTotalPartitionCount( + const std::vector& dimension_partition_counts) { + int64 total_partition_count = 1; + for (int64 dim_partition_count : dimension_partition_counts) { + total_partition_count *= dim_partition_count; + } + return total_partition_count; +} + +ShapePartitionIterator::ShapePartitionIterator( + const Shape& shape, const std::vector& dimension_partition_counts) + : shape_(shape), + dimension_partition_counts_(dimension_partition_counts), + dimensions_(dimension_partition_counts_.size()), + dimension_partition_sizes_(dimension_partition_counts_.size()), + dimension_partition_strides_(dimension_partition_counts_.size()) { + // Store partitioned outer dimensions from 'shape_'. + for (int i = 0; i < dimensions_.size(); ++i) { + dimensions_[i] = shape_.layout().minor_to_major( + shape_.layout().minor_to_major_size() - 1 - i); + } + + // Calculate partition size for each dimension (note that the size of + // the last partition in each dimension may be different if the dimension + // size is not a multiple of partition size). + for (int i = 0; i < dimension_partition_sizes_.size(); ++i) { + const int64 dim_size = shape_.dimensions(dimensions_[i]); + dimension_partition_sizes_[i] = + std::max(1LL, dim_size / dimension_partition_counts_[i]); + } + + // Calculate the partition strides for each dimension. + dimension_partition_strides_[dimension_partition_strides_.size() - 1] = 1; + for (int i = dimension_partition_strides_.size() - 2; i >= 0; --i) { + dimension_partition_strides_[i] = dimension_partition_strides_[i + 1] * + dimension_partition_counts_[i + 1]; + } +} + +std::vector> ShapePartitionIterator::GetPartition( + int64 index) const { + // Calculate and return the partition for 'index'. + // Returns for each dimension: (partition_start, partition_size). + std::vector> partition(dimensions_.size()); + for (int64 i = 0; i < partition.size(); ++i) { + // Calculate the index for dimension 'i'. + const int64 partition_index = index / dimension_partition_strides_[i]; + // Calculate dimension partition start at 'partition_index'. + partition[i].first = partition_index * dimension_partition_sizes_[i]; + // Calculate dimension partition size (note that the last partition size + // may be adjusted if dimension size is not a multiple of partition size). + if (partition_index == dimension_partition_counts_[i] - 1) { + // Last partition in this dimension. + partition[i].second = + shape_.dimensions(dimensions_[i]) - partition[i].first; + } else { + partition[i].second = dimension_partition_sizes_[i]; + } + CHECK_GT(partition[i].second, 0); + // Update index to remove conribution from current dimension. + index -= partition_index * dimension_partition_strides_[i]; + } + return partition; +} + +int64 ShapePartitionIterator::GetTotalPartitionCount() const { + return ShapePartitionAssigner::GetTotalPartitionCount( + dimension_partition_counts_); +} + +} // namespace cpu +} // namespace xla diff --git a/tensorflow/compiler/xla/service/cpu/shape_partition.h b/tensorflow/compiler/xla/service/cpu/shape_partition.h new file mode 100644 index 0000000000000000000000000000000000000000..7a2d00421cfdc8e41ec48698a16665621de16bda --- /dev/null +++ b/tensorflow/compiler/xla/service/cpu/shape_partition.h @@ -0,0 +1,105 @@ +/* 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 THIRD_PARTY_TENSORFLOW_COMPILER_XLA_SERVICE_CPU_SHAPE_PARTITION_H_ +#define THIRD_PARTY_TENSORFLOW_COMPILER_XLA_SERVICE_CPU_SHAPE_PARTITION_H_ + +#include + +#include "tensorflow/compiler/xla/shape_util.h" + +namespace xla { +namespace cpu { + +// ShapePartitionAssigner partitions the most-major dimensions of 'shape' such +// that the total partition count <= 'target_partition_count'. +// +// Example 1: +// +// Let 'shape' = [8, 16, 32] and 'target_partition_count' = 6. +// +// Because the most-major dimension size is <= 'target_partition_count', we +// can generate our target number of partitions by partition the most-major +// dimensions. +// +// This will result in the following partitions of the most-major dimension: +// +// [0, 1), [1, 2), [2, 3), [3, 4), [4, 5) [5, 8) +// +// Note that the last partition has residule because the dimension size is +// not a multiple of the partition count. +// +// +// Example 2: +// +// Let 'shape' = [8, 16, 32] and 'target_partition_count' = 16. +// +// Because the most-major dimension only has size 8, we must also partition +// the next most-major dimension to generate the target of 16 partitions. +// We factor 'target_partition_count' by the number of most-major dimensions +// we need to partition, to get a per-dimension target partition count: +// +// target_dimension_partition_count = 16 ^ (1 / 2) == 4 +// +// This will result in the following partitions of the most-major dimension: +// +// [0, 2), [2, 4), [4, 6), [6, 8) +// +// This will result in the following partitions of the second most-major +// dimension: +// +// [0, 4), [4, 8), [8, 12), [12, 16) +// +class ShapePartitionAssigner { + public: + ShapePartitionAssigner(const Shape& shape) : shape_(shape) {} + + // Returns dimension partition counts (starting at outer-most dimension). + std::vector Run(int64 target_partition_count); + + // Returns the total partition count based on 'dimension_partition_counts'. + static int64 GetTotalPartitionCount( + const std::vector& dimension_partition_counts); + + private: + const Shape& shape_; +}; + +// ShapePartitionIterator iterates through outer-dimension partitions of +// 'shape' as specified by 'dimension_partition_counts'. +class ShapePartitionIterator { + public: + ShapePartitionIterator(const Shape& shape, + const std::vector& dimension_partition_counts); + + // Returns a partition [start, size] for each dimension. + // Partitions are listed starting from outer-most dimension first. + std::vector> GetPartition(int64 index) const; + + int64 GetTotalPartitionCount() const; + + private: + const Shape& shape_; + const std::vector dimension_partition_counts_; + + std::vector dimensions_; + std::vector dimension_partition_sizes_; + std::vector dimension_partition_strides_; +}; + +} // namespace cpu +} // namespace xla + +#endif // THIRD_PARTY_TENSORFLOW_COMPILER_XLA_SERVICE_CPU_SHAPE_PARTITION_H_ diff --git a/tensorflow/compiler/xla/service/cpu/shape_partition_test.cc b/tensorflow/compiler/xla/service/cpu/shape_partition_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..ee0c53fa6d7c41481a53350e57e5844dea2644c1 --- /dev/null +++ b/tensorflow/compiler/xla/service/cpu/shape_partition_test.cc @@ -0,0 +1,248 @@ +/* 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/cpu/shape_partition.h" + +#include +#include + +#include "tensorflow/compiler/xla/test_helpers.h" +#include "tensorflow/compiler/xla/tests/hlo_test_base.h" +#include "tensorflow/compiler/xla/util.h" + +namespace xla { +namespace cpu { +namespace { + +class ShapePartitionAssignerTest : public HloTestBase { + protected: + typedef std::vector Vec; + + void RunR2Test(const Shape& shape, const int64 expected_max_partition_count) { + ShapePartitionAssigner assigner(shape); + // Check all partitions of outer dimension. + for (int64 i = 1; i <= expected_max_partition_count; ++i) { + EXPECT_TRUE(ContainersEqual(Vec({i}), + assigner.Run(/*target_partition_count=*/i))); + } + // Check target_partition_count > outer dimension size. + EXPECT_TRUE(ContainersEqual( + Vec({expected_max_partition_count}), + assigner.Run( + /*target_partition_count=*/expected_max_partition_count + 1))); + } +}; + +TEST_F(ShapePartitionAssignerTest, Shape13WithLayout10) { + RunR2Test(ShapeUtil::MakeShapeWithLayout(F32, {1, 3}, {1, 0}), 1); +} + +TEST_F(ShapePartitionAssignerTest, Shape31WithLayout01) { + RunR2Test(ShapeUtil::MakeShapeWithLayout(F32, {3, 1}, {0, 1}), 1); +} + +TEST_F(ShapePartitionAssignerTest, Shape53WithLayout10) { + RunR2Test(ShapeUtil::MakeShapeWithLayout(F32, {5, 3}, {1, 0}), 5); +} + +TEST_F(ShapePartitionAssignerTest, Shape53WithLayout01) { + RunR2Test(ShapeUtil::MakeShapeWithLayout(F32, {5, 3}, {0, 1}), 3); +} + +TEST_F(ShapePartitionAssignerTest, Shape532WithLayout210) { + Shape shape = ShapeUtil::MakeShapeWithLayout(F32, {5, 3, 2}, {2, 1, 0}); + ShapePartitionAssigner assigner(shape); + + for (int64 i = 1; i <= 5; ++i) { + EXPECT_TRUE(ContainersEqual(Vec({i}), assigner.Run( + /*target_partition_count=*/i))); + } + + EXPECT_TRUE( + ContainersEqual(Vec({3, 2}), assigner.Run(/*target_partition_count=*/6))); + EXPECT_TRUE( + ContainersEqual(Vec({3, 2}), assigner.Run(/*target_partition_count=*/7))); + EXPECT_TRUE( + ContainersEqual(Vec({4, 2}), assigner.Run(/*target_partition_count=*/8))); + EXPECT_TRUE( + ContainersEqual(Vec({3, 3}), assigner.Run(/*target_partition_count=*/9))); + EXPECT_TRUE(ContainersEqual(Vec({3, 3}), + assigner.Run(/*target_partition_count=*/10))); + EXPECT_TRUE(ContainersEqual(Vec({3, 3}), + assigner.Run(/*target_partition_count=*/11))); + EXPECT_TRUE(ContainersEqual(Vec({4, 3}), + assigner.Run(/*target_partition_count=*/12))); + EXPECT_TRUE(ContainersEqual(Vec({4, 3}), + assigner.Run(/*target_partition_count=*/13))); + EXPECT_TRUE(ContainersEqual(Vec({4, 3}), + assigner.Run(/*target_partition_count=*/14))); + EXPECT_TRUE(ContainersEqual(Vec({5, 3}), + assigner.Run(/*target_partition_count=*/15))); + EXPECT_TRUE(ContainersEqual(Vec({5, 3}), + assigner.Run(/*target_partition_count=*/16))); +} + +TEST_F(ShapePartitionAssignerTest, Shape532WithLayout201) { + Shape shape = ShapeUtil::MakeShapeWithLayout(F32, {5, 3, 2}, {2, 0, 1}); + ShapePartitionAssigner assigner(shape); + + for (int64 i = 1; i <= 3; ++i) { + EXPECT_TRUE(ContainersEqual(Vec({i}), assigner.Run( + /*target_partition_count=*/i))); + } + + EXPECT_TRUE( + ContainersEqual(Vec({2, 2}), assigner.Run(/*target_partition_count=*/4))); + EXPECT_TRUE( + ContainersEqual(Vec({2, 2}), assigner.Run(/*target_partition_count=*/5))); + EXPECT_TRUE( + ContainersEqual(Vec({3, 2}), assigner.Run(/*target_partition_count=*/6))); + EXPECT_TRUE( + ContainersEqual(Vec({3, 2}), assigner.Run(/*target_partition_count=*/7))); + EXPECT_TRUE( + ContainersEqual(Vec({3, 2}), assigner.Run(/*target_partition_count=*/8))); + EXPECT_TRUE( + ContainersEqual(Vec({3, 3}), assigner.Run(/*target_partition_count=*/9))); + EXPECT_TRUE(ContainersEqual(Vec({3, 3}), + assigner.Run(/*target_partition_count=*/10))); + EXPECT_TRUE(ContainersEqual(Vec({3, 3}), + assigner.Run(/*target_partition_count=*/11))); + EXPECT_TRUE(ContainersEqual(Vec({3, 4}), + assigner.Run(/*target_partition_count=*/12))); + EXPECT_TRUE(ContainersEqual(Vec({3, 4}), + assigner.Run(/*target_partition_count=*/13))); + EXPECT_TRUE(ContainersEqual(Vec({3, 4}), + assigner.Run(/*target_partition_count=*/14))); + EXPECT_TRUE(ContainersEqual(Vec({3, 5}), + assigner.Run(/*target_partition_count=*/15))); + EXPECT_TRUE(ContainersEqual(Vec({3, 5}), + assigner.Run(/*target_partition_count=*/16))); +} + +class ShapePartitionIteratorTest : public HloTestBase { + protected: + typedef std::vector> Partition; +}; + +TEST_F(ShapePartitionIteratorTest, Shape53WithLayout10) { + Shape shape = ShapeUtil::MakeShapeWithLayout(F32, {5, 3}, {1, 0}); + + { + ShapePartitionIterator iterator(shape, {1}); + EXPECT_EQ(1, iterator.GetTotalPartitionCount()); + EXPECT_TRUE(ContainersEqual(Partition({{0, 5}}), iterator.GetPartition(0))); + } + + { + ShapePartitionIterator iterator(shape, {2}); + EXPECT_EQ(2, iterator.GetTotalPartitionCount()); + EXPECT_TRUE(ContainersEqual(Partition({{0, 2}}), iterator.GetPartition(0))); + EXPECT_TRUE(ContainersEqual(Partition({{2, 3}}), iterator.GetPartition(1))); + } + + { + ShapePartitionIterator iterator(shape, {3}); + EXPECT_EQ(3, iterator.GetTotalPartitionCount()); + EXPECT_TRUE(ContainersEqual(Partition({{0, 1}}), iterator.GetPartition(0))); + EXPECT_TRUE(ContainersEqual(Partition({{1, 1}}), iterator.GetPartition(1))); + EXPECT_TRUE(ContainersEqual(Partition({{2, 3}}), iterator.GetPartition(2))); + } +} + +TEST_F(ShapePartitionIteratorTest, Shape532WithLayout210) { + Shape shape = ShapeUtil::MakeShapeWithLayout(F32, {5, 3, 2}, {2, 1, 0}); + + { + ShapePartitionIterator iterator(shape, {1, 1}); + EXPECT_EQ(1, iterator.GetTotalPartitionCount()); + EXPECT_TRUE( + ContainersEqual(Partition({{0, 5}, {0, 3}}), iterator.GetPartition(0))); + } + + { + ShapePartitionIterator iterator(shape, {2, 2}); + EXPECT_EQ(4, iterator.GetTotalPartitionCount()); + EXPECT_TRUE( + ContainersEqual(Partition({{0, 2}, {0, 1}}), iterator.GetPartition(0))); + EXPECT_TRUE( + ContainersEqual(Partition({{0, 2}, {1, 2}}), iterator.GetPartition(1))); + EXPECT_TRUE( + ContainersEqual(Partition({{2, 3}, {0, 1}}), iterator.GetPartition(2))); + EXPECT_TRUE( + ContainersEqual(Partition({{2, 3}, {1, 2}}), iterator.GetPartition(3))); + } +} + +class RandomShapePartitionIteratorTest : public HloTestBase { + protected: + typedef std::vector> Partition; + RandomShapePartitionIteratorTest() + : generator_(rd_()), distribution_(1, 10) {} + + std::vector RandR4Dims() { return {Rand(), Rand(), Rand(), Rand()}; } + + int64 Rand() { return distribution_(generator_); } + + std::random_device rd_; + std::mt19937 generator_; + std::uniform_int_distribution distribution_; +}; + +TEST_F(RandomShapePartitionIteratorTest, RandomShapeAndPartitions) { + // Choose random dimensions for R4 shape. + Shape shape = ShapeUtil::MakeShapeWithLayout(F32, RandR4Dims(), {3, 2, 1, 0}); + // Choose random number of outer dimensions to partition. + const int num_outer_dims_to_partition = 1 + (Rand() % 3); + // Choose random outer dimension partition counts. + std::vector dim_sizes(num_outer_dims_to_partition); + std::vector dim_partition_counts(num_outer_dims_to_partition); + int64 total_dim_size = 1; + for (int i = 0; i < num_outer_dims_to_partition; ++i) { + const int64 dimension = shape.layout().minor_to_major( + shape.layout().minor_to_major_size() - 1 - i); + dim_sizes[i] = shape.dimensions(dimension); + total_dim_size *= dim_sizes[i]; + // Choose dimension partition count in [1, dim_size] + const int64 dim_partition_count = 1 + Rand() % dim_sizes[i]; + dim_partition_counts[i] = dim_partition_count; + } + // Iterate through all partition: for each partition record covered + // index ranges by dimension. + std::vector> ranges(num_outer_dims_to_partition); + ShapePartitionIterator partition_iterator(shape, dim_partition_counts); + const int64 partition_count = partition_iterator.GetTotalPartitionCount(); + for (int64 i = 0; i < partition_count; ++i) { + const auto& dim_partition = partition_iterator.GetPartition(i); + for (int dim = 0; dim < dim_partition.size(); ++dim) { + ranges[dim].insert( + std::make_pair(dim_partition[dim].first, + dim_partition[dim].first + dim_partition[dim].second)); + } + } + // Check that partitions cover entire dimension size range (for each + // partitioned dimension). + for (int i = 0; i < ranges.size(); ++i) { + int64 expected_index = 0; + for (auto& r : ranges[i]) { + EXPECT_EQ(expected_index, r.first); + expected_index = r.second; + } + EXPECT_EQ(expected_index, dim_sizes[i]); + } +} + +} // namespace +} // namespace cpu +} // namespace xla diff --git a/tensorflow/compiler/xla/service/cpu/simple_orc_jit.cc b/tensorflow/compiler/xla/service/cpu/simple_orc_jit.cc index 8beb565ab3e220f9b9eebac836c8de8c1fc2e8ee..c3c11df090e88c3c24104b66d28b3b16f03baa80 100644 --- a/tensorflow/compiler/xla/service/cpu/simple_orc_jit.cc +++ b/tensorflow/compiler/xla/service/cpu/simple_orc_jit.cc @@ -21,13 +21,15 @@ limitations under the License. #include #include -#include "external/llvm/include/llvm/IR/Mangler.h" -#include "external/llvm/include/llvm/Support/CodeGen.h" -#include "external/llvm/include/llvm/Support/Host.h" +#include "llvm/ExecutionEngine/ExecutionEngine.h" +#include "llvm/ExecutionEngine/SectionMemoryManager.h" +#include "llvm/IR/Mangler.h" +#include "llvm/Support/CodeGen.h" +#include "llvm/Support/Host.h" #include "tensorflow/compiler/xla/ptr_util.h" -#include "tensorflow/compiler/xla/service/cpu/compiler_functor.h" #include "tensorflow/compiler/xla/service/cpu/cpu_runtime.h" #include "tensorflow/compiler/xla/service/cpu/cpu_runtime_avx.h" +#include "tensorflow/compiler/xla/service/cpu/cpu_runtime_neon.h" #include "tensorflow/compiler/xla/service/cpu/cpu_runtime_sse4_1.h" #include "tensorflow/compiler/xla/service/cpu/runtime_conv2d.h" #include "tensorflow/compiler/xla/service/cpu/runtime_matmul.h" @@ -41,7 +43,7 @@ namespace cpu { namespace { // Converts a symbol 'name' into the form expected by dlsym(). -std::string CanonicalizeSymbol(const std::string &name) { +std::string CanonicalizeSymbol(const std::string& name) { #if defined(__APPLE__) // On Mac OS X, dlsym() expects names not to be prefixed with a leading // underscore. @@ -52,47 +54,77 @@ std::string CanonicalizeSymbol(const std::string &name) { return name; } +class JITSymbolTable { + public: + JITSymbolTable() { Populate(); } + + void* Lookup(llvm::StringRef jit_symbol_name) const { + auto it = jit_symbol_table_.find(jit_symbol_name); + return it == jit_symbol_table_.end() ? nullptr : it->getValue(); + } + + static bool MustBeInTable(llvm::StringRef name) { + // In particular, names starting with + // runtime::kXlaCpuRuntimeSymbolNamePrefix should not be dlsym'ed. + return name.startswith(runtime::kXlaCpuRuntimeSymbolNamePrefix); + } + + private: + void AddJITSymbolToTable(llvm::StringRef jit_symbol_name, + llvm::StringRef cpp_symbol_name, + void* jit_symbol_value) { + // The JIT symbol name and the C++ symbol name (with an extern "C" linkage) + // need to match, otherwise AOT links will fail. + CHECK(jit_symbol_name == cpp_symbol_name); + CHECK(jit_symbol_table_.insert({jit_symbol_name, jit_symbol_value}).second); + } + + void Populate() { +#define ADD_JIT_SYMBOL_TO_TABLE(base_name) \ + do { \ + AddJITSymbolToTable( \ + xla::cpu::runtime::k##base_name##SymbolName, \ + "__xla_cpu_runtime_" #base_name, \ + reinterpret_cast(__xla_cpu_runtime_##base_name)); \ + } while (false) + + ADD_JIT_SYMBOL_TO_TABLE(AcquireInfeedBufferForDequeue); + ADD_JIT_SYMBOL_TO_TABLE(ReleaseInfeedBufferAfterDequeue); + ADD_JIT_SYMBOL_TO_TABLE(AcquireOutfeedBufferForPopulation); + ADD_JIT_SYMBOL_TO_TABLE(ReleaseOutfeedBufferAfterPopulation); + ADD_JIT_SYMBOL_TO_TABLE(ExpV8F32AVX); + ADD_JIT_SYMBOL_TO_TABLE(LogV8F32AVX); + ADD_JIT_SYMBOL_TO_TABLE(ExpV4F32SSE); + ADD_JIT_SYMBOL_TO_TABLE(LogV4F32SSE); + ADD_JIT_SYMBOL_TO_TABLE(ExpV4F32NEON); + ADD_JIT_SYMBOL_TO_TABLE(LogV4F32NEON); + ADD_JIT_SYMBOL_TO_TABLE(EigenConvF32); + ADD_JIT_SYMBOL_TO_TABLE(EigenMatMulF32); + ADD_JIT_SYMBOL_TO_TABLE(EigenMatMulF64); + ADD_JIT_SYMBOL_TO_TABLE(EigenSingleThreadedConvF32); + ADD_JIT_SYMBOL_TO_TABLE(EigenSingleThreadedMatMulF32); + ADD_JIT_SYMBOL_TO_TABLE(EigenSingleThreadedMatMulF64); + +#undef ADD_JIT_SYMBOL_TO_TABLE + } + + llvm::StringMap jit_symbol_table_; +}; + +const JITSymbolTable& GetJITSymbolTable() { + static JITSymbolTable* symbol_table = new JITSymbolTable; + return *symbol_table; +} + // A simple SymbolResolver that delegates to the host dynamic linker. struct SimpleResolver : public llvm::JITSymbolResolver { - llvm::JITSymbol findSymbol(const std::string &name) override { - void *func_addr = nullptr; - + llvm::JITSymbol findSymbol(const std::string& name) override { std::string canonical_name = CanonicalizeSymbol(name); - if (canonical_name == runtime::kEigenMatmulF32SymbolName) { - func_addr = reinterpret_cast(__xla_cpu_runtime_EigenMatMulF32); - } else if (canonical_name == - runtime::kEigenSingleThreadedMatmulF32SymbolName) { - func_addr = reinterpret_cast( - __xla_cpu_runtime_EigenSingleThreadedMatMulF32); - } else if (canonical_name == runtime::kEigenConvF32SymbolName) { - func_addr = reinterpret_cast(__xla_cpu_runtime_EigenConvF32); - } else if (canonical_name == - runtime::kEigenSingleThreadedConvF32SymbolName) { - func_addr = reinterpret_cast( - __xla_cpu_runtime_EigenSingleThreadedConvF32); - } else if (canonical_name == - runtime::kAcquireInfeedBufferForDequeueSymbolName) { - func_addr = reinterpret_cast( - __xla_cpu_runtime_AcquireInfeedBufferForDequeue); - } else if (canonical_name == - runtime::kReleaseInfeedBufferAfterDequeueSymbolName) { - func_addr = reinterpret_cast( - __xla_cpu_runtime_ReleaseInfeedBufferAfterDequeue); - } else if (canonical_name == runtime::kExpV4F32) { - func_addr = reinterpret_cast(runtime::ExpV4F32); - } else if (canonical_name == runtime::kExpV8F32) { - func_addr = reinterpret_cast(runtime::ExpV8F32); - } else if (canonical_name == runtime::kLogV4F32) { - func_addr = reinterpret_cast(runtime::LogV4F32); - } else if (canonical_name == runtime::kLogV8F32) { - func_addr = reinterpret_cast(runtime::LogV8F32); - } else if (canonical_name == runtime::kTanhV4F32) { - func_addr = reinterpret_cast(runtime::TanhV4F32); - } else if (canonical_name == runtime::kTanhV8F32) { - func_addr = reinterpret_cast(runtime::TanhV8F32); - } else { - func_addr = dlsym(RTLD_DEFAULT, canonical_name.c_str()); - } + const JITSymbolTable& jit_symbol_table = GetJITSymbolTable(); + + void* func_addr = JITSymbolTable::MustBeInTable(canonical_name) + ? jit_symbol_table.Lookup(canonical_name) + : dlsym(RTLD_DEFAULT, canonical_name.c_str()); if (func_addr == nullptr) { return nullptr; @@ -101,7 +133,7 @@ struct SimpleResolver : public llvm::JITSymbolResolver { llvm::JITSymbolFlags::None); return symbol_info; } - llvm::JITSymbol findSymbolInLogicalDylib(const std::string &name) override { + llvm::JITSymbol findSymbolInLogicalDylib(const std::string& name) override { return nullptr; } }; @@ -110,60 +142,82 @@ llvm::SmallVector DetectMachineAttributes() { llvm::SmallVector result; llvm::StringMap host_features; if (llvm::sys::getHostCPUFeatures(host_features)) { - for (auto &feature : host_features) { + for (auto& feature : host_features) { if (feature.second) { - result.push_back(feature.first()); + llvm::StringRef feature_name = feature.first(); + // Skip avx512 for now, it isn't quite ready in LLVM. + if (feature_name.startswith("avx512")) { + continue; + } + result.push_back(feature_name); } } } return result; } +llvm::StringRef GetHostCpuName() { + auto cpu_name = llvm::sys::getHostCPUName(); + // Skip avx512 for now, it isn't quite ready in LLVM. + cpu_name.consume_back("-avx512"); + return cpu_name; +} + CompilerFunctor::VectorIntrinsics GetAvailableIntrinsics() { CompilerFunctor::VectorIntrinsics intrinsics; - intrinsics.sse_intrinsics = (&runtime::ExpV4F32 != nullptr); - intrinsics.avx_intrinsics = (&runtime::ExpV8F32 != nullptr); + intrinsics.sse_intrinsics = (&__xla_cpu_runtime_ExpV4F32SSE != nullptr); + intrinsics.avx_intrinsics = (&__xla_cpu_runtime_ExpV8F32AVX != nullptr); + intrinsics.neon_intrinsics = (&__xla_cpu_runtime_ExpV4F32NEON != nullptr); return intrinsics; } } // namespace -SimpleOrcJIT::SimpleOrcJIT(const llvm::TargetOptions &target_options, - llvm::CodeGenOpt::Level opt_level) +SimpleOrcJIT::SimpleOrcJIT(const llvm::TargetOptions& target_options, + llvm::CodeGenOpt::Level opt_level, + bool optimize_for_size, bool enable_fast_math, + bool disable_expensive_passes, + LLVMCompiler::ModuleHook pre_optimization_hook, + LLVMCompiler::ModuleHook post_optimization_hook) : target_machine_( CHECK_NOTNULL(llvm::EngineBuilder() .setTargetOptions(target_options) .setOptLevel(opt_level) .selectTarget( /*TargetTriple=*/llvm::Triple(), /*MArch=*/"", - /*MCPU=*/llvm::sys::getHostCPUName(), + /*MCPU=*/GetHostCpuName(), /*MAttrs=*/DetectMachineAttributes()))), disassembler_(*target_machine_), data_layout_(target_machine_->createDataLayout()), - compile_layer_(object_layer_, - CompilerFunctor(target_machine_.get(), &disassembler_, - opt_level, GetAvailableIntrinsics())) {} + object_layer_( + [] { return std::make_shared(); }), + compile_layer_( + object_layer_, + CompilerFunctor(target_machine_.get(), &disassembler_, opt_level, + optimize_for_size, enable_fast_math, + disable_expensive_passes, GetAvailableIntrinsics(), + std::move(pre_optimization_hook), + std::move(post_optimization_hook))) { + VLOG(1) << "CPU target: " << target_machine_->getTargetCPU().str() + << " features: " << target_machine_->getTargetFeatureString().str(); +} SimpleOrcJIT::ModuleHandleT SimpleOrcJIT::AddModule( std::unique_ptr module) { - // The Orc API adds a whole iterable "set" of modules, so we wrap the module - // in a vector. - std::vector> module_set; - module_set.push_back(std::move(module)); - auto handle = compile_layer_.addModuleSet( - std::move(module_set), MakeUnique(), - MakeUnique()); + auto handle = cantFail(compile_layer_.addModule( + std::move(module), MakeUnique())); module_handles_.push_back(handle); return handle; } void SimpleOrcJIT::RemoveModule(SimpleOrcJIT::ModuleHandleT handle) { module_handles_.erase( - std::remove(module_handles_.begin(), module_handles_.end(), handle)); - compile_layer_.removeModuleSet(handle); + std::remove(module_handles_.begin(), module_handles_.end(), handle), + module_handles_.end()); + cantFail(compile_layer_.removeModule(handle)); } -llvm::JITSymbol SimpleOrcJIT::FindSymbol(const std::string &name) { +llvm::JITSymbol SimpleOrcJIT::FindSymbol(const std::string& name) { std::string mangled_name; { llvm::raw_string_ostream mangled_name_stream(mangled_name); @@ -172,7 +226,7 @@ llvm::JITSymbol SimpleOrcJIT::FindSymbol(const std::string &name) { // Resolve symbol from last module to first, allowing later redefinitions of // symbols shadow earlier ones. - for (auto &handle : + for (auto& handle : llvm::make_range(module_handles_.rbegin(), module_handles_.rend())) { if (auto symbol = compile_layer_.findSymbolIn(handle, mangled_name, diff --git a/tensorflow/compiler/xla/service/cpu/simple_orc_jit.h b/tensorflow/compiler/xla/service/cpu/simple_orc_jit.h index 4d8653484a037a345321dbe11c384f650e0142d0..e476c0e3812cc0fb2a2d633832374b3165ca072a 100644 --- a/tensorflow/compiler/xla/service/cpu/simple_orc_jit.h +++ b/tensorflow/compiler/xla/service/cpu/simple_orc_jit.h @@ -20,11 +20,12 @@ limitations under the License. #include #include -#include "external/llvm/include/llvm/ADT/Triple.h" -#include "external/llvm/include/llvm/ExecutionEngine/Orc/IRCompileLayer.h" -#include "external/llvm/include/llvm/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.h" -#include "external/llvm/include/llvm/IR/Module.h" -#include "external/llvm/include/llvm/Target/TargetMachine.h" +#include "llvm/ADT/Triple.h" +#include "llvm/ExecutionEngine/Orc/IRCompileLayer.h" +#include "llvm/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.h" +#include "llvm/IR/Module.h" +#include "llvm/Target/TargetMachine.h" +#include "tensorflow/compiler/xla/service/cpu/compiler_functor.h" #include "tensorflow/compiler/xla/service/cpu/disassembler.h" #include "tensorflow/compiler/xla/types.h" @@ -41,9 +42,12 @@ namespace cpu { // it's added to the JIT. class SimpleOrcJIT { public: - using ObjLayerT = llvm::orc::RTDyldObjectLinkingLayer<>; - using CompileLayerT = llvm::orc::IRCompileLayer; - using ModuleHandleT = CompileLayerT::ModuleSetHandleT; + using ObjLayerT = llvm::orc::RTDyldObjectLinkingLayer; + using CompileFtor = + std::function( + llvm::Module&)>; + using CompileLayerT = llvm::orc::IRCompileLayer; + using ModuleHandleT = CompileLayerT::ModuleHandleT; // Create a new JIT, targeting the host architecture. // The |target_options| parameter allows customization of certain code @@ -51,8 +55,19 @@ class SimpleOrcJIT { // can be reassociated, etc.). // The |opt_level| parameter controls the optimization level of the code // generator. + // The |optimize_for_size| parameter specifies that the code generator should + // optimize to reduce code size, potentially at the cost of performance. + // The |disable_expensive_passes| parameter will disable certain optimization + // passes + // The |pre_optimization_hook| is invoked on the module before any IR + // level optimizations are applied. + // The |post_optimization_hook| is invoked on the module after all IR + // level optimizations are applied. SimpleOrcJIT(const llvm::TargetOptions& target_options, - llvm::CodeGenOpt::Level opt_level); + llvm::CodeGenOpt::Level opt_level, bool optimize_for_size, + bool enable_fast_math, bool disable_expensive_passes, + LLVMCompiler::ModuleHook pre_optimization_hook, + LLVMCompiler::ModuleHook post_optimization_hook); // Data layout this JIT was created with. const llvm::DataLayout& data_layout() const { return data_layout_; } @@ -73,6 +88,8 @@ class SimpleOrcJIT { // nullptr if the symbol cannot be found. llvm::JITSymbol FindSymbol(const std::string& name); + llvm::TargetMachine* target_machine() const { return target_machine_.get(); } + private: std::vector module_handles_; std::unique_ptr target_machine_; diff --git a/tensorflow/compiler/xla/service/cpu/xfeed_manager.cc b/tensorflow/compiler/xla/service/cpu/xfeed_manager.cc new file mode 100644 index 0000000000000000000000000000000000000000..d0f214202908266371639af8f431ad8269ad0e35 --- /dev/null +++ b/tensorflow/compiler/xla/service/cpu/xfeed_manager.cc @@ -0,0 +1,78 @@ +/* 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/cpu/xfeed_manager.h" + +#include "tensorflow/compiler/xla/shape_util.h" +#include "tensorflow/core/platform/logging.h" + +namespace xla { +namespace cpu { +namespace runtime { + +void XfeedManager::Reset() { + infeed()->Reset(); + outfeed()->Reset(); +} + +void XfeedQueueManager::Reset() { + tensorflow::mutex_lock l(mu_); + CHECK(current_buffer_ == nullptr); + for (auto buffer : enqueued_buffers_) { + buffer->Done(ShapeUtil::MakeNil()); + } + enqueued_buffers_.clear(); +} + +void XfeedQueueManager::EnqueueBuffersAtomically( + tensorflow::gtl::ArraySlice buffers) { + tensorflow::mutex_lock l(mu_); + bool was_empty = enqueued_buffers_.empty(); + for (XfeedBuffer* b : buffers) { + enqueued_buffers_.push_back(b); + } + if (was_empty && !buffers.empty()) { + // This has the potential to suffer from the notified thread + // immediately trying and failing to acquire mu_, but seems + // preferable to the alternative of notifying outside the lock + // on every enqueue. + cv_.notify_one(); + } +} + +XfeedBuffer* XfeedQueueManager::BlockingDequeueBuffer() { + tensorflow::mutex_lock l(mu_); + while (enqueued_buffers_.empty()) { + cv_.wait(l); + } + CHECK(current_buffer_ == nullptr); + current_buffer_ = enqueued_buffers_.front(); + enqueued_buffers_.pop_front(); + return current_buffer_; +} + +void XfeedQueueManager::ReleaseCurrentBuffer(int32 length, void* data, + StatusOr shape) { + tensorflow::mutex_lock l(mu_); + CHECK(current_buffer_ != nullptr); + CHECK_EQ(length, current_buffer_->length()); + CHECK_EQ(data, current_buffer_->data()); + current_buffer_->Done(std::move(shape)); + current_buffer_ = nullptr; +} + +} // namespace runtime +} // namespace cpu +} // namespace xla diff --git a/tensorflow/compiler/xla/service/cpu/xfeed_manager.h b/tensorflow/compiler/xla/service/cpu/xfeed_manager.h new file mode 100644 index 0000000000000000000000000000000000000000..6af55700052007a2ca419d52b63dddea2052bd0b --- /dev/null +++ b/tensorflow/compiler/xla/service/cpu/xfeed_manager.h @@ -0,0 +1,123 @@ +/* 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 header declares the abstract class for the infeed manager that +// is used by the CPU runtime to transfer buffers into an executing +// CPU computation, e.g., to feed data into a while loop. + +#ifndef TENSORFLOW_COMPILER_XLA_SERVICE_CPU_XFEED_MANAGER_H_ +#define TENSORFLOW_COMPILER_XLA_SERVICE_CPU_XFEED_MANAGER_H_ + +#include + +#include "tensorflow/compiler/xla/statusor.h" +#include "tensorflow/compiler/xla/types.h" +#include "tensorflow/compiler/xla/xla_data.pb.h" +#include "tensorflow/core/lib/gtl/array_slice.h" +#include "tensorflow/core/platform/mutex.h" + +namespace xla { +namespace cpu { +namespace runtime { + +// Abstract class defining an infeed buffer that is passed to the +// runtime by a client. The client manages the storage of the buffer. +class XfeedBuffer { + public: + virtual ~XfeedBuffer() = default; + + virtual int32 length() = 0; + virtual void* data() = 0; + + // The 'shape' parameter reflects what shape the embedded program was + // expecting / producing with respect to this XfeedBuffer. E.g. this will + // contain information about the layout of an outfed buffer. + virtual void Done(StatusOr shape) = 0; +}; + +// Reusable component for managing the infeed and outfeed queue state. +class XfeedQueueManager { + public: + XfeedQueueManager() = default; + + // Calls the completion callback for any enqueued buffers that have + // not been dequeued by the runtime, and empties the + // queue. Reset may not be called while a runtime computation is + // processing a dequeued buffer. The only safe way to ensure this + // condition is to call Reset when no computation is taking place. + void Reset(); + + // Adds a sequence of buffers to the queue atomically. buffer->Done will be + // called when the buffer will no longer be accessed by the XfeedManager, + // either as a result of a call to Reset or because the runtime has dequeued + // and used the buffer. + void EnqueueBuffersAtomically( + tensorflow::gtl::ArraySlice buffers); + + // Blocks until the queue is non-empty, then returns the buffer at the head of + // the queue. Sets the current buffer to be the returned buffer. It is an + // error to call BlockingDequeueBuffer if there is an unreleased current + // buffer, i.e., ReleaseCurrentBuffer must be called between calls to + // BlockingDequeueBuffer. + XfeedBuffer* BlockingDequeueBuffer(); + + // Releases the current buffer, which is the last buffer returned by + // BlockingDequeuBuffer and not yet released. length and data must + // match the buffer->length() and buffer->data() for the current + // buffer. + // + // 'shape' communicates the shape of the buffer being released. If the program + // passed a value that could not be decoded as a shape, 'shape' will be an + // error status. In the case of outfeed, this indicates the layout of the + // shape that has been outfed. In the case of infeed, this can be used for + // sanity checking purposes. + void ReleaseCurrentBuffer(int32 length, void* data, StatusOr shape); + + private: + tensorflow::mutex mu_; + + // Condition variable that is signaled every time a buffer is + // enqueued to an empty queue. + tensorflow::condition_variable cv_; + + // XfeedBuffer* queue contents are not owned, but buffer->Done must + // be called when the buffer is no longer needed by the runtime. + std::deque enqueued_buffers_; + + // If non-NULL, the buffer that is currently being processed by the + // runtime. Not owned. + XfeedBuffer* current_buffer_ = nullptr; +}; + +// Client-side class used to enqueue infeed buffers. +class XfeedManager { + public: + XfeedManager() = default; + + void Reset(); + + XfeedQueueManager* infeed() { return &infeed_; } + XfeedQueueManager* outfeed() { return &outfeed_; } + + private: + XfeedQueueManager infeed_; + XfeedQueueManager outfeed_; +}; + +} // namespace runtime +} // namespace cpu +} // namespace xla + +#endif // TENSORFLOW_COMPILER_XLA_SERVICE_CPU_XFEED_MANAGER_H_ diff --git a/tensorflow/compiler/xla/service/cpu/xfeed_manager_test.cc b/tensorflow/compiler/xla/service/cpu/xfeed_manager_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..8fe65f488a2f0c4031926fa4c5f02dcf5473568d --- /dev/null +++ b/tensorflow/compiler/xla/service/cpu/xfeed_manager_test.cc @@ -0,0 +1,140 @@ +/* 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/cpu/xfeed_manager.h" + +#include + +#include "tensorflow/compiler/xla/service/cpu/cpu_runtime.h" +#include "tensorflow/compiler/xla/shape_util.h" +#include "tensorflow/core/lib/core/status_test_util.h" +#include "tensorflow/core/lib/core/threadpool.h" +#include "tensorflow/core/platform/env.h" +#include "tensorflow/core/platform/logging.h" +#include "tensorflow/core/platform/test.h" + +namespace xla { +namespace { + +class InfeedManagerTest : public ::testing::Test {}; + +class TestInfeedBuffer : public cpu::runtime::XfeedBuffer { + public: + explicit TestInfeedBuffer(int32 length, bool expect_shape_match = true) + : shape_(ShapeUtil::MakeShape(U8, {length})), + done_called_(false), + length_(length), + expect_shape_match_(expect_shape_match) {} + ~TestInfeedBuffer() override { EXPECT_TRUE(done_called_); } + + int32 length() override { return length_; } + void* data() override { return nullptr; } + void Done(StatusOr shape) override { + CHECK(!done_called_); + done_called_ = true; + TF_ASSERT_OK(shape.status()); + EXPECT_EQ(expect_shape_match_, ShapeUtil::Equal(shape_, shape.ValueOrDie())) + << "want " << ShapeUtil::HumanString(shape_) << " " + << (expect_shape_match_ ? "==" : "!=") << " " + << ShapeUtil::HumanString(shape.ValueOrDie()); + } + + const Shape& shape() const { return shape_; } + + private: + Shape shape_; + bool done_called_; + int32 length_; + bool expect_shape_match_; +}; + +// Performs the acquire/release sequence on the infeed, as the generated CPU +// code would in the process of executing the infeed operation. +void ProcessNextBuffer(int32 length) { + auto shape = ShapeUtil::MakeShape(U8, {length}); + string bytes = shape.SerializeAsString(); + void* buffer = __xla_cpu_runtime_AcquireInfeedBufferForDequeue( + length, bytes.data(), bytes.size()); + __xla_cpu_runtime_ReleaseInfeedBufferAfterDequeue(length, buffer, + bytes.data(), bytes.size()); +} + +// Performs the acquire/release sequence on the outfeed, as the generated CPU +// code would in the process of executing the outfeed operation. +void ProcessNextOutfeedBuffer(int32 length, const Shape& shape) { + string bytes = shape.SerializeAsString(); + void* buffer = __xla_cpu_runtime_AcquireOutfeedBufferForPopulation( + length, bytes.data(), bytes.size()); + __xla_cpu_runtime_ReleaseOutfeedBufferAfterPopulation( + length, buffer, bytes.data(), bytes.size()); +} + +TEST_F(InfeedManagerTest, SingleThreadedSequential) { + TestInfeedBuffer* a = new TestInfeedBuffer(64); + TestInfeedBuffer* b = new TestInfeedBuffer(32); + + cpu::runtime::XfeedManager* xfeed = cpu::runtime::GetXfeedManager(); + + xfeed->infeed()->EnqueueBuffersAtomically({a}); + xfeed->infeed()->EnqueueBuffersAtomically({b}); + ProcessNextBuffer(a->length()); + ProcessNextBuffer(b->length()); +} + +TEST_F(InfeedManagerTest, SingleThreadedInterleaved) { + TestInfeedBuffer* a = new TestInfeedBuffer(64); + TestInfeedBuffer* b = new TestInfeedBuffer(32); + + cpu::runtime::XfeedManager* xfeed = cpu::runtime::GetXfeedManager(); + + xfeed->infeed()->EnqueueBuffersAtomically({a}); + ProcessNextBuffer(a->length()); + xfeed->infeed()->EnqueueBuffersAtomically({b}); + ProcessNextBuffer(b->length()); +} + +TEST_F(InfeedManagerTest, MultiThreaded) { + tensorflow::thread::ThreadPool pool(tensorflow::Env::Default(), "test", 2); + + cpu::runtime::XfeedManager* xfeed = cpu::runtime::GetXfeedManager(); + + const int32 length = 64; + + pool.Schedule([xfeed]() { + // Spin for 100 milliseconds + int64 start_micros = tensorflow::Env::Default()->NowMicros(); + while (true) { + int64 end_micros = tensorflow::Env::Default()->NowMicros(); + if ((end_micros - start_micros) >= 100000) { // 100 ms + break; + } + } + TestInfeedBuffer* a = new TestInfeedBuffer(length); + xfeed->infeed()->EnqueueBuffersAtomically({a}); + }); + + ProcessNextBuffer(length); +} + +TEST_F(InfeedManagerTest, OutfeedWrongShape) { + TestInfeedBuffer* b = new TestInfeedBuffer(32, /*expect_shape_match=*/false); + cpu::runtime::XfeedManager* xfeed = cpu::runtime::GetXfeedManager(); + xfeed->outfeed()->EnqueueBuffersAtomically({b}); + + ProcessNextOutfeedBuffer(32, ShapeUtil::MakeShape(U8, {33})); +} + +} // namespace +} // namespace xla diff --git a/tensorflow/compiler/xla/service/cpu_transfer_manager.cc b/tensorflow/compiler/xla/service/cpu_transfer_manager.cc index 423ec29fdc9a6b6b32c7ce94ea7f5fb3a275ba4c..bf43c04ae2ac5f4b846ffd34e4c90d6765f8ba15 100644 --- a/tensorflow/compiler/xla/service/cpu_transfer_manager.cc +++ b/tensorflow/compiler/xla/service/cpu_transfer_manager.cc @@ -21,15 +21,17 @@ limitations under the License. #include "tensorflow/compiler/xla/literal_util.h" #include "tensorflow/compiler/xla/service/cpu/cpu_runtime.h" -#include "tensorflow/compiler/xla/service/cpu/infeed_manager.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/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" +#include "tensorflow/core/platform/notification.h" #include "tensorflow/core/platform/stream_executor_no_cuda.h" namespace se = ::perftools::gputools; @@ -38,7 +40,7 @@ namespace xla { namespace { -class CpuInfeedBuffer : public cpu::runtime::InfeedBuffer { +class CpuInfeedBuffer : public cpu::runtime::XfeedBuffer { public: explicit CpuInfeedBuffer(int32 length) : length_(length), @@ -48,7 +50,7 @@ class CpuInfeedBuffer : public cpu::runtime::InfeedBuffer { int32 length() override { return length_; } void* data() override { return buffer_; } - void Done() override { delete this; } + void Done(StatusOr /*shape*/) override { delete this; } se::DeviceMemoryBase* device_memory() { return &device_memory_; } @@ -58,6 +60,30 @@ class CpuInfeedBuffer : public cpu::runtime::InfeedBuffer { se::DeviceMemoryBase device_memory_; }; +class CpuOutfeedBuffer : public cpu::runtime::XfeedBuffer { + public: + CpuOutfeedBuffer(void* destination, int32 length) + : destination_(destination), length_(length) {} + + StatusOr WaitForNotification() { + done_.WaitForNotification(); + return status_; + } + + int32 length() override { return length_; } + void* data() override { return destination_; } + void Done(StatusOr shape) override { + status_ = std::move(shape); + done_.Notify(); + } + + private: + void* destination_; + int32 length_; + StatusOr status_; + tensorflow::Notification done_; +}; + } // namespace CpuTransferManager::CpuTransferManager() @@ -66,38 +92,218 @@ CpuTransferManager::CpuTransferManager() Status CpuTransferManager::TransferLiteralToInfeed(se::StreamExecutor* executor, const Literal& literal) { const Shape& shape = literal.shape(); - VLOG(2) << "transferring literal shape to infeed: " + VLOG(2) << "Transferring literal to infeed with shape: " << ShapeUtil::HumanString(shape); - // TODO(b/31381668) handle tuples. - if (ShapeUtil::IsTuple(shape)) { - return Unimplemented("Infeed with a tuple shape is not supported: %s", - ShapeUtil::HumanString(literal.shape()).c_str()); + if (!ShapeUtil::IsTuple(shape)) { + int64 size = GetByteSizeRequirement(shape); + return TransferBufferToInfeed(executor, size, literal.InternalData()); + } + + if (ShapeUtil::IsNestedTuple(shape)) { + return Unimplemented( + "Infeed with a nested tuple shape is not supported: %s", + ShapeUtil::HumanString(literal.shape()).c_str()); } - cpu::runtime::InfeedManager* infeed_manager = - cpu::runtime::GetInfeedManager(); + // For a tuple, we transfer each of its elements to the device and + // enqueue the resulting destination device addresses with the + // infeed manager. + std::vector buffers; + buffers.reserve(literal.tuple_literals_size()); + auto cleanup = tensorflow::gtl::MakeCleanup([&buffers]() { + for (cpu::runtime::XfeedBuffer* b : buffers) { + b->Done(Cancelled("Failed to infeed buffer to device.")); + } + }); - int64 size = GetByteSizeRequirement(shape); + for (const auto& tuple_element : literal.tuple_literals()) { + const Shape& tuple_element_shape = tuple_element.shape(); + int64 tuple_element_size = GetByteSizeRequirement(tuple_element_shape); + TF_ASSIGN_OR_RETURN( + cpu::runtime::XfeedBuffer * buffer, + TransferBufferToInfeedInternal(executor, tuple_element_size, + tuple_element.InternalData())); + buffers.push_back(buffer); + } + + cpu::runtime::XfeedManager* xfeed_manager = cpu::runtime::GetXfeedManager(); + xfeed_manager->infeed()->EnqueueBuffersAtomically(buffers); + + cleanup.release(); + return Status::OK(); +} + +Status CpuTransferManager::TransferBufferToInfeed(se::StreamExecutor* executor, + int64 size, + const void* source) { + TF_ASSIGN_OR_RETURN(cpu::runtime::XfeedBuffer * buffer, + TransferBufferToInfeedInternal(executor, size, source)); + + cpu::runtime::XfeedManager* xfeed_manager = cpu::runtime::GetXfeedManager(); + xfeed_manager->infeed()->EnqueueBuffersAtomically({buffer}); + + return Status::OK(); +} + +StatusOr +CpuTransferManager::TransferBufferToInfeedInternal(se::StreamExecutor* executor, + int64 size, + const void* source) { if (size > std::numeric_limits::max()) { - return Unimplemented("Infeed shape is too large: %s needs %lld bytes", - ShapeUtil::HumanString(literal.shape()).c_str(), size); + return InvalidArgument("Infeed shape is too large: needs %lld bytes", size); + } + + if (size <= 0) { + return InvalidArgument("Infeed shape must have positive size; got %lld", + size); } + int32 size_32 = static_cast(size); CpuInfeedBuffer* queued_buffer = new CpuInfeedBuffer(size_32); - TF_RETURN_IF_ERROR(TransferBufferToDevice( - executor, /*size=*/size, /*source=*/LiteralUtil::InternalData(literal), - queued_buffer->device_memory())); + Status s = + TransferBufferToDevice(executor, /*size=*/size, + /*source=*/source, queued_buffer->device_memory()); + + if (!s.ok()) { + queued_buffer->Done(s); + return s; + } + return queued_buffer; +} - infeed_manager->EnqueueBuffer(queued_buffer); +Status CpuTransferManager::TransferLiteralFromOutfeed( + se::StreamExecutor* executor, const Shape& literal_shape, + Literal* literal) { + if (!ShapeUtil::IsTuple(literal_shape)) { + 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. + tensorflow::gtl::ArraySlice dimensions( + tensorflow::bit_cast(literal_shape.dimensions().data()), + literal_shape.dimensions().size()); + auto empty = + Literal::CreateFromDimensions(literal_shape.element_type(), dimensions); + literal->Swap(empty.get()); + TF_ASSIGN_OR_RETURN(Shape received_shape, + TransferArrayBufferFromOutfeed( + executor, literal->MutableInternalData(), size)); + TF_RET_CHECK(ShapeUtil::Compatible(received_shape, literal->shape())) + << "Shape received from outfeed " + << ShapeUtil::HumanString(received_shape) + << " did not match the shape that was requested for outfeed: " + << ShapeUtil::HumanString(literal_shape); + TF_RET_CHECK(size == GetByteSizeRequirement(received_shape)); + *literal->mutable_shape() = received_shape; + return Status::OK(); + } + + if (ShapeUtil::IsNestedTuple(literal_shape)) { + return Unimplemented( + "Nested tuple outfeeds are not yet implemented on CPU."); + } + std::vector> elements; + std::vector> buffer_data; + for (int64 i = 0; i < literal_shape.tuple_shapes_size(); ++i) { + const Shape& tuple_element_shape = + ShapeUtil::GetTupleElementShape(literal_shape, i); + // Note: OSS build didn't like implicit conversion from + // literal_shape.dimensions() to the array slice on 2017-07-10. + tensorflow::gtl::ArraySlice dimensions( + tensorflow::bit_cast( + tuple_element_shape.dimensions().data()), + tuple_element_shape.dimensions().size()); + auto empty = Literal::CreateFromDimensions( + tuple_element_shape.element_type(), dimensions); + int64 size = GetByteSizeRequirement(tuple_element_shape); + buffer_data.push_back({empty->MutableInternalData(), size}); + elements.push_back(std::move(empty)); + } + + TF_ASSIGN_OR_RETURN(Shape received_shape, + TransferTupleBuffersFromOutfeed(executor, buffer_data)); + + TF_RET_CHECK(ShapeUtil::Compatible(received_shape, literal_shape)) + << "Shape received from outfeed " + << ShapeUtil::HumanString(received_shape) + << " did not match the shape that was requested for outfeed: " + << ShapeUtil::HumanString(literal_shape); + TF_RET_CHECK(GetByteSizeRequirement(literal_shape) == + GetByteSizeRequirement(received_shape)); + + for (int64 i = 0; i < literal_shape.tuple_shapes_size(); ++i) { + *elements[i]->mutable_shape() = received_shape.tuple_shapes(i); + } + auto result = Literal::MakeTupleOwned(std::move(elements)); + literal->Swap(result.get()); + TF_RET_CHECK(ShapeUtil::Equal(literal->shape(), literal_shape)); return Status::OK(); } +StatusOr CpuTransferManager::TransferTupleBuffersFromOutfeed( + perftools::gputools::StreamExecutor* executor, + tensorflow::gtl::ArraySlice> buffer_data) { + return TransferBuffersFromOutfeedInternal(executor, buffer_data, + /*is_tuple=*/true); +} + +StatusOr CpuTransferManager::TransferArrayBufferFromOutfeed( + perftools::gputools::StreamExecutor* executor, void* destination, + int64 size_bytes) { + return TransferBuffersFromOutfeedInternal( + executor, {{destination, size_bytes}}, /*is_tuple=*/false); +} + +StatusOr CpuTransferManager::TransferBuffersFromOutfeedInternal( + perftools::gputools::StreamExecutor* executor, + tensorflow::gtl::ArraySlice> buffer_data, + bool is_tuple) { + std::vector> buffers; + for (auto b : buffer_data) { + int64 size = b.second; + if (size > std::numeric_limits::max()) { + return InvalidArgument("Outfeed shape is too large: needs %lld bytes", + size); + } + + if (size <= 0) { + return InvalidArgument("Outfeed shape must have positive size; got %lld", + size); + } + + int32 size_32 = static_cast(size); + VLOG(2) + << "Enqueueing outfeed buffer (for the device to populate) of length " + << size_32 << "B"; + buffers.emplace_back(MakeUnique(b.first, size_32)); + } + + std::vector buffer_pointers; + buffer_pointers.reserve(buffers.size()); + for (auto& b : buffers) { + buffer_pointers.push_back(b.get()); + } + + cpu::runtime::XfeedManager* xfeed_manager = cpu::runtime::GetXfeedManager(); + xfeed_manager->outfeed()->EnqueueBuffersAtomically(buffer_pointers); + VLOG(2) << "Waiting for buffer to be notified as populated."; + std::vector outfed_shapes; + for (auto& buffer : buffers) { + TF_ASSIGN_OR_RETURN(Shape outfed_shape, buffer->WaitForNotification()); + outfed_shapes.push_back(std::move(outfed_shape)); + } + if (is_tuple) { + return ShapeUtil::MakeTupleShape(outfed_shapes); + } + TF_RET_CHECK(outfed_shapes.size() == 1); + return std::move(outfed_shapes[0]); +} + } // namespace xla -static xla::TransferManager* CreateCpuTransferManager() { - return new xla::CpuTransferManager(); +static std::unique_ptr CreateCpuTransferManager() { + return xla::MakeUnique(); } static bool InitModule() { diff --git a/tensorflow/compiler/xla/service/cpu_transfer_manager.h b/tensorflow/compiler/xla/service/cpu_transfer_manager.h index 727462252d7291959fd09c05c87e36411eb3ddab..6c7524d94716464218ba18ad9950f702d2759f89 100644 --- a/tensorflow/compiler/xla/service/cpu_transfer_manager.h +++ b/tensorflow/compiler/xla/service/cpu_transfer_manager.h @@ -18,10 +18,12 @@ limitations under the License. #include +#include "tensorflow/compiler/xla/service/cpu/xfeed_manager.h" #include "tensorflow/compiler/xla/service/generic_transfer_manager.h" #include "tensorflow/compiler/xla/service/transfer_manager.h" #include "tensorflow/compiler/xla/statusor.h" #include "tensorflow/compiler/xla/xla_data.pb.h" +#include "tensorflow/core/lib/gtl/array_slice.h" #include "tensorflow/core/platform/macros.h" #include "tensorflow/core/platform/stream_executor_no_cuda.h" #include "tensorflow/core/platform/types.h" @@ -37,8 +39,37 @@ class CpuTransferManager : public GenericTransferManager { Status TransferLiteralToInfeed(perftools::gputools::StreamExecutor* executor, const Literal& literal) override; + Status TransferBufferToInfeed(perftools::gputools::StreamExecutor* executor, + int64 size, const void* source) override; + Status TransferLiteralFromOutfeed( + perftools::gputools::StreamExecutor* executor, const Shape& literal_shape, + Literal* literal) override; private: + // Transfers infeed data to device. InfeedBuffer->Done() must be + // called to clean up the memory allocated for InfeedBuffer. + StatusOr TransferBufferToInfeedInternal( + perftools::gputools::StreamExecutor* executor, int64 size, + const void* source); + + // Helper that transfers a tuple of element buffers from the device's outfeed. + StatusOr TransferTupleBuffersFromOutfeed( + perftools::gputools::StreamExecutor* executor, + tensorflow::gtl::ArraySlice> buffer_data); + + // Helper that transfers an array buffer from the device's outfeed. + StatusOr TransferArrayBufferFromOutfeed( + perftools::gputools::StreamExecutor* executor, void* destination, + int64 size_bytes); + + // On success, returns the shape that was transferred from the outfeed -- if + // is_tuple is true, the returned shape will be a tuple of the returned shapes + // for the given buffers. + StatusOr TransferBuffersFromOutfeedInternal( + perftools::gputools::StreamExecutor* executor, + tensorflow::gtl::ArraySlice> buffer_data, + bool is_tuple); + TF_DISALLOW_COPY_AND_ASSIGN(CpuTransferManager); }; diff --git a/tensorflow/compiler/xla/service/device_memory_allocator.cc b/tensorflow/compiler/xla/service/device_memory_allocator.cc index c13c86741cc4291d5ae76cb4b3d7913927c565ea..2e4b0a5230516b5308aeed892de9a49565a09f2e 100644 --- a/tensorflow/compiler/xla/service/device_memory_allocator.cc +++ b/tensorflow/compiler/xla/service/device_memory_allocator.cc @@ -35,7 +35,15 @@ StreamExecutorMemoryAllocator::Allocate(int device_ordinal, uint64 size, bool retry_on_failure) { TF_ASSIGN_OR_RETURN(perftools::gputools::StreamExecutor * stream_executor, GetStreamExecutor(device_ordinal)); - return stream_executor->AllocateArray(size); + perftools::gputools::DeviceMemoryBase result = + stream_executor->AllocateArray(size); + if (size > 0 && result == nullptr) { + return ResourceExhausted( + "Failed to allocate request for %s (%lluB) on device ordinal %d", + tensorflow::strings::HumanReadableNumBytes(size).c_str(), size, + device_ordinal); + } + return result; } tensorflow::Status StreamExecutorMemoryAllocator::Deallocate( diff --git a/tensorflow/compiler/xla/service/dfs_hlo_visitor.cc b/tensorflow/compiler/xla/service/dfs_hlo_visitor.cc index 5b296861006923f438df1ad4fb5898f82f11b9e0..6efd0bcee58d19b355b6c2afa6d9497f75ef4b3c 100644 --- a/tensorflow/compiler/xla/service/dfs_hlo_visitor.cc +++ b/tensorflow/compiler/xla/service/dfs_hlo_visitor.cc @@ -17,56 +17,38 @@ limitations under the License. #include +#include "tensorflow/compiler/xla/service/hlo_instruction.h" #include "tensorflow/compiler/xla/types.h" #include "tensorflow/compiler/xla/util.h" #include "tensorflow/core/platform/logging.h" namespace xla { -Status DfsHloVisitor::HandleElementwiseUnary(HloInstruction* hlo, - HloOpcode opcode, - HloInstruction* operand) { +Status DfsHloVisitor::HandleElementwiseUnary(HloInstruction* hlo) { return Unimplemented("DfsHloVisitor::HandleElementwiseUnary: %s", - HloOpcodeString(opcode).c_str()); + HloOpcodeString(hlo->opcode()).c_str()); } -Status DfsHloVisitor::HandleElementwiseBinary(HloInstruction* hlo, - HloOpcode opcode, - HloInstruction* lhs, - HloInstruction* rhs) { +Status DfsHloVisitor::HandleElementwiseBinary(HloInstruction* hlo) { return Unimplemented("DfsHloVisitor::HandleElementwiseBinary: %s", - HloOpcodeString(opcode).c_str()); + HloOpcodeString(hlo->opcode()).c_str()); +} + +DfsHloVisitor::VisitState DfsHloVisitor::GetVisitState( + const HloInstruction& instruction) { + return GetVisitState(instruction.unique_id()); } void DfsHloVisitor::SetVisiting(const HloInstruction& instruction) { VLOG(3) << "marking HLO " << &instruction << " as visiting: "; - CHECK(NotVisited(instruction)); - visit_state_[&instruction] = VisitState::kVisiting; + DCHECK(NotVisited(instruction)); + visit_state_.SetState(instruction.unique_id(), VisitState::kVisiting); } void DfsHloVisitor::SetVisited(const HloInstruction& instruction) { VLOG(3) << "marking HLO " << &instruction << " as visited: "; - CHECK(NotVisited(instruction) || IsVisiting(instruction)); - visit_state_[&instruction] = VisitState::kVisited; -} - -bool DfsHloVisitor::IsVisiting(const HloInstruction& instruction) { - if (visit_state_.count(&instruction) == 0) { - return false; - } - return visit_state_[&instruction] == VisitState::kVisiting; -} - -bool DfsHloVisitor::DidVisit(const HloInstruction& instruction) { - if (visit_state_.count(&instruction) == 0) { - return false; - } - return visit_state_[&instruction] == VisitState::kVisited; -} - -bool DfsHloVisitor::NotVisited(const HloInstruction& instruction) { - return visit_state_.count(&instruction) == 0 || - visit_state_[&instruction] == VisitState::kNotVisited; + DCHECK(NotVisited(instruction) || IsVisiting(instruction)); + visit_state_.SetState(instruction.unique_id(), VisitState::kVisited); } Status DfsHloVisitor::Preprocess(HloInstruction* hlo) { return Status::OK(); } diff --git a/tensorflow/compiler/xla/service/dfs_hlo_visitor.h b/tensorflow/compiler/xla/service/dfs_hlo_visitor.h index 351efa82dd21dd9f618ed38cdb54bd2e26fcd5d5..2f21043a1d341aecd14c0476fb61a8ff511656ea 100644 --- a/tensorflow/compiler/xla/service/dfs_hlo_visitor.h +++ b/tensorflow/compiler/xla/service/dfs_hlo_visitor.h @@ -18,6 +18,7 @@ limitations under the License. #include +#include "tensorflow/compiler/xla/literal_util.h" #include "tensorflow/compiler/xla/service/hlo_opcode.h" #include "tensorflow/compiler/xla/status.h" #include "tensorflow/compiler/xla/types.h" @@ -56,51 +57,43 @@ class HloInstruction; // instruction that is accessible from the instruction object itself. class DfsHloVisitor { public: - DfsHloVisitor() - : visit_state_(32) // Start the hash table a bit larger to avoid resizes - {} + DfsHloVisitor() {} virtual ~DfsHloVisitor() {} // These routines are self-descriptive, see class comment for usage // information. - virtual Status HandleElementwiseUnary(HloInstruction* hlo, HloOpcode opcode, - HloInstruction* operand); - virtual Status HandleElementwiseBinary(HloInstruction* hlo, HloOpcode opcode, - HloInstruction* lhs, - HloInstruction* rhs); + virtual Status HandleElementwiseUnary(HloInstruction* hlo); + virtual Status HandleElementwiseBinary(HloInstruction* hlo); virtual Status HandleClamp(HloInstruction* clamp, HloInstruction* min, HloInstruction* arg, HloInstruction* max) = 0; virtual Status HandleSelect(HloInstruction* select, HloInstruction* pred, HloInstruction* on_true, HloInstruction* on_false) = 0; - virtual Status HandleMaximum(HloInstruction* maximum, HloInstruction* lhs, - HloInstruction* rhs) { - return HandleElementwiseBinary(maximum, HloOpcode::kMaximum, lhs, rhs); + virtual Status HandleMaximum(HloInstruction* maximum) { + return HandleElementwiseBinary(maximum); } - virtual Status HandleMinimum(HloInstruction* minimum, HloInstruction* lhs, - HloInstruction* rhs) { - return HandleElementwiseBinary(minimum, HloOpcode::kMinimum, lhs, rhs); + virtual Status HandleMinimum(HloInstruction* minimum) { + return HandleElementwiseBinary(minimum); } virtual Status HandleConcatenate( HloInstruction* concatenate, tensorflow::gtl::ArraySlice operands) = 0; - virtual Status HandleConvert(HloInstruction* convert, - HloInstruction* operand) { - return HandleElementwiseUnary(convert, HloOpcode::kConvert, operand); + virtual Status HandleConvert(HloInstruction* convert) { + return HandleElementwiseUnary(convert); } - virtual Status HandleCopy(HloInstruction* copy, HloInstruction* operand) { - return HandleElementwiseUnary(copy, HloOpcode::kCopy, operand); + virtual Status HandleCopy(HloInstruction* copy) { + return HandleElementwiseUnary(copy); } virtual Status HandleMultiply(HloInstruction* multiply, HloInstruction* lhs, HloInstruction* rhs) { - return HandleElementwiseBinary(multiply, HloOpcode::kMultiply, lhs, rhs); + return HandleElementwiseBinary(multiply); } virtual Status HandleDot(HloInstruction* dot, HloInstruction* lhs, HloInstruction* rhs) = 0; virtual Status HandlePower(HloInstruction* power, HloInstruction* lhs, HloInstruction* rhs) { - return HandleElementwiseBinary(power, HloOpcode::kPower, lhs, rhs); + return HandleElementwiseBinary(power); } virtual Status HandleConvolution(HloInstruction* convolution, HloInstruction* lhs, HloInstruction* rhs, @@ -108,64 +101,72 @@ class DfsHloVisitor { virtual Status HandleCrossReplicaSum(HloInstruction* crs) = 0; virtual Status HandleCompare(HloInstruction* compare, HloOpcode opcode, HloInstruction* lhs, HloInstruction* rhs) { - return HandleElementwiseBinary(compare, opcode, lhs, rhs); + return HandleElementwiseBinary(compare); } virtual Status HandleAdd(HloInstruction* add, HloInstruction* lhs, HloInstruction* rhs) { - return HandleElementwiseBinary(add, HloOpcode::kAdd, lhs, rhs); + return HandleElementwiseBinary(add); } virtual Status HandleDivide(HloInstruction* divide, HloInstruction* lhs, HloInstruction* rhs) { - return HandleElementwiseBinary(divide, HloOpcode::kDivide, lhs, rhs); + return HandleElementwiseBinary(divide); } virtual Status HandleRemainder(HloInstruction* remainder, HloInstruction* lhs, HloInstruction* rhs) { - return HandleElementwiseBinary(remainder, HloOpcode::kRemainder, lhs, rhs); + return HandleElementwiseBinary(remainder); } virtual Status HandleSubtract(HloInstruction* subtract, HloInstruction* lhs, HloInstruction* rhs) { - return HandleElementwiseBinary(subtract, HloOpcode::kSubtract, lhs, rhs); + return HandleElementwiseBinary(subtract); } virtual Status HandleAbs(HloInstruction* abs, HloInstruction* operand) { - return HandleElementwiseUnary(abs, HloOpcode::kAbs, operand); + return HandleElementwiseUnary(abs); } virtual Status HandleSign(HloInstruction* sign, HloInstruction* operand) { - return HandleElementwiseUnary(sign, HloOpcode::kSign, operand); + return HandleElementwiseUnary(sign); } virtual Status HandleNegate(HloInstruction* negate, HloInstruction* operand) { - return HandleElementwiseUnary(negate, HloOpcode::kNegate, operand); + return HandleElementwiseUnary(negate); } virtual Status HandleExp(HloInstruction* exp, HloInstruction* operand) { - return HandleElementwiseUnary(exp, HloOpcode::kExp, operand); + return HandleElementwiseUnary(exp); } virtual Status HandleFloor(HloInstruction* floor, HloInstruction* operand) { - return HandleElementwiseUnary(floor, HloOpcode::kFloor, operand); + return HandleElementwiseUnary(floor); } virtual Status HandleCeil(HloInstruction* ceil, HloInstruction* operand) { - return HandleElementwiseUnary(ceil, HloOpcode::kCeil, operand); + return HandleElementwiseUnary(ceil); } virtual Status HandleLog(HloInstruction* log, HloInstruction* operand) { - return HandleElementwiseUnary(log, HloOpcode::kLog, operand); + return HandleElementwiseUnary(log); + } + virtual Status HandleCos(HloInstruction* cos, HloInstruction* operand) { + return HandleElementwiseUnary(cos); + } + virtual Status HandleSin(HloInstruction* sin, HloInstruction* operand) { + return HandleElementwiseUnary(sin); } virtual Status HandleTanh(HloInstruction* tanh, HloInstruction* operand) { - return HandleElementwiseUnary(tanh, HloOpcode::kTanh, operand); + return HandleElementwiseUnary(tanh); } virtual Status HandleIsFinite(HloInstruction* is_finite, HloInstruction* operand) { - return HandleElementwiseUnary(is_finite, HloOpcode::kIsFinite, operand); + return HandleElementwiseUnary(is_finite); } virtual Status HandleLogicalAnd(HloInstruction* logical_and, HloInstruction* lhs, HloInstruction* rhs) { - return HandleElementwiseBinary(logical_and, HloOpcode::kLogicalAnd, lhs, - rhs); + return HandleElementwiseBinary(logical_and); } virtual Status HandleLogicalNot(HloInstruction* logical_not, HloInstruction* operand) { - return HandleElementwiseUnary(logical_not, HloOpcode::kLogicalNot, operand); + return HandleElementwiseUnary(logical_not); } virtual Status HandleLogicalOr(HloInstruction* logical_or, HloInstruction* lhs, HloInstruction* rhs) { - return HandleElementwiseBinary(logical_or, HloOpcode::kLogicalOr, lhs, rhs); + return HandleElementwiseBinary(logical_or); + } + virtual Status HandleReducePrecision(HloInstruction* reduce_precision) { + return HandleElementwiseUnary(reduce_precision); } virtual Status HandleInfeed(HloInstruction* infeed) = 0; @@ -189,19 +190,16 @@ class DfsHloVisitor { virtual Status HandleTranspose(HloInstruction* transpose) = 0; virtual Status HandleParameter(HloInstruction* parameter) = 0; virtual Status HandleFusion(HloInstruction* fusion) = 0; - virtual Status HandleCall( - HloInstruction* call, - tensorflow::gtl::ArraySlice operands, - HloComputation* computation) = 0; + virtual Status HandleCall(HloInstruction* call) = 0; virtual Status HandleCustomCall( HloInstruction* custom_call, tensorflow::gtl::ArraySlice operands, tensorflow::StringPiece custom_call_target) = 0; virtual Status HandleSlice(HloInstruction* slice, HloInstruction* operand) = 0; - virtual Status HandleDynamicSlice( - HloInstruction* slice, - tensorflow::gtl::ArraySlice operands) = 0; + virtual Status HandleDynamicSlice(HloInstruction* dynamic_slice, + HloInstruction* operand, + HloInstruction* start_indices) = 0; virtual Status HandleDynamicUpdateSlice(HloInstruction* dynamic_update_slice, HloInstruction* operand, HloInstruction* update, @@ -219,9 +217,7 @@ class DfsHloVisitor { const Window& window, HloComputation* function) = 0; virtual Status HandleSelectAndScatter(HloInstruction* instruction) = 0; - virtual Status HandleWhile(HloInstruction* xla_while, HloInstruction* init, - HloComputation* condition, - HloComputation* body) = 0; + virtual Status HandleWhile(HloInstruction* xla_while) = 0; virtual Status HandlePad(HloInstruction* pad) = 0; @@ -229,6 +225,13 @@ class DfsHloVisitor { virtual Status HandleRecv(HloInstruction* recv) = 0; + virtual Status HandleBatchNormTraining(HloInstruction* batchNormTraining) = 0; + + virtual Status HandleBatchNormInference( + HloInstruction* batchNormInference) = 0; + + virtual Status HandleBatchNormGrad(HloInstruction* batchNormGrad) = 0; + // Invoked to inform the visitor that the traversal has completed, and that // the root was "root". virtual Status FinishVisit(HloInstruction* root) = 0; @@ -236,11 +239,23 @@ class DfsHloVisitor { // 3 possible visitation states of HLO instructions. Each instruction's // state only flows one way: kNotVisited -> kVisiting -> kVisited. enum VisitState { - kNotVisited, - kVisiting, - kVisited, + kNotVisited = 0, + kVisiting = 1, + kVisited = 2, }; + VisitState GetVisitState(int id) { return visit_state_.GetState(id); } + VisitState GetVisitState(const HloInstruction& instruction); + + // Resize internal state if necessary to hold state for ids <= num. + // This call is purely a performance hint and can be omitted without + // affecting correctness. + void ReserveVisitStates(int num) { visit_state_.Reserve(num); } + + void SetVisitState(int id, VisitState state) { + visit_state_.SetState(id, state); + } + // Sets the visitation state of the given instruction as kVisiting. // // Precondition: current state must be kNotVisited. @@ -252,13 +267,19 @@ class DfsHloVisitor { void SetVisited(const HloInstruction& instruction); // Returns whether the state of the given instruction is kVisiting. - bool IsVisiting(const HloInstruction& instruction); + bool IsVisiting(const HloInstruction& instruction) { + return GetVisitState(instruction) == kVisiting; + } // Returns whether the state of the given instruction is kVisited. - bool DidVisit(const HloInstruction& instruction); + bool DidVisit(const HloInstruction& instruction) { + return GetVisitState(instruction) == kVisited; + } // Returns whether the state of the given instruction is kNotVisited. - bool NotVisited(const HloInstruction& instruction); + bool NotVisited(const HloInstruction& instruction) { + return GetVisitState(instruction) == kNotVisited; + } // This method should be overridden by subclasses that wish to run some // operation on an op before its Handle* visitor method is called. @@ -282,9 +303,43 @@ class DfsHloVisitor { virtual Status Postprocess(HloInstruction* visited); private: - // Tracks the visitation state of each instruction. Any instructions that are - // not found from the map are considered as VisitState::kNotVisited. - tensorflow::gtl::FlatMap visit_state_; + class DFSVisitStates { + public: + DFSVisitStates() {} + void Reserve(uint64 num) { + states_.reserve((num + kStatesPerWord - 1) / kStatesPerWord); + } + VisitState GetState(uint64 id) { + uint64 word_index = id / kStatesPerWord; + if (word_index >= states_.size()) { + return VisitState::kNotVisited; + } + static_assert(static_cast(VisitState::kVisited) < 3, + "VisitState must fit in two bits"); + uint64 w = states_[word_index]; + uint32 shift = 2 * (id % kStatesPerWord); // 2 bits per state + return static_cast((w >> shift) & 0x3); + } + void SetState(uint64 id, VisitState state) { + uint64 word_index = id / kStatesPerWord; + if (word_index >= states_.size()) { + states_.resize(word_index + 1, 0); + } + uint64* w = &states_[word_index]; + uint32 shift = 2 * (id % kStatesPerWord); // 2 bits per state + uint64 mask = 0x3ull << shift; + *w = (*w & ~mask) | (static_cast(state) << shift); + DCHECK_EQ(GetState(id), state); + } + + private: + static const uint32 kStatesPerWord = sizeof(uint64) / 2 /*bits per entry*/; + // Map from id to two-bit states. We store 32 such states per 64-bit + // value + std::vector states_; + }; + + DFSVisitStates visit_state_; TF_DISALLOW_COPY_AND_ASSIGN(DfsHloVisitor); }; 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 18cfaf83e1cd558928c9fc65452524567f3cbb49..a5fe120598416235dff2af9d8a5c0ae64ac9edcc 100644 --- a/tensorflow/compiler/xla/service/dfs_hlo_visitor_with_default.h +++ b/tensorflow/compiler/xla/service/dfs_hlo_visitor_with_default.h @@ -16,6 +16,7 @@ limitations under the License. #ifndef TENSORFLOW_COMPILER_XLA_SERVICE_DFS_HLO_VISITOR_WITH_DEFAULT_H_ #define TENSORFLOW_COMPILER_XLA_SERVICE_DFS_HLO_VISITOR_WITH_DEFAULT_H_ +#include "tensorflow/compiler/xla/literal_util.h" #include "tensorflow/compiler/xla/service/dfs_hlo_visitor.h" #include "tensorflow/compiler/xla/service/hlo_opcode.h" #include "tensorflow/compiler/xla/types.h" @@ -40,15 +41,25 @@ class DfsHloVisitorWithDefault : public DfsHloVisitor { // Default action performed on HloInstruction. virtual Status DefaultAction(HloInstruction* hlo_instruction) = 0; - Status HandleElementwiseUnary(HloInstruction* hlo, HloOpcode opcode, - HloInstruction* operand) override { + Status HandleElementwiseUnary(HloInstruction* hlo) override { return DefaultAction(hlo); } - Status HandleElementwiseBinary(HloInstruction* hlo, HloOpcode opcode, - HloInstruction* lhs, - HloInstruction* rhs) override { + Status HandleElementwiseBinary(HloInstruction* hlo) override { return DefaultAction(hlo); } + + Status HandleBatchNormTraining(HloInstruction* hlo) override { + return DefaultAction(hlo); + } + + Status HandleBatchNormInference(HloInstruction* hlo) override { + return DefaultAction(hlo); + } + + Status HandleBatchNormGrad(HloInstruction* hlo) override { + return DefaultAction(hlo); + } + Status HandleClamp(HloInstruction* clamp, HloInstruction* /*min*/, HloInstruction* /*arg*/, HloInstruction* /*max*/) override { @@ -59,12 +70,10 @@ class DfsHloVisitorWithDefault : public DfsHloVisitor { tensorflow::gtl::ArraySlice /*operands*/) override { return DefaultAction(concatenate); } - Status HandleConvert(HloInstruction* convert, - HloInstruction* /*operand*/) override { + Status HandleConvert(HloInstruction* convert) override { return DefaultAction(convert); } - Status HandleCopy(HloInstruction* copy, - HloInstruction* /*operand*/) override { + Status HandleCopy(HloInstruction* copy) override { return DefaultAction(copy); } Status HandleSelect(HloInstruction* select, HloInstruction* /*pred*/, @@ -121,9 +130,7 @@ class DfsHloVisitorWithDefault : public DfsHloVisitor { Status HandleFusion(HloInstruction* fusion) override { return DefaultAction(fusion); } - Status HandleCall(HloInstruction* call, - tensorflow::gtl::ArraySlice /*operands*/, - HloComputation* /*computation*/) override { + Status HandleCall(HloInstruction* call) override { return DefaultAction(call); } Status HandleCustomCall( @@ -136,10 +143,10 @@ class DfsHloVisitorWithDefault : public DfsHloVisitor { HloInstruction* /*operand*/) override { return DefaultAction(slice); } - Status HandleDynamicSlice( - HloInstruction* slice, - tensorflow::gtl::ArraySlice /*operands*/) override { - return DefaultAction(slice); + Status HandleDynamicSlice(HloInstruction* dynamic_slice, + HloInstruction* /*operand*/, + HloInstruction* /*start_indices*/) override { + return DefaultAction(dynamic_slice); } Status HandleDynamicUpdateSlice(HloInstruction* dynamic_update_slice, HloInstruction* /*operand*/, @@ -188,9 +195,7 @@ class DfsHloVisitorWithDefault : public DfsHloVisitor { Status HandleTranspose(HloInstruction* transpose) override { return DefaultAction(transpose); } - Status HandleWhile(HloInstruction* xla_while, HloInstruction* /*init*/, - HloComputation* /*condition*/, - HloComputation* /*body*/) override { + Status HandleWhile(HloInstruction* xla_while) override { return DefaultAction(xla_while); } Status HandleSend(HloInstruction* send) override { diff --git a/tensorflow/compiler/xla/service/elemental_ir_emitter.cc b/tensorflow/compiler/xla/service/elemental_ir_emitter.cc index a04815dad94484a6f01ebd27d3ec73f547086722..350dbc321fb2234912d2143adfe70b75b48d0e27 100644 --- a/tensorflow/compiler/xla/service/elemental_ir_emitter.cc +++ b/tensorflow/compiler/xla/service/elemental_ir_emitter.cc @@ -21,14 +21,15 @@ limitations under the License. #include // IWYU pragma: no_include "llvm/IR/Intrinsics.gen.inc" -#include "external/llvm/include/llvm/IR/BasicBlock.h" -#include "external/llvm/include/llvm/IR/Instructions.h" -#include "external/llvm/include/llvm/IR/Intrinsics.h" -#include "external/llvm/include/llvm/Transforms/Utils/BasicBlockUtils.h" +#include "llvm/IR/BasicBlock.h" +#include "llvm/IR/Instructions.h" +#include "llvm/IR/Intrinsics.h" +#include "llvm/Transforms/Utils/BasicBlockUtils.h" #include "tensorflow/compiler/xla/primitive_util.h" #include "tensorflow/compiler/xla/service/hlo_module.h" #include "tensorflow/compiler/xla/service/hlo_opcode.h" #include "tensorflow/compiler/xla/service/llvm_ir/ir_array.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/shape_util.h" #include "tensorflow/compiler/xla/status_macros.h" @@ -63,7 +64,7 @@ StatusOr ElementalIrEmitter::EmitIntegerUnaryOp( case HloOpcode::kConvert: { PrimitiveType from_type = op->operand(0)->shape().element_type(); PrimitiveType to_type = op->shape().element_type(); - CHECK(primitive_util::IsIntegralType(from_type)); + CHECK(primitive_util::IsIntegralType(from_type) || from_type == PRED); if (from_type == to_type) { return operand_value; } @@ -78,7 +79,8 @@ StatusOr ElementalIrEmitter::EmitIntegerUnaryOp( operand_value, llvm_ir::PrimitiveTypeToIrType(to_type, ir_builder_)); } - if (primitive_util::IsUnsignedIntegralType(from_type)) { + if (primitive_util::IsUnsignedIntegralType(from_type) || + from_type == PRED) { return ir_builder_->CreateUIToFP( operand_value, llvm_ir::PrimitiveTypeToIrType(to_type, ir_builder_)); @@ -172,6 +174,14 @@ StatusOr ElementalIrEmitter::EmitFloatUnaryOp( return llvm_ir::EmitCallToIntrinsic(llvm::Intrinsic::log, {operand_value}, {operand_value->getType()}, ir_builder_); + case HloOpcode::kCos: + return llvm_ir::EmitCallToIntrinsic(llvm::Intrinsic::cos, {operand_value}, + {operand_value->getType()}, + ir_builder_); + case HloOpcode::kSin: + return llvm_ir::EmitCallToIntrinsic(llvm::Intrinsic::sin, {operand_value}, + {operand_value->getType()}, + ir_builder_); case HloOpcode::kFloor: return llvm_ir::EmitCallToIntrinsic( llvm::Intrinsic::floor, {operand_value}, {operand_value->getType()}, @@ -240,14 +250,18 @@ StatusOr ElementalIrEmitter::EmitFloatBinaryOp( return ir_builder_->CreateFDiv(lhs_value, rhs_value); case HloOpcode::kRemainder: return ir_builder_->CreateFRem(lhs_value, rhs_value); - - // The 'O' prefix on the LLVM ops means "ordered" compare where comparisons - // with NAN always return false. + // LLVM comparisons can be "unordered" (U) or "ordered" (O) -- ordered + // comparisons always return false when one of the operands is NaN, whereas + // unordered comparisons return true. + // + // We use ordered comparisons for everything except kNe, where we use an + // unordered comparison. This makes x != y equivalent to !(x == y), and + // matches C++'s semantics. case HloOpcode::kEq: return llvm_ir::EmitComparison(llvm::CmpInst::FCMP_OEQ, lhs_value, rhs_value, ir_builder_); case HloOpcode::kNe: - return llvm_ir::EmitComparison(llvm::CmpInst::FCMP_ONE, lhs_value, + return llvm_ir::EmitComparison(llvm::CmpInst::FCMP_UNE, lhs_value, rhs_value, ir_builder_); case HloOpcode::kLt: return llvm_ir::EmitComparison(llvm::CmpInst::FCMP_OLT, lhs_value, @@ -279,6 +293,7 @@ StatusOr ElementalIrEmitter::EmitFloatBinaryOp( llvm::Value* ElementalIrEmitter::EmitFloatMax(llvm::Value* lhs_value, llvm::Value* rhs_value) const { + // TODO(b/64580527): We can do better here if fast-math is enabled. return llvm_ir::EmitCallToIntrinsic(llvm::Intrinsic::maxnum, {lhs_value, rhs_value}, {lhs_value->getType()}, ir_builder_); @@ -286,6 +301,7 @@ llvm::Value* ElementalIrEmitter::EmitFloatMax(llvm::Value* lhs_value, llvm::Value* ElementalIrEmitter::EmitFloatMin(llvm::Value* lhs_value, llvm::Value* rhs_value) const { + // TODO(b/64580527): We can do better here if fast-math is enabled. return llvm_ir::EmitCallToIntrinsic(llvm::Intrinsic::minnum, {lhs_value, rhs_value}, {lhs_value->getType()}, ir_builder_); @@ -377,6 +393,118 @@ StatusOr ElementalIrEmitter::EmitErfcInv( return EmitErfInv(prim_type, ir_builder_->CreateFSub(one, value)); } +StatusOr ElementalIrEmitter::EmitReducePrecision( + const HloInstruction* hlo, llvm::Value* x) const { + if (hlo->operand(0)->shape().element_type() != F32) { + return Unimplemented("reduce-precision only implemented for F32"); + } + + // Integer and float types for casting and constant generation. + llvm::Type* float_type = x->getType(); + llvm::IntegerType* int_type = ir_builder_->getInt32Ty(); + + // Cast the input value to an integer for bitwise manipulation. + llvm::Value* x_as_int = ir_builder_->CreateBitCast(x, int_type); + + if (hlo->mantissa_bits() < 23) { + // Last remaining mantissa bit. + const uint32_t last_mantissa_bit_mask = 1u << (23 - hlo->mantissa_bits()); + + // Compute rounding bias for round-to-nearest with ties to even. This is + // equal to a base value of 0111... plus one bit if the last remaining + // mantissa bit is 1. + const uint32_t base_rounding_bias = (last_mantissa_bit_mask >> 1) - 1; + llvm::Value* x_last_mantissa_bit = ir_builder_->CreateLShr( + ir_builder_->CreateAnd( + x_as_int, llvm::ConstantInt::get(int_type, last_mantissa_bit_mask)), + (23 - hlo->mantissa_bits())); + llvm::Value* x_rounding_bias = ir_builder_->CreateAdd( + x_last_mantissa_bit, + llvm::ConstantInt::get(int_type, base_rounding_bias)); + + // Add rounding bias, and mask out truncated bits. Note that the case + // where adding the rounding bias overflows into the exponent bits is + // correct; the non-masked mantissa bits will all be zero, and the + // exponent will be incremented by one. + const uint32_t truncation_mask = ~(last_mantissa_bit_mask - 1); + x_as_int = ir_builder_->CreateAdd(x_as_int, x_rounding_bias); + x_as_int = ir_builder_->CreateAnd( + x_as_int, llvm::ConstantInt::get(int_type, truncation_mask)); + } + + if (hlo->exponent_bits() < 8) { + // Masks for f32 values. + const uint32_t f32_sign_bit_mask = 1u << 31; + const uint32_t f32_exp_bits_mask = 0xffu << 23; + + // An exponent of 2^(n-1)-1 -- that is, 0111... with the zero in the most- + // significant bit -- is equal to 1.0f for all exponent sizes. Adding + // 2^(n-1)-1 to this gives us the highest non-infinite exponent for a bit- + // size of n, and subtracting 2^(n-1)-1 from this gives us the lowest' + // exponent (corresponding to 0.0f). + // + // Thus, the f32 exponent corresponding to the highest non-infinite + // exponent for a bit size of n is (2^7-1) + 2^(n-1)-1, and the f32 + // exponent corresponding to the lowest exponent for a bit size of n is + // (2^7-1) - 2^(n-1)-1. + // + // Note that we have already checked that exponents_bits >= 1. + const uint32_t f32_exponent_bias = (1 << 7) - 1; + const uint32_t reduced_exponent_bias = + (1 << (hlo->exponent_bits() - 1)) - 1; + const uint32_t reduced_max_exponent = + f32_exponent_bias + reduced_exponent_bias; + const uint32_t reduced_min_exponent = + f32_exponent_bias - reduced_exponent_bias; + + // Do we overflow or underflow? + llvm::Value* x_exponent = ir_builder_->CreateAnd( + x_as_int, llvm::ConstantInt::get(int_type, f32_exp_bits_mask)); + llvm::Value* x_overflows = ir_builder_->CreateICmpUGT( + x_exponent, + llvm::ConstantInt::get(int_type, reduced_max_exponent << 23)); + llvm::Value* x_underflows = ir_builder_->CreateICmpULE( + x_exponent, + llvm::ConstantInt::get(int_type, reduced_min_exponent << 23)); + + // Compute appropriately-signed values of zero and infinity. + llvm::Value* x_signed_zero = ir_builder_->CreateAnd( + x_as_int, llvm::ConstantInt::get(int_type, f32_sign_bit_mask)); + llvm::Value* x_signed_inf = ir_builder_->CreateOr( + x_signed_zero, llvm::ConstantInt::get(int_type, f32_exp_bits_mask)); + + // Force to zero or infinity if overflow or underflow. (Note that this + // truncates all denormal values to zero, rather than rounding them.) + x_as_int = ir_builder_->CreateSelect(x_overflows, x_signed_inf, x_as_int); + x_as_int = ir_builder_->CreateSelect(x_underflows, x_signed_zero, x_as_int); + } + + // Cast the result back to a floating-point type. + llvm::Value* result = ir_builder_->CreateBitCast(x_as_int, float_type); + + // Correct result for NaN inputs. + // + // The exponent handling will "normalize" NaN values to infinities, which is + // undesirable (except in the case with no mantissa bits, in which case it + // is mandatory). This logic also handles cases where mantissa-rounding + // causes a NaN's mantissa to overflow into the exponent bits, which would + // otherwise create an erroneous zero value. + // + // If the fast-math flags are set to assume no NaNs, the comparison is likely + // to be optimized away, so there's no point in even emitting it. + if (!ir_builder_->getFastMathFlags().noNaNs()) { + llvm::Value* x_is_nan = ir_builder_->CreateFCmpUNO(x, x); + + if (hlo->mantissa_bits() > 0) { + result = ir_builder_->CreateSelect(x_is_nan, x, result); + } else { + result = ir_builder_->CreateSelect( + x_is_nan, llvm::ConstantFP::getInfinity(float_type), result); + } + } + return result; +} + StatusOr ElementalIrEmitter::EmitIntegerBinaryOp( const HloInstruction* op, llvm::Value* lhs_value, llvm::Value* rhs_value, bool is_signed) const { @@ -488,7 +616,7 @@ llvm_ir::ElementGenerator ElementalIrEmitter::MakeRngElementGenerator( auto random_value = [hlo]() { const HloModule* module = - hlo->IsFused() ? hlo->fusion_instruction()->parent()->parent() + hlo->IsFused() ? hlo->parent()->FusionInstruction()->parent()->parent() : hlo->parent()->parent(); return module->RandomNew64(); }; @@ -581,23 +709,40 @@ llvm_ir::ElementGenerator ElementalIrEmitter::MakeRngElementGenerator( } else { auto r = ir_builder_->CreateSub(q, p); auto leading_zeros = llvm_ir::EmitCallToIntrinsic( - llvm::Intrinsic::ctlz, {r, ir_builder_->getInt1(1)}, + llvm::Intrinsic::ctlz, {r, ir_builder_->getInt1(true)}, {param_ir_type}, ir_builder_); auto in_block = ir_builder_->GetInsertBlock(); - auto body_block = in_block->splitBasicBlock( - ir_builder_->GetInsertPoint(), "rng_body"); - SetToFirstInsertPoint(body_block, ir_builder_); - auto out_block = body_block->splitBasicBlock( - ir_builder_->GetInsertPoint(), "rng_out"); + + // A terminator should be present iff we're emitting code + // into the middle (as opposed to the end) of a basic block. + CHECK_EQ(ir_builder_->GetInsertPoint() == in_block->end(), + in_block->getTerminator() == nullptr); + + llvm::BasicBlock* body_block; + llvm::BasicBlock* out_block; + + if (ir_builder_->GetInsertPoint() == in_block->end()) { + body_block = + llvm_ir::CreateBasicBlock(nullptr, "rng_body", ir_builder_); + out_block = + llvm_ir::CreateBasicBlock(nullptr, "rng_out", ir_builder_); + llvm::BranchInst::Create(body_block, in_block); + } else { + body_block = in_block->splitBasicBlock( + ir_builder_->GetInsertPoint(), "rng_body"); + out_block = body_block->splitBasicBlock( + ir_builder_->GetInsertPoint(), "rng_out"); + body_block->getTerminator()->eraseFromParent(); + } + SetToFirstInsertPoint(body_block, ir_builder_); auto random = ir_builder_->CreateAnd( ir_builder_->CreateZExtOrTrunc(get_next_i64(), param_ir_type), ir_builder_->CreateLShr(llvm::ConstantInt::get(param_ir_type, ~0), leading_zeros)); - llvm::ReplaceInstWithInst( - body_block->getTerminator(), - llvm::BranchInst::Create(out_block, body_block, - ir_builder_->CreateICmpULT(random, r))); + llvm::BranchInst::Create(out_block, body_block, + ir_builder_->CreateICmpULT(random, r), + body_block); SetToFirstInsertPoint(out_block, ir_builder_); return ir_builder_->CreateAdd( p, ir_builder_->CreateSelect( @@ -643,12 +788,14 @@ llvm_ir::ElementGenerator ElementalIrEmitter::MakeElementGenerator( case HloOpcode::kCeil: case HloOpcode::kConvert: case HloOpcode::kCopy: + case HloOpcode::kCos: case HloOpcode::kExp: case HloOpcode::kFloor: case HloOpcode::kIsFinite: case HloOpcode::kLog: case HloOpcode::kNegate: case HloOpcode::kSign: + case HloOpcode::kSin: case HloOpcode::kTanh: case HloOpcode::kLogicalNot: return [this, hlo, &operand_to_generator]( @@ -716,6 +863,14 @@ llvm_ir::ElementGenerator ElementalIrEmitter::MakeElementGenerator( ElementwiseSourceIndex(index, *hlo, 2))); return EmitFloatMin(max_value, EmitFloatMax(min_value, arg_value)); }; + case HloOpcode::kReducePrecision: + return [this, hlo, &operand_to_generator]( + const IrArray::Index& index) -> StatusOr { + TF_ASSIGN_OR_RETURN(llvm::Value * operand_value, + operand_to_generator.at(hlo->operand(0))( + ElementwiseSourceIndex(index, *hlo, 0))); + return EmitReducePrecision(hlo, operand_value); + }; case HloOpcode::kConcatenate: return [this, hlo, &operand_to_generator]( const IrArray::Index target_index) -> StatusOr { @@ -739,11 +894,11 @@ llvm_ir::ElementGenerator ElementalIrEmitter::MakeElementGenerator( const HloInstruction* operand = hlo->operand(operand_idx); auto true_block = llvm_ir::CreateBasicBlock( exit_block, tensorflow::strings::StrCat( - "concat_index_from_operand", operand_idx), + "concat_index_from_operand", operand_idx), ir_builder_); auto false_block = llvm_ir::CreateBasicBlock( exit_block, tensorflow::strings::StrCat( - "concat_index_not_from_operand", operand_idx), + "concat_index_not_from_operand", operand_idx), ir_builder_); auto concat_dim_size = llvm::ConstantInt::get(source_index[concat_dim]->getType(), @@ -801,12 +956,9 @@ llvm_ir::ElementGenerator ElementalIrEmitter::MakeElementGenerator( case HloOpcode::kSlice: return [this, hlo, &operand_to_generator]( const IrArray::Index& index) -> StatusOr { - IrArray::Index sliced_index(index.size()); - for (int i = 0; i < index.size(); ++i) { - sliced_index[i] = ir_builder_->CreateAdd( - index[i], llvm::ConstantInt::get(index[i]->getType(), - hlo->slice_starts(i))); - } + IrArray::Index sliced_index = index.SourceIndexOfSlice( + /*shape=*/hlo->shape(), /*starts=*/hlo->slice_starts(), + /*strides=*/hlo->slice_strides(), /*builder=*/ir_builder_); return operand_to_generator.at(hlo->operand(0))(sliced_index); }; case HloOpcode::kDynamicSlice: @@ -999,6 +1151,74 @@ llvm_ir::ElementGenerator ElementalIrEmitter::MakeElementGenerator( // if_data.after_block. return ir_builder_->CreateLoad(ret_value_addr); }; + + case HloOpcode::kDot: + return [=, &operand_to_generator](const IrArray::Index& dot_result_index) + -> StatusOr { + auto lhs_generator = operand_to_generator.at(hlo->operand(0)); + auto rhs_generator = operand_to_generator.at(hlo->operand(1)); + int64 contracted_dim_size = hlo->operand(0)->shape().dimensions( + hlo->operand(0)->shape().dimensions_size() - 1); + int64 lhs_dims = hlo->operand(0)->shape().dimensions_size(); + int64 rhs_dims = hlo->operand(1)->shape().dimensions_size(); + + std::unique_ptr inner_loop = + llvm_ir::ForLoop::EmitForLoop( + "dot.inner", ir_builder_->getInt64(0), + ir_builder_->getInt64(contracted_dim_size), + ir_builder_->getInt64(1), ir_builder_); + + SetToFirstInsertPoint(inner_loop->GetPreheaderBasicBlock(), + ir_builder_); + PrimitiveType primitive_type = hlo->shape().element_type(); + llvm::Type* primitive_type_llvm = + llvm_ir::PrimitiveTypeToIrType(primitive_type, ir_builder_); + llvm::Value* accumulator_alloca = llvm_ir::EmitAllocaAtFunctionEntry( + primitive_type_llvm, "dot_acc", ir_builder_); + ir_builder_->CreateStore( + llvm::Constant::getNullValue(primitive_type_llvm), + accumulator_alloca); + + SetToFirstInsertPoint(inner_loop->GetBodyBasicBlock(), ir_builder_); + + // This is the inner reduction loop for a dot operation that produces + // one element in the output. If the operands to the dot operation have + // shapes [A,B,C,T] and [D,T,E], the result has a shape [A,B,C,D,E]. + // Given an output index [a,b,c,d,e] in the result, we compute: + // sum(lhs[a,b,c,t]*rhs[d,t,e] for t in [0, T)) + + IrArray::Index lhs_index, rhs_index; + + for (int64 i = 0; i < lhs_dims - 1; i++) { + lhs_index.push_back(dot_result_index[i]); + } + lhs_index.push_back(inner_loop->GetIndVarValue()); + + for (int64 i = 0; i < rhs_dims - 2; i++) { + rhs_index.push_back(dot_result_index[lhs_dims - 1 + i]); + } + rhs_index.push_back(inner_loop->GetIndVarValue()); + rhs_index.push_back(dot_result_index.back()); + + llvm::Value* current_accumulator = + ir_builder_->CreateLoad(accumulator_alloca); + TF_ASSIGN_OR_RETURN(llvm::Value * lhs_value, lhs_generator(lhs_index)); + TF_ASSIGN_OR_RETURN(llvm::Value * rhs_value, rhs_generator(rhs_index)); + llvm::Value* next_accumulator; + if (primitive_util::IsFloatingPointType(primitive_type)) { + next_accumulator = ir_builder_->CreateFAdd( + current_accumulator, + ir_builder_->CreateFMul(lhs_value, rhs_value)); + } else { + next_accumulator = ir_builder_->CreateAdd( + current_accumulator, + ir_builder_->CreateMul(lhs_value, rhs_value)); + } + ir_builder_->CreateStore(next_accumulator, accumulator_alloca); + + SetToFirstInsertPoint(inner_loop->GetExitBasicBlock(), ir_builder_); + return ir_builder_->CreateLoad(accumulator_alloca); + }; default: return [this, hlo, &operand_to_generator](const IrArray::Index& index) { return Unimplemented("%s", HloOpcodeString(hlo->opcode()).c_str()); diff --git a/tensorflow/compiler/xla/service/elemental_ir_emitter.h b/tensorflow/compiler/xla/service/elemental_ir_emitter.h index 2576d3823e06ed3050554b38766dbd6c6a48ca5c..35dfa88e9b02e3ec7686dc7fdded8cf4e88201fb 100644 --- a/tensorflow/compiler/xla/service/elemental_ir_emitter.h +++ b/tensorflow/compiler/xla/service/elemental_ir_emitter.h @@ -18,9 +18,9 @@ limitations under the License. #include -#include "external/llvm/include/llvm/IR/IRBuilder.h" -#include "external/llvm/include/llvm/IR/Module.h" -#include "external/llvm/include/llvm/IR/Value.h" +#include "llvm/IR/IRBuilder.h" +#include "llvm/IR/Module.h" +#include "llvm/IR/Value.h" #include "tensorflow/compiler/xla/service/hlo_instruction.h" #include "tensorflow/compiler/xla/service/hlo_module_config.h" #include "tensorflow/compiler/xla/service/llvm_ir/loop_emitter.h" @@ -84,6 +84,9 @@ class ElementalIrEmitter { virtual StatusOr EmitErfcInv(PrimitiveType prim_type, llvm::Value* value) const; + virtual StatusOr EmitReducePrecision(const HloInstruction* hlo, + llvm::Value* x) const; + // A helper method for MakeElementGenerator. Given an elementwise op `hlo` and // the target array index, computes the source array index of its // `operand_no`-th operand. diff --git a/tensorflow/compiler/xla/service/executable.cc b/tensorflow/compiler/xla/service/executable.cc index ef973676ea4adfd233e20579929506fe9f46d412..79fedb61c971862fc0e3a59e01e55825f09c587d 100644 --- a/tensorflow/compiler/xla/service/executable.cc +++ b/tensorflow/compiler/xla/service/executable.cc @@ -15,7 +15,8 @@ limitations under the License. #include "tensorflow/compiler/xla/service/executable.h" -#include "tensorflow/compiler/xla/legacy_flags/service_flags.h" +#include "tensorflow/compiler/xla/legacy_flags/debug_options_flags.h" +#include "tensorflow/compiler/xla/service/hlo_graph_dumper.h" #include "tensorflow/compiler/xla/status_macros.h" #include "tensorflow/core/lib/io/path.h" #include "tensorflow/core/lib/strings/stringprintf.h" @@ -56,8 +57,8 @@ Executable::ExecuteOnStreams( Status Executable::DumpSessionModule() { TF_RET_CHECK(dumping()); - legacy_flags::ServiceFlags* flags = legacy_flags::GetServiceFlags(); - const string& directory_path = flags->xla_dump_executions_to; + const string& directory_path = + module_config().debug_options().xla_dump_executions_to(); VersionedComputationHandle versioned_handle = entry_computation_handle(); // This filename does not include the version number because the computation // is only ever executed at one version. @@ -71,7 +72,7 @@ Status Executable::DumpSessionModule() { // Removes illegal characters from filenames. static void SanitizeFilename(string* name) { for (char& c : *name) { - if (c == '/' || c == '\\') { + if (c == '/' || c == '\\' || c == '[' || c == ']') { c = '_'; } } @@ -82,7 +83,11 @@ static void SanitizeFilename(string* name) { const SessionModule& session_module) { tensorflow::Env* env = tensorflow::Env::Default(); if (!env->IsDirectory(directory_path).ok()) { - TF_RETURN_IF_ERROR(env->CreateDir(directory_path)); + // NB! CreateDir does not work reliably with multiple XLA threads -- two + // threads can race to observe the absence of the dump directory and + // simultaneously try to create it, causing the "losing" thread to get a + // "directory already exists" error. + TF_RETURN_IF_ERROR(env->RecursivelyCreateDir(directory_path)); } SanitizeFilename(&filename); string file_path = tensorflow::io::JoinPath(directory_path, filename); diff --git a/tensorflow/compiler/xla/service/executable.h b/tensorflow/compiler/xla/service/executable.h index eb36aba33a7694c43985b5e5636e7e0fa2ad4794..b58dee9c20a6431968358fe90babf2fa813e7e11 100644 --- a/tensorflow/compiler/xla/service/executable.h +++ b/tensorflow/compiler/xla/service/executable.h @@ -19,16 +19,19 @@ limitations under the License. #include #include +#include "tensorflow/compiler/xla/legacy_flags/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_cost_analysis.h" #include "tensorflow/compiler/xla/service/hlo_execution_profile.h" +#include "tensorflow/compiler/xla/service/hlo_graph_dumper.h" #include "tensorflow/compiler/xla/service/hlo_module.h" -#include "tensorflow/compiler/xla/service/hlo_module_config.h" #include "tensorflow/compiler/xla/service/service_executable_run_options.h" #include "tensorflow/compiler/xla/service/session.pb.h" #include "tensorflow/compiler/xla/service/shaped_buffer.h" #include "tensorflow/compiler/xla/service/versioned_computation_handle.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/gtl/array_slice.h" #include "tensorflow/core/platform/mutex.h" @@ -41,10 +44,8 @@ namespace xla { // interface that is used for launching compiled programs across platforms. class Executable { public: - explicit Executable(std::unique_ptr hlo_module, - std::unique_ptr module_config) - : hlo_module_(std::move(hlo_module)), - module_config_(std::move(module_config)) {} + explicit Executable(std::unique_ptr hlo_module) + : hlo_module_(std::move(hlo_module)) {} virtual ~Executable() {} // Enqueues the compilation result on the provided stream, passing the given @@ -87,6 +88,16 @@ class Executable { tensorflow::gtl::ArraySlice> arguments); + // Convenience wrapper for calling Executable::ExecuteOnStream. Sets up a + // timer for the execution, sets up HLO profiling if enabled, and fills in the + // given ExecutionProfile if non-null. The ExecuteOnStream overloads have + // different argument types and return types, so this method is templated on + // argument type and return type of the execute function. + template + StatusOr ExecuteOnStreamWrapper( + const ServiceExecutableRunOptions* run_options, ExecutionProfile* profile, + const ArgT& arguments); + // Returns the ExecutionProfile from executing on the device. This includes // the number of cycles taken for the computation or the compilation time. ExecutionProfile execution_profile() const { @@ -94,19 +105,26 @@ class Executable { return execution_profile_; } + // Returns Status::ok() if the two executables are equal to each other. + // + // An error status is returned otherwise. + virtual const Status EqualOrFail(const Executable& executable) { + return Unimplemented( + "Equality test on this executable is not implemented."); + } + // Returns whether this executable was compiled with HLO profilings support // enabled. If not, the caller should not expect an hlo_execution_profile // passed to ExecuteOnStream above to be populated during execution. bool hlo_profiling_enabled() const { - return module_config_->hlo_profiling_enabled(); + return hlo_module_->config().hlo_profiling_enabled(); } const HloModule& module() const { return *hlo_module_; } - const HloModuleConfig& module_config() const { return *module_config_; } + const bool has_module() const { return hlo_module_ != nullptr; } - // Returns whether this executable has an associated HloModuleConfig. - bool has_module_config() const { return module_config_ != nullptr; } + const HloModuleConfig& module_config() const { return hlo_module_->config(); } // Returns the versioned computation handle of the computation computed by // this executable. @@ -117,7 +135,7 @@ class Executable { // The shape (including layout) that results from this execution. This is the // shape of the DeviceMemoryBase result value in ExecuteOnStream above. const Shape& result_shape() const { - return module_config_->entry_computation_layout().result_shape(); + return hlo_module_->config().entry_computation_layout().result_shape(); } // Dumping helpers. @@ -132,6 +150,10 @@ class Executable { static Status DumpToDirectory(const string& directory_path, string filename, const SessionModule& session_module); + // Returns a cost analysis object appropriate for the platform on which this + // executable can run. + virtual std::unique_ptr CreateCostAnalysis() const = 0; + protected: mutable tensorflow::mutex mutex_; @@ -143,10 +165,6 @@ class Executable { // around. std::unique_ptr hlo_module_; - // The configuration used to build this executable (parameter layouts, result - // layout, profiling enabled, etc). - std::unique_ptr module_config_; - // SessionModule this was compiled from. Null if not dumping executions. std::unique_ptr session_module_; @@ -155,6 +173,78 @@ class Executable { int64 execution_count_ = 0; }; +template +StatusOr Executable::ExecuteOnStreamWrapper( + const ServiceExecutableRunOptions* run_options, ExecutionProfile* profile, + const ArgT& arguments) { + perftools::gputools::Stream* stream = run_options->stream(); + std::unique_ptr timer; + if (profile != nullptr) { + timer.reset(new perftools::gputools::Timer(stream->parent())); + stream->InitTimer(timer.get()).ThenStartTimer(timer.get()); + } + + VLOG(1) << "enqueueing executable on stream..."; + // If the profiling flag isn't enabled, we pass nullptr as the profile to + // indicate profiling is not requested. + HloExecutionProfile hlo_execution_profile; + HloExecutionProfile* profile_ptr = + module_config().debug_options().xla_hlo_profile() && + hlo_profiling_enabled() + ? &hlo_execution_profile + : nullptr; + + auto return_value = ExecuteOnStream(run_options, arguments, profile_ptr); + + if (profile != nullptr) { + VLOG(1) << "enqueueing 'stop timer' and blocking host until done..."; + stream->ThenStopTimer(timer.get()).BlockHostUntilDone(); + VLOG(1) << "done with block-host-until-done"; + + // Merge in run-time profile information from execution_profile. + profile->MergeFrom(execution_profile()); + + // Overall execution time (in nanoseconds) from the executor timer. + profile->set_compute_and_transfer_time_ns(timer->Nanoseconds()); + + // TODO(b/28123297): On GPU we end up including transfer time in + // the compute time this way. Instead, we should get the correct + // value by measuring it. Setting the field here at least lets + // benchmarks provide *some* value for GPU computations. + // + // TODO(b/28447609): The value in compute_and_transfer_time_ns is actually + // the compute time without the transfer time, so this way we get the + // correct compute time. We should instead have the correct value for + // compute_and_transfer_time and set compute_time to the compute time. + if (profile->compute_time_ns() == 0) { + profile->set_compute_time_ns(profile->compute_and_transfer_time_ns()); + } + } + + if (profile_ptr != nullptr) { + std::unordered_set profiled_computations = + profile_ptr->profiled_computations(); + // To ensure we have print the profiles in a stable order, iterate over the + // computations in post order. + std::list all_computations = + module().MakeComputationPostOrder(); + for (xla::HloComputation* computation : all_computations) { + if (profiled_computations.count(computation) > 0) { + string profile_string = profile_ptr->ToString( + *computation, stream->parent()->GetDeviceDescription(), + CreateCostAnalysis().get()); + if (!profile_string.empty()) { + XLA_LOG_LINES(tensorflow::INFO, profile_string); + } + } + } + hlo_graph_dumper::MaybeDumpHloModule(module(), "Service::Execute", + profile_ptr); + } + + return return_value; +} + } // namespace xla #endif // TENSORFLOW_COMPILER_XLA_SERVICE_EXECUTABLE_H_ diff --git a/tensorflow/compiler/xla/service/execution_tracker.cc b/tensorflow/compiler/xla/service/execution_tracker.cc index 8d79d07f9424f58dfc43a0c0595520a441fd3f82..c225e62e3e11d2d01251b0f92272b0949eff8af1 100644 --- a/tensorflow/compiler/xla/service/execution_tracker.cc +++ b/tensorflow/compiler/xla/service/execution_tracker.cc @@ -31,7 +31,7 @@ AsyncExecution::AsyncExecution(Backend* backend, : backend_(CHECK_NOTNULL(backend)), streams_(std::move(streams)), profile_(profile), - result_(result) { + result_(std::move(result)) { for (const auto& stream : streams_) { CHECK(stream != nullptr); } diff --git a/tensorflow/compiler/xla/service/flatten_call_graph.cc b/tensorflow/compiler/xla/service/flatten_call_graph.cc new file mode 100644 index 0000000000000000000000000000000000000000..297a4f7599f9c127386b2f53f7ffb987befc456e --- /dev/null +++ b/tensorflow/compiler/xla/service/flatten_call_graph.cc @@ -0,0 +1,112 @@ +/* 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/flatten_call_graph.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_module.h" +#include "tensorflow/compiler/xla/util.h" +#include "tensorflow/core/lib/core/errors.h" + +namespace xla { + +namespace { + +// Helper to replace the called computation at a while- or call-instruction. +void ReplaceCalledComputation(HloInstruction* instruction, + HloComputation* computation, + HloComputation* new_computation) { + switch (instruction->opcode()) { + case HloOpcode::kWhile: { + if (computation == instruction->while_condition()) { + instruction->set_while_condition(new_computation); + } else { + CHECK_EQ(computation, instruction->while_body()); + instruction->set_while_body(new_computation); + } + break; + } + case HloOpcode::kCall: { + CHECK_EQ(instruction->to_apply(), computation); + instruction->set_to_apply(new_computation); + break; + } + default: + LOG(FATAL) << "unexpected opcode: " + << HloOpcodeString(instruction->opcode()); + } +} + +// Flatten a single call graph node. Expects to visit nodes in postorder. +Status FlattenNode(const CallGraphNode& node) { + HloComputation* computation = node.computation(); + HloModule* module = computation->parent(); + // Clone callee for all call-sites except the first one. + for (int i = 0; i < node.caller_callsites().size(); ++i) { + CallSite call_site = node.caller_callsites()[i]; + // Only consider sequential call contexts. + if (call_site.context() == CallContext::kParallel) { + continue; + } + CHECK_EQ(call_site.context(), CallContext::kSequential); + + // Skip first element if this computation is only called from a sequential + // context. + if (node.context() != CallContext::kBoth && i == 0) { + continue; + } + + // Clone computation for the remaining sequential context call sites. + HloComputation* clone = + module->AddEmbeddedComputation(computation->Clone()); + ReplaceCalledComputation(call_site.instruction(), computation, clone); + // Clone the sub-tree of all computations called from this node. + std::vector worklist; + worklist.push_back(clone); + while (!worklist.empty()) { + auto current = worklist.back(); + worklist.pop_back(); + for (auto& instruction : current->instructions()) { + if (GetInstructionCallContext(instruction.get()) != + CallContext::kSequential) { + continue; + } + for (auto callee : instruction->called_computations()) { + HloComputation* callee_clone = + module->AddEmbeddedComputation(callee->Clone()); + ReplaceCalledComputation(instruction.get(), callee, callee_clone); + worklist.push_back(callee_clone); + } + } + } + } + return Status::OK(); +} + +} // namespace + +StatusOr FlattenCallGraph::Run(HloModule* module) { + XLA_VLOG_LINES(3, "Before flatten call graph:\n" + module->ToString()); + + std::unique_ptr call_graph = CallGraph::Build(module); + TF_RETURN_IF_ERROR(call_graph->VisitNodes(FlattenNode)); + + XLA_VLOG_LINES(3, "After flatten call graph:\n" + module->ToString()); + return true; +} + +} // namespace xla diff --git a/tensorflow/compiler/xla/service/flatten_call_graph.h b/tensorflow/compiler/xla/service/flatten_call_graph.h new file mode 100644 index 0000000000000000000000000000000000000000..d3efab3614912e4b0c2c8aa3b80277c326382ed0 --- /dev/null +++ b/tensorflow/compiler/xla/service/flatten_call_graph.h @@ -0,0 +1,40 @@ +/* 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. +==============================================================================*/ + +// Flatten the call graph for an HLO module into a tree. + +#ifndef TENSORFLOW_COMPILER_XLA_SERVICE_FLATTEN_CALL_GRAPH_H_ +#define TENSORFLOW_COMPILER_XLA_SERVICE_FLATTEN_CALL_GRAPH_H_ + +#include "tensorflow/compiler/xla/service/hlo_pass_interface.h" +#include "tensorflow/compiler/xla/statusor.h" + +namespace xla { + +// Flattening associates each call site with a unique computation (for +// sequential calling contexts) This simplifies buffer assignment and +// points-to analysis (see b/36865746 for details). +class FlattenCallGraph : public HloPassInterface { + public: + tensorflow::StringPiece name() const override { return "flatten-call-graph"; } + + // Duplicates computations called from multiple call- or while-nodes to + // flatten the call graph. + StatusOr Run(HloModule* module) override; +}; + +} // namespace xla + +#endif // TENSORFLOW_COMPILER_XLA_SERVICE_FLATTEN_CALL_GRAPH_H_ diff --git a/tensorflow/compiler/xla/service/flatten_call_graph_test.cc b/tensorflow/compiler/xla/service/flatten_call_graph_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..12a6794ac177deb54dd66822a5f830ff213c7b40 --- /dev/null +++ b/tensorflow/compiler/xla/service/flatten_call_graph_test.cc @@ -0,0 +1,231 @@ +/* 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/flatten_call_graph.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/shape_util.h" +#include "tensorflow/compiler/xla/status_macros.h" +#include "tensorflow/compiler/xla/test.h" +#include "tensorflow/compiler/xla/test_helpers.h" +#include "tensorflow/compiler/xla/tests/hlo_test_base.h" +#include "tensorflow/compiler/xla/util.h" +#include "tensorflow/compiler/xla/xla_data.pb.h" +#include "tensorflow/core/lib/core/status_test_util.h" + +namespace xla { +namespace { + +class FlattenCallGraphTest : public HloTestBase { + protected: + // Build and return a trivial computation taking and returning a scalar. + std::unique_ptr MakeScalarComputation() { + HloComputation::Builder builder(TestName() + ".ScalarComputation"); + HloInstruction* param0 = builder.AddInstruction( + HloInstruction::CreateParameter(0, kScalarShape, "param0")); + builder.AddInstruction( + HloInstruction::CreateUnary(kScalarShape, HloOpcode::kNegate, param0)); + return builder.Build(); + } + + // Build and return a computation which takes a scalar and maps (kMap) the + // given computation to the value 'callsites' number of times. + std::unique_ptr MakeMappingComputation( + HloComputation* map_computation, int64 callsites) { + HloComputation::Builder builder(TestName() + ".MappingComputation"); + HloInstruction* param0 = builder.AddInstruction( + HloInstruction::CreateParameter(0, kScalarShape, "param0")); + HloInstruction* last_value = param0; + for (int64 i = 0; i < callsites; ++i) { + last_value = builder.AddInstruction(HloInstruction::CreateMap( + kScalarShape, {last_value}, map_computation)); + } + return builder.Build(); + } + + // Build and return a computation which takes a scalar and calls (kCall) the + // given computation with value 'callsites' number of times. + std::unique_ptr MakeCallingComputation( + HloComputation* callee_computation, int64 callsites, + const string& suffix = ".CallingComputation") { + HloComputation::Builder builder(TestName() + suffix); + HloInstruction* param0 = builder.AddInstruction( + HloInstruction::CreateParameter(0, kScalarShape, "param0")); + HloInstruction* last_value = param0; + for (int64 i = 0; i < callsites; ++i) { + last_value = builder.AddInstruction(HloInstruction::CreateCall( + kScalarShape, {last_value}, callee_computation)); + } + return builder.Build(); + } + + // Build and return a computation which takes a scalar and returns a PRED + // value. + std::unique_ptr MakeConditionComputation() { + HloComputation::Builder builder(TestName() + ".ConditionComputation"); + HloInstruction* param0 = builder.AddInstruction( + HloInstruction::CreateParameter(0, kScalarShape, "param0")); + HloInstruction* zero = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(0.0f))); + builder.AddInstruction(HloInstruction::CreateBinary( + ShapeUtil::MakeShape(PRED, {}), HloOpcode::kGt, param0, zero)); + return builder.Build(); + } + + StatusOr RunFlattenCallGraph(HloModule* module) { + FlattenCallGraph flatten; + TF_ASSIGN_OR_RETURN(bool result, flatten.Run(module)); + return result; + } + + const Shape kScalarShape = ShapeUtil::MakeShape(F32, {}); +}; + +TEST_F(FlattenCallGraphTest, ComplexGraph) { + // Test a call graph of a module with several computation called in various + // contexts. The call graph looks like: + // + // entry + // / | + // a | + // / | \ | + // b | cond + // \ | + // c + // + // Calls are made via kCall, kWhile, and kMap instructions. + auto module = CreateNewModule(); + HloComputation* cond_computation = + module->AddEmbeddedComputation(MakeConditionComputation()); + HloComputation* c_computation = + module->AddEmbeddedComputation(MakeScalarComputation()); + HloComputation* b_computation = module->AddEmbeddedComputation( + MakeMappingComputation(c_computation, /*callsites=*/1)); + + HloComputation* a_computation; + { + HloComputation::Builder builder(TestName() + ".a"); + HloInstruction* param0 = builder.AddInstruction( + HloInstruction::CreateParameter(0, kScalarShape, "param0")); + HloInstruction* call = builder.AddInstruction( + HloInstruction::CreateCall(kScalarShape, {param0}, c_computation)); + builder.AddInstruction(HloInstruction::CreateWhile( + kScalarShape, cond_computation, b_computation, call)); + a_computation = module->AddEmbeddedComputation(builder.Build()); + } + + HloComputation* entry_computation; + { + HloComputation::Builder builder(TestName() + ".entry"); + HloInstruction* param0 = builder.AddInstruction( + HloInstruction::CreateParameter(0, kScalarShape, "param0")); + builder.AddInstruction(HloInstruction::CreateWhile( + kScalarShape, cond_computation, a_computation, param0)); + entry_computation = module->AddEntryComputation(builder.Build()); + } + + { + TF_ASSERT_OK_AND_ASSIGN(bool result, RunFlattenCallGraph(module.get())); + EXPECT_TRUE(result); + 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()); + } +} + +// Test corner case of a computation used as a body and a loop condition. +TEST_F(FlattenCallGraphTest, SharedWhileConditionAndBody) { + auto module = CreateNewModule(); + HloComputation* cond_computation; + { + HloComputation::Builder builder(TestName() + ".cond"); + HloInstruction* param0 = + builder.AddInstruction(HloInstruction::CreateParameter( + 0, ShapeUtil::MakeShape(PRED, {}), "param0")); + HloInstruction* false_constant = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(false))); + builder.AddInstruction( + HloInstruction::CreateBinary(ShapeUtil::MakeShape(PRED, {}), + HloOpcode::kEq, param0, false_constant)); + cond_computation = module->AddEmbeddedComputation(builder.Build()); + } + + HloComputation* entry_computation; + { + HloComputation::Builder builder(TestName() + ".entry"); + HloInstruction* false_constant = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(false))); + builder.AddInstruction(HloInstruction::CreateWhile( + ShapeUtil::MakeShape(PRED, {}), cond_computation, cond_computation, + false_constant)); + entry_computation = module->AddEntryComputation(builder.Build()); + } + + { + 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.get())); + EXPECT_TRUE(result); + 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()); + } +} + +// Test flattening of a nested calling computations. +// +// Entry +// / \ +// \ / +// B +// / \ +// \ / +// C +// +TEST_F(FlattenCallGraphTest, FlattenCalls) { + auto module = CreateNewModule(); + HloComputation* c_computation = + module->AddEmbeddedComputation(MakeScalarComputation()); + + HloComputation* b_computation = module->AddEmbeddedComputation( + MakeCallingComputation(c_computation, /*callsites=*/2, ".B")); + + module->AddEntryComputation( + MakeCallingComputation(b_computation, /*callsites=*/2, ".Entry")); + + TF_ASSERT_OK_AND_ASSIGN(bool result, RunFlattenCallGraph(module.get())); + EXPECT_TRUE(result); + std::unique_ptr call_graph = CallGraph::Build(module.get()); + EXPECT_EQ(7, module->computations().size()); + + const CallGraphNode& c_node = call_graph->GetNode(c_computation); + EXPECT_EQ(1, c_node.caller_callsites().size()); + + const CallGraphNode& b_node = call_graph->GetNode(b_computation); + EXPECT_EQ(1, b_node.caller_callsites().size()); +} + +} // namespace +} // namespace xla + +int main(int argc, char** argv) { + return xla::ParseDebugOptionsFlagsAndRunTests(argc, argv); +} diff --git a/tensorflow/compiler/xla/service/generic_transfer_manager.cc b/tensorflow/compiler/xla/service/generic_transfer_manager.cc index 7b87ac6da1d1aa4efef94feabc93f36c58751605..69195c45ed33bbb689a0633471686a03bb6d2654 100644 --- a/tensorflow/compiler/xla/service/generic_transfer_manager.cc +++ b/tensorflow/compiler/xla/service/generic_transfer_manager.cc @@ -82,13 +82,12 @@ Status GenericTransferManager::TransferLiteralFromDevice( } *literal->mutable_shape() = device_shape; - LiteralUtil::Reserve(ShapeUtil::ElementsIn(device_shape), literal); + literal->Reserve(ShapeUtil::ElementsIn(device_shape)); TF_RETURN_IF_ERROR(TransferBufferFromDevice( executor, source, /*size=*/ShapeUtil::ByteSizeOf(device_shape), - /*destination=*/LiteralUtil::MutableInternalData(literal))); + /*destination=*/literal->MutableInternalData())); if (!ShapeUtil::Equal(literal_shape, device_shape)) { - literal->Swap( - LiteralUtil::Relayout(*literal, literal_shape.layout()).get()); + literal->Swap(literal->Relayout(literal_shape.layout()).get()); } TF_RET_CHECK(ShapeUtil::Equal(literal_shape, literal->shape())); return Status::OK(); @@ -152,27 +151,34 @@ Status GenericTransferManager::TransferLiteralToDevice( tuple_elements_on_device.data(), destination); } - return TransferBufferToDevice( - executor, /*size=*/GetByteSizeRequirement(shape), - /*source=*/LiteralUtil::InternalData(literal), destination); + return TransferBufferToDevice(executor, + /*size=*/GetByteSizeRequirement(shape), + /*source=*/literal.InternalData(), destination); } Status GenericTransferManager::TransferLiteralToInfeed( se::StreamExecutor* executor, const Literal& literal) { - return Unimplemented("Infeed is not supported on GPU (b/30467474)"); + return Unimplemented("Generic transfer to Infeed"); +} + +Status GenericTransferManager::TransferBufferToInfeed( + perftools::gputools::StreamExecutor* executor, int64 size, + const void* source) { + return Unimplemented("Generic transfer to Infeed"); } Status GenericTransferManager::TransferLiteralFromOutfeed( perftools::gputools::StreamExecutor* executor, const Shape& literal_shape, Literal* literal) { - return Unimplemented("Outfeed is not supported on CPU/GPU (b/30467474)"); + return Unimplemented( + "Outfeed is not supported on this platform (b/30467474)"); } Status GenericTransferManager::ResetDevices( tensorflow::gtl::ArraySlice - executors) { + /*executors*/) { return Unimplemented( - "Device reset is not yet supported on CPU and GPU (b/30481585)"); + "Device reset is not yet supported on this platform (b/30481585)"); } int64 GenericTransferManager::GetByteSizeRequirement(const Shape& shape) { @@ -180,14 +186,3 @@ int64 GenericTransferManager::GetByteSizeRequirement(const Shape& shape) { } } // namespace xla - -static xla::TransferManager* CreateGenericTransferManager() { - return new xla::GenericTransferManager(se::cuda::kCudaPlatformId); -} - -static bool InitModule() { - xla::TransferManager::RegisterTransferManager(se::cuda::kCudaPlatformId, - CreateGenericTransferManager); - return true; -} -static bool module_initialized = InitModule(); diff --git a/tensorflow/compiler/xla/service/generic_transfer_manager.h b/tensorflow/compiler/xla/service/generic_transfer_manager.h index 2fbdb94f06f1b12763571dc2aa9b0d770f420406..48c061d28e5967f903e9ea665fdaeb02fab7e02e 100644 --- a/tensorflow/compiler/xla/service/generic_transfer_manager.h +++ b/tensorflow/compiler/xla/service/generic_transfer_manager.h @@ -54,6 +54,8 @@ class GenericTransferManager : public TransferManager { Status TransferLiteralToInfeed(perftools::gputools::StreamExecutor* executor, const Literal& literal) override; + Status TransferBufferToInfeed(perftools::gputools::StreamExecutor* executor, + int64 size, const void* source) override; Status TransferLiteralFromOutfeed( perftools::gputools::StreamExecutor* executor, const Shape& literal_shape, diff --git a/tensorflow/compiler/xla/service/gpu/BUILD b/tensorflow/compiler/xla/service/gpu/BUILD index 9de6d65a27bfcb6747d59eac75f8b13debba0ebd..28cc43bb5ae947153c70ae98cd71e194c9ca3dad 100644 --- a/tensorflow/compiler/xla/service/gpu/BUILD +++ b/tensorflow/compiler/xla/service/gpu/BUILD @@ -68,14 +68,15 @@ cc_library( deps = [ ":ir_emission_utils", "//tensorflow/compiler/xla:util", - "//tensorflow/compiler/xla/legacy_flags:stream_assignment_flags", "//tensorflow/compiler/xla/service:hlo", + "//tensorflow/compiler/xla/service:hlo_reachability", "//tensorflow/core:lib", ], ) cc_test( name = "stream_assignment_test", + size = "small", srcs = [ "stream_assignment_test.cc", ], @@ -86,7 +87,6 @@ cc_test( "//tensorflow/compiler/xla/service:hlo", "//tensorflow/compiler/xla/tests:hlo_test_base", "//tensorflow/core:lib", - "//tensorflow/core:test_main", ], ) @@ -219,6 +219,7 @@ cc_library( "for_thunk.cc", "gemm_thunk.cc", "gpu_executable.cc", + "infeed_thunk.cc", "kernel_thunk.cc", "sequential_thunk.cc", "thunk_schedule.cc", @@ -231,6 +232,7 @@ cc_library( "for_thunk.h", "gemm_thunk.h", "gpu_executable.h", + "infeed_thunk.h", "kernel_thunk.h", "sequential_thunk.h", "thunk.h", @@ -240,6 +242,7 @@ cc_library( ], deps = [ ":buffer_allocations", + ":infeed_manager", ":partition_assignment", ":stream_assignment", "//tensorflow/compiler/xla:array2d", @@ -250,21 +253,20 @@ cc_library( "//tensorflow/compiler/xla:types", "//tensorflow/compiler/xla:util", "//tensorflow/compiler/xla:xla_data_proto", - "//tensorflow/compiler/xla/legacy_flags:convolution_thunk_flags", "//tensorflow/compiler/xla/service:buffer_assignment", "//tensorflow/compiler/xla/service:device_memory_allocator", "//tensorflow/compiler/xla/service:executable", "//tensorflow/compiler/xla/service:hlo", "//tensorflow/compiler/xla/service:hlo_execution_profile", - "//tensorflow/compiler/xla/service:hlo_module_config", "//tensorflow/compiler/xla/service:logical_buffer", - "//tensorflow/compiler/xla/service:pool", "//tensorflow/compiler/xla/service:shaped_buffer", "//tensorflow/compiler/xla/service:transfer_manager", "//tensorflow/compiler/xla/service:tuple_points_to_analysis", "//tensorflow/core:lib", "//tensorflow/core:stream_executor_no_cuda", - "//tensorflow/core/platform/default/build_config:stream_executor_cuda", + "//tensorflow/core/platform/default/build_config:cublas_plugin", + "//tensorflow/core/platform/default/build_config:cudnn_plugin", + "//tensorflow/core/platform/default/build_config:stream_executor_cuda", # build_cleaner: keep ], ) @@ -302,6 +304,7 @@ cc_library( cc_test( name = "convolution_folding_test", + size = "small", srcs = ["convolution_folding_test.cc"], deps = [ ":convolution_folding", @@ -310,7 +313,6 @@ cc_test( "//tensorflow/compiler/xla/service:shape_inference", "//tensorflow/compiler/xla/tests:hlo_test_base", "//tensorflow/core:test", - "//tensorflow/core:test_main", ], ) @@ -328,12 +330,11 @@ cc_library( cc_test( name = "instruction_fusion_test", + size = "small", srcs = ["instruction_fusion_test.cc"], deps = [ ":instruction_fusion", "//tensorflow/compiler/xla/tests:hlo_test_base", - "//tensorflow/core:test", - "//tensorflow/core:test_main", ], ) @@ -368,18 +369,13 @@ cc_library( cc_test( name = "fusion_merger_test", + size = "small", srcs = ["fusion_merger_test.cc"], deps = [ ":fusion_merger", ":instruction_fusion", - "//tensorflow/compiler/xla:literal_util", - "//tensorflow/compiler/xla:shape_util", "//tensorflow/compiler/xla:test_helpers", - "//tensorflow/compiler/xla:util", - "//tensorflow/compiler/xla:xla_data_proto", - "//tensorflow/compiler/xla/service:hlo", "//tensorflow/compiler/xla/tests:hlo_test_base", - "//tensorflow/core:test_main", ], ) @@ -415,25 +411,29 @@ cc_library( ":pad_insertion", ":partition_assignment", ":stream_assignment", + "//tensorflow/compiler/xla:protobuf_util", "//tensorflow/compiler/xla:status_macros", "//tensorflow/compiler/xla:statusor", "//tensorflow/compiler/xla:types", "//tensorflow/compiler/xla:util", - "//tensorflow/compiler/xla/legacy_flags:gpu_compiler_flags", "//tensorflow/compiler/xla/service:algebraic_simplifier", + "//tensorflow/compiler/xla/service:batchnorm_rewriter", "//tensorflow/compiler/xla/service:buffer_assignment", "//tensorflow/compiler/xla/service:buffer_liveness", - "//tensorflow/compiler/xla/service:compiler", "//tensorflow/compiler/xla/service:executable", + "//tensorflow/compiler/xla/service:flatten_call_graph", "//tensorflow/compiler/xla/service:hlo", "//tensorflow/compiler/xla/service:hlo_constant_folding", "//tensorflow/compiler/xla/service:hlo_cse", "//tensorflow/compiler/xla/service:hlo_dce", - "//tensorflow/compiler/xla/service:hlo_module_config", "//tensorflow/compiler/xla/service:hlo_pass", "//tensorflow/compiler/xla/service:hlo_pass_pipeline", + "//tensorflow/compiler/xla/service:hlo_proto", + "//tensorflow/compiler/xla/service:hlo_proto_util", "//tensorflow/compiler/xla/service:hlo_subcomputation_unification", "//tensorflow/compiler/xla/service:hlo_verifier", + "//tensorflow/compiler/xla/service:llvm_compiler", + "//tensorflow/compiler/xla/service:reduce_precision_insertion", "//tensorflow/compiler/xla/service:reshape_mover", "//tensorflow/compiler/xla/service:transpose_folding", "//tensorflow/compiler/xla/service/gpu/llvm_gpu_backend", @@ -447,6 +447,18 @@ cc_library( alwayslink = True, # Contains compiler registration ) +cc_library( + name = "infeed_manager", + srcs = ["infeed_manager.cc"], + hdrs = ["infeed_manager.h"], + deps = [ + "//tensorflow/compiler/xla:types", + "//tensorflow/compiler/xla:util", + "//tensorflow/core:lib", + "//tensorflow/core:stream_executor_no_cuda", + ], +) + cc_library( name = "layout_assignment", srcs = ["layout_assignment.cc"], @@ -465,6 +477,7 @@ cc_library( cc_test( name = "layout_assignment_test", + size = "small", srcs = ["layout_assignment_test.cc"], deps = [ ":layout_assignment", @@ -474,7 +487,6 @@ cc_test( "//tensorflow/compiler/xla/service:computation_layout", "//tensorflow/compiler/xla/service:hlo", "//tensorflow/compiler/xla/tests:hlo_test_base", - "//tensorflow/core:test_main", ], ) @@ -487,13 +499,16 @@ cc_library( "//tensorflow/compiler/xla:statusor", "//tensorflow/compiler/xla:types", "//tensorflow/compiler/xla:util", - "//tensorflow/compiler/xla/service:buffer_liveness", "//tensorflow/compiler/xla/service:hlo", + "//tensorflow/compiler/xla/service:hlo_ordering", + "//tensorflow/compiler/xla/service:hlo_reachability", + "//tensorflow/compiler/xla/service:hlo_scheduling", ], ) cc_test( name = "hlo_schedule_test", + size = "small", srcs = [ "hlo_schedule_test.cc", ], @@ -504,7 +519,6 @@ cc_test( "//tensorflow/compiler/xla:types", "//tensorflow/compiler/xla/service:hlo", "//tensorflow/compiler/xla/tests:hlo_test_base", - "//tensorflow/core:test_main", ], ) @@ -525,19 +539,15 @@ cc_library( cc_test( name = "while_transformer_test", + size = "small", srcs = ["while_transformer_test.cc"], deps = [ ":instruction_fusion", ":while_transformer", - "//tensorflow/compiler/xla:literal_util", - "//tensorflow/compiler/xla:shape_util", + "//tensorflow/compiler/xla:test", "//tensorflow/compiler/xla:test_helpers", "//tensorflow/compiler/xla/service:copy_insertion", - "//tensorflow/compiler/xla/service:hlo", "//tensorflow/compiler/xla/tests:hlo_test_base", - "//tensorflow/core:lib", - "//tensorflow/core:test", - "//tensorflow/core:test_main", ], ) diff --git a/tensorflow/compiler/xla/service/gpu/convolution_folding.cc b/tensorflow/compiler/xla/service/gpu/convolution_folding.cc index 16febea14de233c554045a1fe95221d802c0882c..c598025b5e8f3ff72656ff370068bb0ff3a80f2a 100644 --- a/tensorflow/compiler/xla/service/gpu/convolution_folding.cc +++ b/tensorflow/compiler/xla/service/gpu/convolution_folding.cc @@ -43,7 +43,7 @@ MatchBackwardFilter(HloInstruction* conv) { ConvolutionDimensionNumbers()); // Step 1: match the instruction pattern without considering the paddings and // dimension numbers just yet. We may need some generic pattern matcher - // similar to external/llvm/include/llvm/IR/PatternMatch.h + // similar to third_party/llvm/llvm/include/llvm/IR/PatternMatch.h // // Backward filter convolution is implemented in XLA as the forward // convolution of padded activations and dilated gradients. Padding on diff --git a/tensorflow/compiler/xla/service/gpu/convolution_folding_test.cc b/tensorflow/compiler/xla/service/gpu/convolution_folding_test.cc index 83922cbe14af1f5a2a8c0a9a6a678a181cee1fca..ba9c70ded36d8e3f6d25f6a7b95daa17683bca99 100644 --- a/tensorflow/compiler/xla/service/gpu/convolution_folding_test.cc +++ b/tensorflow/compiler/xla/service/gpu/convolution_folding_test.cc @@ -97,10 +97,10 @@ TEST_F(ConvolutionFoldingTest, BackwardFilterConvolveWithoutTranspose) { activations, gradients, conv_window, tf_default_dnums_for_backward_filter_)); - HloModule module(TestName()); + auto module = CreateNewModule(); HloComputation* entry_computation = - module.AddEntryComputation(builder.Build()); - EXPECT_TRUE(FoldConvolution(&module)); + module->AddEntryComputation(builder.Build()); + EXPECT_TRUE(FoldConvolution(module.get())); EXPECT_EQ(HloOpcode::kFusion, entry_computation->root_instruction()->opcode()); EXPECT_TRUE(HloInstruction::FusionKind::kConvBackwardFilter == @@ -126,9 +126,9 @@ TEST_F(ConvolutionFoldingTest, activations, gradients, conv_window, tf_default_dnums_for_backward_filter_)); - HloModule module(TestName()); - module.AddEntryComputation(builder.Build()); - EXPECT_FALSE(FoldConvolution(&module)); + auto module = CreateNewModule(); + module->AddEntryComputation(builder.Build()); + EXPECT_FALSE(FoldConvolution(module.get())); } // Extracted from block35 training. @@ -155,10 +155,10 @@ TEST_F(ConvolutionFoldingTest, BackwardFilterConvolveWithPaddedActivations) { builder.AddInstruction(HloInstruction::CreateTranspose( ShapeUtil::MakeShape(F32, {3, 3, 32, 32}), convolution, {1, 2, 3, 0})); - HloModule module(TestName()); + auto module = CreateNewModule(); HloComputation* entry_computation = - module.AddEntryComputation(builder.Build()); - EXPECT_TRUE(FoldConvolution(&module)); + module->AddEntryComputation(builder.Build()); + EXPECT_TRUE(FoldConvolution(module.get())); EXPECT_EQ(HloOpcode::kFusion, entry_computation->root_instruction()->opcode()); EXPECT_TRUE(HloInstruction::FusionKind::kConvBackwardFilter == @@ -189,10 +189,10 @@ TEST_F(ConvolutionFoldingTest, BackwardFilterConvolveWithPaddedGradients) { builder.AddInstruction(HloInstruction::CreateTranspose( ShapeUtil::MakeShape(F32, {3, 3, 192, 320}), convolution, {1, 2, 3, 0})); - HloModule module(TestName()); + auto module = CreateNewModule(); HloComputation* entry_computation = - module.AddEntryComputation(builder.Build()); - EXPECT_TRUE(FoldConvolution(&module)); + module->AddEntryComputation(builder.Build()); + EXPECT_TRUE(FoldConvolution(module.get())); EXPECT_EQ(HloOpcode::kFusion, entry_computation->root_instruction()->opcode()); EXPECT_TRUE(HloInstruction::FusionKind::kConvBackwardFilter == @@ -222,10 +222,10 @@ TEST_F(ConvolutionFoldingTest, BackwardFilterConvolveWithUnevenPadding) { builder.AddInstruction(HloInstruction::CreateTranspose( ShapeUtil::MakeShape(F32, {2, 2, 32, 32}), convolution, {1, 2, 3, 0})); - HloModule module(TestName()); + auto module = CreateNewModule(); HloComputation* entry_computation = - module.AddEntryComputation(builder.Build()); - EXPECT_TRUE(FoldConvolution(&module)); + module->AddEntryComputation(builder.Build()); + EXPECT_TRUE(FoldConvolution(module.get())); EXPECT_EQ(HloOpcode::kFusion, entry_computation->root_instruction()->opcode()); EXPECT_TRUE(HloInstruction::FusionKind::kConvBackwardFilter == @@ -269,10 +269,10 @@ TEST_F(ConvolutionFoldingTest, BackwardInputConvolveEvenPadding) { output->shape(), reverse_kernel->shape(), conv_window, conv_dnums) .ValueOrDie())); - HloModule module(TestName()); + auto module = CreateNewModule(); HloComputation* entry_computation = - module.AddEntryComputation(builder.Build()); - EXPECT_TRUE(FoldConvolution(&module)); + module->AddEntryComputation(builder.Build()); + EXPECT_TRUE(FoldConvolution(module.get())); EXPECT_EQ(HloOpcode::kFusion, entry_computation->root_instruction()->opcode()); EXPECT_TRUE(HloInstruction::FusionKind::kConvBackwardInput == @@ -313,10 +313,10 @@ TEST_F(ConvolutionFoldingTest, BackwardInputConvolve1x1Filter) { /*lhs=*/output, /*rhs=*/kernel, conv_window, tf_default_dnums_for_backward_input_)); - HloModule module(TestName()); + auto module = CreateNewModule(); HloComputation* entry_computation = - module.AddEntryComputation(builder.Build()); - EXPECT_TRUE(FoldConvolution(&module)); + module->AddEntryComputation(builder.Build()); + EXPECT_TRUE(FoldConvolution(module.get())); EXPECT_EQ(HloOpcode::kFusion, entry_computation->root_instruction()->opcode()); EXPECT_TRUE(HloInstruction::FusionKind::kConvBackwardInput == @@ -346,9 +346,9 @@ TEST_F(ConvolutionFoldingTest, /*lhs=*/output, /*rhs=*/kernel, default_conv_window_, tf_default_dnums_for_backward_input_)); - HloModule module(TestName()); - module.AddEntryComputation(builder.Build()); - EXPECT_FALSE(FoldConvolution(&module)); + auto module = CreateNewModule(); + module->AddEntryComputation(builder.Build()); + EXPECT_FALSE(FoldConvolution(module.get())); } // Extracted from Inception V3 training. @@ -394,10 +394,10 @@ TEST_F(ConvolutionFoldingTest, BackwardInputConvolveUnevenPaddingOnGradients) { tf_default_dnums_for_backward_input_) .ValueOrDie())); - HloModule module(TestName()); + auto module = CreateNewModule(); HloComputation* entry_computation = - module.AddEntryComputation(builder.Build()); - EXPECT_TRUE(FoldConvolution(&module)); + module->AddEntryComputation(builder.Build()); + EXPECT_TRUE(FoldConvolution(module.get())); EXPECT_EQ(HloOpcode::kFusion, entry_computation->root_instruction()->opcode()); EXPECT_TRUE(HloInstruction::FusionKind::kConvBackwardInput == @@ -441,9 +441,9 @@ TEST_F(ConvolutionFoldingTest, BackwardInputConvolveLowPaddingTooLarge) { tf_default_dnums_for_backward_input_) .ValueOrDie())); - HloModule module(TestName()); - module.AddEntryComputation(builder.Build()); - EXPECT_FALSE(FoldConvolution(&module)); + auto module = CreateNewModule(); + module->AddEntryComputation(builder.Build()); + EXPECT_FALSE(FoldConvolution(module.get())); } // Extracted from //learning/brain/google/xla/benchmarks/resnet.py @@ -490,10 +490,10 @@ TEST_F(ConvolutionFoldingTest, tf_default_dnums_for_backward_input_) .ValueOrDie())); - HloModule module(TestName()); + auto module = CreateNewModule(); const HloComputation* entry_computation = - module.AddEntryComputation(builder.Build()); - EXPECT_TRUE(FoldConvolution(&module)); + module->AddEntryComputation(builder.Build()); + EXPECT_TRUE(FoldConvolution(module.get())); const HloInstruction* backward_conv = entry_computation->root_instruction(); EXPECT_EQ(HloOpcode::kFusion, backward_conv->opcode()); EXPECT_TRUE(HloInstruction::FusionKind::kConvBackwardInput == @@ -543,10 +543,14 @@ TEST_F(ConvolutionFoldingTest, tf_default_dnums_for_backward_input_) .ValueOrDie())); - HloModule module(TestName()); - module.AddEntryComputation(builder.Build()); - EXPECT_FALSE(FoldConvolution(&module)); + auto module = CreateNewModule(); + module->AddEntryComputation(builder.Build()); + EXPECT_FALSE(FoldConvolution(module.get())); } } // namespace gpu } // namespace xla + +int main(int argc, char** argv) { + return xla::ParseDebugOptionsFlagsAndRunTests(argc, argv); +} diff --git a/tensorflow/compiler/xla/service/gpu/convolution_thunk.cc b/tensorflow/compiler/xla/service/gpu/convolution_thunk.cc index f6b7fe1e8ef10e4e66018d887707e587ecfa3465..20e0d8eb785daa07b3fcc5339efe950aac0dacad 100644 --- a/tensorflow/compiler/xla/service/gpu/convolution_thunk.cc +++ b/tensorflow/compiler/xla/service/gpu/convolution_thunk.cc @@ -17,7 +17,6 @@ limitations under the License. #include -#include "tensorflow/compiler/xla/legacy_flags/convolution_thunk_flags.h" #include "tensorflow/compiler/xla/types.h" #include "tensorflow/compiler/xla/util.h" #include "tensorflow/core/lib/strings/stringprintf.h" @@ -125,7 +124,7 @@ tensorflow::Status ConvolutionThunk::ExecuteOnStream( CHECK_LE(num_dimensions, 3); // cuDNN does not support 1D convolutions. We therefore express 1D // convolutions as 2D convolutions where the first spatial dimension is 1. - // This matches the behaviour of TF (see definition of conv1d in + // This matches the behavior of TF (see definition of conv1d in // tensorflow/python/ops/nn_ops.py). const int effective_num_dimensions = std::max(2, num_dimensions); @@ -258,15 +257,21 @@ tensorflow::Status ConvolutionThunk::Convolve( std::vector ConvolutionThunk::GetAlgorithms( se::StreamExecutor* stream_exec) const { std::vector algorithms; + // TODO(yangzihao): Currently disable the use of winograd nonfused in XLA + // by default. Should send in conv parameters and enable it when + // ShouldIncludeWinogradNonfusedAlgo() returns true. switch (convolution_kind_) { case ConvolutionKind::kBackwardFilter: - CHECK(stream_exec->GetConvolveBackwardFilterAlgorithms(&algorithms)); + CHECK(stream_exec->GetConvolveBackwardFilterAlgorithms( + /*with_winograd_nonfused=*/false, &algorithms)); break; case ConvolutionKind::kBackwardInput: - CHECK(stream_exec->GetConvolveBackwardDataAlgorithms(&algorithms)); + CHECK(stream_exec->GetConvolveBackwardDataAlgorithms( + /*with_winograd_nonfused=*/false, &algorithms)); break; case ConvolutionKind::kForward: - CHECK(stream_exec->GetConvolveAlgorithms(&algorithms)); + CHECK(stream_exec->GetConvolveAlgorithms(/*with_winograd_nonfused=*/false, + &algorithms)); break; } return algorithms; @@ -281,10 +286,7 @@ tensorflow::Status ConvolutionThunk::ConvolveWithTune( const ConvolutionDescriptor& convolution_descriptor, const BufferAllocations& buffer_allocations, se::Stream* stream) { // TODO(b/29126320): Try cudnn v5's new auto-tuner when it's rolled out. - legacy_flags::ConvolutionThunkFlags* flags = - legacy_flags::GetConvolutionThunkFlags(); - if (flags->xla_gpu_autotune_convolution_algorithm && - best_algorithm_.algorithm() == se::dnn::kDefaultAlgorithm) { + if (best_algorithm_.algorithm() == se::dnn::kDefaultAlgorithm) { // Auto-tuning either is disabled or only happens in the first run of this // function. VLOG(2) << "Profiling for best convolution algorithm used for " diff --git a/tensorflow/compiler/xla/service/gpu/convolution_thunk.h b/tensorflow/compiler/xla/service/gpu/convolution_thunk.h index aaf72935e61ee8b8da7df410ba3aaed63800cfd9..91d6df299da2686d6d836445d391c4b0eaf4ed00 100644 --- a/tensorflow/compiler/xla/service/gpu/convolution_thunk.h +++ b/tensorflow/compiler/xla/service/gpu/convolution_thunk.h @@ -81,9 +81,8 @@ class ConvolutionThunk : public Thunk { ConvolutionThunk(const ConvolutionThunk&) = delete; ConvolutionThunk& operator=(const ConvolutionThunk&) = delete; - // Does the convolution for the thunk on "stream". If the - // xla_gpu_autotune_convolution_algorithm is turned on, auto-tuning happens on - // the first run of this function. + // Does the convolution for the thunk on "stream". Auto-tuning happens on the + // first run of this function. tensorflow::Status ExecuteOnStream( const BufferAllocations& buffer_allocations, perftools::gputools::Stream* stream) override; diff --git a/tensorflow/compiler/xla/service/gpu/elemental_ir_emitter.cc b/tensorflow/compiler/xla/service/gpu/elemental_ir_emitter.cc index 1667ab36792c91cbbf3c6396a673bedff2208045..5edaaba3ebe482126c800059968d0e430076f950 100644 --- a/tensorflow/compiler/xla/service/gpu/elemental_ir_emitter.cc +++ b/tensorflow/compiler/xla/service/gpu/elemental_ir_emitter.cc @@ -23,12 +23,12 @@ limitations under the License. #include "tensorflow/core/platform/types.h" // IWYU pragma: no_include "llvm/IR/Attributes.gen.inc" // IWYU pragma: no_include "llvm/IR/Intrinsics.gen.inc" -#include "external/llvm/include/llvm/ADT/APInt.h" -#include "external/llvm/include/llvm/IR/BasicBlock.h" -#include "external/llvm/include/llvm/IR/Instructions.h" -#include "external/llvm/include/llvm/IR/Intrinsics.h" -#include "external/llvm/include/llvm/IR/Module.h" -#include "external/llvm/include/llvm/IR/Type.h" +#include "llvm/ADT/APInt.h" +#include "llvm/IR/BasicBlock.h" +#include "llvm/IR/Instructions.h" +#include "llvm/IR/Intrinsics.h" +#include "llvm/IR/Module.h" +#include "llvm/IR/Type.h" #include "tensorflow/compiler/xla/literal_util.h" #include "tensorflow/compiler/xla/primitive_util.h" #include "tensorflow/compiler/xla/service/hlo_opcode.h" @@ -55,7 +55,7 @@ using tensorflow::strings::StrAppend; // Returns whether operand is a floating-point literal with the given value. bool IsFPLiteralWithValue(const HloInstruction* operand, float value) { return operand->opcode() == HloOpcode::kConstant && - LiteralUtil::IsAllFloat(operand->literal(), value); + operand->literal().IsAllFloat(value); } GpuElementalIrEmitter::GpuElementalIrEmitter( @@ -113,7 +113,7 @@ StatusOr GpuElementalIrEmitter::EmitMathCall( tensorflow::gtl::ArraySlice operands, tensorflow::gtl::ArraySlice input_types, PrimitiveType output_type) const { - // Binary math functions tranform are of type [T] -> T. + // Binary math functions transform are of type [T] -> T. for (PrimitiveType input_type : input_types) { if (output_type != input_type) { return Unimplemented("Input type ≠ output type: %s ≠ %s", @@ -175,7 +175,7 @@ StatusOr GpuElementalIrEmitter::EmitPowerOp( return make_sqrt(); } - if (!hlo_module_config_.fast_math_disabled() && + if (hlo_module_config_.debug_options().xla_enable_fast_math() && IsFPLiteralWithValue(rhs, -.5)) { VLOG(10) << "emitting pow(A, -.5) as 1/sqrt(A): " << op->ToString(); // LLVM's NVPTX backend knows how to transform 1/sqrt(A) into the NVPTX @@ -211,6 +211,12 @@ StatusOr GpuElementalIrEmitter::EmitFloatUnaryOp( case HloOpcode::kLog: return EmitLibdeviceMathCall("__nv_log", {operand_value}, {input_type}, output_type); + case HloOpcode::kCos: + return EmitLibdeviceMathCall("__nv_cos", {operand_value}, {input_type}, + output_type); + case HloOpcode::kSin: + return EmitLibdeviceMathCall("__nv_sin", {operand_value}, {input_type}, + output_type); case HloOpcode::kTanh: return EmitLibdeviceMathCall("__nv_tanh", {operand_value}, {input_type}, output_type); @@ -328,7 +334,7 @@ llvm_ir::ElementGenerator GpuElementalIrEmitter::MakeElementGenerator( SetToFirstInsertPoint(loops.GetInnerLoopBodyBasicBlock(), ir_builder_); IrArray::Index input_index(index.size()); - llvm::Value* in_bounds = ir_builder_->getInt1(1); + llvm::Value* in_bounds = ir_builder_->getInt1(true); for (size_t i = 0; i < index.size(); ++i) { llvm::Value* stridden_index = ir_builder_->CreateNSWMul( index[i], ir_builder_->getInt64(window.dimensions(i).stride())); diff --git a/tensorflow/compiler/xla/service/gpu/elemental_ir_emitter.h b/tensorflow/compiler/xla/service/gpu/elemental_ir_emitter.h index 3b484a0f1a2cb7e88eaa14969f4eda82decac971..6ddfc3710c56a4e129f050f862812a3d78d8dba0 100644 --- a/tensorflow/compiler/xla/service/gpu/elemental_ir_emitter.h +++ b/tensorflow/compiler/xla/service/gpu/elemental_ir_emitter.h @@ -20,8 +20,8 @@ limitations under the License. #include #include -#include "external/llvm/include/llvm/IR/IRBuilder.h" -#include "external/llvm/include/llvm/IR/Value.h" +#include "llvm/IR/IRBuilder.h" +#include "llvm/IR/Value.h" #include "tensorflow/compiler/xla/service/elemental_ir_emitter.h" #include "tensorflow/compiler/xla/service/hlo_computation.h" #include "tensorflow/compiler/xla/service/hlo_instruction.h" diff --git a/tensorflow/compiler/xla/service/gpu/fusion_merger.cc b/tensorflow/compiler/xla/service/gpu/fusion_merger.cc index afb78b8300b457ba9384bd66f789d333630b51e4..a9ef204b46facafabcf16d1d38d69c14e6aab497 100644 --- a/tensorflow/compiler/xla/service/gpu/fusion_merger.cc +++ b/tensorflow/compiler/xla/service/gpu/fusion_merger.cc @@ -98,7 +98,13 @@ double CalculateFlopsToBytesRatio(HloInstruction* fusion) { // Calculate total bytes transferred in/out. double bytes = CalculateBytesReadByFusionInstruction(fusion); // Add bytes written to root instructions buffer. - bytes += ShapeUtil::ByteSizeOf(fusion->fused_expression_root()->shape()); + if (fusion->IsMultiOutputFusion()) { + for (auto& operand : fusion->fused_expression_root()->operands()) { + bytes += ShapeUtil::ByteSizeOf(operand->shape()); + } + } else { + bytes += ShapeUtil::ByteSizeOf(fusion->fused_expression_root()->shape()); + } // Calculate flops for all fused instructions. Use a null shape size function // because we don't care about bytes accessed by the ops. HloCostAnalysis analysis([](const Shape& shape) { return 0; }); @@ -112,8 +118,15 @@ double CalculateFlopsToBytesRatio(HloInstruction* fusion) { double GetCurrentBytesTransferred(HloInstruction* fusion) { CHECK_EQ(HloOpcode::kFusion, fusion->opcode()); const double bytes_read = CalculateBytesReadByFusionInstruction(fusion); - const double bytes_written = - ShapeUtil::ByteSizeOf(fusion->fused_expression_root()->shape()); + double bytes_written = 0; + if (fusion->IsMultiOutputFusion()) { + for (auto& operand : fusion->fused_expression_root()->operands()) { + bytes_written += ShapeUtil::ByteSizeOf(operand->shape()); + } + } else { + bytes_written = + ShapeUtil::ByteSizeOf(fusion->fused_expression_root()->shape()); + } // Current bytes transferred (ignoring non 'fusion' user operands) is bytes // read and written by 'fusion', plus reads of size 'bytes_written' for each // user. @@ -198,6 +211,12 @@ Status FusionInstructionMerger::HandleFusion(HloInstruction* fusion) { ++num_fail_not_loop_fusion_; return Status::OK(); } + + // Skip multiple output fusion. It's not yet supported. + if (fusion->IsMultiOutputFusion()) { + ++num_fail_not_loop_fusion_; + return Status::OK(); + } // Skip 'fusion' instruction if we cannot merge into all of its users. // Merging into all users enables the removal of 'fusion' from the // computation. @@ -274,12 +293,19 @@ Status FusionInstructionMerger::HandleFusion(HloInstruction* fusion) { StatusOr FusionMerger::Run(HloModule* module) { bool changed = false; VLOG(2) << "FusionMerger for module: " << module->name(); + std::vector computations; for (auto& computation : module->computations()) { + if (computation->IsFusionComputation()) { + continue; + } + computations.push_back(computation.get()); + } + for (auto& computation : computations) { VLOG(1) << "Before running FusionInstructionMerger for computation: " << computation->name(); XLA_VLOG_LINES(3, computation->ToString()); - FusionInstructionMerger fusion_merger(computation.get()); + FusionInstructionMerger fusion_merger(computation); TF_RETURN_IF_ERROR(fusion_merger.Run()); changed |= fusion_merger.changed(); diff --git a/tensorflow/compiler/xla/service/gpu/fusion_merger.h b/tensorflow/compiler/xla/service/gpu/fusion_merger.h index 9a989d26f93a4abd59eecc54d56e049f84a54155..bd720f8584f6254c43a3e2a1a5399aa919eebbc0 100644 --- a/tensorflow/compiler/xla/service/gpu/fusion_merger.h +++ b/tensorflow/compiler/xla/service/gpu/fusion_merger.h @@ -25,7 +25,7 @@ namespace gpu { // An HLO pass that attempts to merge fusion instructions to reduce kernel // launch overhead and improve data locality. // -// Fusion instructions are merged into their users if two conditons are met: +// Fusion instructions are merged into their users if two conditions are met: // // 1) The flops_to_bytes ratio of the fusion instruction is below the threshold // value of 1.0. diff --git a/tensorflow/compiler/xla/service/gpu/fusion_merger_test.cc b/tensorflow/compiler/xla/service/gpu/fusion_merger_test.cc index a87e66ca869e249c560c9e477b1406d09c2886ba..242c32936d31d0cb578825cade5f35979077a44e 100644 --- a/tensorflow/compiler/xla/service/gpu/fusion_merger_test.cc +++ b/tensorflow/compiler/xla/service/gpu/fusion_merger_test.cc @@ -25,7 +25,7 @@ namespace { class FusionMergerTest : public HloTestBase { protected: - FusionMergerTest() : module_(TestName()) {} + FusionMergerTest() : module_(CreateNewModule()) {} // Builds the following computation: // @@ -59,7 +59,7 @@ class FusionMergerTest : public HloTestBase { // Create const vector of ones to be used in element-wise computations. auto one_vec = builder.AddInstruction(HloInstruction::CreateConstant( - LiteralUtil::CreateR1({1.f, 1.f, 1.f, 1.f}))); + Literal::CreateR1({1.f, 1.f, 1.f, 1.f}))); // Create simple fusable computation for tuple element 0 (wont get merged). auto out0 = builder.AddInstruction(HloInstruction::CreateBinary( @@ -86,7 +86,7 @@ class FusionMergerTest : public HloTestBase { // Create output Tuple. builder.AddInstruction(HloInstruction::CreateTuple({out0, out1, out2})); - return module_.AddEntryComputation(builder.Build()); + return module_->AddEntryComputation(builder.Build()); } // Builds the following computation: @@ -138,7 +138,7 @@ class FusionMergerTest : public HloTestBase { // Create two sub-computations, both of which are users of 'mul0'. auto one_vec = builder.AddInstruction(HloInstruction::CreateConstant( - LiteralUtil::CreateR1({1.f, 1.f, 1.f, 1.f}))); + Literal::CreateR1({1.f, 1.f, 1.f, 1.f}))); // First sub-computation: out0 = Mul(Add(mul0, one_vec), one_vec) auto add0 = builder.AddInstruction(HloInstruction::CreateBinary( @@ -154,7 +154,7 @@ class FusionMergerTest : public HloTestBase { // Create output Tuple. builder.AddInstruction(HloInstruction::CreateTuple({out0, out1})); - return module_.AddEntryComputation(builder.Build()); + return module_->AddEntryComputation(builder.Build()); } // Builds the following computation: @@ -209,7 +209,7 @@ class FusionMergerTest : public HloTestBase { // Create two fusable sub-computations which are dependent on shared // computation 'reduce_out'. auto one_vec = builder.AddInstruction(HloInstruction::CreateConstant( - LiteralUtil::CreateR1({1.f, 1.f, 1.f, 1.f}))); + Literal::CreateR1({1.f, 1.f, 1.f, 1.f}))); // First sub-computation: out0 = Mul(Add(reduce_out, one_vec), one_vec) auto add2 = builder.AddInstruction(HloInstruction::CreateBinary( @@ -225,7 +225,7 @@ class FusionMergerTest : public HloTestBase { // Create output Tuple. builder.AddInstruction(HloInstruction::CreateTuple({out0, out1})); - return module_.AddEntryComputation(builder.Build()); + return module_->AddEntryComputation(builder.Build()); } Shape data_shape_ = ShapeUtil::MakeShape(F32, {4}); @@ -235,7 +235,7 @@ class FusionMergerTest : public HloTestBase { Shape tuple_shape4_ = ShapeUtil::MakeTupleShape( {data_shape_, data_shape_, data_shape_, data_shape_}); - HloModule module_; + std::unique_ptr module_; }; // Tests that we can merge a fusion instruction that is below threshold. @@ -278,13 +278,15 @@ class FusionMergerTest : public HloTestBase { TEST_F(FusionMergerTest, MergeSharedFusionInstruction) { auto computation = BuildComputation0(); // Run standard fusion passes. - EXPECT_TRUE( - GpuInstructionFusion(/*may_duplicate=*/false).Run(&module_).ValueOrDie()); - EXPECT_FALSE( - GpuInstructionFusion(/*may_duplicate=*/true).Run(&module_).ValueOrDie()); + EXPECT_TRUE(GpuInstructionFusion(/*may_duplicate=*/false) + .Run(module_.get()) + .ValueOrDie()); + EXPECT_FALSE(GpuInstructionFusion(/*may_duplicate=*/true) + .Run(module_.get()) + .ValueOrDie()); // Run fusion merger pass, which should merge the shared fusion instruction // into its two users. - EXPECT_TRUE(FusionMerger().Run(&module_).ValueOrDie()); + EXPECT_TRUE(FusionMerger().Run(module_.get()).ValueOrDie()); auto* root = computation->root_instruction(); EXPECT_EQ(HloOpcode::kTuple, root->opcode()); @@ -338,14 +340,16 @@ TEST_F(FusionMergerTest, MergeSharedFusionInstruction) { TEST_F(FusionMergerTest, FlopsToBytesRatioThresholdExceeded) { BuildComputation1(); // Run standard fusion passes. - EXPECT_TRUE( - GpuInstructionFusion(/*may_duplicate=*/false).Run(&module_).ValueOrDie()); - EXPECT_FALSE( - GpuInstructionFusion(/*may_duplicate=*/true).Run(&module_).ValueOrDie()); + EXPECT_TRUE(GpuInstructionFusion(/*may_duplicate=*/false) + .Run(module_.get()) + .ValueOrDie()); + EXPECT_FALSE(GpuInstructionFusion(/*may_duplicate=*/true) + .Run(module_.get()) + .ValueOrDie()); // Run fusion merger pass, which should detect that the flops/bytes of the // shared fusion instruction exceeds the threshold ratio, and therefore // cannot be merged with other fusion instructions. - EXPECT_FALSE(FusionMerger().Run(&module_).ValueOrDie()); + EXPECT_FALSE(FusionMerger().Run(module_.get()).ValueOrDie()); } // Tests that threshold for bytes transferred if merged is exceeded. @@ -388,13 +392,15 @@ TEST_F(FusionMergerTest, FlopsToBytesRatioThresholdExceeded) { TEST_F(FusionMergerTest, BytesTransferredThresholdExeceeded) { BuildComputation2(/*add_extra_input=*/true); // Run standard fusion passes. - EXPECT_TRUE( - GpuInstructionFusion(/*may_duplicate=*/false).Run(&module_).ValueOrDie()); - EXPECT_FALSE( - GpuInstructionFusion(/*may_duplicate=*/true).Run(&module_).ValueOrDie()); + EXPECT_TRUE(GpuInstructionFusion(/*may_duplicate=*/false) + .Run(module_.get()) + .ValueOrDie()); + EXPECT_FALSE(GpuInstructionFusion(/*may_duplicate=*/true) + .Run(module_.get()) + .ValueOrDie()); // Run fusion merger pass, which should detect that the net bytes transferred // (if merged) would increase. - EXPECT_FALSE(FusionMerger().Run(&module_).ValueOrDie()); + EXPECT_FALSE(FusionMerger().Run(module_.get()).ValueOrDie()); } // Tests that threshold for bytes transferred if merged is not exceeded. @@ -442,15 +448,21 @@ TEST_F(FusionMergerTest, BytesTransferredThresholdExeceeded) { TEST_F(FusionMergerTest, BytesTransferredThresholdNotExeceeded) { BuildComputation2(/*add_extra_input=*/false); // Run standard fusion passes. - EXPECT_TRUE( - GpuInstructionFusion(/*may_duplicate=*/false).Run(&module_).ValueOrDie()); - EXPECT_FALSE( - GpuInstructionFusion(/*may_duplicate=*/true).Run(&module_).ValueOrDie()); + EXPECT_TRUE(GpuInstructionFusion(/*may_duplicate=*/false) + .Run(module_.get()) + .ValueOrDie()); + EXPECT_FALSE(GpuInstructionFusion(/*may_duplicate=*/true) + .Run(module_.get()) + .ValueOrDie()); // Run fusion merger pass, which should detect that the net bytes transferred // (if merged) would not increase. - EXPECT_TRUE(FusionMerger().Run(&module_).ValueOrDie()); + EXPECT_TRUE(FusionMerger().Run(module_.get()).ValueOrDie()); } } // namespace } // namespace gpu } // namespace xla + +int main(int argc, char** argv) { + return xla::ParseDebugOptionsFlagsAndRunTests(argc, argv); +} diff --git a/tensorflow/compiler/xla/service/gpu/gemm_thunk.cc b/tensorflow/compiler/xla/service/gpu/gemm_thunk.cc index a80f969b9ded2e4a50e77ce35f807014ee521b2a..e784046450ed1cca088770c65c786e80adda869f 100644 --- a/tensorflow/compiler/xla/service/gpu/gemm_thunk.cc +++ b/tensorflow/compiler/xla/service/gpu/gemm_thunk.cc @@ -245,7 +245,7 @@ tensorflow::Status GemmThunk::ExecuteOnStream( // Therefore, we need to convert dot between row-major matrices to that // between column-major matrices. The key insight for the conversion is that, // in linear storage, matrix M in column-major order is identical to the - // tranpose of M in row-major order. In other words, + // transpose of M in row-major order. In other words, // // column-major(M) = row-major(M^T). // diff --git a/tensorflow/compiler/xla/service/gpu/gpu_compiler.cc b/tensorflow/compiler/xla/service/gpu/gpu_compiler.cc index f692f28bd9858ab809732389fcc2908b8fa66a42..7f5be602beb1c6b337ee7ef86ec19271d1f73cb5 100644 --- a/tensorflow/compiler/xla/service/gpu/gpu_compiler.cc +++ b/tensorflow/compiler/xla/service/gpu/gpu_compiler.cc @@ -19,15 +19,17 @@ limitations under the License. #include #include -#include "external/llvm/include/llvm/IR/DiagnosticInfo.h" -#include "external/llvm/include/llvm/IR/DiagnosticPrinter.h" -#include "external/llvm/include/llvm/IR/LLVMContext.h" -#include "external/llvm/include/llvm/IR/Module.h" -#include "tensorflow/compiler/xla/legacy_flags/gpu_compiler_flags.h" +#include "llvm/IR/DiagnosticInfo.h" +#include "llvm/IR/DiagnosticPrinter.h" +#include "llvm/IR/LLVMContext.h" +#include "llvm/IR/Module.h" +#include "tensorflow/compiler/xla/protobuf_util.h" #include "tensorflow/compiler/xla/ptr_util.h" #include "tensorflow/compiler/xla/service/algebraic_simplifier.h" +#include "tensorflow/compiler/xla/service/batchnorm_rewriter.h" #include "tensorflow/compiler/xla/service/buffer_assignment.h" #include "tensorflow/compiler/xla/service/buffer_liveness.h" +#include "tensorflow/compiler/xla/service/flatten_call_graph.h" #include "tensorflow/compiler/xla/service/gpu/convolution_folding.h" #include "tensorflow/compiler/xla/service/gpu/copy_insertion.h" #include "tensorflow/compiler/xla/service/gpu/fusion_merger.h" @@ -43,6 +45,7 @@ limitations under the License. #include "tensorflow/compiler/xla/service/gpu/partition_assignment.h" #include "tensorflow/compiler/xla/service/gpu/stream_assignment.h" #include "tensorflow/compiler/xla/service/gpu/thunk_schedule.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" @@ -50,9 +53,11 @@ limitations under the License. #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" +#include "tensorflow/compiler/xla/service/hlo_proto_util.h" #include "tensorflow/compiler/xla/service/hlo_subcomputation_unification.h" #include "tensorflow/compiler/xla/service/hlo_verifier.h" #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/transpose_folding.h" #include "tensorflow/compiler/xla/status_macros.h" @@ -78,7 +83,7 @@ const char* kTargetTriple = "nvptx64-nvidia-cuda"; // The data layout of the emitted module. Copied from computeDataLayout in // NVPTXTargetMachine.cpp. -const char* kDataLayout = "e-i64:64-v16:16-v32:32-n16:32:64"; +const char* kDataLayout = "e-i64:64-i128:128-v16:16-v32:32-n16:32:64"; // Any address of a variable residing in global memory or returned by one of the // memory allocation routines from the driver or runtime API is always aligned @@ -91,11 +96,9 @@ constexpr int64 kMemoryAlignment = 256; // called in GpuCompiler's constructor, so can't return an error. But // GpuCompiler::Compile will return an error when the wanted libdevice file // doesn't exist in the folder this function returns. -string GetLibdeviceDir() { +string GetLibdeviceDir(const HloModuleConfig& config) { std::vector potential_libdevice_dirs; - // Flag xla_cuda_data_dir specified by the user. - legacy_flags::GpuCompilerFlags* flags = legacy_flags::GetGpuCompilerFlags(); - const string datadir = flags->xla_cuda_data_dir; + const string datadir = config.debug_options().xla_gpu_cuda_data_dir(); if (!datadir.empty()) { potential_libdevice_dirs.push_back(datadir); } @@ -117,15 +120,27 @@ string GetLibdeviceDir() { } // Runs optimization passes on the given HLO module. -tensorflow::Status OptimizeHloModule(HloModule* hlo_module, - const Compiler::HloDumper& dump_hlo, - const se::DeviceDescription& device_desc) { +tensorflow::Status OptimizeHloModule( + HloModule* hlo_module, const se::DeviceDescription& device_desc, + const HloCostAnalysis::ShapeSizeFunction& shape_size_function) { { - HloPassPipeline pipeline("optimization", dump_hlo); - pipeline.AddInvariantChecker(); + HloPassPipeline pipeline("optimization"); + pipeline.AddInvariantChecker(shape_size_function); + ReducePrecisionInsertion::AddPasses( + &pipeline, hlo_module->config().debug_options(), + ReducePrecisionInsertion::PassTiming::BEFORE_OPTIMIZATION); { - auto& pass = pipeline.AddPass>( - "simplification", dump_hlo); + auto& pass = + pipeline.AddPass>("simplification"); + pass.AddInvariantChecker(shape_size_function); + + // TODO(b/62764704): Do not rewrite on GPU, use cuDNN's BatchNorm APIs + // instead. + pass.AddPass( + /*rewrite_training_op=*/true, + /*rewrite_inference_op=*/true, + /*rewrite_grad_op=*/true, + /*use_fusion=*/false); pass.AddPass( /*is_layout_sensitive=*/false, [](const Shape&, const Shape&) { return false; }); @@ -133,36 +148,57 @@ tensorflow::Status OptimizeHloModule(HloModule* hlo_module, pass.AddPass(); } pipeline.AddPass(); - pipeline.AddPass(ImplementedAsGemm); - pipeline.AddPass(); + pipeline.AddPass( + [](const HloInstruction& dot, + const TransposeFolding::OperandIndices& candidate_operands) { + return ImplementedAsGemm(dot) ? candidate_operands + : TransposeFolding::OperandIndices{}; + }, + TransposeFolding::NeverFoldTranspose); pipeline.AddPass(/*is_layout_sensitive=*/false); pipeline.AddPass(); TF_RETURN_IF_ERROR(pipeline.Run(hlo_module).status()); } { - HloPassFix fusion("fusion", dump_hlo); + HloPassFix fusion("fusion"); + fusion.AddInvariantChecker(shape_size_function); fusion.AddPass(/*may_duplicate=*/false); fusion.AddPass(/*may_duplicate=*/true); fusion.AddPass(); - return fusion.Run(hlo_module).status(); + TF_RETURN_IF_ERROR(fusion.Run(hlo_module).status()); + + HloPassPipeline reduce_pipeline("reduce-precision"); + reduce_pipeline.AddInvariantChecker(shape_size_function); + ReducePrecisionInsertion::AddPasses( + &reduce_pipeline, hlo_module->config().debug_options(), + ReducePrecisionInsertion::PassTiming::AFTER_FUSION); + StatusOr reduce_result = reduce_pipeline.Run(hlo_module); + TF_RETURN_IF_ERROR(reduce_result.status()); + + if (reduce_result.ValueOrDie()) { + // Do another fusion pass, with the expectation that we may be able to + // fuse the new ReducePrecision operations. + TF_RETURN_IF_ERROR(fusion.Run(hlo_module).status()); + } } + return tensorflow::Status::OK(); } // Modifies the given HLO module so that it will be accepted by IrEmitter. // Unlike optimization passes, the passes are necessary for correctness. tensorflow::Status PrepareHloModuleForIrEmitting( - const Compiler::HloDumper& dump_hlo, HloModule* hlo_module, - HloModuleConfig* module_config) { + HloModule* hlo_module, + const HloCostAnalysis::ShapeSizeFunction& shape_size_function) { // In some cases, we have to place the result of an instruction in a temporary // buffer. For instance, the buffer that holds an external parameter is // assumed immutable at this point, and should not be reused for output // (b/27180329). Therefore, in that case, we set the output to be a copy of // the parameter. - HloPassPipeline pipeline("GPU-ir-emit-prepare", dump_hlo); - pipeline.AddInvariantChecker(); + HloPassPipeline pipeline("GPU-ir-emit-prepare"); + pipeline.AddInvariantChecker(shape_size_function); pipeline.AddPass(); pipeline.AddPass( - module_config->mutable_entry_computation_layout()); + hlo_module->mutable_entry_computation_layout()); // The LayoutAssignment pass may leave behind kCopy instructions which are // duplicate or NOPs, so remove them with algebraic simplification and CSE. pipeline.AddPass>( @@ -172,16 +208,20 @@ tensorflow::Status PrepareHloModuleForIrEmitting( // Copy insertion should be performed immediately before IR emission to avoid // inserting unnecessary copies (later pass adds an instruction which // materializes the value) or missing a necessary copy (later pass removes an - // instruction which materializes a value). + // instruction which materializes a value). DCE must be run immediately before + // (and sometime after) copy insertion, to avoid dead code from interfering + // with the rewrites. + pipeline.AddPass(); pipeline.AddPass(); pipeline.AddPass(); + pipeline.AddPass(); return pipeline.Run(hlo_module).status(); } // Invokes the ptxas tool on the given PTX string, and dumps its output. void DumpPtxasInfo(const string& ptx) { - legacy_flags::GpuCompilerFlags* flags = legacy_flags::GetGpuCompilerFlags(); - const string ptxas_path = flags->xla_ptxas_path; + const string ptxas_path = + tensorflow::io::JoinPath(tensorflow::CudaRoot(), "bin/ptxas"); // Do not log PTX stats if ptxas is not found at the given path. if (!tensorflow::Env::Default()->FileExists(ptxas_path).ok()) { LOG(WARNING) @@ -218,19 +258,17 @@ void DumpPtxasInfo(const string& ptx) { } // namespace GpuCompiler::GpuCompiler() - : libdevice_dir_(GetLibdeviceDir()), - pointer_size_(llvm::DataLayout(kDataLayout).getPointerSize()) {} + : pointer_size_(llvm::DataLayout(kDataLayout).getPointerSize()) {} StatusOr> GpuCompiler::Compile( - std::unique_ptr hlo_module, - std::unique_ptr module_config, HloDumper dump_hlo, - se::StreamExecutor* stream_exec) { + std::unique_ptr module, se::StreamExecutor* stream_exec) { TF_RET_CHECK(stream_exec != nullptr); - TF_RETURN_IF_ERROR(OptimizeHloModule(hlo_module.get(), dump_hlo, - stream_exec->GetDeviceDescription())); - TF_RETURN_IF_ERROR(PrepareHloModuleForIrEmitting(dump_hlo, hlo_module.get(), - module_config.get())); + TF_RETURN_IF_ERROR(OptimizeHloModule(module.get(), + stream_exec->GetDeviceDescription(), + ShapeSizeBytesFunction())); + TF_RETURN_IF_ERROR( + PrepareHloModuleForIrEmitting(module.get(), ShapeSizeBytesFunction())); llvm::LLVMContext llvm_context; std::string buffer; @@ -243,7 +281,7 @@ StatusOr> GpuCompiler::Compile( }; llvm_context.setDiagnosticHandler(DiagnosticHandler, &printer); - llvm::Module llvm_module(hlo_module->name().c_str(), llvm_context); + llvm::Module llvm_module(module->name().c_str(), llvm_context); // Set the target triple and the data layout. llvm_module.setTargetTriple(kTargetTriple); llvm_module.setDataLayout(kDataLayout); @@ -251,36 +289,46 @@ StatusOr> GpuCompiler::Compile( // Determine the HLO schedule, which is an ordering of HLO instructions. This // is used by buffer assignment to enable buffer reuse, and the same ordering // must also be used to determine the thunk launch schedule. - std::unique_ptr stream_assignment = - AssignStreams(*hlo_module); + std::unique_ptr stream_assignment = AssignStreams(*module); TF_ASSIGN_OR_RETURN( std::unique_ptr hlo_schedule, - HloSchedule::Build(*hlo_module, *stream_assignment, pointer_size_)); + HloSchedule::Build(*module, *stream_assignment, pointer_size_)); // Run buffer analysis on the HLO graph. This analysis figures out which // temporary buffers are required to run the computation. TF_ASSIGN_OR_RETURN( std::unique_ptr buffer_assignment, - BufferAssigner::Run(hlo_module.get(), hlo_schedule->ConsumeHloOrdering(), - [this](const LogicalBuffer& buffer) { - return ShapeSizeBytes(buffer.shape()); - }, - kMemoryAlignment)); + BufferAssigner::Run(module.get(), hlo_schedule->ConsumeHloOrdering(), + BufferSizeBytesFunction(), [](LogicalBuffer::Color) { + return kMemoryAlignment; + })); + + const string dump_debug_json_to = + module->config().debug_options().xla_dump_debug_json_to(); + if (!dump_debug_json_to.empty()) { + HloProto proto = MakeHloProto(*module, *buffer_assignment); + TF_RETURN_IF_ERROR(protobuf_util::DumpJsonToDirectory( + proto, dump_debug_json_to, module->name())); + } - IrEmitterContext ir_emitter_context(hlo_module.get(), buffer_assignment.get(), + IrEmitterContext ir_emitter_context(module.get(), buffer_assignment.get(), &stream_exec->GetDeviceDescription(), &llvm_module); - HloComputation* entry_computation = hlo_module->entry_computation(); - IrEmitterUnnested ir_emitter(*module_config, entry_computation, - module_config->has_hybrid_result(), + HloComputation* entry_computation = module->entry_computation(); + IrEmitterUnnested ir_emitter(module->config(), entry_computation, + module->config().has_hybrid_result(), &ir_emitter_context); TF_RETURN_IF_ERROR( entry_computation->root_instruction()->Accept(&ir_emitter)); + if (user_pre_optimization_hook_) { + TF_CHECK_OK(user_pre_optimization_hook_(llvm_module)); + } string ir_module_string_before_opt; - legacy_flags::GpuCompilerFlags* flags = legacy_flags::GetGpuCompilerFlags(); - if (VLOG_IS_ON(2) || flags->xla_gpu_embed_ir) { + const bool embed_ir_in_executable = + module->config().debug_options().xla_embed_ir_in_executable(); + if (VLOG_IS_ON(2) || 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); @@ -301,9 +349,16 @@ StatusOr> GpuCompiler::Compile( cc_major = 2; cc_minor = 0; } + if (libdevice_dir_.empty()) { + // Compute libdevice_dir_ just once and cache it in this member. + libdevice_dir_ = GetLibdeviceDir(module->config()); + } TF_ASSIGN_OR_RETURN(*ptx, CompileToPtx(&llvm_module, {cc_major, cc_minor}, - *module_config, libdevice_dir_)); + module->config(), libdevice_dir_)); + 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:"; @@ -319,9 +374,9 @@ StatusOr> GpuCompiler::Compile( XLA_VLOG_LINES(2, thunk_schedule->ToString()); auto* gpu_executable = - new GpuExecutable(*ptx, std::move(thunk_schedule), std::move(hlo_module), - std::move(module_config), std::move(buffer_assignment)); - if (flags->xla_gpu_embed_ir) { + new GpuExecutable(*ptx, std::move(thunk_schedule), std::move(module), + std::move(buffer_assignment), ShapeSizeBytesFunction()); + if (embed_ir_in_executable) { DCHECK_NE("", ir_module_string_before_opt); gpu_executable->set_ir_module_string(ir_module_string_before_opt); } @@ -329,18 +384,15 @@ StatusOr> GpuCompiler::Compile( } StatusOr>> GpuCompiler::Compile( - std::vector> hlo_modules, - std::vector> module_configs, - HloDumper dump_hlos, std::vector stream_execs) { + std::vector> modules, + std::vector stream_execs) { return Unimplemented( "Compilation of multiple HLO modules is not yet supported on GPU."); } StatusOr>> -GpuCompiler::CompileAheadOfTime( - std::vector> module, - std::vector> module_config, - HloDumper dump_hlo, const AotCompilationOptions& options) { +GpuCompiler::CompileAheadOfTime(std::vector> module, + const AotCompilationOptions& options) { return Unimplemented("not yet implemented: GpuCompiler::CompileAheadOfTime"); } @@ -348,10 +400,6 @@ se::Platform::Id GpuCompiler::PlatformId() const { return se::cuda::kCudaPlatformId; } -int64 GpuCompiler::ShapeSizeBytes(const Shape& shape) const { - return ShapeUtil::ByteSizeOf(shape, pointer_size_); -} - } // namespace gpu } // namespace xla diff --git a/tensorflow/compiler/xla/service/gpu/gpu_compiler.h b/tensorflow/compiler/xla/service/gpu/gpu_compiler.h index 22f492b42294838bf323b70f492d83fa9c7b4ce2..e8073935990938c9ecf0d835066c4c490c7cc2c4 100644 --- a/tensorflow/compiler/xla/service/gpu/gpu_compiler.h +++ b/tensorflow/compiler/xla/service/gpu/gpu_compiler.h @@ -20,10 +20,9 @@ limitations under the License. #include #include -#include "tensorflow/compiler/xla/service/compiler.h" #include "tensorflow/compiler/xla/service/executable.h" #include "tensorflow/compiler/xla/service/hlo_module.h" -#include "tensorflow/compiler/xla/service/hlo_module_config.h" +#include "tensorflow/compiler/xla/service/llvm_compiler.h" #include "tensorflow/compiler/xla/statusor.h" #include "tensorflow/compiler/xla/types.h" #include "tensorflow/core/lib/gtl/array_slice.h" @@ -36,35 +35,36 @@ namespace xla { namespace gpu { // The GPU compiler generates efficient GPU executables. -class GpuCompiler : public Compiler { +class GpuCompiler : public LLVMCompiler { public: GpuCompiler(); ~GpuCompiler() override {} StatusOr> Compile( - std::unique_ptr hlo_module, - std::unique_ptr module_config, HloDumper dump_hlo, + std::unique_ptr module, perftools::gputools::StreamExecutor* stream_exec) override; StatusOr>> Compile( - std::vector> hlo_module, - std::vector> module_config, - HloDumper dump_hlo, + std::vector> modules, std::vector stream_exec) override; StatusOr>> - CompileAheadOfTime( - std::vector> module, - std::vector> module_config, - HloDumper dump_hlo, AotCompilationOptions const& options) override; + CompileAheadOfTime(std::vector> module, + AotCompilationOptions const& options) override; perftools::gputools::Platform::Id PlatformId() const override; - int64 ShapeSizeBytes(const Shape& shape) const override; + HloCostAnalysis::ShapeSizeFunction ShapeSizeBytesFunction() const override { + // Capture just the pointer size, not the entire GpuCompiler object. + int64 pointer_size = pointer_size_; + return [pointer_size](const Shape& shape) { + return ShapeUtil::ByteSizeOf(shape, pointer_size); + }; + } private: // The parent directory of libdevice IR libraries. - const string libdevice_dir_; + string libdevice_dir_; // The list of PTX strings generated by this GpuCompiler. We let GpuCompiler // to own them because they need to be alive across the life span of the @@ -72,7 +72,7 @@ class GpuCompiler : public Compiler { tensorflow::mutex mutex_; std::vector> generated_ptxes_ GUARDED_BY(mutex_); - // The size in bytes of a pointer. Used for computing ShapeSizeBytes. + // The size in bytes of a pointer. Used by ShapeSizeBytesFunction. int64 pointer_size_; TF_DISALLOW_COPY_AND_ASSIGN(GpuCompiler); diff --git a/tensorflow/compiler/xla/service/gpu/gpu_executable.cc b/tensorflow/compiler/xla/service/gpu/gpu_executable.cc index 32f0368b4bc523d3d81147a8cbbde745387c21d4..db7f9826d798181b55b5fd6cef4ea749d4fe7d53 100644 --- a/tensorflow/compiler/xla/service/gpu/gpu_executable.cc +++ b/tensorflow/compiler/xla/service/gpu/gpu_executable.cc @@ -107,15 +107,16 @@ class HloExecutionProfiler { // Implementation note: HLO profiling is always enabled for GPU executables, // since we can use timers around thunks. -GpuExecutable::GpuExecutable(tensorflow::StringPiece ptx, - std::unique_ptr thunk_schedule, - std::unique_ptr hlo_module, - std::unique_ptr module_config, - std::unique_ptr assignment) - : Executable(std::move(hlo_module), std::move(module_config)), +GpuExecutable::GpuExecutable( + tensorflow::StringPiece ptx, std::unique_ptr thunk_schedule, + std::unique_ptr hlo_module, + std::unique_ptr assignment, + HloCostAnalysis::ShapeSizeFunction shape_size_function) + : Executable(std::move(hlo_module)), ptx_(ptx), thunk_schedule_(std::move(thunk_schedule)), - assignment_(std::move(assignment)) {} + assignment_(std::move(assignment)), + shape_size_function_(std::move(shape_size_function)) {} Status GpuExecutable::ExecuteThunks( const ServiceExecutableRunOptions* run_options, @@ -228,10 +229,10 @@ StatusOr GpuExecutable::ExecuteOnStream( // The points-to set of the root is unambiguous so it's known statically // which buffers are in the result. Gather these buffers using the root's // points-to set. - TF_RETURN_IF_ERROR(GetRootPointsToSet().ForEachElement( + TF_RETURN_IF_ERROR(GetRootPointsToSet().ForEachElementWithStatus( [&referred_by_output, &buffer_allocations, this]( - const ShapeIndex& /*index*/, bool /*is_leaf*/, - const std::vector& buffers) { + const ShapeIndex& /*index*/, + const PointsToSet::BufferList& buffers) { // The points to set is unambiguous so the set should be a // singleton. That is, we know exactly which instruction produced // the array at this element. @@ -306,11 +307,11 @@ StatusOr> GpuExecutable::ExecuteOnStream( std::set buffers_in_result; TF_RETURN_IF_ERROR( shaped_buffer->mutable_shape_index_to_buffer_entry() - ->ForEachMutableElement( + ->ForEachMutableElementWithStatus( [&buffer_allocations, &buffers_in_result, &shaped_buffer, this]( - const ShapeIndex& index, bool is_leaf, size_t* buffer_entry) { - if (is_leaf) { - const std::vector& sources = + const ShapeIndex& index, size_t* buffer_entry) { + if (ShapeUtil::IsLeafIndex(shaped_buffer->shape(), index)) { + const auto& sources = this->GetRootPointsToSet().element(index); // The points to set is unambiguous so the set should be a // singleton. That is, we know exactly which instruction @@ -356,5 +357,9 @@ const PointsToSet& GpuExecutable::GetRootPointsToSet() const { module().entry_computation()->root_instruction()); } +std::unique_ptr GpuExecutable::CreateCostAnalysis() const { + return MakeUnique(shape_size_function_); +} + } // namespace gpu } // namespace xla diff --git a/tensorflow/compiler/xla/service/gpu/gpu_executable.h b/tensorflow/compiler/xla/service/gpu/gpu_executable.h index e308de79ba582d3497e7f217285ae4b1ed0be1a7..bbf8549fdbcd1017c95b2a6485319f72e91df5c5 100644 --- a/tensorflow/compiler/xla/service/gpu/gpu_executable.h +++ b/tensorflow/compiler/xla/service/gpu/gpu_executable.h @@ -28,7 +28,6 @@ limitations under the License. #include "tensorflow/compiler/xla/service/gpu/thunk_schedule.h" #include "tensorflow/compiler/xla/service/hlo_execution_profile.h" #include "tensorflow/compiler/xla/service/hlo_module.h" -#include "tensorflow/compiler/xla/service/hlo_module_config.h" #include "tensorflow/compiler/xla/service/shaped_buffer.h" #include "tensorflow/compiler/xla/service/tuple_points_to_analysis.h" #include "tensorflow/compiler/xla/statusor.h" @@ -51,8 +50,8 @@ class GpuExecutable : public Executable { GpuExecutable(tensorflow::StringPiece ptx, std::unique_ptr thunk_schedule, std::unique_ptr hlo_module, - std::unique_ptr module_config, - std::unique_ptr assignment); + std::unique_ptr assignment, + HloCostAnalysis::ShapeSizeFunction shape_size_function); // This should be called after set_ir_module_string. const string& ir_module_string() const { return ir_module_string_; } @@ -81,6 +80,13 @@ class GpuExecutable : public Executable { tensorflow::gtl::ArraySlice arguments) override; + const Status EqualOrFail(const Executable& executable) { + // TODO(b/62952745) Implement equality test on GPU executable. + return Unimplemented("Equality test on GPU executable is not implemented."); + } + + std::unique_ptr CreateCostAnalysis() const override; + private: // If `block_host_until_done` is false, execution will not block the host // until the kernels have completed. This is used as an optimization for @@ -115,6 +121,9 @@ class GpuExecutable : public Executable { // memory for every output/temp buffers. const std::unique_ptr assignment_; + // Function to compute the size of a given Shape, in bytes. + HloCostAnalysis::ShapeSizeFunction shape_size_function_; + TF_DISALLOW_COPY_AND_ASSIGN(GpuExecutable); }; diff --git a/tensorflow/compiler/xla/service/gpu/hlo_schedule.cc b/tensorflow/compiler/xla/service/gpu/hlo_schedule.cc index d16a1d4ee5be00e685fc181f19c1a3cfda253f6a..81e905a06665436875b17991a8635e7bb31600de 100644 --- a/tensorflow/compiler/xla/service/gpu/hlo_schedule.cc +++ b/tensorflow/compiler/xla/service/gpu/hlo_schedule.cc @@ -20,6 +20,8 @@ limitations under the License. #include "tensorflow/compiler/xla/service/gpu/hlo_schedule.h" #include "tensorflow/compiler/xla/ptr_util.h" +#include "tensorflow/compiler/xla/service/hlo_reachability.h" +#include "tensorflow/compiler/xla/service/hlo_scheduling.h" #include "tensorflow/compiler/xla/types.h" namespace xla { @@ -67,38 +69,38 @@ GpuHloOrdering::GpuHloOrdering( // waits for its operands before executing. // // The predecessor map is built incrementally, in thunk launch order. We - // record the instructions already visited per stream in - // 'instructions_per_stream'. This lets us quickly determine the same-stream - // predecessors of each instruction. To capture cross-stream dependency edges, - // we use the predecessor map to insert each operand as well as its transitive - // closure of dependencies. - - // Compute the set of all instructions we will want to set reachability on - auto predecessor_map = MakeUnique( + // record the most-recently seen instructions per stream in + // 'last_instruction_per_stream'. This lets us quickly determine the + // same-stream predecessors of each instruction. + + // Compute the set of all instructions we will want to set reachability on. + auto predecessor_map = MakeUnique( module->entry_computation()->MakeInstructionPostOrder()); - std::vector> instructions_per_stream( - stream_assignment.StreamCount()); + // The most recently visited instruction per stream. + std::vector last_instruction_per_stream( + stream_assignment.StreamCount(), nullptr); for (const HloInstruction* hlo : thunk_launch_order) { + predecessor_map->SetReachable(hlo, hlo); if (stream_assignment.HasStreamAssigned(*hlo)) { + // Gather all instruction which are immediate predecessors of 'hlo' in the + // reachability graph. + std::vector immediate_preds; + immediate_preds.insert(immediate_preds.end(), hlo->operands().begin(), + hlo->operands().end()); + immediate_preds.insert(immediate_preds.end(), + hlo->control_predecessors().begin(), + hlo->control_predecessors().end()); + // All ops already queued on the same instruction stream, and their - // transitive predecessors, are predecessors. Since the relation is - // transitive, we just set the transitive closure of the previous op. + // transitive predecessors, are predecessors. const int stream_no = stream_assignment.StreamNumberForHlo(*hlo); - std::vector* instructions = - &instructions_per_stream[stream_no]; - if (!instructions->empty()) { - const HloInstruction* back = instructions->back(); - predecessor_map->SetReachableAndTransitiveClosure(hlo, back); - } - // All operands and their transitive predecessors are predecessors. Each - // operand must already exist in 'predecessor_map', since we're iterating - // in thunk launch order. - for (const HloInstruction* operand : hlo->operands()) { - predecessor_map->SetReachableAndTransitiveClosure(hlo, operand); + if (last_instruction_per_stream[stream_no] != nullptr) { + immediate_preds.push_back(last_instruction_per_stream[stream_no]); } - instructions->push_back(hlo); + predecessor_map->SetReachabilityToUnion(immediate_preds, hlo); + last_instruction_per_stream[stream_no] = hlo; } else { // Only parameters and constants don't have an assigned stream, since they // don't require a thunk. These ops don't have any predecessors. @@ -107,21 +109,21 @@ GpuHloOrdering::GpuHloOrdering( CHECK_EQ(hlo->operand_count(), 0); } } - strict_predecessors_.emplace(module->entry_computation(), - std::move(predecessor_map)); + predecessors_.emplace(module->entry_computation(), + std::move(predecessor_map)); - // The ordering of instructions in subcomputations is based solely on data - // dependencies. I.e. the strict predecessors of each subcomputation - // instruction is its transitive operands. + // The ordering of instructions in subcomputations is based solely on control + // and data dependencies. // // TODO(toddw): Each subcomputation is actually emitted as a function in DFS // postorder, so we can do better and establish the total order here. We don't // do that yet since it's hard to ensure that the order here is the order used // by IrEmitterNested. And mismatched ordering bugs would be hard to find. for (auto& computation : module->computations()) { - if (computation.get() != module->entry_computation()) { - strict_predecessors_.emplace(computation.get(), - computation->ComputeTransitiveOperands()); + if (computation.get() != module->entry_computation() && + !computation->IsFusionComputation()) { + predecessors_.emplace(computation.get(), + computation->ComputeReachability()); } } } diff --git a/tensorflow/compiler/xla/service/gpu/hlo_schedule.h b/tensorflow/compiler/xla/service/gpu/hlo_schedule.h index 773973010a46bb4a2af1f536c43201ba8c0be5d8..1ce7a48ac8fcbbad0b3697845681582fe806b322 100644 --- a/tensorflow/compiler/xla/service/gpu/hlo_schedule.h +++ b/tensorflow/compiler/xla/service/gpu/hlo_schedule.h @@ -19,9 +19,9 @@ limitations under the License. #include #include -#include "tensorflow/compiler/xla/service/buffer_liveness.h" #include "tensorflow/compiler/xla/service/gpu/stream_assignment.h" #include "tensorflow/compiler/xla/service/hlo_module.h" +#include "tensorflow/compiler/xla/service/hlo_ordering.h" #include "tensorflow/compiler/xla/statusor.h" namespace xla { diff --git a/tensorflow/compiler/xla/service/gpu/hlo_schedule_test.cc b/tensorflow/compiler/xla/service/gpu/hlo_schedule_test.cc index c6749851dbb3de56f71e6a700c2ef4a4fec6add6..118ef18c44b85f8298391160398e542329673f97 100644 --- a/tensorflow/compiler/xla/service/gpu/hlo_schedule_test.cc +++ b/tensorflow/compiler/xla/service/gpu/hlo_schedule_test.cc @@ -69,14 +69,14 @@ TEST_F(HloScheduleTest, SequentialMatMul) { HloInstruction* dot2 = builder.AddInstruction( HloInstruction::CreateBinary(f32_2x2_, HloOpcode::kDot, dot1, z)); - HloModule module(TestName()); - module.AddEntryComputation(builder.Build(dot2)); + auto module = CreateNewModule(); + module->AddEntryComputation(builder.Build(dot2)); - std::unique_ptr streams = AssignStreams(module); + std::unique_ptr streams = AssignStreams(*module); EXPECT_EQ(streams->StreamNumberForHlo(*dot1), streams->StreamNumberForHlo(*dot2)); - auto schedule = BuildHloSchedule(module, *streams); + auto schedule = BuildHloSchedule(*module, *streams); // Remove parameters, which are unordered. EXPECT_EQ(RemoveHlo(schedule->ThunkLaunchOrder(), {x, y, z}), HloVec({dot1, dot2})); @@ -129,16 +129,16 @@ TEST_F(HloScheduleTest, SequentialAdd) { HloInstruction* add3 = builder.AddInstruction( HloInstruction::CreateBinary(f32_2x2_, HloOpcode::kAdd, add1, add2)); - HloModule module(TestName()); - module.AddEntryComputation(builder.Build(add3)); + auto module = CreateNewModule(); + module->AddEntryComputation(builder.Build(add3)); - std::unique_ptr streams = AssignStreams(module); + std::unique_ptr streams = AssignStreams(*module); EXPECT_EQ(streams->StreamNumberForHlo(*add1), streams->StreamNumberForHlo(*add2)); EXPECT_EQ(streams->StreamNumberForHlo(*add1), streams->StreamNumberForHlo(*add3)); - auto schedule = BuildHloSchedule(module, *streams); + auto schedule = BuildHloSchedule(*module, *streams); // Remove parameters, which are unordered. EXPECT_EQ(RemoveHlo(schedule->ThunkLaunchOrder(), {x, y, z}), HloVec({add1, add2, add3})); @@ -199,14 +199,14 @@ TEST_F(HloScheduleTest, ConcurrentMatMul) { HloInstruction* add = builder.AddInstruction( HloInstruction::CreateBinary(f32_2x2_, HloOpcode::kAdd, dot1, dot2)); - HloModule module(TestName()); - module.AddEntryComputation(builder.Build(add)); + auto module = CreateNewModule(); + module->AddEntryComputation(builder.Build(add)); - std::unique_ptr streams = AssignStreams(module); + std::unique_ptr streams = AssignStreams(*module); EXPECT_NE(streams->StreamNumberForHlo(*dot1), streams->StreamNumberForHlo(*dot2)); - auto schedule = BuildHloSchedule(module, *streams); + auto schedule = BuildHloSchedule(*module, *streams); // Remove parameters, which are unordered. HloVec thunk_launch_order = RemoveHlo(schedule->ThunkLaunchOrder(), {x, y}); EXPECT_TRUE(thunk_launch_order == HloVec({dot1, dot2, add}) || @@ -254,6 +254,7 @@ TEST_F(HloScheduleTest, LatticeMatMul) { // d40 -- layer 4 HloComputation::Builder builder("entry_computation"); std::vector params; + params.reserve(6); for (int i = 0; i < 6; ++i) { params.push_back(builder.AddInstruction(HloInstruction::CreateParameter( i, f32_2x2_, /*name=*/tensorflow::strings::Printf("param%d", i)))); @@ -277,10 +278,10 @@ TEST_F(HloScheduleTest, LatticeMatMul) { HloInstruction* d40 = builder.AddInstruction( HloInstruction::CreateBinary(f32_2x2_, HloOpcode::kDot, d30, d31)); - HloModule module(TestName()); - module.AddEntryComputation(builder.Build(d40)); + auto module = CreateNewModule(); + module->AddEntryComputation(builder.Build(d40)); - std::unique_ptr streams = AssignStreams(module); + std::unique_ptr streams = AssignStreams(*module); // The two dots on layer 1 are concurrent. EXPECT_NE(streams->StreamNumberForHlo(*d10), streams->StreamNumberForHlo(*d11)); @@ -297,7 +298,7 @@ TEST_F(HloScheduleTest, LatticeMatMul) { // We don't check the thunk launch order, since there are many valid total // orders, and it's annoying to express. - auto schedule = BuildHloSchedule(module, *streams); + auto schedule = BuildHloSchedule(*module, *streams); auto order = schedule->ConsumeHloOrdering(); const HloVec all_params( @@ -392,3 +393,7 @@ TEST_F(HloScheduleTest, LatticeMatMul) { } // namespace gpu } // namespace xla + +int main(int argc, char** argv) { + return xla::ParseDebugOptionsFlagsAndRunTests(argc, argv); +} 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 1a61eec353740202065c1ce98e8c91274facfd19..715c7001c8a5efc564061d419eb1ebfc81bebc60 100644 --- a/tensorflow/compiler/xla/service/gpu/hlo_to_ir_bindings.cc +++ b/tensorflow/compiler/xla/service/gpu/hlo_to_ir_bindings.cc @@ -15,9 +15,9 @@ limitations under the License. #include "tensorflow/compiler/xla/service/gpu/hlo_to_ir_bindings.h" -#include "external/llvm/include/llvm/IR/BasicBlock.h" -#include "external/llvm/include/llvm/IR/Function.h" -#include "external/llvm/include/llvm/IR/Instructions.h" +#include "llvm/IR/BasicBlock.h" +#include "llvm/IR/Function.h" +#include "llvm/IR/Instructions.h" #include "tensorflow/compiler/xla/service/gpu/ir_emission_utils.h" #include "tensorflow/compiler/xla/service/hlo_opcode.h" #include "tensorflow/compiler/xla/service/llvm_ir/llvm_util.h" @@ -86,23 +86,35 @@ void HloToIrBindings::EmitBasePointersForHlos( continue; } - // A non-IO HLO with a buffer is bound to - // (1) an alloca if it is thread-local, or - // (2) an internal pointer in temp_buffer_base according to its offset. - const BufferAllocation::Slice slice = - buffer_assignment_->GetUniqueTopLevelSlice(non_io_hlo) - .ConsumeValueOrDie(); - if (slice.allocation()->is_thread_local()) { - llvm::Type* pointee_type = - llvm_ir::ShapeToIrType(non_io_hlo->shape(), ir_builder_); - BindHloToIrValue(*non_io_hlo, ir_builder_->CreateAlloca(pointee_type)); - } else { - const int64 offset = slice.offset(); - CHECK_NE(nullptr, temp_buffer_base_); - BindHloToIrValue(*non_io_hlo, - ir_builder_->CreateInBoundsGEP( - temp_buffer_base_, ir_builder_->getInt64(offset))); - } + ShapeUtil::ForEachSubshape( + non_io_hlo->shape(), + [&](const Shape& /*subshape*/, const ShapeIndex& index) { + // A non-IO HLO with a buffer is bound to + // (1) an alloca if it is thread-local, or + // (2) an internal pointer in temp_buffer_base according to its + // offset. + auto slice_result = + buffer_assignment_->GetUniqueSlice(non_io_hlo, index); + if (!slice_result.ok()) { + return; + } + const BufferAllocation::Slice slice = + slice_result.ConsumeValueOrDie(); + if (slice.allocation()->is_thread_local()) { + llvm::Type* pointee_type = + llvm_ir::ShapeToIrType(non_io_hlo->shape(), ir_builder_); + BindHloToIrValue(*non_io_hlo, + ir_builder_->CreateAlloca(pointee_type), index); + } else { + const int64 offset = slice.offset(); + CHECK_NE(nullptr, temp_buffer_base_); + BindHloToIrValue( + *non_io_hlo, + ir_builder_->CreateInBoundsGEP(temp_buffer_base_, + ir_builder_->getInt64(offset)), + index); + } + }); } } @@ -112,7 +124,7 @@ llvm::Value* HloToIrBindings::EmitGetTupleElement(const HloInstruction* gte, if (gte->operand(0)->opcode() != HloOpcode::kGetTupleElement) { return llvm_ir::EmitGetTupleElement( gte->shape(), gte->tuple_index(), /*alignment=*/1, - GetTypedIrValue(*gte->operand(0), base_ptr), ir_builder_); + GetTypedIrValue(*gte->operand(0), {}, base_ptr), ir_builder_); } return llvm_ir::EmitGetTupleElement( gte->shape(), gte->tuple_index(), /*alignment=*/1, @@ -120,8 +132,10 @@ llvm::Value* HloToIrBindings::EmitGetTupleElement(const HloInstruction* gte, } llvm::Value* HloToIrBindings::GetTypedIrValue(const HloInstruction& hlo, + const ShapeIndex& shape_index, llvm::Value* ir_value) { - llvm::Type* pointee_type = llvm_ir::ShapeToIrType(hlo.shape(), ir_builder_); + llvm::Type* pointee_type = llvm_ir::ShapeToIrType( + ShapeUtil::GetSubshape(hlo.shape(), shape_index), ir_builder_); llvm::Type* dest_type = pointee_type->getPointerTo(); llvm::Value* typed_ir_value; @@ -139,13 +153,24 @@ llvm::Value* HloToIrBindings::GetTypedIrValue(const HloInstruction& hlo, } void HloToIrBindings::BindHloToIrValue(const HloInstruction& hlo, - llvm::Value* ir_value) { + llvm::Value* ir_value, + const ShapeIndex& shape_index) { VLOG(2) << "Binding " << hlo.ToString(); - InsertOrDie(&base_ptrs_, &hlo, GetTypedIrValue(hlo, ir_value)); + + const Shape& hlo_shape = hlo.shape(); + llvm::Value* typed_ir_value = GetTypedIrValue(hlo, shape_index, ir_value); + + if (!BoundToIrValue(hlo)) { + // Set the root of ShapeTree first before assigning the element ir value. + InsertOrDie(&base_ptrs_, &hlo, ShapeTree(hlo_shape, nullptr)); + } + *(base_ptrs_[&hlo].mutable_element(shape_index)) = typed_ir_value; } -llvm_ir::IrArray HloToIrBindings::GetIrArray(const HloInstruction& hlo) { - llvm_ir::IrArray ir_array(GetBasePointer(hlo), hlo.shape()); +llvm_ir::IrArray HloToIrBindings::GetIrArray(const HloInstruction& hlo, + const ShapeIndex& shape_index) { + llvm_ir::IrArray ir_array(GetBasePointer(hlo, shape_index), + ShapeUtil::GetSubshape(hlo.shape(), shape_index)); alias_analysis_.AddAliasingInformationToIrArray(hlo, &ir_array); return ir_array; } @@ -154,7 +179,7 @@ void HloToIrBindings::UnbindAllLocalIrValues() { std::vector hlos_to_unbind; for (auto& key_value : base_ptrs_) { if (!llvm::isa( - key_value.second->stripPointerCasts())) { + (key_value.second.element({}))->stripPointerCasts())) { hlos_to_unbind.push_back(key_value.first); } } 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 5be2150801fbd2a3a624d9c87513d5cee7288bbd..d43e09e8a8c5cc2efcd8e1fbf9a7c0697e24d73c 100644 --- a/tensorflow/compiler/xla/service/gpu/hlo_to_ir_bindings.h +++ b/tensorflow/compiler/xla/service/gpu/hlo_to_ir_bindings.h @@ -18,8 +18,8 @@ limitations under the License. #include -#include "external/llvm/include/llvm/IR/IRBuilder.h" -#include "external/llvm/include/llvm/IR/Value.h" +#include "llvm/IR/IRBuilder.h" +#include "llvm/IR/Value.h" #include "tensorflow/compiler/xla/map_util.h" #include "tensorflow/compiler/xla/service/buffer_assignment.h" #include "tensorflow/compiler/xla/service/hlo_instruction.h" @@ -48,7 +48,8 @@ class HloToIrBindings { tensorflow::gtl::ArraySlice non_io_hlos); // Rebinds the given HLO to the LLVM IR value that represent its address. - void BindHloToIrValue(const HloInstruction& hlo, llvm::Value* ir_value); + void BindHloToIrValue(const HloInstruction& hlo, llvm::Value* ir_value, + const ShapeIndex& shape_index = {}); // Unbinds all IR values that's defined in an LLVM function, e.g., function // arguments and stack variables. Global variables will be kept in bindings_. @@ -64,15 +65,18 @@ class HloToIrBindings { llvm::Value* GetTempBufferBase() const { return temp_buffer_base_; } - // A helper method that returns the base pointer of the IrArray for "inst". - llvm::Value* GetBasePointer(const HloInstruction& hlo) const { + // A helper method that returns the base pointer of the IrArray containing the + // output of "inst".at the given ShapeIndex. + llvm::Value* GetBasePointer(const HloInstruction& hlo, + const ShapeIndex& shape_index = {}) const { auto it = base_ptrs_.find(&hlo); CHECK(it != base_ptrs_.end()); - return it->second; + return it->second.element(shape_index); } // Return the underlying IrArray of the output of the given instruction. - llvm_ir::IrArray GetIrArray(const HloInstruction& hlo); + llvm_ir::IrArray GetIrArray(const HloInstruction& hlo, + const ShapeIndex& shape_index = {}); private: // Emits IR to resolve (possibly) recursive GetTupleElement instructions. @@ -81,6 +85,7 @@ class HloToIrBindings { // Returns an llvm typed ir representation of 'ir_value' based on 'hlo' shape. llvm::Value* GetTypedIrValue(const HloInstruction& hlo, + const ShapeIndex& shape_index, llvm::Value* ir_value); const BufferAssignment* buffer_assignment_; @@ -90,7 +95,10 @@ class HloToIrBindings { llvm::IRBuilder<>* ir_builder_; // Stores the underlying llvm::IrArray for each HloInstruction. - std::unordered_map base_ptrs_; + // 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_; // The address of the memory block that contains all temporary buffers. llvm::Value* temp_buffer_base_; diff --git a/tensorflow/compiler/xla/service/gpu/infeed_manager.cc b/tensorflow/compiler/xla/service/gpu/infeed_manager.cc new file mode 100644 index 0000000000000000000000000000000000000000..ee5b447c9cd0b1fde4d3a0943d5d4cb8cc5b3376 --- /dev/null +++ b/tensorflow/compiler/xla/service/gpu/infeed_manager.cc @@ -0,0 +1,99 @@ +/* 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/gpu/infeed_manager.h" + +#include "tensorflow/compiler/xla/map_util.h" +#include "tensorflow/compiler/xla/ptr_util.h" +#include "tensorflow/core/platform/logging.h" + +namespace se = ::perftools::gputools; + +namespace xla { +namespace gpu { + +InfeedManager::InfeedManager() : host_to_device_executor_(nullptr) {} + +void InfeedManager::Reset() { + tensorflow::mutex_lock l(mu_); + CHECK(dequeued_buffer_.empty()); + for (auto buffer : enqueued_buffer_) { + buffer->Done(); + } + enqueued_buffer_.clear(); +} + +void InfeedManager::EnqueueBuffers(const std::vector& buffers) { + tensorflow::mutex_lock l(mu_); + bool was_empty = enqueued_buffer_.empty(); + for (gpu::InfeedBuffer* b : buffers) { + enqueued_buffer_.push_back(b); + } + if (was_empty) { + // This has the potential to suffer from the notified thread + // immediately trying and failing to acquire mu_, but seems + // preferable to the alternative of notifying outside the lock + // on every enqueue. + cv_.notify_one(); + } +} + +InfeedBuffer* InfeedManager::BlockingDequeueBuffer() { + tensorflow::mutex_lock l(mu_); + while (enqueued_buffer_.empty()) { + cv_.wait(l); + } + InfeedBuffer* current_buffer = enqueued_buffer_.front(); + enqueued_buffer_.pop_front(); + dequeued_buffer_.insert(current_buffer); + return current_buffer; +} + +void InfeedManager::ReleaseBuffers(const std::vector& buffers) { + { + tensorflow::mutex_lock l(mu_); + for (gpu::InfeedBuffer* b : buffers) { + CHECK(ContainsKey(dequeued_buffer_, b)); + dequeued_buffer_.erase(b); + } + } + for (gpu::InfeedBuffer* b : buffers) { + b->Done(); + } +} + +se::Stream* InfeedManager::GetStream(se::StreamExecutor* executor) { + if (host_to_device_executor_ == nullptr) { + host_to_device_executor_ = executor; + host_to_device_stream_ = MakeUnique(executor); + host_to_device_stream_->Init(); + } + + if (executor != host_to_device_executor_) { + // The requested executor must be the same as the one for which + // the stream is cached. + return nullptr; + } + + return host_to_device_stream_.get(); +} + +InfeedManager* GetOrCreateInfeedManager() { + static InfeedManager* manager = new InfeedManager; + return manager; +} + +} // namespace gpu +} // namespace xla diff --git a/tensorflow/compiler/xla/service/gpu/infeed_manager.h b/tensorflow/compiler/xla/service/gpu/infeed_manager.h new file mode 100644 index 0000000000000000000000000000000000000000..73d5a5ce35497f156a181371bfb97fc37a8eb09e --- /dev/null +++ b/tensorflow/compiler/xla/service/gpu/infeed_manager.h @@ -0,0 +1,136 @@ +/* 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 header declares classes for the infeed manager and the infeed +// buffer that are used by the GPU runtime to transfer buffers into an +// executing GPU computation, e.g., to feed data into a while loop. + +#ifndef TENSORFLOW_COMPILER_XLA_SERVICE_GPU_INFEED_MANAGER_H_ +#define TENSORFLOW_COMPILER_XLA_SERVICE_GPU_INFEED_MANAGER_H_ + +#include + +#include "tensorflow/compiler/xla/types.h" +#include "tensorflow/core/lib/gtl/flatset.h" +#include "tensorflow/core/platform/mutex.h" +#include "tensorflow/core/platform/stream_executor_no_cuda.h" + +namespace xla { +namespace gpu { + +// TODO(b/30467474) Once GPU infeed implementation settles, consider +// folding back the cpu and gpu infeed implementations into a generic +// one if possible. +// +// Current limitations: +// * Does not handle multiple devices/replicas. +// +// * Buffer space on GPU is allocated on every infeed enqueue request, +// and it does not handle the case when it runs out of +// memory. Potential solution is to pre-allocate a fixed amount of +// memory and block when that memory is full. + +// Defines an infeed buffer that is passed to the runtime by +// the client. The client manages the memory of the buffer. +class InfeedBuffer { + public: + InfeedBuffer(perftools::gputools::StreamExecutor* executor, int64 length) + : executor_(executor), length_(length) { + device_memory_ = executor_->AllocateArray(length); + CHECK(!device_memory_.is_null()); + } + + ~InfeedBuffer() { executor_->Deallocate(&device_memory_); } + + int64 length() const { return length_; } + + // Callback to signal that this buffer is consumed. This helps the + // client to manage memory for the infeed buffers. + void Done() { delete this; } + + perftools::gputools::DeviceMemoryBase* device_memory() { + return &device_memory_; + } + + private: + perftools::gputools::StreamExecutor* executor_; // Not owned. + const int64 length_; + perftools::gputools::DeviceMemoryBase device_memory_; +}; + +// Client-side class used to enqueue infeed buffers. +class InfeedManager { + public: + InfeedManager(); + + // Calls the completion callback for any enqueued buffers that have + // not been dequeued by the runtime, and empties the infeed + // queue. Reset may not be called while a runtime computation is + // processing a dequeued buffer. The only safe way to ensure this + // condition is to call Reset when no computation is taking place. + void Reset(); + + // Adds a set of buffers to the infeed queue atomically. buffer->Done + // will be called when the buffer will no longer be accessed by the + // InfeedManager, either as a result of a call to Reset or because the + // runtime has dequeued and used the buffer. + void EnqueueBuffers(const std::vector& buffers); + + // Blocks until the infeed queue is non-empty, then returns the + // buffer at the head of the queue. Adds the current buffer to the + // to-be released set. + InfeedBuffer* BlockingDequeueBuffer(); + + // Releases a set of buffers from the to-be released set. + void ReleaseBuffers(const std::vector& buffers); + + // Returns a cached stream associated with an executor. Allocates a + // new stream on the first invocation. On subsequent invocations, if + // the cached executor is not the same as the requested executor, + // returns null. + perftools::gputools::Stream* GetStream( + perftools::gputools::StreamExecutor* executor); + + private: + // TODO(b/30467474): Revisit if this mutex becomes a point of + // contention. + tensorflow::mutex mu_; + + // Condition variable that is signaled every time a buffer is + // enqueued to an empty queue. + tensorflow::condition_variable cv_; + + // InfeedBuffer* queue contents are not owned, but buffer->Done must + // be called when the buffer is no longer needed by the runtime. + std::deque enqueued_buffer_; + + // Buffers that are dequeued and currently being processed by the + // runtime. Not owned. + tensorflow::gtl::FlatSet dequeued_buffer_; + + // Cached host to device stream for queuing infeed data. + std::unique_ptr host_to_device_stream_; + + // Executor that the host_to_device_stream belongs to. Not owned. + perftools::gputools::StreamExecutor* host_to_device_executor_; +}; + +// Singleton creator-or-accessor: Returns the GPU infeed manager. +InfeedManager* GetOrCreateInfeedManager(); + +} // namespace gpu +} // namespace xla + +#endif // TENSORFLOW_COMPILER_XLA_SERVICE_GPU_INFEED_MANAGER_H_ diff --git a/tensorflow/compiler/xla/service/gpu/infeed_thunk.cc b/tensorflow/compiler/xla/service/gpu/infeed_thunk.cc new file mode 100644 index 0000000000000000000000000000000000000000..e33e904692ca5ad41e17d2e165dbb40b6bd4aa33 --- /dev/null +++ b/tensorflow/compiler/xla/service/gpu/infeed_thunk.cc @@ -0,0 +1,81 @@ +/* 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/gpu/infeed_manager.h" +#include "tensorflow/compiler/xla/service/gpu/infeed_thunk.h" +#include "tensorflow/compiler/xla/util.h" +#include "tensorflow/core/platform/stream_executor_no_cuda.h" + +namespace xla { +namespace gpu { + +InfeedThunk::InfeedThunk( + tensorflow::gtl::ArraySlice tuple_element_buffers, + const BufferAllocation::Slice& destination_buffer, + const HloInstruction* hlo_instruction) + : Thunk(Kind::kInfeed, hlo_instruction), + tuple_element_buffers_(tuple_element_buffers.begin(), + tuple_element_buffers.end()), + destination_buffer_(destination_buffer) {} + +tensorflow::Status InfeedThunk::ExecuteOnStream( + const BufferAllocations& buffer_allocations, + perftools::gputools::Stream* stream) { + VLOG(2) << "Infeeding to GPU "; + + perftools::gputools::DeviceMemoryBase destination_address = + buffer_allocations.GetDeviceAddress(destination_buffer_); + + InfeedManager* infeed_manager = GetOrCreateInfeedManager(); + std::vector infeed_buffers; + if (ShapeUtil::IsTuple(hlo_instruction()->shape())) { + CHECK(!ShapeUtil::IsNestedTuple(hlo_instruction()->shape())); + // Transfer the tuple elements first. + std::vector tuple_element_addresses; + for (BufferAllocation::Slice tuple_element_buffer : + tuple_element_buffers_) { + perftools::gputools::DeviceMemoryBase tuple_element_address = + buffer_allocations.GetDeviceAddress(tuple_element_buffer); + + InfeedBuffer* buffer = infeed_manager->BlockingDequeueBuffer(); + infeed_buffers.push_back(buffer); + stream->ThenMemcpy(&tuple_element_address, *(buffer->device_memory()), + buffer->length()); + tuple_element_addresses.push_back(tuple_element_address.opaque()); + } + // Transfer the tuple outer buffer. + auto host_size = tuple_element_addresses.size() * sizeof(void*); + stream->ThenMemcpy(&destination_address, tuple_element_addresses.data(), + host_size); + } else { + InfeedBuffer* buffer = infeed_manager->BlockingDequeueBuffer(); + infeed_buffers.push_back(buffer); + stream->ThenMemcpy(&destination_address, *(buffer->device_memory()), + buffer->length()); + } + + if (!stream->BlockHostUntilDone()) { + return InternalError("Failed to complete data transfer on stream %p", + stream); + } + + infeed_manager->ReleaseBuffers(infeed_buffers); + + VLOG(2) << "Infeeding to GPU complete"; + return tensorflow::Status::OK(); +} + +} // namespace gpu +} // namespace xla diff --git a/tensorflow/compiler/xla/service/gpu/infeed_thunk.h b/tensorflow/compiler/xla/service/gpu/infeed_thunk.h new file mode 100644 index 0000000000000000000000000000000000000000..371d71f9dbdd21cb5f36cc3108c8f398a4a91c29 --- /dev/null +++ b/tensorflow/compiler/xla/service/gpu/infeed_thunk.h @@ -0,0 +1,58 @@ +/* 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_GPU_INFEED_THUNK_H_ +#define TENSORFLOW_COMPILER_XLA_SERVICE_GPU_INFEED_THUNK_H_ + +#include "tensorflow/compiler/xla/service/buffer_assignment.h" +#include "tensorflow/compiler/xla/service/gpu/buffer_allocations.h" +#include "tensorflow/compiler/xla/service/gpu/thunk.h" +#include "tensorflow/compiler/xla/service/hlo_instruction.h" +#include "tensorflow/core/platform/stream_executor_no_cuda.h" +#include "tensorflow/core/platform/types.h" + +namespace xla { +namespace gpu { + +// A thunk that infeeds data. Data must be already resident on the +// device. This thunk performs an intra-device copy from that location +// to the buffer allocated for the infeed op. +class InfeedThunk : public Thunk { + public: + // Constructs a InfeedThunk that copies data from the on-device + // infeed queue to the device buffer + // `destination_buffer`. `mem_size` is the size of the data in + // bytes. + InfeedThunk(tensorflow::gtl::ArraySlice + tuple_element_buffers, + const BufferAllocation::Slice& destination_buffer, + const HloInstruction* hlo_instruction); + + InfeedThunk(const InfeedThunk&) = delete; + InfeedThunk& operator=(const InfeedThunk&) = delete; + + tensorflow::Status ExecuteOnStream( + const BufferAllocations& buffer_allocations, + perftools::gputools::Stream* stream) override; + + private: + const std::vector tuple_element_buffers_; + const BufferAllocation::Slice destination_buffer_; +}; + +} // namespace gpu +} // namespace xla + +#endif // TENSORFLOW_COMPILER_XLA_SERVICE_GPU_INFEED_THUNK_H_ diff --git a/tensorflow/compiler/xla/service/gpu/instruction_fusion.cc b/tensorflow/compiler/xla/service/gpu/instruction_fusion.cc index 34a44ad40548272a0c2a87efadfa1ab2aca7b979..a36dcbbd2faf3258ec2790f51bb2aec3ce834a6c 100644 --- a/tensorflow/compiler/xla/service/gpu/instruction_fusion.cc +++ b/tensorflow/compiler/xla/service/gpu/instruction_fusion.cc @@ -46,6 +46,11 @@ bool GpuInstructionFusion::ShouldFuse(HloInstruction* consumer, int64 operand_index) { HloInstruction* producer = consumer->mutable_operand(operand_index); + // Output fusion is not currently supported on GPUs. + if (producer->opcode() == HloOpcode::kFusion) { + return false; + } + // RNG operations are not currently parallel-friendly on GPU. if (producer->opcode() == HloOpcode::kRng) { return false; diff --git a/tensorflow/compiler/xla/service/gpu/instruction_fusion_test.cc b/tensorflow/compiler/xla/service/gpu/instruction_fusion_test.cc index c58af04bad081060d5d4f6a426d07ff207f9fef3..896f6ea84252a34b254f4468bd55284ec0432243 100644 --- a/tensorflow/compiler/xla/service/gpu/instruction_fusion_test.cc +++ b/tensorflow/compiler/xla/service/gpu/instruction_fusion_test.cc @@ -16,7 +16,6 @@ limitations under the License. #include "tensorflow/compiler/xla/service/gpu/instruction_fusion.h" #include "tensorflow/compiler/xla/tests/hlo_test_base.h" -#include "tensorflow/core/platform/test.h" namespace xla { namespace gpu { @@ -32,7 +31,7 @@ TEST_F(InstructionFusionTest, PotentialBitcastReshapeOfDotUnfused) { auto reshape2 = builder.AddInstruction(HloInstruction::CreateReshape( ShapeUtil::MakeShape(S32, {1, 1, 1}), dot1)); - auto module = MakeUnique(TestName()); + auto module = CreateNewModule(); auto computation = module->AddEntryComputation(builder.Build()); EXPECT_EQ(reshape2, computation->root_instruction()); EXPECT_FALSE(GpuInstructionFusion(/*may_duplicate=*/true) @@ -49,7 +48,7 @@ TEST_F(InstructionFusionTest, PotentialBitcastTransposeOfDotUnfused) { auto transpose2 = builder.AddInstruction(HloInstruction::CreateTranspose( ShapeUtil::MakeShape(S32, {1, 1}), dot1, {0, 1})); - auto module = MakeUnique(TestName()); + auto module = CreateNewModule(); auto computation = module->AddEntryComputation(builder.Build()); EXPECT_EQ(transpose2, computation->root_instruction()); EXPECT_FALSE(GpuInstructionFusion(/*may_duplicate=*/true) @@ -89,7 +88,7 @@ TEST_F(InstructionFusionTest, PotentialBitcastTransposeOfConvolutionUnfused) { builder.AddInstruction( HloInstruction::CreateReshape(ShapeUtil::MakeShape(F32, {3}), transpose)); - auto module = MakeUnique(TestName()); + auto module = CreateNewModule(); module->AddEntryComputation(builder.Build()); EXPECT_FALSE(GpuInstructionFusion(/*may_duplicate=*/true) .Run(module.get()) @@ -108,7 +107,7 @@ TEST_F(InstructionFusionTest, GetTupleElementFused) { HloInstruction::CreateGetTupleElement(data_shape, param, 1)); builder.AddInstruction( HloInstruction::CreateBinary(data_shape, HloOpcode::kAdd, gte0, gte1)); - auto module = MakeUnique(TestName()); + auto module = CreateNewModule(); auto computation = module->AddEntryComputation(builder.Build()); EXPECT_TRUE(GpuInstructionFusion(/*may_duplicate=*/true) .Run(module.get()) @@ -124,3 +123,7 @@ TEST_F(InstructionFusionTest, GetTupleElementFused) { } // namespace gpu } // namespace xla + +int main(int argc, char** argv) { + return xla::ParseDebugOptionsFlagsAndRunTests(argc, argv); +} diff --git a/tensorflow/compiler/xla/service/gpu/ir_emission_utils.cc b/tensorflow/compiler/xla/service/gpu/ir_emission_utils.cc index e8378a7f447cebf8d491e98595188d2391333c58..6be26dde8f957040c73db6a7e52f050e44d44c06 100644 --- a/tensorflow/compiler/xla/service/gpu/ir_emission_utils.cc +++ b/tensorflow/compiler/xla/service/gpu/ir_emission_utils.cc @@ -18,10 +18,11 @@ limitations under the License. #include #include -#include "external/llvm/include/llvm/IR/Module.h" +#include "llvm/IR/Module.h" #include "tensorflow/compiler/xla/layout_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/service/hlo_opcode.h" #include "tensorflow/compiler/xla/service/llvm_ir/llvm_util.h" #include "tensorflow/compiler/xla/shape_util.h" @@ -59,6 +60,11 @@ bool AreValidGemmShapes(const Shape& lhs_shape, const Shape& rhs_shape, } // namespace bool ImplementedAsGemm(const HloInstruction& hlo) { + // We can only do this if the HLO is unnested. + if (hlo.parent() != hlo.GetModule()->entry_computation()) { + return false; + } + // For certain types of Dot, we can call pre-canned BLAS gemm. if (hlo.opcode() == HloOpcode::kDot) { const Shape& lhs_shape = hlo.operand(0)->shape(); @@ -85,6 +91,11 @@ bool ImplementedAsGemm(const HloInstruction& hlo) { } bool ImplementedAsDnnConvolution(const HloInstruction& hlo) { + // We can only do this if the HLO is unnested. + if (hlo.parent() != hlo.GetModule()->entry_computation()) { + return false; + } + // Forward convolution. if (hlo.opcode() == HloOpcode::kConvolution) { const ConvolutionDimensionNumbers& dnums = diff --git a/tensorflow/compiler/xla/service/gpu/ir_emission_utils.h b/tensorflow/compiler/xla/service/gpu/ir_emission_utils.h index 4d3e9b10b2e69b083d74cf7b56edc5b781991b55..422972762ee3da793852429a71b4cee76e41e2bc 100644 --- a/tensorflow/compiler/xla/service/gpu/ir_emission_utils.h +++ b/tensorflow/compiler/xla/service/gpu/ir_emission_utils.h @@ -18,23 +18,14 @@ limitations under the License. #include -#include "external/llvm/include/llvm/IR/IRBuilder.h" -#include "external/llvm/include/llvm/IR/Value.h" +#include "llvm/IR/IRBuilder.h" +#include "llvm/IR/Value.h" #include "tensorflow/compiler/xla/service/hlo_instruction.h" namespace xla { namespace gpu { -const int64 kWarpSize = 32; - -// Precondition: "hlo" is an operand of a Dot instruction. -// -// Returns whether "hlo" is foldable to its user. -bool IsOperandFoldableToDot(const HloInstruction& hlo); - -// Returns true if GpuCompiler can fold any operands of "dot" into "dot" for -// better performance. -bool CanFoldOperandsIntoDot(const HloInstruction& dot); +constexpr int64 kWarpSize = 32; // Returns true if `hlo` will be implemented as a call to BLAS gemm. bool ImplementedAsGemm(const HloInstruction& hlo); diff --git a/tensorflow/compiler/xla/service/gpu/ir_emitter.cc b/tensorflow/compiler/xla/service/gpu/ir_emitter.cc index 5f3ce85f857a96ca0cca6b0bea4bf1e86b971827..2f9675f3a758eb80678eb509f65a41cf6d0fbde8 100644 --- a/tensorflow/compiler/xla/service/gpu/ir_emitter.cc +++ b/tensorflow/compiler/xla/service/gpu/ir_emitter.cc @@ -20,10 +20,10 @@ limitations under the License. #include "tensorflow/core/platform/logging.h" // IWYU pragma: no_include "llvm/IR/Intrinsics.gen.inc" -#include "external/llvm/include/llvm/IR/BasicBlock.h" -#include "external/llvm/include/llvm/IR/Constants.h" -#include "external/llvm/include/llvm/IR/Instructions.h" -#include "external/llvm/include/llvm/IR/Module.h" +#include "llvm/IR/BasicBlock.h" +#include "llvm/IR/Constants.h" +#include "llvm/IR/Instructions.h" +#include "llvm/IR/Module.h" #include "tensorflow/compiler/xla/primitive_util.h" #include "tensorflow/compiler/xla/service/elemental_ir_emitter.h" #include "tensorflow/compiler/xla/service/gpu/elemental_ir_emitter.h" @@ -57,7 +57,9 @@ IrEmitter::IrEmitter(const HloModuleConfig& hlo_module_config, &ir_emitter_context->buffer_assignment(), &ir_builder_, is_nested), hlo_module_config_(hlo_module_config) { - ir_builder_.setFastMathFlags(llvm_ir::GetFastMathFlags(hlo_module_config)); + ir_builder_.setFastMathFlags(llvm_ir::GetFastMathFlags( + /*fast_math_enabled=*/hlo_module_config.debug_options() + .xla_enable_fast_math())); } Status IrEmitter::DefaultAction(HloInstruction* hlo) { @@ -200,18 +202,22 @@ bool IrEmitter::MaybeEmitSpecialAtomicOperation( // NVPTX supports atomicMax and atomicMin on only integer types. if (root_opcode == HloOpcode::kMaximum && primitive_util::IsIntegralType(element_type)) { - // min(integral, integral) - ir_builder_.CreateAtomicRMW(llvm::AtomicRMWInst::Max, output_address, - source, + // max(integral, integral) + auto opcode = primitive_util::IsSignedIntegralType(element_type) + ? llvm::AtomicRMWInst::Max + : llvm::AtomicRMWInst::UMax; + ir_builder_.CreateAtomicRMW(opcode, output_address, source, llvm::AtomicOrdering::SequentiallyConsistent); return true; } if (root_opcode == HloOpcode::kMinimum && primitive_util::IsIntegralType(element_type)) { - // max(integral, integral) - ir_builder_.CreateAtomicRMW(llvm::AtomicRMWInst::Min, output_address, - source, + // min(integral, integral) + auto opcode = primitive_util::IsSignedIntegralType(element_type) + ? llvm::AtomicRMWInst::Min + : llvm::AtomicRMWInst::UMin; + ir_builder_.CreateAtomicRMW(opcode, output_address, source, llvm::AtomicOrdering::SequentiallyConsistent); return true; } @@ -399,7 +405,7 @@ Status IrEmitter::HandleDot(HloInstruction* dot, llvm::Type* accum_type = target_array.GetElementLlvmType(); llvm::Value* accum_address = llvm_ir::EmitAllocaAtFunctionEntry( accum_type, // The pointee type of the alloca instruction. - "accum_address", // The name of the alloca instuction. + "accum_address", // The name of the alloca instruction. &ir_builder_); // Initialize the accumulator in the preheader to zero. @@ -549,14 +555,12 @@ Status IrEmitter::HandleFusion(HloInstruction* fusion) { return EmitTargetElementLoop(*fusion, fused_emitter.GetRootGenerator()); } -Status IrEmitter::HandleCall( - HloInstruction* call, tensorflow::gtl::ArraySlice operands, - HloComputation* computation) { +Status IrEmitter::HandleCall(HloInstruction* call) { std::vector operand_addresses; - for (HloInstruction* operand : operands) { + for (HloInstruction* operand : call->operands()) { operand_addresses.push_back(GetBasePointer(*operand)); } - return EmitCallToNestedComputation(*computation, operand_addresses, + return EmitCallToNestedComputation(*call->to_apply(), operand_addresses, GetBasePointer(*call)); } diff --git a/tensorflow/compiler/xla/service/gpu/ir_emitter.h b/tensorflow/compiler/xla/service/gpu/ir_emitter.h index 1aefee2739978ec05f4094f79acaece39e221bea..2f6b3514497bff386d9f3e6f0d6c9737e8da4871 100644 --- a/tensorflow/compiler/xla/service/gpu/ir_emitter.h +++ b/tensorflow/compiler/xla/service/gpu/ir_emitter.h @@ -35,9 +35,9 @@ limitations under the License. #include #include -#include "external/llvm/include/llvm/IR/Function.h" -#include "external/llvm/include/llvm/IR/IRBuilder.h" -#include "external/llvm/include/llvm/IR/Value.h" +#include "llvm/IR/Function.h" +#include "llvm/IR/IRBuilder.h" +#include "llvm/IR/Value.h" #include "tensorflow/compiler/xla/service/buffer_assignment.h" #include "tensorflow/compiler/xla/service/dfs_hlo_visitor_with_default.h" #include "tensorflow/compiler/xla/service/gpu/elemental_ir_emitter.h" @@ -101,9 +101,7 @@ class IrEmitter : public DfsHloVisitorWithDefault { HloInstruction* on_true, HloInstruction* on_false) override; Status HandleFusion(HloInstruction* fusion) override; - Status HandleCall(HloInstruction* call, - tensorflow::gtl::ArraySlice operands, - HloComputation* computation) override; + Status HandleCall(HloInstruction* call) override; Status HandleCustomCall(HloInstruction* custom_call, tensorflow::gtl::ArraySlice operands, tensorflow::StringPiece custom_call_target) override; @@ -120,8 +118,9 @@ class IrEmitter : public DfsHloVisitorWithDefault { IrEmitterContext* ir_emitter_context, bool is_nested); // A convenient helper for calling HloToIrBindings::GetIrArray. - llvm_ir::IrArray GetIrArray(const HloInstruction& inst) { - return bindings_.GetIrArray(inst); + llvm_ir::IrArray GetIrArray(const HloInstruction& inst, + const ShapeIndex& shape_index = {}) { + return bindings_.GetIrArray(inst, shape_index); } // A convenient helper for calling HloToIrBindings::GetBasePointer. llvm::Value* GetBasePointer(const HloInstruction& inst) const { @@ -233,7 +232,7 @@ class IrEmitterUnnested : public IrEmitter { // IrEmitterUnnested handles the following instructions differently from // IrEmitter. - Status HandleCopy(HloInstruction* copy, HloInstruction* operand) override; + Status HandleCopy(HloInstruction* copy) override; Status HandleConvolution(HloInstruction* convolution, HloInstruction* lhs, HloInstruction* rhs, const Window& window) override; Status HandleDot(HloInstruction* dot, HloInstruction* lhs_instruction, @@ -249,8 +248,8 @@ class IrEmitterUnnested : public IrEmitter { Status HandleTuple( HloInstruction* tuple, tensorflow::gtl::ArraySlice operands) override; - Status HandleWhile(HloInstruction* xla_while, HloInstruction* init, - HloComputation* condition, HloComputation* body) override; + Status HandleWhile(HloInstruction* xla_while) override; + Status HandleInfeed(HloInstruction* xla_infeed) override; Status HandleRng(HloInstruction* random, RandomDistribution distribution) override; Status HandleSelect(HloInstruction* select, HloInstruction* pred, @@ -344,6 +343,10 @@ class IrEmitterUnnested : public IrEmitter { // Returns a CopyThunk that calls host-to-device cuMemcpy to implement `inst`. std::unique_ptr BuildCopyThunk(const HloInstruction* inst); + // Returns an InfeedThunk that performs device-to-device memcpy to implement + // `inst`. + std::unique_ptr BuildInfeedThunk(const HloInstruction* inst); + // Returns a WhileThunk that invokes thunk sequences for 'condition' and // 'body' sub-computations of while instruction 'hlo'. std::unique_ptr BuildWhileThunk(const HloInstruction* hlo); diff --git a/tensorflow/compiler/xla/service/gpu/ir_emitter_context.h b/tensorflow/compiler/xla/service/gpu/ir_emitter_context.h index 454c3f9ab2df00114341b6b59e6133950d260940..3790ed313b9d0e167185a8b12c812132ee78811f 100644 --- a/tensorflow/compiler/xla/service/gpu/ir_emitter_context.h +++ b/tensorflow/compiler/xla/service/gpu/ir_emitter_context.h @@ -16,7 +16,7 @@ limitations under the License. #ifndef TENSORFLOW_COMPILER_XLA_SERVICE_GPU_IR_EMITTER_CONTEXT_H_ #define TENSORFLOW_COMPILER_XLA_SERVICE_GPU_IR_EMITTER_CONTEXT_H_ -#include "external/llvm/include/llvm/IR/Module.h" +#include "llvm/IR/Module.h" #include "tensorflow/compiler/xla/service/buffer_assignment.h" #include "tensorflow/compiler/xla/service/gpu/partition_assignment.h" #include "tensorflow/compiler/xla/service/name_uniquer.h" diff --git a/tensorflow/compiler/xla/service/gpu/ir_emitter_nested.cc b/tensorflow/compiler/xla/service/gpu/ir_emitter_nested.cc index dc5e2d8f0241320aa0f6a781eed0b0355d8df8fd..202a0171dbe742d205f86afc79b663cfabe1c706 100644 --- a/tensorflow/compiler/xla/service/gpu/ir_emitter_nested.cc +++ b/tensorflow/compiler/xla/service/gpu/ir_emitter_nested.cc @@ -16,10 +16,10 @@ limitations under the License. #include #include -#include "external/llvm/include/llvm/IR/BasicBlock.h" -#include "external/llvm/include/llvm/IR/Function.h" -#include "external/llvm/include/llvm/IR/IRBuilder.h" -#include "external/llvm/include/llvm/IR/Instructions.h" +#include "llvm/IR/BasicBlock.h" +#include "llvm/IR/Function.h" +#include "llvm/IR/IRBuilder.h" +#include "llvm/IR/Instructions.h" #include "tensorflow/compiler/xla/service/gpu/hlo_to_ir_bindings.h" #include "tensorflow/compiler/xla/service/gpu/ir_emitter.h" #include "tensorflow/compiler/xla/service/gpu/ir_emitter_context.h" diff --git a/tensorflow/compiler/xla/service/gpu/ir_emitter_unnested.cc b/tensorflow/compiler/xla/service/gpu/ir_emitter_unnested.cc index 9b7aa7c860b14e03c238bd7037f0df832eacfef3..749badf3f235c6a56a624c9ee1decf95c510e957 100644 --- a/tensorflow/compiler/xla/service/gpu/ir_emitter_unnested.cc +++ b/tensorflow/compiler/xla/service/gpu/ir_emitter_unnested.cc @@ -17,13 +17,13 @@ limitations under the License. #include #include -#include "external/llvm/include/llvm/ADT/StringRef.h" -#include "external/llvm/include/llvm/IR/BasicBlock.h" -#include "external/llvm/include/llvm/IR/Function.h" -#include "external/llvm/include/llvm/IR/IRBuilder.h" -#include "external/llvm/include/llvm/IR/Instructions.h" -#include "external/llvm/include/llvm/IR/LLVMContext.h" -#include "external/llvm/include/llvm/IR/Module.h" +#include "llvm/ADT/StringRef.h" +#include "llvm/IR/BasicBlock.h" +#include "llvm/IR/Function.h" +#include "llvm/IR/IRBuilder.h" +#include "llvm/IR/Instructions.h" +#include "llvm/IR/LLVMContext.h" +#include "llvm/IR/Module.h" #include "tensorflow/compiler/xla/literal_util.h" #include "tensorflow/compiler/xla/ptr_util.h" #include "tensorflow/compiler/xla/service/buffer_assignment.h" @@ -33,6 +33,7 @@ limitations under the License. #include "tensorflow/compiler/xla/service/gpu/for_thunk.h" #include "tensorflow/compiler/xla/service/gpu/gemm_thunk.h" #include "tensorflow/compiler/xla/service/gpu/hlo_to_ir_bindings.h" +#include "tensorflow/compiler/xla/service/gpu/infeed_thunk.h" #include "tensorflow/compiler/xla/service/gpu/ir_emission_utils.h" #include "tensorflow/compiler/xla/service/gpu/ir_emitter.h" #include "tensorflow/compiler/xla/service/gpu/ir_emitter_context.h" @@ -192,11 +193,13 @@ llvm::Function* IrEmitterUnnested::BuildKernelPrototype( // the result). We also know how many bytes can be dereferenced in it. const llvm::Argument& temp_buffer = *std::prev(kernel->arg_end()); int64 temp_buffer_arg_no = temp_buffer.getArgNo(); - if (const BufferAllocation* allocation = - ir_emitter_context_->buffer_assignment().GetTempAllocation()) { - kernel->addDereferenceableAttr(temp_buffer_arg_no + 1, allocation->size()); + int64 temp_allocation_total_size = + ir_emitter_context_->buffer_assignment().temp_allocation_total_size(); + if (temp_allocation_total_size != 0) { + kernel->addDereferenceableAttr(temp_buffer_arg_no + 1, + temp_allocation_total_size); } - kernel->setDoesNotAlias(temp_buffer_arg_no + 1); + kernel->addAttribute(temp_buffer_arg_no + 1, llvm::Attribute::NoAlias); // Add the declaration of this kernel to llvm.nvvm.annotations so that NVPTX // treats it as a CUDA kernel. @@ -719,8 +722,7 @@ int64 EmitTranspose021Tiled(llvm_ir::IrArray input, llvm_ir::IrArray output, } // namespace -Status IrEmitterUnnested::HandleCopy(HloInstruction* copy, - HloInstruction* operand) { +Status IrEmitterUnnested::HandleCopy(HloInstruction* copy) { if (ImplementedAsMemcpy(*copy)) { thunk_sequence_->emplace_back(BuildCopyThunk(copy)); return Status::OK(); @@ -728,7 +730,7 @@ Status IrEmitterUnnested::HandleCopy(HloInstruction* copy, bool is_transpose_021; Shape reduced_input_shape, reduced_output_shape; std::tie(is_transpose_021, reduced_input_shape, reduced_output_shape) = - IsTranspose021(operand->shape(), copy->shape()); + IsTranspose021(copy->operand(0)->shape(), copy->shape()); if (is_transpose_021 && reduced_input_shape.dimensions(1) >= kMinDimensionToTransposeTiled && reduced_input_shape.dimensions(2) >= kMinDimensionToTransposeTiled) { @@ -736,7 +738,8 @@ Status IrEmitterUnnested::HandleCopy(HloInstruction* copy, VLOG(3) << "Emitting tiled 0-2-1 transposition"; constexpr int64 tile_size = 32; int64 num_tiles = EmitTranspose021Tiled( - GetIrArray(*operand).CastToShape(reduced_input_shape, &ir_builder_), + GetIrArray(*(copy->operand(0))) + .CastToShape(reduced_input_shape, &ir_builder_), GetIrArray(*copy).CastToShape(reduced_output_shape, &ir_builder_), tile_size, &ir_builder_); UpdateLaunchDimensions(LaunchDimensions(num_tiles, tile_size), LastThunk(), @@ -744,7 +747,7 @@ Status IrEmitterUnnested::HandleCopy(HloInstruction* copy, return Status::OK(); } - return IrEmitter::HandleCopy(copy, operand); + return IrEmitter::HandleCopy(copy); } Status IrEmitterUnnested::EmitColumnReduction( @@ -891,7 +894,7 @@ Status IrEmitterUnnested::EmitColumnReduction( llvm_ir::SetToFirstInsertPoint(if_tile_in_bounds_data.after_block, &ir_builder_); const HloInstruction* output = - reduce->IsFused() ? reduce->fusion_instruction() : reduce; + reduce->IsFused() ? reduce->parent()->FusionInstruction() : reduce; llvm::Value* output_address = GetIrArray(*output).EmitArrayElementAddress( llvm_ir::IrArray::Index(x, output->shape(), &ir_builder_), &ir_builder_, "output_element_address"); @@ -1139,7 +1142,7 @@ Status IrEmitterUnnested::EmitRowReduction( } const HloInstruction* output = - reduce->IsFused() ? reduce->fusion_instruction() : reduce; + 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 @@ -1540,10 +1543,8 @@ Status IrEmitterUnnested::HandleSelectAndScatter( .EmitLoop(); } -Status IrEmitterUnnested::HandleWhile(HloInstruction* xla_while, - HloInstruction* init, - HloComputation* condition, - HloComputation* body) { +Status IrEmitterUnnested::HandleWhile(HloInstruction* xla_while) { + HloComputation* condition = xla_while->while_condition(); TF_RET_CHECK(ShapeUtil::IsScalar(condition->root_instruction()->shape()) && condition->root_instruction()->shape().element_type() == PRED) << "While condition computation must return bool"; @@ -1579,6 +1580,11 @@ Status IrEmitterUnnested::HandleSelect(HloInstruction* select, return IrEmitter::HandleSelect(select, pred, on_true, on_false); } +Status IrEmitterUnnested::HandleInfeed(HloInstruction* infeed) { + thunk_sequence_->emplace_back(BuildInfeedThunk(infeed)); + return Status::OK(); +} + llvm::Function* IrEmitterUnnested::EmitBasePointersForHloAndItsOperands( const HloInstruction& hlo, std::vector* io_hlos) { const BufferAssignment& buffer_assignment = @@ -1627,6 +1633,7 @@ std::unique_ptr IrEmitterUnnested::BuildKernelThunk( // Compute the input buffer indices. std::vector io_buffers; + io_buffers.reserve(io_hlos.size()); for (const HloInstruction* io_hlo : io_hlos) { io_buffers.push_back(GetAllocationSlice(*LatestNonGteAncestor(io_hlo))); } @@ -1641,7 +1648,7 @@ std::unique_ptr IrEmitterUnnested::BuildCopyThunk( const HloInstruction* operand = inst->operand(0); CHECK_EQ(HloOpcode::kConstant, operand->opcode()); return MakeUnique( - /*source_address=*/LiteralUtil::InternalData(operand->literal()), + /*source_address=*/operand->literal().InternalData(), /*destination_buffer=*/GetAllocationSlice(*inst), /*mem_size=*/ llvm_ir::ByteSizeOf(operand->shape(), @@ -1649,6 +1656,23 @@ std::unique_ptr IrEmitterUnnested::BuildCopyThunk( inst); } +std::unique_ptr IrEmitterUnnested::BuildInfeedThunk( + const HloInstruction* inst) { + CHECK_EQ(HloOpcode::kInfeed, inst->opcode()); + + std::vector tuple_element_buffers; + for (int64 i = 0; i < inst->shape().tuple_shapes_size(); ++i) { + BufferAllocation::Slice buffer = ir_emitter_context_->buffer_assignment() + .GetUniqueSlice(inst, {i}) + .ConsumeValueOrDie(); + tuple_element_buffers.push_back(buffer); + } + + return MakeUnique( + tuple_element_buffers, + /*destination_buffer=*/GetAllocationSlice(*inst), inst); +} + std::unique_ptr IrEmitterUnnested::BuildGemmThunk( const HloInstruction* inst) { if (inst->opcode() == HloOpcode::kDot) { @@ -1781,7 +1805,7 @@ namespace { Status CheckWhileBuffersShareAllocation( const HloInstruction* xla_while, const BufferAssignment& buffer_assignment) { - return ShapeUtil::ForEachSubshape( + return ShapeUtil::ForEachSubshapeWithStatus( xla_while->shape(), [&buffer_assignment, &xla_while](const Shape& /*subshape*/, const ShapeIndex& index) -> Status { @@ -1862,15 +1886,35 @@ std::unique_ptr IrEmitterUnnested::BuildForThunk( Status IrEmitterUnnested::EmitTargetElementLoopInThunk( const HloInstruction& hlo, const llvm_ir::ElementGenerator& element_generator, KernelThunk* thunk) { + const Shape& element_shape = hlo.IsMultiOutputFusion() + ? ShapeUtil::GetSubshape(hlo.shape(), {0}) + : hlo.shape(); LaunchDimensions launch_dimensions = CalculateLaunchDimensions( - hlo.shape(), ir_emitter_context_->device_description()); + element_shape, ir_emitter_context_->device_description()); UpdateLaunchDimensions(launch_dimensions, thunk, ir_emitter_context_->llvm_module()); - // Otherwise, emit a parallel loop that computes the partition that each - // thread is in charge of. - return ParallelLoopEmitter(element_generator, GetIrArray(hlo), - launch_dimensions, &ir_builder_) - .EmitLoop(); + if (!hlo.IsMultiOutputFusion()) { + return ParallelLoopEmitter(element_generator, GetIrArray(hlo), + launch_dimensions, &ir_builder_) + .EmitLoop(); + } + + // For multiple outputs fusion, we need to emit each operand and the root. + std::vector output_arrays; + for (int64 i = 0; i < ShapeUtil::TupleElementCount(hlo.shape()); ++i) { + output_arrays.push_back(GetIrArray(hlo, {i})); + } + TF_RETURN_IF_ERROR(ParallelLoopEmitter(element_generator, output_arrays, + launch_dimensions, &ir_builder_) + .EmitLoop()); + + std::vector tuple_operand_ptrs; + for (int64 i = 0; i < output_arrays.size(); ++i) { + tuple_operand_ptrs.push_back(output_arrays[i].GetBasePointer()); + } + ir_builder_.SetInsertPoint(ir_builder_.GetInsertBlock()->getTerminator()); + llvm_ir::EmitTuple(GetIrArray(hlo), tuple_operand_ptrs, &ir_builder_); + return Status::OK(); } Status IrEmitterUnnested::EmitTargetElementLoop( diff --git a/tensorflow/compiler/xla/service/gpu/layout_assignment_test.cc b/tensorflow/compiler/xla/service/gpu/layout_assignment_test.cc index 692ec8147d3345d6fd1f5f7dca40a2e878ab2cfc..fa258b6e567e474ba8240f02ad54bc7a7f3c3258 100644 --- a/tensorflow/compiler/xla/service/gpu/layout_assignment_test.cc +++ b/tensorflow/compiler/xla/service/gpu/layout_assignment_test.cc @@ -55,9 +55,9 @@ TEST_F(LayoutAssignmentTest, Elementwise) { HloInstruction::CreateParameter(1, ashape, "y")); auto add = builder.AddInstruction( HloInstruction::CreateBinary(ashape, HloOpcode::kAdd, x, y)); - HloModule module(TestName()); + auto module = CreateNewModule(); HloComputation* computation = - module.AddEntryComputation(builder.Build(add)); + module->AddEntryComputation(builder.Build(add)); ComputationLayout computation_layout( computation->ComputeProgramShape()); @@ -69,7 +69,7 @@ TEST_F(LayoutAssignmentTest, Elementwise) { ShapeLayout(result_shape_with_layout); GpuLayoutAssignment layout_assignment(&computation_layout); - EXPECT_TRUE(layout_assignment.Run(&module).ValueOrDie()); + EXPECT_TRUE(layout_assignment.Run(module.get()).ValueOrDie()); for (const HloInstruction* operand : add->operands()) { EXPECT_TRUE(LayoutUtil::Equal(add->shape().layout(), @@ -83,3 +83,7 @@ TEST_F(LayoutAssignmentTest, Elementwise) { } // namespace } // namespace gpu } // namespace xla + +int main(int argc, char** argv) { + return xla::ParseDebugOptionsFlagsAndRunTests(argc, argv); +} diff --git a/tensorflow/compiler/xla/service/gpu/llvm_gpu_backend/BUILD b/tensorflow/compiler/xla/service/gpu/llvm_gpu_backend/BUILD index 724549c0c4ef46e7526953f41439ea8eff71a779..876f14f5c44f5d65185a9bfb4e6c7de86e3633b9 100644 --- a/tensorflow/compiler/xla/service/gpu/llvm_gpu_backend/BUILD +++ b/tensorflow/compiler/xla/service/gpu/llvm_gpu_backend/BUILD @@ -2,10 +2,6 @@ licenses(["notice"]) # Apache 2.0 package( default_visibility = [":friends"], - features = [ - "-parse_headers", - "no_layering_check", - ], ) package_group( @@ -28,27 +24,25 @@ cc_library( "utils.h", ], deps = [ + "//tensorflow/compiler/xla:status_macros", "//tensorflow/compiler/xla:statusor", "//tensorflow/compiler/xla:types", "//tensorflow/compiler/xla:util", - "//tensorflow/compiler/xla/legacy_flags:gpu_backend_lib_flags", "//tensorflow/compiler/xla/service:hlo_module_config", "//tensorflow/compiler/xla/service/llvm_ir:llvm_util", "//tensorflow/core:lib", "//tensorflow/core:lib_internal", "@llvm//:analysis", - "@llvm//:asm_printer", "@llvm//:bit_reader", "@llvm//:bit_writer", "@llvm//:code_gen", "@llvm//:core", - "@llvm//:instrumentation", "@llvm//:ipo", "@llvm//:ir_reader", "@llvm//:linker", - "@llvm//:mc", - "@llvm//:nvptx_code_gen", - "@llvm//:objc_arc", + "@llvm//:nvptx_code_gen", # buildcleaner: keep + "@llvm//:objc_arc", # buildcleaner: keep + "@llvm//:scalar", "@llvm//:support", "@llvm//:target", "@llvm//:transform_utils", diff --git a/tensorflow/compiler/xla/service/gpu/llvm_gpu_backend/dump_ir_pass.cc b/tensorflow/compiler/xla/service/gpu/llvm_gpu_backend/dump_ir_pass.cc index aeec3a03ca9d594f9ac26cfb8ff0292671e5b8bb..12a8a59488bfdd6ce55f762926cd63ba56bf9d7f 100644 --- a/tensorflow/compiler/xla/service/gpu/llvm_gpu_backend/dump_ir_pass.cc +++ b/tensorflow/compiler/xla/service/gpu/llvm_gpu_backend/dump_ir_pass.cc @@ -15,9 +15,9 @@ limitations under the License. #include "tensorflow/compiler/xla/service/gpu/llvm_gpu_backend/dump_ir_pass.h" -#include "external/llvm/include/llvm/IR/Module.h" -#include "external/llvm/include/llvm/Support/FileSystem.h" -#include "external/llvm/include/llvm/Support/raw_ostream.h" +#include "llvm/IR/Module.h" +#include "llvm/Support/FileSystem.h" +#include "llvm/Support/raw_ostream.h" #include "tensorflow/compiler/xla/service/gpu/llvm_gpu_backend/utils.h" #include "tensorflow/compiler/xla/types.h" #include "tensorflow/core/lib/core/stringpiece.h" diff --git a/tensorflow/compiler/xla/service/gpu/llvm_gpu_backend/dump_ir_pass.h b/tensorflow/compiler/xla/service/gpu/llvm_gpu_backend/dump_ir_pass.h index 1d515a0f28ca66d1ecbc50f90832376dd691f2a5..d0a863499f0a033767768426e523d4589b61da9c 100644 --- a/tensorflow/compiler/xla/service/gpu/llvm_gpu_backend/dump_ir_pass.h +++ b/tensorflow/compiler/xla/service/gpu/llvm_gpu_backend/dump_ir_pass.h @@ -18,8 +18,8 @@ limitations under the License. #include -#include "external/llvm/include/llvm/IR/LegacyPassManager.h" -#include "external/llvm/include/llvm/Pass.h" +#include "llvm/IR/LegacyPassManager.h" +#include "llvm/Pass.h" #include "tensorflow/compiler/xla/types.h" namespace xla { diff --git a/tensorflow/compiler/xla/service/gpu/llvm_gpu_backend/gpu_backend_lib.cc b/tensorflow/compiler/xla/service/gpu/llvm_gpu_backend/gpu_backend_lib.cc index 485216837dc727bfe8565ff22678dd2fa470bc40..2e7765c4c61a18f482fcc659dc1de8408a9d37b8 100644 --- a/tensorflow/compiler/xla/service/gpu/llvm_gpu_backend/gpu_backend_lib.cc +++ b/tensorflow/compiler/xla/service/gpu/llvm_gpu_backend/gpu_backend_lib.cc @@ -20,40 +20,39 @@ limitations under the License. #include #include -#include "tensorflow/compiler/xla/legacy_flags/gpu_backend_lib_flags.h" #include "tensorflow/compiler/xla/ptr_util.h" #include "tensorflow/compiler/xla/service/gpu/llvm_gpu_backend/dump_ir_pass.h" #include "tensorflow/compiler/xla/service/gpu/llvm_gpu_backend/utils.h" #include "tensorflow/compiler/xla/service/llvm_ir/llvm_util.h" +#include "tensorflow/compiler/xla/status_macros.h" #include "tensorflow/compiler/xla/util.h" -#include "external/llvm/include/llvm/ADT/STLExtras.h" -#include "external/llvm/include/llvm/ADT/StringMap.h" -#include "external/llvm/include/llvm/ADT/StringSet.h" -#include "external/llvm/include/llvm/Analysis/TargetLibraryInfo.h" -#include "external/llvm/include/llvm/Analysis/TargetTransformInfo.h" -#include "external/llvm/include/llvm/Bitcode/BitcodeReader.h" -#include "external/llvm/include/llvm/Bitcode/BitcodeWriter.h" -#include "external/llvm/include/llvm/CodeGen/CommandFlags.h" -#include "external/llvm/include/llvm/IR/LLVMContext.h" -#include "external/llvm/include/llvm/IR/LegacyPassManager.h" -#include "external/llvm/include/llvm/IR/Module.h" -#include "external/llvm/include/llvm/LinkAllIR.h" -#include "external/llvm/include/llvm/LinkAllPasses.h" -#include "external/llvm/include/llvm/Linker/Linker.h" -#include "external/llvm/include/llvm/PassRegistry.h" -#include "external/llvm/include/llvm/Support/CommandLine.h" -#include "external/llvm/include/llvm/Support/FileSystem.h" -#include "external/llvm/include/llvm/Support/FormattedStream.h" -#include "external/llvm/include/llvm/Support/TargetRegistry.h" -#include "external/llvm/include/llvm/Support/TargetSelect.h" -#include "external/llvm/include/llvm/Support/ToolOutputFile.h" -#include "external/llvm/include/llvm/Target/TargetMachine.h" -#include "external/llvm/include/llvm/Transforms/IPO.h" -#include "external/llvm/include/llvm/Transforms/IPO/AlwaysInliner.h" -#include "external/llvm/include/llvm/Transforms/IPO/PassManagerBuilder.h" - -#include "external/llvm/include/llvm/Transforms/IPO/Internalize.h" +#include "llvm/ADT/STLExtras.h" +#include "llvm/ADT/StringMap.h" +#include "llvm/ADT/StringSet.h" +#include "llvm/Analysis/TargetLibraryInfo.h" +#include "llvm/Analysis/TargetTransformInfo.h" +#include "llvm/Bitcode/BitcodeReader.h" +#include "llvm/Bitcode/BitcodeWriter.h" +#include "llvm/CodeGen/CommandFlags.h" +#include "llvm/IR/LLVMContext.h" +#include "llvm/IR/LegacyPassManager.h" +#include "llvm/IR/Module.h" +#include "llvm/IR/Verifier.h" +#include "llvm/Linker/Linker.h" +#include "llvm/PassRegistry.h" +#include "llvm/Support/CommandLine.h" +#include "llvm/Support/FileSystem.h" +#include "llvm/Support/FormattedStream.h" +#include "llvm/Support/TargetRegistry.h" +#include "llvm/Support/TargetSelect.h" +#include "llvm/Support/ToolOutputFile.h" +#include "llvm/Target/TargetMachine.h" +#include "llvm/Transforms/IPO.h" +#include "llvm/Transforms/IPO/AlwaysInliner.h" +#include "llvm/Transforms/IPO/Internalize.h" +#include "llvm/Transforms/IPO/PassManagerBuilder.h" +#include "llvm/Transforms/Scalar.h" #include "tensorflow/compiler/xla/types.h" #include "tensorflow/core/lib/core/stringpiece.h" #include "tensorflow/core/lib/io/path.h" @@ -133,13 +132,8 @@ static string GetSmName(std::pair compute_capability) { // from the input filename. string MakeNameForTempProduct(const std::string& input_filename, tensorflow::StringPiece extension) { - legacy_flags::GpuBackendLibFlags* flags = - legacy_flags::GetGpuBackendLibFlags(); - return tensorflow::io::JoinPath( - flags->dump_temp_products_to, - ReplaceFilenameExtension( - tensorflow::io::Basename(llvm_ir::AsString(input_filename)), - extension)); + return ReplaceFilenameExtension( + tensorflow::io::Basename(llvm_ir::AsString(input_filename)), extension); } // Initializes LLVM passes. Uses the PassRegistry mechanism. @@ -171,23 +165,21 @@ std::unique_ptr GetTargetMachine( } TargetOptions target_options = InitTargetOptionsFromCodeGenFlags(); - // Set options from hlo_module_config (specifically, fast-math flags). - llvm_ir::SetTargetOptions(hlo_module_config, &target_options); - - // Enable FMA synthesis if desired. - legacy_flags::GpuBackendLibFlags* flags = - legacy_flags::GetGpuBackendLibFlags(); - if (flags->fma) { - target_options.AllowFPOpFusion = FPOpFusion::Fast; - } + llvm_ir::SetTargetOptions( + /*fast_math_enabled=*/hlo_module_config.debug_options() + .xla_enable_fast_math(), + &target_options); + + // Enable FMA synthesis. + target_options.AllowFPOpFusion = FPOpFusion::Fast; // Set the verbose assembly options. - target_options.MCOptions.AsmVerbose = flags->verbose_ptx_asm; + target_options.MCOptions.AsmVerbose = false; // The selection of codegen optimization level is copied from function - // GetCodeGenOptLevel in //external/llvm/tools/opt/opt.cpp. + // GetCodeGenOptLevel in //third_party/llvm/llvm/tools/opt/opt.cpp. CodeGenOpt::Level codegen_opt_level; - switch (flags->opt_level) { + switch (hlo_module_config.debug_options().xla_backend_optimization_level()) { case 1: codegen_opt_level = CodeGenOpt::Less; break; @@ -202,7 +194,8 @@ std::unique_ptr GetTargetMachine( } return WrapUnique(target->createTargetMachine( triple.str(), llvm_ir::AsStringRef(cpu_name), "+ptx42", target_options, - Optional(RelocModel), CMModel, codegen_opt_level)); + Optional(RelocModel), Optional(CMModel), + codegen_opt_level)); } // Adds the standard LLVM optimization passes, based on the speed optimization @@ -259,12 +252,10 @@ string EmitModuleToPTX(Module* module, llvm::TargetMachine* target_machine) { // The extension is stripped by IrDumpingPassManager, so we need to // get creative to add a suffix. string module_id(llvm_ir::AsString(module->getModuleIdentifier())); - legacy_flags::GpuBackendLibFlags* flags = - legacy_flags::GetGpuBackendLibFlags(); IrDumpingPassManager codegen_passes( ReplaceFilenameExtension(tensorflow::io::Basename(module_id), "-nvptx.dummy"), - flags->dump_temp_products_to, flags->dump_ir_before_passes); + "", false); codegen_passes.add(new llvm::TargetLibraryInfoWrapperPass( llvm::Triple(module->getTargetTriple()))); @@ -342,36 +333,19 @@ StatusOr CompileModuleToPtx(llvm::Module* module, TF_RETURN_IF_ERROR( LinkLibdeviceIfNecessary(module, compute_capability, libdevice_dir_path)); - legacy_flags::GpuBackendLibFlags* flags = - legacy_flags::GetGpuBackendLibFlags(); - if (!flags->dump_temp_products_to.empty()) { - string linked_filename = - MakeNameForTempProduct(module->getModuleIdentifier(), "linked.bc"); - LOG(INFO) << "dumping bitcode after linking libdevice to: " - << linked_filename; - EmitBitcodeToFile(*module, linked_filename); - } - // Set the flush-denormals-to-zero flag on the module so the NVVM reflect pass // can access it. - module->addModuleFlag(llvm::Module::Override, "nvvm-reflect-ftz", flags->ftz); + module->addModuleFlag(llvm::Module::Override, "nvvm-reflect-ftz", + hlo_module_config.debug_options().xla_gpu_ftz()); // If ftz is enabled, set it as an attribute on every function in the module. - if (flags->ftz) { + if (hlo_module_config.debug_options().xla_gpu_ftz()) { for (llvm::Function& fn : *module) { fn.addFnAttr("nvptx-f32ftz", "true"); } } - // Run IR-level optimizations. - if (flags->dump_ir_before_passes && flags->dump_temp_products_to.empty()) { - LOG(FATAL) << "--dump_ir_before_passes must be specified with " - "--dump_temp_products_to"; - } - - IrDumpingPassManager module_passes(module->getModuleIdentifier(), - flags->dump_temp_products_to, - flags->dump_ir_before_passes); + IrDumpingPassManager module_passes(module->getModuleIdentifier(), "", false); // Add an appropriate TargetLibraryInfo pass for the module's triple. llvm::TargetLibraryInfoWrapperPass* tliwp = @@ -396,23 +370,31 @@ StatusOr CompileModuleToPtx(llvm::Module* module, // The LLVM IR verifier performs sanity checking on the IR. This helps // discover problems and report them in a meaningful manner, rather than let - // later passes report obscure assertions becasue of unfulfilled invariants. + // later passes report obscure assertions because of unfulfilled invariants. module_passes.add(llvm::createVerifierPass()); // Create the function-level pass manager. It needs data layout information // too. llvm::legacy::FunctionPassManager function_passes(module); - AddOptimizationPasses(flags->opt_level, /*size_level=*/0, - target_machine.get(), &module_passes, &function_passes); - // Loop unrolling exposes more opportunites for SROA. Therefore, we run SROA + 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"; + + AddOptimizationPasses(opt_level, + /*size_level=*/0, target_machine.get(), &module_passes, + &function_passes); + + // Loop unrolling exposes more opportunities for SROA. Therefore, we run SROA // again after the standard optimization passes [http://b/13329423]. - // TODO(jingyue): SROA may further expose more optimization opportunites, such + // TODO(jingyue): SROA may further expose more optimization opportunities such // as more precise alias analysis and more function inlining (SROA may change // the inlining cost of a function). For now, running SROA already emits good // enough code for the evaluated benchmarks. We may want to run more // optimizations later. - if (flags->opt_level > 0) { + if (opt_level > 0) { // LLVM's optimizer turns on SROA when the optimization level is greater // than 0. We mimic this behavior here. module_passes.add(llvm::createSROAPass()); @@ -430,14 +412,6 @@ StatusOr CompileModuleToPtx(llvm::Module* module, function_passes.doFinalization(); module_passes.run(*module); - if (!flags->dump_temp_products_to.empty()) { - string optimized_filename = - MakeNameForTempProduct(module->getModuleIdentifier(), "optimized.bc"); - LOG(INFO) << "dumping bitcode after optimizations to: " - << optimized_filename; - EmitBitcodeToFile(*module, optimized_filename); - } - // Finally, produce PTX. return EmitModuleToPTX(module, target_machine.get()); } @@ -470,22 +444,6 @@ void GPUBackendInit() { // between those loads. FeedLLVMWithFlags({"-memdep-block-scan-limit=500"}); - legacy_flags::GpuBackendLibFlags* flags = - legacy_flags::GetGpuBackendLibFlags(); - if (!flags->llvm_cl_opts.empty()) { - std::vector opts = - tensorflow::str_util::Split(flags->llvm_cl_opts, ','); - FeedLLVMWithFlags(opts); - } - - if (flags->llvm_dump_passes) { - // Enable LLVM pass debugging dump. LLVM dumps this information when a pass - // manager is initialized for execution. It's done to stderr (this is - // hardcoded within LLVM to the dbgs() stream, we can't change it from the - // outside). - FeedLLVMWithFlags({"-debug-pass=Arguments"}); - } - // Initialize the NVPTX target; it's the only target we link with, so call its // specific initialization functions instead of the catch-all InitializeAll*. LLVMInitializeNVPTXTarget(); diff --git a/tensorflow/compiler/xla/service/gpu/llvm_gpu_backend/gpu_backend_lib.h b/tensorflow/compiler/xla/service/gpu/llvm_gpu_backend/gpu_backend_lib.h index fd8940721701eed49427cea2f39c23cd8b8b1d9c..0a345191d34e6f40db043c559a67a44a6748321c 100644 --- a/tensorflow/compiler/xla/service/gpu/llvm_gpu_backend/gpu_backend_lib.h +++ b/tensorflow/compiler/xla/service/gpu/llvm_gpu_backend/gpu_backend_lib.h @@ -20,7 +20,7 @@ limitations under the License. #include #include -#include "external/llvm/include/llvm/IR/Module.h" +#include "llvm/IR/Module.h" #include "tensorflow/compiler/xla/service/hlo_module_config.h" #include "tensorflow/compiler/xla/statusor.h" #include "tensorflow/compiler/xla/types.h" diff --git a/tensorflow/compiler/xla/service/gpu/llvm_gpu_backend/utils.cc b/tensorflow/compiler/xla/service/gpu/llvm_gpu_backend/utils.cc index c10346bbc235d8949525eb2008bac5312395381d..9ef9bc3a50fc76f83f05e19163ab339f2da6ef3c 100644 --- a/tensorflow/compiler/xla/service/gpu/llvm_gpu_backend/utils.cc +++ b/tensorflow/compiler/xla/service/gpu/llvm_gpu_backend/utils.cc @@ -17,10 +17,10 @@ limitations under the License. #include "tensorflow/core/platform/logging.h" -#include "external/llvm/include/llvm/IR/LLVMContext.h" -#include "external/llvm/include/llvm/IR/Module.h" -#include "external/llvm/include/llvm/IRReader/IRReader.h" -#include "external/llvm/include/llvm/Support/SourceMgr.h" +#include "llvm/IR/LLVMContext.h" +#include "llvm/IR/Module.h" +#include "llvm/IRReader/IRReader.h" +#include "llvm/Support/SourceMgr.h" #include "tensorflow/compiler/xla/types.h" #include "tensorflow/core/lib/core/stringpiece.h" #include "tensorflow/core/lib/strings/strcat.h" @@ -28,7 +28,8 @@ limitations under the License. namespace { static void DieWithSMDiagnosticError(llvm::SMDiagnostic* diagnostic) { - LOG(FATAL) << diagnostic->getLineNo() << ":" << diagnostic->getColumnNo() + LOG(FATAL) << diagnostic->getFilename().str() << ":" + << diagnostic->getLineNo() << ":" << diagnostic->getColumnNo() << ": " << diagnostic->getMessage().str(); } diff --git a/tensorflow/compiler/xla/service/gpu/llvm_gpu_backend/utils_test.cc b/tensorflow/compiler/xla/service/gpu/llvm_gpu_backend/utils_test.cc index 3848e58b0d3c29901da46e92efa1a3d425a3ff18..8c7f70ebcfbd7e2d84aeceeff2259b02d67ceb24 100644 --- a/tensorflow/compiler/xla/service/gpu/llvm_gpu_backend/utils_test.cc +++ b/tensorflow/compiler/xla/service/gpu/llvm_gpu_backend/utils_test.cc @@ -19,8 +19,8 @@ limitations under the License. #include "tensorflow/core/lib/io/path.h" -#include "external/llvm/include/llvm/IR/LLVMContext.h" -#include "external/llvm/include/llvm/IR/Module.h" +#include "llvm/IR/LLVMContext.h" +#include "llvm/IR/Module.h" #include "tensorflow/compiler/xla/types.h" #include "tensorflow/core/platform/test.h" diff --git a/tensorflow/compiler/xla/service/gpu/pad_insertion.cc b/tensorflow/compiler/xla/service/gpu/pad_insertion.cc index c645e84aa4ff32ac6f1af890b3ec72460ef2b385..b8c61620845a1434cc79dc9a8b00f89944e2ae95 100644 --- a/tensorflow/compiler/xla/service/gpu/pad_insertion.cc +++ b/tensorflow/compiler/xla/service/gpu/pad_insertion.cc @@ -61,7 +61,7 @@ HloInstruction* MaybePaddedAndSlicedInput( PrimitiveType element_type = input->shape().element_type(); HloInstruction* padding = computation->AddInstruction(HloInstruction::CreateConstant( - MakeUnique(LiteralUtil::Zero(element_type)))); + MakeUnique(Literal::Zero(element_type)))); input = computation->AddInstruction(HloInstruction::CreatePad( ShapeInference::InferPadShape( /*operand_shape=*/input->shape(), @@ -80,6 +80,7 @@ HloInstruction* MaybePaddedAndSlicedInput( std::vector start_indices(input->shape().dimensions_size(), 0); std::vector limit_indices(input->shape().dimensions().begin(), input->shape().dimensions().end()); + std::vector strides(input->shape().dimensions_size(), 1); for (size_t i = 0; i < conv_dnums.spatial_dimensions().size(); ++i) { int64 dim = conv_dnums.spatial_dimensions(i); // If dimension "dim" has negative padding, increase the start index or @@ -92,9 +93,9 @@ HloInstruction* MaybePaddedAndSlicedInput( input = computation->AddInstruction(HloInstruction::CreateSlice( ShapeInference::InferSliceShape(input->shape(), start_indices, - limit_indices) + limit_indices, strides) .ConsumeValueOrDie(), - input, start_indices, limit_indices)); + input, start_indices, limit_indices, strides)); } return input; @@ -126,7 +127,7 @@ HloInstruction* MaybePaddedKernel(const Window& conv_window, PrimitiveType element_type = kernel->shape().element_type(); HloInstruction* padding = computation->AddInstruction(HloInstruction::CreateConstant( - MakeUnique(LiteralUtil::Zero(element_type)))); + MakeUnique(Literal::Zero(element_type)))); return computation->AddInstruction(HloInstruction::CreatePad( ShapeInference::InferPadShape( /*operand_shape=*/kernel->shape(), @@ -241,9 +242,9 @@ bool PadInsertion::CanonicalizeBackwardFilterConvolution( // Create a new backward convolution replacing the old one. HloComputation* computation = backward_conv->parent(); HloInstruction* output = backward_conv->mutable_operand(1); - HloInstruction* padding = computation->AddInstruction( - HloInstruction::CreateConstant(MakeUnique( - LiteralUtil::Zero(input->shape().element_type())))); + HloInstruction* padding = + computation->AddInstruction(HloInstruction::CreateConstant( + MakeUnique(Literal::Zero(input->shape().element_type())))); HloInstruction* padded_input = computation->AddInstruction(HloInstruction::CreatePad( ShapeInference::InferPadShape(input->shape(), padding->shape(), @@ -354,6 +355,8 @@ bool PadInsertion::CanonicalizeBackwardInputConvolution( std::vector limit_indices( new_backward_conv->shape().dimensions().begin(), new_backward_conv->shape().dimensions().end()); + std::vector strides(new_backward_conv->shape().dimensions_size(), + 1LL); for (size_t i = 0; i < backward_conv->window().dimensions_size(); ++i) { int64 padding_low = backward_conv->window().dimensions(i).padding_low(); int64 padding_high = backward_conv->window().dimensions(i).padding_high(); @@ -373,13 +376,13 @@ bool PadInsertion::CanonicalizeBackwardInputConvolution( // Replace the old backward convolution with the slice. CHECK(ShapeUtil::Compatible( ShapeInference::InferSliceShape(new_backward_conv->shape(), start_indices, - limit_indices) + limit_indices, strides) .ConsumeValueOrDie(), backward_conv->shape())); TF_CHECK_OK(computation->ReplaceWithNewInstruction( backward_conv, HloInstruction::CreateSlice(backward_conv->shape(), new_backward_conv, - start_indices, limit_indices))); + start_indices, limit_indices, strides))); return true; } diff --git a/tensorflow/compiler/xla/service/gpu/parallel_loop_emitter.cc b/tensorflow/compiler/xla/service/gpu/parallel_loop_emitter.cc index 65610b0995c512cc4a611ac650c581d0180d258d..03ecb6f635200c67ef0d4fe1eea888ca94d68fd6 100644 --- a/tensorflow/compiler/xla/service/gpu/parallel_loop_emitter.cc +++ b/tensorflow/compiler/xla/service/gpu/parallel_loop_emitter.cc @@ -20,8 +20,8 @@ limitations under the License. #include "tensorflow/core/platform/logging.h" #include "tensorflow/core/platform/types.h" // IWYU pragma: no_include "llvm/IR/Intrinsics.gen.inc" -#include "external/llvm/include/llvm/IR/Intrinsics.h" -#include "external/llvm/include/llvm/IR/Value.h" +#include "llvm/IR/Intrinsics.h" +#include "llvm/IR/Value.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/shape_util.h" @@ -36,6 +36,13 @@ ParallelLoopEmitter::ParallelLoopEmitter( : LoopEmitter(body_emitter, shape, ir_builder), launch_dimensions_(launch_dimensions) {} +ParallelLoopEmitter::ParallelLoopEmitter( + const llvm_ir::ElementGenerator& target_element_generator, + tensorflow::gtl::ArraySlice target_arrays, + const LaunchDimensions& launch_dimensions, llvm::IRBuilder<>* ir_builder) + : LoopEmitter(target_element_generator, target_arrays, ir_builder), + launch_dimensions_(launch_dimensions) {} + ParallelLoopEmitter::ParallelLoopEmitter( const llvm_ir::ElementGenerator& target_element_generator, const llvm_ir::IrArray& target_array, diff --git a/tensorflow/compiler/xla/service/gpu/parallel_loop_emitter.h b/tensorflow/compiler/xla/service/gpu/parallel_loop_emitter.h index 73ca28cd842fe350ecd10885d983907e7288a350..8855043236ca931a5c494bb1a71491a8e9d945d2 100644 --- a/tensorflow/compiler/xla/service/gpu/parallel_loop_emitter.h +++ b/tensorflow/compiler/xla/service/gpu/parallel_loop_emitter.h @@ -16,7 +16,7 @@ limitations under the License. #ifndef TENSORFLOW_COMPILER_XLA_SERVICE_GPU_PARALLEL_LOOP_EMITTER_H_ #define TENSORFLOW_COMPILER_XLA_SERVICE_GPU_PARALLEL_LOOP_EMITTER_H_ -#include "external/llvm/include/llvm/IR/IRBuilder.h" +#include "llvm/IR/IRBuilder.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/loop_emitter.h" @@ -41,6 +41,12 @@ class ParallelLoopEmitter : public llvm_ir::LoopEmitter { const llvm_ir::IrArray& target_array, const LaunchDimensions& launch_dimensions, llvm::IRBuilder<>* ir_builder); + + ParallelLoopEmitter( + const llvm_ir::ElementGenerator& target_element_generator, + tensorflow::gtl::ArraySlice target_arrays, + const LaunchDimensions& launch_dimensions, llvm::IRBuilder<>* ir_builder); + ParallelLoopEmitter(const ParallelLoopEmitter&) = delete; ParallelLoopEmitter& operator=(const ParallelLoopEmitter&) = delete; ~ParallelLoopEmitter() override = default; diff --git a/tensorflow/compiler/xla/service/gpu/partition_assignment.h b/tensorflow/compiler/xla/service/gpu/partition_assignment.h index 8ac4c5996632587fe4518df5560a1a74d9e8caa6..8f7fce884acc93fd39510ad0826b819a6d9731a7 100644 --- a/tensorflow/compiler/xla/service/gpu/partition_assignment.h +++ b/tensorflow/compiler/xla/service/gpu/partition_assignment.h @@ -33,7 +33,7 @@ namespace gpu { enum class PartitionStrategy { // Optimized for latency by allowing maximum number of registers per thread. kLatency, - // Optimized for throughtput. This may limit registers per thread and cause + // Optimized for throughput. This may limit registers per thread and cause // longer latency. kThroughput }; diff --git a/tensorflow/compiler/xla/service/gpu/stream_assignment.cc b/tensorflow/compiler/xla/service/gpu/stream_assignment.cc index 5065e7aedd08c591f33c152c6709823948db54f0..e4cfc6999f2da04dd7e7a34d854fdb3d75b8bfc6 100644 --- a/tensorflow/compiler/xla/service/gpu/stream_assignment.cc +++ b/tensorflow/compiler/xla/service/gpu/stream_assignment.cc @@ -15,11 +15,11 @@ limitations under the License. #include "tensorflow/compiler/xla/service/gpu/stream_assignment.h" -#include "tensorflow/compiler/xla/legacy_flags/stream_assignment_flags.h" #include "tensorflow/compiler/xla/map_util.h" #include "tensorflow/compiler/xla/ptr_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_reachability.h" namespace xla { namespace gpu { @@ -46,10 +46,9 @@ namespace { // Returns whether the two HLOs can run concurrently, i.e., neither is a // transitive consumer of the other. -bool CanRunConcurrently( - const HloInstruction& a, const HloInstruction& b, - const HloComputation::ReachabilityMap& transitive_operands) { - return !transitive_operands.IsConnected(&a, &b); +bool CanRunConcurrently(const HloInstruction& a, const HloInstruction& b, + const HloReachabilityMap& reachability) { + return !reachability.IsConnected(&a, &b); } // Returns which existing stream to assign to `hlo`, or -1 if a stream is not @@ -58,7 +57,7 @@ bool CanRunConcurrently( // are topologically before `hlo`. int ComputeStreamToAssign( const HloInstruction& hlo, const StreamAssignment& stream_assignment, - const HloComputation::ReachabilityMap& transitive_operands, + const HloReachabilityMap& reachability, const std::vector& seen_gemms) { if (hlo.opcode() == HloOpcode::kParameter || hlo.opcode() == HloOpcode::kConstant) { @@ -66,9 +65,10 @@ int ComputeStreamToAssign( return -1; } - legacy_flags::StreamAssignmentFlags* flags = - legacy_flags::GetStreamAssignmentFlags(); - if (flags->xla_gpu_disable_multi_streaming) { + if (hlo.GetModule() + ->config() + .debug_options() + .xla_gpu_disable_multi_streaming()) { return 0; } @@ -96,7 +96,7 @@ int ComputeStreamToAssign( for (const auto* seen_gemm : seen_gemms) { int stream_no = stream_assignment.StreamNumberForHlo(*seen_gemm); if (!forbidden_stream_numbers.count(stream_no) && - CanRunConcurrently(*seen_gemm, hlo, transitive_operands)) { + CanRunConcurrently(*seen_gemm, hlo, reachability)) { forbidden_stream_numbers.insert(stream_no); } } @@ -115,12 +115,12 @@ int ComputeStreamToAssign( std::unique_ptr AssignStreams(const HloModule& module) { auto stream_assignment = MakeUnique(); const HloComputation& computation = *module.entry_computation(); - std::unique_ptr transitive_operands = - computation.ComputeTransitiveOperands(); + std::unique_ptr reachability = + computation.ComputeReachability(); std::vector seen_gemms; for (const auto* hlo : computation.MakeInstructionPostOrder()) { int stream_no = ComputeStreamToAssign(*hlo, *stream_assignment, - *transitive_operands, seen_gemms); + *reachability, seen_gemms); if (stream_no != -1) { stream_assignment->AssignStreamToHlo(hlo, stream_no); } diff --git a/tensorflow/compiler/xla/service/gpu/stream_assignment_test.cc b/tensorflow/compiler/xla/service/gpu/stream_assignment_test.cc index 28d47d2b0f8b79e86954ca0b9bd82e6ea978879f..a5230b3e8e993d09e6ae01f927619da3f199a0c6 100644 --- a/tensorflow/compiler/xla/service/gpu/stream_assignment_test.cc +++ b/tensorflow/compiler/xla/service/gpu/stream_assignment_test.cc @@ -45,10 +45,10 @@ TEST_F(StreamAssignmentTest, SequentialMatMul) { HloInstruction* dot2 = builder.AddInstruction( HloInstruction::CreateBinary(f32_2x2_, HloOpcode::kDot, dot1, z)); - HloModule module(TestName()); - module.AddEntryComputation(builder.Build(dot2)); + auto module = CreateNewModule(); + module->AddEntryComputation(builder.Build(dot2)); - std::unique_ptr assignment = AssignStreams(module); + std::unique_ptr assignment = AssignStreams(*module); EXPECT_EQ(assignment->StreamNumberForHlo(*dot1), assignment->StreamNumberForHlo(*dot2)); } @@ -66,10 +66,10 @@ TEST_F(StreamAssignmentTest, ConcurrentMatMul) { HloInstruction* add = builder.AddInstruction( HloInstruction::CreateBinary(f32_2x2_, HloOpcode::kAdd, dot1, dot2)); - HloModule module(TestName()); - module.AddEntryComputation(builder.Build(add)); + auto module = CreateNewModule(); + module->AddEntryComputation(builder.Build(add)); - std::unique_ptr assignment = AssignStreams(module); + std::unique_ptr assignment = AssignStreams(*module); EXPECT_NE(assignment->StreamNumberForHlo(*dot1), assignment->StreamNumberForHlo(*dot2)); } @@ -86,6 +86,7 @@ TEST_F(StreamAssignmentTest, LatticeMatMul) { // d40 -- layer 4 HloComputation::Builder builder("entry_computation"); std::vector params; + params.reserve(6); for (int i = 0; i < 6; ++i) { params.push_back(builder.AddInstruction(HloInstruction::CreateParameter( i, f32_2x2_, /*name=*/tensorflow::strings::Printf("param%d", i)))); @@ -109,10 +110,10 @@ TEST_F(StreamAssignmentTest, LatticeMatMul) { HloInstruction* d40 = builder.AddInstruction( HloInstruction::CreateBinary(f32_2x2_, HloOpcode::kDot, d30, d31)); - HloModule module(TestName()); - module.AddEntryComputation(builder.Build(d40)); + auto module = CreateNewModule(); + module->AddEntryComputation(builder.Build(d40)); - std::unique_ptr assignment = AssignStreams(module); + std::unique_ptr assignment = AssignStreams(*module); // The two dots on layer 1 are concurrent. EXPECT_NE(assignment->StreamNumberForHlo(*d10), assignment->StreamNumberForHlo(*d11)); @@ -130,3 +131,7 @@ TEST_F(StreamAssignmentTest, LatticeMatMul) { } // namespace gpu } // namespace xla + +int main(int argc, char** argv) { + return xla::ParseDebugOptionsFlagsAndRunTests(argc, argv); +} diff --git a/tensorflow/compiler/xla/service/gpu/thunk.h b/tensorflow/compiler/xla/service/gpu/thunk.h index 3ced3484007ac288ee7bbe39d7c7bd6bf77c9d45..0ff27888ad72f8190400c22a9086d1965448662c 100644 --- a/tensorflow/compiler/xla/service/gpu/thunk.h +++ b/tensorflow/compiler/xla/service/gpu/thunk.h @@ -44,6 +44,7 @@ class Thunk { kConvolution, kCopy, kGemm, + kInfeed, kKernel, kSequential, kTuple, diff --git a/tensorflow/compiler/xla/service/gpu/while_transformer.cc b/tensorflow/compiler/xla/service/gpu/while_transformer.cc index ec75e1358142764d80152a6d8abbc6d5b72acb9a..ccdd1717593e4fa7c1d1deb3f0f9ebfab1bf7209 100644 --- a/tensorflow/compiler/xla/service/gpu/while_transformer.cc +++ b/tensorflow/compiler/xla/service/gpu/while_transformer.cc @@ -37,8 +37,8 @@ namespace { // patterns to match. // // Each ExprTree node is comprised of an HloOpcode, and a set of operands (each -// of type ExprTree). Operands can be added by specifing the index and HloOpcode -// of the operand. +// of type ExprTree). Operands can be added by specifying the index and +// HloOpcode of the operand. // // For example, the following computation: // @@ -122,10 +122,12 @@ class ExprTree { Status Match(const HloInstruction* instruction, TaggedInstructionMap* tagged_instructions) const { if (opcode_ != instruction->opcode()) { - return InvalidArgument("Unexpected opcode: %s", - HloOpcodeString(instruction->opcode()).c_str()); + return InvalidArgument("got opcode %s, want %s", + HloOpcodeString(instruction->opcode()).c_str(), + HloOpcodeString(opcode_).c_str()); } + VLOG(2) << "Matched " << HloOpcodeString(opcode_) << ": " << tag_; if (!tag_.empty()) { tagged_instructions->insert({tag_, instruction}); } @@ -166,7 +168,7 @@ class MatcherBase { virtual ~MatcherBase() {} // Attempts to match each ExprTree in 'expr_trees_'. - // Returns OK on the first succesful match, error status otherwise. + // Returns OK on the first successful match, error status otherwise. virtual tensorflow::Status Run() { Status status; for (const ExprTree& expr_tree : expr_trees_) { @@ -195,10 +197,9 @@ class MatcherBase { return InvalidArgument("Must use S32 or S64 integral types."); } if (type == S32) { - *const_value = - static_cast(LiteralUtil::GetFirstElement(literal)); + *const_value = static_cast(literal.GetFirstElement()); } else if (type == S64) { - *const_value = LiteralUtil::GetFirstElement(literal); + *const_value = literal.GetFirstElement(); } return tensorflow::Status::OK(); } @@ -221,7 +222,7 @@ class MatcherBase { TF_DISALLOW_COPY_AND_ASSIGN(MatcherBase); }; -// WhileConditionComputationMatcher attempst to match a target computation +// WhileConditionComputationMatcher attempts to match a target computation // pattern in the while condition sub-computation. // If the target pattern is matched, two pieces of information are extracted // from 'tagged' instructions returned by the matcher: @@ -238,7 +239,7 @@ class MatcherBase { // class WhileConditionComputationMatcher : public MatcherBase { public: - WhileConditionComputationMatcher(const HloComputation* computation) + explicit WhileConditionComputationMatcher(const HloComputation* computation) : computation_(computation) { expr_trees_.emplace_back(BuildCondExprTree()); } @@ -275,6 +276,7 @@ class WhileConditionComputationMatcher : public MatcherBase { } Status MatchExprTree(const ExprTree& expr_tree) override { + VLOG(2) << "MATCHING while condition"; ExprTree::TaggedInstructionMap tagged_instructions; TF_RETURN_IF_ERROR(expr_tree.Match(computation_->root_instruction(), &tagged_instructions)); @@ -306,7 +308,7 @@ class WhileConditionComputationMatcher : public MatcherBase { GetTaggedInstruction("gte.fusion_param.param0", tagged_instructions)); CHECK_EQ(HloOpcode::kParameter, gte_fusion_param0->opcode()); CHECK(gte_fusion_param0->IsFused()); - if (gte_fusion_param0->fusion_instruction()->operand( + if (gte_fusion_param0->parent()->FusionInstruction()->operand( gte_fusion_param0->parameter_number()) != computation_->parameter_instruction(0)) { return InvalidArgument("Could not match fusion param: %s", @@ -344,10 +346,6 @@ class WhileInitOperandMatcher : public MatcherBase { // // Const // | - // Tuple1 - // | - // GTE0 - // | // Copy // | // Tuple0 @@ -355,15 +353,15 @@ class WhileInitOperandMatcher : public MatcherBase { // While // ExprTree BuildInitExprTree() { - ExprTree gte0(HloOpcode::kGetTupleElement, "gte", - ExprTree(HloOpcode::kTuple, tuple_index_, - ExprTree(HloOpcode::kConstant, "loop_start"))); - return ExprTree(HloOpcode::kWhile, "while", - ExprTree(HloOpcode::kTuple, tuple_index_, - ExprTree(HloOpcode::kCopy, gte0))); + return ExprTree( + HloOpcode::kWhile, "while", + ExprTree(HloOpcode::kTuple, tuple_index_, + ExprTree(HloOpcode::kCopy, + ExprTree(HloOpcode::kConstant, "loop_start")))); } Status MatchExprTree(const ExprTree& expr_tree) override { + VLOG(2) << "MATCHING while init"; ExprTree::TaggedInstructionMap tagged_instructions; TF_RETURN_IF_ERROR(expr_tree.Match(while_hlo_, &tagged_instructions)); @@ -375,14 +373,6 @@ class WhileInitOperandMatcher : public MatcherBase { while_hlo->name().c_str()); } - // Get tagged GTE instruction and check 'tuple_index_'. - TF_ASSIGN_OR_RETURN(const HloInstruction* gte, - GetTaggedInstruction("gte", tagged_instructions)); - if (gte->tuple_index() != tuple_index_) { - return InvalidArgument("Unexpected tuple index instruction : %s", - gte->name().c_str()); - } - // Get tagged Constant instruction and parse 'loop_start_'. TF_ASSIGN_OR_RETURN( const HloInstruction* const_hlo, @@ -427,10 +417,6 @@ class WhileBodyComputationMatcher : public MatcherBase { // \ / \ / // Fusion -----------> Add // | - // Tuple1 - // | - // GTE0 - // | // Copy // | // Tuple0 @@ -450,15 +436,13 @@ class WhileBodyComputationMatcher : public MatcherBase { fusion.SetFusedRoot(fused_root); // Build top-level computation. - ExprTree tuple0( - HloOpcode::kTuple, tuple_index_, - ExprTree(HloOpcode::kCopy, - ExprTree(HloOpcode::kGetTupleElement, "gte", - ExprTree(HloOpcode::kTuple, tuple_index_, fusion)))); + ExprTree tuple0(HloOpcode::kTuple, tuple_index_, + ExprTree(HloOpcode::kCopy, fusion)); return tuple0; } Status MatchExprTree(const ExprTree& expr_tree) override { + VLOG(2) << "MATCHING while body"; ExprTree::TaggedInstructionMap tagged_instructions; TF_RETURN_IF_ERROR(expr_tree.Match(computation_->root_instruction(), &tagged_instructions)); @@ -485,7 +469,8 @@ class WhileBodyComputationMatcher : public MatcherBase { // Fusion parameter: lookup and compare with associated fusion operand. CHECK_EQ(HloOpcode::kParameter, inst->opcode()); CHECK(inst->IsFused()); - if (inst->fusion_instruction()->operand(inst->parameter_number()) != + if (inst->parent()->FusionInstruction()->operand( + inst->parameter_number()) != computation_->parameter_instruction(0)) { return InvalidArgument("Could not match fusion param: %s", inst->name().c_str()); diff --git a/tensorflow/compiler/xla/service/gpu/while_transformer_test.cc b/tensorflow/compiler/xla/service/gpu/while_transformer_test.cc index ddf9676e378c5445418d30ae767d19ef2fb74be8..51d38f84212b01c08c33f1b648c579c5672769ba 100644 --- a/tensorflow/compiler/xla/service/gpu/while_transformer_test.cc +++ b/tensorflow/compiler/xla/service/gpu/while_transformer_test.cc @@ -17,16 +17,20 @@ limitations under the License. #include "tensorflow/compiler/xla/service/copy_insertion.h" #include "tensorflow/compiler/xla/service/gpu/instruction_fusion.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 { +using ::testing::Eq; +using ::testing::HasSubstr; + class WhileTransformerTest : public HloTestBase { protected: WhileTransformerTest() - : module_(TestName()), + : module_(CreateNewModule()), induction_variable_shape_(ShapeUtil::MakeShape(S32, {})), data_shape_(ShapeUtil::MakeShape(F32, {8})), loop_state_shape_(ShapeUtil::MakeTupleShape( @@ -37,7 +41,7 @@ class WhileTransformerTest : public HloTestBase { const int64 tuple_index, const int64 limit) { auto builder = HloComputation::Builder(TestName() + ".Condition"); auto limit_const = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(limit))); + HloInstruction::CreateConstant(Literal::CreateR0(limit))); auto loop_state = builder.AddInstruction( HloInstruction::CreateParameter(0, loop_state_shape_, "loop_state")); auto induction_variable = @@ -60,8 +64,8 @@ class WhileTransformerTest : public HloTestBase { auto induction_variable = builder.AddInstruction(HloInstruction::CreateGetTupleElement( induction_variable_shape_, loop_state, ind_var_tuple_index)); - auto inc = builder.AddInstruction(HloInstruction::CreateConstant( - LiteralUtil::CreateR0(increment))); + auto inc = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(increment))); auto add0 = builder.AddInstruction(HloInstruction::CreateBinary( induction_variable->shape(), HloOpcode::kAdd, induction_variable, inc)); // Update data GTE(data_tuple_index). @@ -84,12 +88,10 @@ class WhileTransformerTest : public HloTestBase { const int64 ind_var_tuple_index, const int64 ind_var_init) { auto builder = HloComputation::Builder(TestName() + ".While"); - auto induction_var_init = - builder.AddInstruction(HloInstruction::CreateConstant( - LiteralUtil::CreateR0(ind_var_init))); - auto data_init = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR1( - {0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f}))); + auto induction_var_init = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(ind_var_init))); + auto data_init = builder.AddInstruction(HloInstruction::CreateConstant( + Literal::CreateR1({0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f}))); auto loop_state_init = ind_var_tuple_index == 0 ? builder.AddInstruction( @@ -98,26 +100,26 @@ class WhileTransformerTest : public HloTestBase { HloInstruction::CreateTuple({data_init, induction_var_init})); auto while_hlo = builder.AddInstruction(HloInstruction::CreateWhile( loop_state_shape_, condition, body, loop_state_init)); - module_.AddEntryComputation(builder.Build()); + module_->AddEntryComputation(builder.Build()); return while_hlo; } void RunFusionPasses() { // Run standard fusion passes. EXPECT_TRUE(gpu::GpuInstructionFusion(/*may_duplicate=*/false) - .Run(&module_) + .Run(module_.get()) .ValueOrDie()); EXPECT_TRUE(gpu::GpuInstructionFusion(/*may_duplicate=*/true) - .Run(&module_) + .Run(module_.get()) .ValueOrDie()); } void RunCopyInsertionPass() { CopyInsertion copy_insertion; - EXPECT_IS_OK(copy_insertion.Run(&module_).status()); + EXPECT_IS_OK(copy_insertion.Run(module_.get()).status()); } - HloModule module_; + std::unique_ptr module_; Shape induction_variable_shape_; Shape data_shape_; Shape loop_state_shape_; @@ -127,74 +129,72 @@ class WhileTransformerTest : public HloTestBase { TEST_F(WhileTransformerTest, InductionVariableAtTupleElement0) { // Build computation with induction variable at tuple element 0. auto condition = - module_.AddEmbeddedComputation(BuildConditionComputation(0, 10)); - auto body = module_.AddEmbeddedComputation(BuildBodyComputation(0, 1, 1)); + module_->AddEmbeddedComputation(BuildConditionComputation(0, 10)); + auto body = module_->AddEmbeddedComputation(BuildBodyComputation(0, 1, 1)); auto while_hlo = BuildWhileInstruction(condition, body, 0, 0); // Run HLO Optimization passes. RunFusionPasses(); RunCopyInsertionPass(); // Run WhileTransformer. auto result = gpu::CanTransformWhileToFor(while_hlo); - EXPECT_TRUE(result.ok()); + ASSERT_TRUE(result.ok()); // Check results. - auto tuple = result.ConsumeValueOrDie(); - EXPECT_EQ(0, std::get<0>(tuple)); - EXPECT_EQ(10, std::get<1>(tuple)); - EXPECT_EQ(1, std::get<2>(tuple)); + EXPECT_THAT(result.ConsumeValueOrDie(), + Eq(std::tuple(0, 10, 1))); } TEST_F(WhileTransformerTest, InductionVariableAtTupleElement1) { // Build computation with induction variable at tuple element 1. auto condition = - module_.AddEmbeddedComputation(BuildConditionComputation(1, 10)); - auto body = module_.AddEmbeddedComputation(BuildBodyComputation(1, 0, 1)); + module_->AddEmbeddedComputation(BuildConditionComputation(1, 10)); + auto body = module_->AddEmbeddedComputation(BuildBodyComputation(1, 0, 1)); auto while_hlo = BuildWhileInstruction(condition, body, 1, 0); // Run HLO Optimization passes. RunFusionPasses(); RunCopyInsertionPass(); // Run WhileTransformer. auto result = gpu::CanTransformWhileToFor(while_hlo); - EXPECT_TRUE(result.ok()); + ASSERT_TRUE(result.ok()); // Check results. - auto tuple = result.ConsumeValueOrDie(); - EXPECT_EQ(0, std::get<0>(tuple)); - EXPECT_EQ(10, std::get<1>(tuple)); - EXPECT_EQ(1, std::get<2>(tuple)); + EXPECT_THAT(result.ConsumeValueOrDie(), + Eq(std::tuple(0, 10, 1))); } TEST_F(WhileTransformerTest, InvalidLoopLimit) { // Build computation with invalid loop limit. auto condition = - module_.AddEmbeddedComputation(BuildConditionComputation(0, 5)); - auto body = module_.AddEmbeddedComputation(BuildBodyComputation(0, 1, 1)); + module_->AddEmbeddedComputation(BuildConditionComputation(0, 5)); + auto body = module_->AddEmbeddedComputation(BuildBodyComputation(0, 1, 1)); auto while_hlo = BuildWhileInstruction(condition, body, 0, 10); // Run HLO Optimization passes. RunFusionPasses(); RunCopyInsertionPass(); // Run WhileTransformer. auto result = gpu::CanTransformWhileToFor(while_hlo); - EXPECT_FALSE(result.ok()); - EXPECT_MATCH( - result.status().error_message(), - testing::ContainsRegex("Loop start must be less than loop limit.")); + ASSERT_FALSE(result.ok()); + EXPECT_THAT(result.status().error_message(), + HasSubstr("Loop start must be less than loop limit.")); } TEST_F(WhileTransformerTest, InvalidLoopIncrement) { // Build computation with invalid loop increment. auto condition = - module_.AddEmbeddedComputation(BuildConditionComputation(0, 10)); - auto body = module_.AddEmbeddedComputation(BuildBodyComputation(0, 1, -1)); + module_->AddEmbeddedComputation(BuildConditionComputation(0, 10)); + auto body = module_->AddEmbeddedComputation(BuildBodyComputation(0, 1, -1)); auto while_hlo = BuildWhileInstruction(condition, body, 0, 0); // Run HLO Optimization passes. RunFusionPasses(); RunCopyInsertionPass(); // Run WhileTransformer. auto result = gpu::CanTransformWhileToFor(while_hlo); - EXPECT_FALSE(result.ok()); - EXPECT_MATCH( - result.status().error_message(), - testing::ContainsRegex("Loop increment must greater than zero.")); + ASSERT_FALSE(result.ok()); + EXPECT_THAT(result.status().error_message(), + HasSubstr("Loop increment must greater than zero.")); } } // namespace } // namespace xla + +int main(int argc, char** argv) { + return xla::ParseDebugOptionsFlagsAndRunTests(argc, argv); +} diff --git a/tensorflow/compiler/xla/service/gpu_transfer_manager.cc b/tensorflow/compiler/xla/service/gpu_transfer_manager.cc new file mode 100644 index 0000000000000000000000000000000000000000..74f0bdb7db1847119c5bd75cc9fd9d921c6e162a --- /dev/null +++ b/tensorflow/compiler/xla/service/gpu_transfer_manager.cc @@ -0,0 +1,153 @@ +/* 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/gpu_transfer_manager.h" + +#include +#include +#include + +#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/compiler/xla/types.h" +#include "tensorflow/compiler/xla/util.h" +#include "tensorflow/compiler/xla/xla_data.pb.h" +#include "tensorflow/core/lib/core/errors.h" +#include "tensorflow/core/lib/gtl/cleanup.h" +#include "tensorflow/core/platform/logging.h" +#include "tensorflow/core/platform/stream_executor_no_cuda.h" + +namespace se = ::perftools::gputools; + +namespace xla { + +// TODO(b/30467474) Once GPU infeed implementation settles, consider +// folding back the cpu and gpu infeed implementations into a generic +// one if possible. +GpuTransferManager::GpuTransferManager() + : GenericTransferManager(se::cuda::kCudaPlatformId) {} + +Status GpuTransferManager::TransferLiteralToInfeed(se::StreamExecutor* executor, + const Literal& literal) { + const Shape& shape = literal.shape(); + VLOG(2) << "Transferring literal to infeed with shape: " + << ShapeUtil::HumanString(shape); + + if (!ShapeUtil::IsTuple(shape)) { + int64 size = GetByteSizeRequirement(shape); + return TransferBufferToInfeed(executor, size, literal.InternalData()); + } + + if (ShapeUtil::IsNestedTuple(shape)) { + return Unimplemented( + "Infeed with a nested tuple shape is not supported: %s", + ShapeUtil::HumanString(literal.shape()).c_str()); + } + + // For a tuple, we transfer each of its elements to the device and + // enqueue the resulting destination device addresses with the + // infeed manager. + std::vector buffers; + buffers.reserve(literal.tuple_literals_size()); + auto cleanup = tensorflow::gtl::MakeCleanup([buffers]() { + for (gpu::InfeedBuffer* b : buffers) { + b->Done(); + } + }); + + for (const auto& tuple_element : literal.tuple_literals()) { + const Shape& tuple_element_shape = tuple_element.shape(); + int64 tuple_element_size = GetByteSizeRequirement(tuple_element_shape); + TF_ASSIGN_OR_RETURN( + gpu::InfeedBuffer * buffer, + TransferBufferToInfeedInternal(executor, tuple_element_size, + tuple_element.InternalData())); + buffers.push_back(buffer); + } + + cleanup.release(); + return EnqueueBuffersToInfeed(executor, buffers); +} + +Status GpuTransferManager::TransferBufferToInfeed(se::StreamExecutor* executor, + int64 size, + const void* source) { + TF_ASSIGN_OR_RETURN(gpu::InfeedBuffer * buffer, + TransferBufferToInfeedInternal(executor, size, source)); + return EnqueueBuffersToInfeed(executor, {buffer}); +} + +Status GpuTransferManager::EnqueueBuffersToInfeed( + se::StreamExecutor* executor, std::vector buffers) { + gpu::InfeedManager* infeed_manager = gpu::GetOrCreateInfeedManager(); + se::Stream* stream = infeed_manager->GetStream(executor); + + // TODO(b/30467474): Since this stream is shared across different + // infeed requests, blocking on the stream might be + // heavy-handed. Figure out if finer-grained acknowledgement is + // possible. + if (!stream->BlockHostUntilDone()) { + for (gpu::InfeedBuffer* b : buffers) { + b->Done(); + } + return InternalError("Failed to complete data transfer on stream %p", + stream); + } + + infeed_manager->EnqueueBuffers(buffers); + + VLOG(2) << "Infeed data transferred"; + + return Status::OK(); +} + +StatusOr GpuTransferManager::TransferBufferToInfeedInternal( + se::StreamExecutor* executor, int64 size, const void* source) { + if (size > std::numeric_limits::max()) { + return InvalidArgument("Infeed shape is too large: needs %lld bytes", size); + } + + if (size == 0) { + return InvalidArgument("Infeed shape needs 0 bytes"); + } + + gpu::InfeedManager* infeed_manager = gpu::GetOrCreateInfeedManager(); + se::Stream* stream = infeed_manager->GetStream(executor); + if (stream == nullptr) { + return InternalError("Failed to obtain a stream"); + } + + gpu::InfeedBuffer* buffer = new gpu::InfeedBuffer(executor, size); + stream->ThenMemcpy(buffer->device_memory(), source, size); + + VLOG(2) << "Queued infeed data on stream " << stream; + + return buffer; +} + +} // namespace xla + +static std::unique_ptr CreateGpuTransferManager() { + return xla::MakeUnique(); +} + +static bool InitModule() { + xla::TransferManager::RegisterTransferManager(se::cuda::kCudaPlatformId, + &CreateGpuTransferManager); + return true; +} +static bool module_initialized = InitModule(); diff --git a/tensorflow/compiler/xla/service/gpu_transfer_manager.h b/tensorflow/compiler/xla/service/gpu_transfer_manager.h new file mode 100644 index 0000000000000000000000000000000000000000..9aa369c668364079504ead3491903e2590a142cc --- /dev/null +++ b/tensorflow/compiler/xla/service/gpu_transfer_manager.h @@ -0,0 +1,61 @@ +/* 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_GPU_TRANSFER_MANAGER_H_ +#define TENSORFLOW_COMPILER_XLA_SERVICE_GPU_TRANSFER_MANAGER_H_ + +#include + +#include "tensorflow/compiler/xla/service/generic_transfer_manager.h" +#include "tensorflow/compiler/xla/service/gpu/infeed_manager.h" +#include "tensorflow/compiler/xla/service/transfer_manager.h" +#include "tensorflow/compiler/xla/statusor.h" +#include "tensorflow/compiler/xla/xla_data.pb.h" +#include "tensorflow/core/platform/macros.h" +#include "tensorflow/core/platform/stream_executor_no_cuda.h" +#include "tensorflow/core/platform/types.h" + +namespace xla { + +// An implementation of the XLA GenericTransferManager that +// handles GPU-specific infeed. +class GpuTransferManager : public GenericTransferManager { + public: + GpuTransferManager(); + ~GpuTransferManager() override {} + + Status TransferLiteralToInfeed(perftools::gputools::StreamExecutor* executor, + const Literal& literal) override; + Status TransferBufferToInfeed(perftools::gputools::StreamExecutor* executor, + int64 size, const void* source) override; + + private: + // Initiates the infeed data transfers. InfeedBuffer->Done() must be + // called to clean up the memory allocated for InfeedBuffer. + StatusOr TransferBufferToInfeedInternal( + perftools::gputools::StreamExecutor* executor, int64 size, + const void* source); + + // Enqueues infeed data buffers with the infeed manager after their + // transfer completes. + Status EnqueueBuffersToInfeed(perftools::gputools::StreamExecutor* executor, + std::vector buffers); + + TF_DISALLOW_COPY_AND_ASSIGN(GpuTransferManager); +}; + +} // namespace xla + +#endif // TENSORFLOW_COMPILER_XLA_SERVICE_GPU_TRANSFER_MANAGER_H_ diff --git a/tensorflow/compiler/xla/service/graphviz_example.cc b/tensorflow/compiler/xla/service/graphviz_example.cc index cd00a41a03718502fcfa63e035639390b6fe6e07..049e8d80d80c835bca4a4d38592564ba82a3ecf9 100644 --- a/tensorflow/compiler/xla/service/graphviz_example.cc +++ b/tensorflow/compiler/xla/service/graphviz_example.cc @@ -47,7 +47,7 @@ HloComputation* AddScalarConstantComputation(int64 addend, HloModule* module) { auto x_value = builder.AddInstruction(HloInstruction::CreateParameter( 0, ShapeUtil::MakeShape(F32, {}), "x_value")); auto half = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(0.5))); + HloInstruction::CreateConstant(Literal::CreateR0(0.5))); builder.AddInstruction(HloInstruction::CreateBinary( half->shape(), HloOpcode::kAdd, x_value, half)); return module->AddEmbeddedComputation(builder.Build()); @@ -118,7 +118,7 @@ std::unique_ptr MakeBigGraph() { auto rng = builder.AddInstruction( HloInstruction::CreateRng(vshape, RNG_UNIFORM, {param_m, param_m})); auto one = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(1.0))); + HloInstruction::CreateConstant(Literal::CreateR0(1.0))); auto add_computation = ScalarSumComputation(module.get()); builder.AddInstruction( HloInstruction::CreateReduce(vshape, rng, one, {1}, add_computation)); @@ -156,10 +156,9 @@ int main(int argc, char** argv) { auto module = xla::MakeBigGraph(); - printf("Graph URL: %s\n", - xla::hlo_graph_dumper::DumpGraph( - *module->entry_computation(), "Example computation", - /*show_addresses=*/false, /*show_layouts=*/false) - .c_str()); + printf("Graph URL: %s\n", xla::hlo_graph_dumper::DumpGraph( + *module->entry_computation(), + "Example computation", xla::DebugOptions()) + .c_str()); return 0; } diff --git a/tensorflow/compiler/xla/service/heap_simulator.cc b/tensorflow/compiler/xla/service/heap_simulator.cc index 46c0d8edead1eaba518fd1040b7dd7d0d6c79159..c85e97b691c04a5e2181ce39732b42bdf59ee679 100644 --- a/tensorflow/compiler/xla/service/heap_simulator.cc +++ b/tensorflow/compiler/xla/service/heap_simulator.cc @@ -35,30 +35,67 @@ namespace { std::vector UniqueOperandSourceBuffers( const HloInstruction* instruction, const TuplePointsToAnalysis& points_to_analysis) { - FlatSet buffers; + std::vector buffers; for (const HloInstruction* operand : instruction->operands()) { - FlatSet sources = - points_to_analysis.GetPointsToSet(operand).CreateFlattenedSet(); - buffers.insert(sources.begin(), sources.end()); + points_to_analysis.GetPointsToSet(operand).ForEachElement( + [&](const ShapeIndex& /*index*/, + const PointsToSet::BufferList& points_to) { + buffers.insert(buffers.end(), points_to.begin(), points_to.end()); + }); } - std::vector sorted(buffers.begin(), buffers.end()); - std::sort(sorted.begin(), sorted.end(), + + // Sort and then remove duplicates from buffers. + std::sort(buffers.begin(), buffers.end(), [](const LogicalBuffer* a, const LogicalBuffer* b) { return a->id() < b->id(); }); - return sorted; + buffers.erase(std::unique(buffers.begin(), buffers.end(), + [](const LogicalBuffer* a, const LogicalBuffer* b) { + return a->id() == b->id(); + }), + buffers.end()); + return buffers; } } // namespace /*static*/ StatusOr HeapSimulator::Run( - std::unique_ptr algorithm, + std::unique_ptr algorithm, const HloModule& module, + const SequentialHloOrdering::HloModuleSequence& module_sequence, + const TuplePointsToAnalysis& points_to_analysis, + const LogicalBuffer::SizeFunction& size_fn, + const FlatSet* buffers_to_assign) { + HeapSimulator heap(std::move(algorithm), size_fn, buffers_to_assign, + &module_sequence); + const HloComputation* entry_computation = module.entry_computation(); + const std::vector& instruction_sequence = + FindOrDie(module_sequence, entry_computation); + TF_RETURN_IF_ERROR(heap.RunComputation( + *entry_computation, instruction_sequence, points_to_analysis)); + return heap.Finish(); +} + +/*static*/ +StatusOr HeapSimulator::Run( + std::unique_ptr algorithm, const HloComputation& computation, const std::vector& instruction_sequence, - const HloComputation& computation, const TuplePointsToAnalysis& points_to_analysis, const LogicalBuffer::SizeFunction& size_fn, const FlatSet* buffers_to_assign) { + HeapSimulator heap(std::move(algorithm), size_fn, buffers_to_assign, + /*module_sequence=*/nullptr); + TF_RETURN_IF_ERROR(heap.RunComputation(computation, instruction_sequence, + points_to_analysis)); + return heap.Finish(); +} + +// Runs a heap simulation for the given 'computation', assuming the given +// 'instruction_sequence'. +Status HeapSimulator::RunComputation( + const HloComputation& computation, + const std::vector& instruction_sequence, + const TuplePointsToAnalysis& points_to_analysis) { // The goal here is to minimize memory usage, assuming the given sequential // ordering of instructions. The strategy is to walk through the instruction // sequence, calling Alloc and Free on the underlying heap algorithm. The @@ -67,33 +104,44 @@ StatusOr HeapSimulator::Run( // 'live_buffers' tracks the liveness of each buffer that we assign, by // associating it with a set of HloInstructions that need to be visited. When // the set becomes empty, the buffer is no longer used, and can be freed. - HeapSimulator heap(std::move(algorithm), size_fn, buffers_to_assign); FlatMap> live_buffers; - for (const HloInstruction* instruction : instruction_sequence) { - const std::vector& buffers_defined_by_instruction = - points_to_analysis.GetBuffersDefinedByInstruction(instruction); + const HloInstruction* root = computation.root_instruction(); + auto output_source_buffers = + points_to_analysis.GetPointsToSet(root).CreateFlattenedSet(); - const HloInstruction* root = computation.root_instruction(); - FlatSet output_source_buffers = - points_to_analysis.GetPointsToSet(root).CreateFlattenedSet(); + std::vector dead_buffers_to_free; + std::vector operand_buffers_to_free; + for (const HloInstruction* instruction : instruction_sequence) { + const TuplePointsToAnalysis::BufferDefinitionVector& + buffers_defined_by_instruction = + points_to_analysis.GetBuffersDefinedByInstruction(instruction); // Initialize live_buffers for each buffer that we're going to assign. The // set of instructions that need to be visited contains all users of all // aliases. The alias itself is not necessary; if it has users, the users // are necessarily scheduled after the alias. And if it has no users, it is // either a dead value or an output, both of which are handled below. - std::vector dead_buffers_to_free; + // + // We ignore control dependencies here. The reasoning is that the control + // dependencies have already been accounted for in the ordering of the given + // 'instruction_sequence', and should not otherwise artificially extend the + // lifetime of buffers that aren't already connected by a data dependency. + dead_buffers_to_free.clear(); for (const LogicalBuffer* buffer : buffers_defined_by_instruction) { - if (heap.IgnoreBuffer(buffer)) { + if (IgnoreBuffer(buffer)) { continue; } + FlatSet* live_set = nullptr; for (const BufferAlias& alias : points_to_analysis.GetBufferAliases(*buffer)) { const std::vector& users = alias.instruction()->users(); if (!users.empty()) { - live_buffers[buffer].insert(users.begin(), users.end()); + if (live_set == nullptr) { + live_set = &live_buffers[buffer]; + } + live_set->insert(users.begin(), users.end()); } } @@ -119,15 +167,17 @@ StatusOr HeapSimulator::Run( // all source buffers of all operands of this instruction. Buffers that // have no instructions left to visit are moved from live_buffers to // operand_buffers_to_free. - std::vector operand_buffers_to_free; + operand_buffers_to_free.clear(); for (const LogicalBuffer* operand_buffer : UniqueOperandSourceBuffers(instruction, points_to_analysis)) { - if (heap.IgnoreBuffer(operand_buffer)) { + if (IgnoreBuffer(operand_buffer)) { continue; } - live_buffers[operand_buffer].erase(instruction); - if (live_buffers[operand_buffer].empty()) { - live_buffers.erase(operand_buffer); + auto it = live_buffers.find(operand_buffer); + FlatSet* live_set = &it->second; + live_set->erase(instruction); + if (live_set->empty()) { + live_buffers.erase(it); operand_buffers_to_free.push_back(operand_buffer); } } @@ -137,10 +187,10 @@ StatusOr HeapSimulator::Run( // happen before dead or operand buffers are freed; the instruction reads // the operand buffers to produce its output. // - // INVARIANT: Either heap.Alloc or heap.ShareBuffer will be called for each - // buffer that we should assign. + // INVARIANT: Either Alloc or ShareBuffer will be called for each buffer + // that we should assign. for (const LogicalBuffer* buffer : buffers_defined_by_instruction) { - if (heap.IgnoreBuffer(buffer)) { + if (IgnoreBuffer(buffer)) { continue; } @@ -151,27 +201,53 @@ StatusOr HeapSimulator::Run( bool shared = false; for (const LogicalBuffer* operand_buffer : operand_buffers_to_free) { if (buffer->instruction()->IsUserOf(operand_buffer->instruction()) && + buffer->instruction()->opcode() != HloOpcode::kCopy && CanShareOperandBufferWithUser( operand_buffer->instruction(), operand_buffer->index(), - buffer->instruction(), buffer->index(), points_to_analysis)) { - heap.ShareBuffer(buffer, operand_buffer); + buffer->instruction(), buffer->index(), &points_to_analysis)) { + ShareBuffer(buffer, operand_buffer, instruction); shared = true; break; } } if (!shared) { - heap.Alloc(buffer); + Alloc(buffer, instruction); } } + // If the whole module is sequential, we can save memory by running the + // heap-simulation for sub-computations inline. E.g. the buffers for the + // condition and body of a kWhile instruction are only live for the duration + // of the instruction itself. + // + // The order that the sub-computations are simulated does not affect + // correctness; since the whole module is sequential, we know that the + // sub-computations will never be run concurrently. + if (module_sequence_ != nullptr) { + if (instruction->opcode() == HloOpcode::kCall || + instruction->opcode() == HloOpcode::kWhile) { + for (const HloComputation* called_computation : + instruction->called_computations()) { + const std::vector& called_sequence = + FindOrDie(*module_sequence_, called_computation); + TF_RETURN_IF_ERROR(RunComputation( + *called_computation, called_sequence, points_to_analysis)); + } + } + + // Other sub-computations (e.g. Map, Reduce, ...) are skipped; they are + // assigned "thread-local" allocations, meaning their buffers are not + // allocated up-front at the beginning of the computation. + } + // Free buffers that are no longer live. This is the earliest point that we // can de-allocate; right after the last use of the buffer. for (const LogicalBuffer* buffer : dead_buffers_to_free) { - heap.Free(buffer); + Free(buffer, instruction); } for (const LogicalBuffer* buffer : operand_buffers_to_free) { - heap.Free(buffer); + Free(buffer, instruction); } } @@ -182,20 +258,24 @@ StatusOr HeapSimulator::Run( const FlatSet& pending = buffer_pending.second; CHECK_EQ(pending.size(), 1) << *buffer; CHECK(*pending.begin() == nullptr) << *buffer; - heap.Free(buffer); + Free(buffer, root); } - return heap.Finish(); + return Status::OK(); } HeapSimulator::HeapSimulator( std::unique_ptr algorithm, const LogicalBuffer::SizeFunction& size_fn, - const FlatSet* buffers_to_assign) + const FlatSet* buffers_to_assign, + const SequentialHloOrdering::HloModuleSequence* module_sequence) : no_fragmentation_stats_(MakeUnique()), algorithm_(std::move(algorithm)), size_fn_(size_fn), - buffers_to_assign_(buffers_to_assign) {} + buffers_to_assign_(buffers_to_assign), + module_sequence_(module_sequence) { + debug_trace_.set_whole_module_simulation(module_sequence_ != nullptr); +} HeapSimulator::~HeapSimulator() {} @@ -210,7 +290,8 @@ bool HeapSimulator::IgnoreBuffer(const LogicalBuffer* buffer) const { } // Alloc always calls the underlying heap algorithm. -void HeapSimulator::Alloc(const LogicalBuffer* buffer) { +void HeapSimulator::Alloc(const LogicalBuffer* buffer, + const HloInstruction* instruction) { CHECK(allocated_buffers_.count(buffer) == 0) << "Alloc called on allocated buffer: " << *buffer; CHECK(freed_buffers_.count(buffer) == 0) @@ -220,13 +301,17 @@ void HeapSimulator::Alloc(const LogicalBuffer* buffer) { const int64 size = size_fn_(*buffer); algorithm_->Alloc(buffer, size); no_fragmentation_stats_->Alloc(buffer, size); + + FillDebugTrace(HeapSimulatorTrace::Event::ALLOC, buffer, instruction, + nullptr); } // Free calls the underlying algorithm for non-shared buffers, and for shared // buffers whose group liveness has expired. Shared group liveness is tracked // by maintaining a refcount; the Free call on the last buffer in the group // causes Free to be called on the underlying algorithm. -void HeapSimulator::Free(const LogicalBuffer* buffer) { +void HeapSimulator::Free(const LogicalBuffer* buffer, + const HloInstruction* instruction) { auto shared_it = shared_buffers_.find(buffer); if (shared_it != shared_buffers_.end()) { std::shared_ptr group = shared_it->second; @@ -248,6 +333,8 @@ void HeapSimulator::Free(const LogicalBuffer* buffer) { const int64 size = size_fn_(*buffer); algorithm_->Free(buffer, size); no_fragmentation_stats_->Free(buffer, size); + + FillDebugTrace(HeapSimulatorTrace::Event::FREE, buffer, instruction, nullptr); } // ShareBuffer associates buffers with their SharedGroup in shared_buffers_. @@ -255,7 +342,8 @@ void HeapSimulator::Free(const LogicalBuffer* buffer) { // Alloc. The 'shared' buffer must be a previously allocated or shared buffer. // Both 'buffer' and 'shared' will be associated with the same SharedGroup. void HeapSimulator::ShareBuffer(const LogicalBuffer* buffer, - const LogicalBuffer* shared) { + const LogicalBuffer* shared, + const HloInstruction* instruction) { CHECK_LE(size_fn_(*buffer), size_fn_(*shared)) << "ShareBuffer oversized buffer" << *buffer << " shared: " << *shared; CHECK(allocated_buffers_.count(buffer) == 0) @@ -265,11 +353,13 @@ void HeapSimulator::ShareBuffer(const LogicalBuffer* buffer, CHECK(freed_buffers_.count(shared) == 0) << "ShareBuffer called on freed shared buffer: " << *shared; + const LogicalBuffer* canonical = nullptr; auto shared_it = shared_buffers_.find(shared); if (shared_it != shared_buffers_.end()) { // The 'shared' buffer already has a group; it might be the canonical, but // also might not be. Just add 'buffer' to the existing group. std::shared_ptr group = shared_it->second; + canonical = group->canonical; ++group->refcount; shared_buffers_.emplace(buffer, group); } else { @@ -278,11 +368,15 @@ void HeapSimulator::ShareBuffer(const LogicalBuffer* buffer, CHECK(allocated_buffers_.count(shared) > 0) << "ShareBuffer called on non-allocated shared buffer: " << *shared; auto group = std::make_shared(); - group->canonical = shared; + canonical = shared; + group->canonical = canonical; group->refcount = 2; shared_buffers_.emplace(buffer, group); shared_buffers_.emplace(shared, group); } + + FillDebugTrace(HeapSimulatorTrace::Event::SHARE_WITH, buffer, instruction, + canonical); } HeapSimulator::Result HeapSimulator::Finish() { @@ -304,15 +398,40 @@ HeapSimulator::Result HeapSimulator::Finish() { result.chunk_map.emplace(buffer, chunk); } } + // If we were told to assign specific buffers, make sure we've assigned + // exactly that many buffers. + if (buffers_to_assign_ != nullptr) { + CHECK_EQ(buffers_to_assign_->size(), result.chunk_map.size()); + } } // Fragmentation is the difference between the actual and ideal sizes. const Result no_frag_result = no_fragmentation_stats_->Finish(); result.fragmentation_size = result.heap_size - no_frag_result.heap_size; + // Copy the debug trace we collected to the final result. + result.debug_trace.Swap(&debug_trace_); + return result; } +void HeapSimulator::FillDebugTrace(HeapSimulatorTrace::Event::Kind kind, + const LogicalBuffer* buffer, + const HloInstruction* instruction, + const LogicalBuffer* share_with_canonical) { + HeapSimulatorTrace::Event* event = debug_trace_.add_events(); + event->set_kind(kind); + event->set_buffer_id(buffer->id()); + event->set_computation_name(instruction->parent()->name()); + event->set_instruction_name(instruction->name()); + if (kind == HeapSimulatorTrace::Event::SHARE_WITH) { + CHECK(share_with_canonical != nullptr); + event->set_share_with_canonical_id(share_with_canonical->id()); + } else { + CHECK(share_with_canonical == nullptr); + } +} + void NoFragmentationStatsHeap::Alloc(const LogicalBuffer* buffer, int64 size) { current_heap_size_ += size; if (current_heap_size_ > max_heap_size_) { diff --git a/tensorflow/compiler/xla/service/heap_simulator.h b/tensorflow/compiler/xla/service/heap_simulator.h index 0ce2906767898bcace45e296d76f958c50a2b3a7..a03ad2f37cf5ede35275ea019ab3d5998fb85d0a 100644 --- a/tensorflow/compiler/xla/service/heap_simulator.h +++ b/tensorflow/compiler/xla/service/heap_simulator.h @@ -21,8 +21,10 @@ limitations under the License. #include #include +#include "tensorflow/compiler/xla/service/hlo.pb.h" #include "tensorflow/compiler/xla/service/hlo_computation.h" #include "tensorflow/compiler/xla/service/hlo_instruction.h" +#include "tensorflow/compiler/xla/service/hlo_ordering.h" #include "tensorflow/compiler/xla/service/logical_buffer.h" #include "tensorflow/compiler/xla/service/tuple_points_to_analysis.h" #include "tensorflow/compiler/xla/statusor.h" @@ -60,42 +62,76 @@ class HeapSimulator { // The total size in bytes of heap fragmentation. int64 fragmentation_size = 0; + + // A trace of heap simulation events. + HeapSimulatorTrace debug_trace; }; // Run the heap simulation with the given algorithm, assuming the given - // sequential ordering of instructions. The 'instruction_sequence' must - // contain a topologically-consistent total ordering of all instructions in - // the computation. The result is invalid if instructions are not run in - // exactly this sequence. + // module_sequence, which must contain a topologically-consistent total + // ordering of all instructions within each computation. The result is invalid + // if instructions are not run in exactly this sequence. + // + // Running heap simulation on the whole module tends to save memory, compared + // to running on a per-computation basis, since we can re-use buffer space for + // called sub-computations. // // If 'buffers_to_assign' is provided, only those buffers are assigned // offsets, otherwise all buffers defined by the instructions are assigned. + static StatusOr Run( + std::unique_ptr algorithm, const HloModule& module, + const SequentialHloOrdering::HloModuleSequence& module_sequence, + const TuplePointsToAnalysis& points_to_analysis, + const LogicalBuffer::SizeFunction& size_fn, + const tensorflow::gtl::FlatSet* buffers_to_assign = + nullptr); + + // Same as above, but runs on a single computation. The 'instruction_sequence' + // must contain a topologically-consistent total ordering of all instructions + // in the computation. The result is invalid if instructions are not run in + // exactly this sequence. static StatusOr Run( std::unique_ptr algorithm, - const std::vector& instruction_sequence, const HloComputation& computation, + const std::vector& instruction_sequence, const TuplePointsToAnalysis& points_to_analysis, const LogicalBuffer::SizeFunction& size_fn, const tensorflow::gtl::FlatSet* buffers_to_assign = nullptr); private: + // If 'module_sequence' is non-null, it is used to find kCall and kWhile + // sub-computations, and the heap simulation for those sub-computations will + // be run recursively. I.e. the simulation is run over the whole module. HeapSimulator( std::unique_ptr algorithm, const LogicalBuffer::SizeFunction& size_fn, - const tensorflow::gtl::FlatSet* buffers_to_assign); + const tensorflow::gtl::FlatSet* buffers_to_assign, + const SequentialHloOrdering::HloModuleSequence* module_sequence); ~HeapSimulator(); + Status RunComputation( + const HloComputation& computation, + const std::vector& instruction_sequence, + const TuplePointsToAnalysis& points_to_analysis); + bool IgnoreBuffer(const LogicalBuffer* buffer) const; - void Alloc(const LogicalBuffer* buffer); - void Free(const LogicalBuffer* buffer); - void ShareBuffer(const LogicalBuffer* buffer, const LogicalBuffer* shared); + void Alloc(const LogicalBuffer* buffer, const HloInstruction* instruction); + void Free(const LogicalBuffer* buffer, const HloInstruction* instruction); + void ShareBuffer(const LogicalBuffer* buffer, const LogicalBuffer* shared, + const HloInstruction* instruction); Result Finish(); + void FillDebugTrace(HeapSimulatorTrace::Event::Kind kind, + const LogicalBuffer* buffer, + const HloInstruction* instruction, + const LogicalBuffer* shared_with_canonical); + const std::unique_ptr no_fragmentation_stats_; const std::unique_ptr algorithm_; const LogicalBuffer::SizeFunction size_fn_; const tensorflow::gtl::FlatSet* buffers_to_assign_; + const SequentialHloOrdering::HloModuleSequence* module_sequence_; // In addition to Alloc and Free, the heap simulator exposes a concept of // buffer sharing. When ShareBuffer is called, instead of allocating new @@ -121,6 +157,9 @@ class HeapSimulator { // Hold some sets for error-checking the sequence of Alloc and Free calls. tensorflow::gtl::FlatSet allocated_buffers_; tensorflow::gtl::FlatSet freed_buffers_; + + // Debugging information filled in while the heap simulator runs. + HeapSimulatorTrace debug_trace_; }; // Abstract base class describing a heap simulation algorithm that assigns diff --git a/tensorflow/compiler/xla/service/heap_simulator_test.cc b/tensorflow/compiler/xla/service/heap_simulator_test.cc index 874bd5f1060c179d5547510c351909069aa935b8..17b926c8748e45b55f380e7595711b9e7a748f64 100644 --- a/tensorflow/compiler/xla/service/heap_simulator_test.cc +++ b/tensorflow/compiler/xla/service/heap_simulator_test.cc @@ -19,13 +19,16 @@ limitations under the License. #include #include +#include "tensorflow/compiler/xla/literal_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/service/hlo_ordering.h" #include "tensorflow/compiler/xla/service/logical_buffer.h" #include "tensorflow/compiler/xla/service/tuple_points_to_analysis.h" #include "tensorflow/compiler/xla/status_macros.h" #include "tensorflow/compiler/xla/tests/hlo_test_base.h" +#include "tensorflow/core/lib/gtl/flatmap.h" namespace xla { namespace { @@ -69,6 +72,7 @@ class HeapCallRecorder : public HeapAlgorithm { // sequence against an expected sequence. class HeapSimulatorTracker { public: + // Constructor for testing a single entry computation. HeapSimulatorTracker( const string& name, std::unique_ptr computation, const std::vector& instruction_sequence) { @@ -83,12 +87,48 @@ class HeapSimulatorTracker { auto zero_size = [](const LogicalBuffer& buffer) { return 0; }; auto algorithm = MakeUnique( MakeUnique(&actual_calls_)); - result_ = HeapSimulator::Run(std::move(algorithm), instruction_sequence, - *module_->entry_computation(), - *points_to_analysis_, zero_size) + result_ = HeapSimulator::Run( + std::move(algorithm), *module_->entry_computation(), + instruction_sequence, *points_to_analysis_, zero_size) .ConsumeValueOrDie(); } + explicit HeapSimulatorTracker(const string& name) { + module_ = MakeUnique(name); + } + + // Similar to the single entry computation constructor above, but runs the + // simulation over the entire module. + void RunWholeModule( + const std::vector& full_module_sequence) { + points_to_analysis_ = + TuplePointsToAnalysis::Run(module_.get()).ConsumeValueOrDie(); + + // Construct the module sequence grouped by computation. + SequentialHloOrdering::HloModuleSequence module_sequence; + tensorflow::gtl::FlatMap reverse_position; + for (int i = 0; i < full_module_sequence.size(); ++i) { + const HloInstruction* instruction = full_module_sequence[i]; + module_sequence[instruction->parent()].push_back(instruction); + reverse_position[instruction] = full_module_sequence.size() - i; + } + + // Hack the size_fn so that it returns a decreasing value as we step through + // the sequence. This lets us ensure the Alloc calls are in the sequence + // order. The Free calls are sorted by LogicalBuffer.id, which is at least + // deterministic. + auto size_fn = [&reverse_position](const LogicalBuffer& buffer) { + return reverse_position[buffer.instruction()]; + }; + auto algorithm = MakeUnique( + MakeUnique(&actual_calls_)); + result_ = HeapSimulator::Run(std::move(algorithm), *module_, + module_sequence, *points_to_analysis_, size_fn) + .ConsumeValueOrDie(); + } + + HloModule* module() { return module_.get(); } + // Returns the buffer defined at the given instruction and index. const LogicalBuffer* BufferAt(const HloInstruction* instruction, const ShapeIndex& index) const { @@ -133,7 +173,7 @@ class HeapSimulatorTest : public HloTestBase { TEST_F(HeapSimulatorTest, ScalarConstant) { auto builder = HloComputation::Builder(TestName()); auto const0 = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(1.0))); + HloInstruction::CreateConstant(Literal::CreateR0(1.0))); // Constants aren't assigned. See b/32248867 HeapSimulatorTracker tracker(TestName(), builder.Build(), {const0}); @@ -358,10 +398,90 @@ TEST_F(HeapSimulatorTest, MultiplyDotDotTuple) { }); } +TEST_F(HeapSimulatorTest, WholeModule) { + HeapSimulatorTracker tracker(TestName()); + + const Shape scalar_shape = ShapeUtil::MakeShape(xla::F32, {}); + const Shape tuple_shape = + ShapeUtil::MakeTupleShape({scalar_shape, scalar_shape}); + + auto cond_builder = HloComputation::Builder("WhileCond"); + HloInstruction* cond_param = cond_builder.AddInstruction( + HloInstruction::CreateParameter(0, tuple_shape, "cond_param")); + HloInstruction* cond_iter = cond_builder.AddInstruction( + HloInstruction::CreateGetTupleElement(scalar_shape, cond_param, 0)); + HloInstruction* cond_data = cond_builder.AddInstruction( + HloInstruction::CreateGetTupleElement(scalar_shape, cond_param, 1)); + HloInstruction* cond_lt = cond_builder.AddInstruction( + HloInstruction::CreateBinary(ShapeUtil::MakeShape(PRED, {}), + HloOpcode::kLt, cond_iter, cond_data)); + HloComputation* cond_computation = + tracker.module()->AddEmbeddedComputation(cond_builder.Build()); + + auto body_builder = HloComputation::Builder("WhileBody"); + HloInstruction* body_param = body_builder.AddInstruction( + HloInstruction::CreateParameter(0, tuple_shape, "body_param")); + HloComputation* body_computation = + tracker.module()->AddEmbeddedComputation(body_builder.Build()); + + auto builder = HloComputation::Builder(TestName()); + HloInstruction* param = builder.AddInstruction( + HloInstruction::CreateParameter(0, tuple_shape, "param")); + HloInstruction* while_op = builder.AddInstruction(HloInstruction::CreateWhile( + tuple_shape, cond_computation, body_computation, param)); + tracker.module()->AddEntryComputation(builder.Build()); + + tracker.RunWholeModule( + {param, while_op, body_param, cond_param, cond_iter, cond_data, cond_lt}); + tracker.ExpectCallSequence({ + // The entry computation param and while_op are allocated first. + {kAlloc, tracker.BufferAt(param, {})}, + {kAlloc, tracker.BufferAt(param, {0})}, + {kAlloc, tracker.BufferAt(param, {1})}, + {kAlloc, tracker.BufferAt(while_op, {})}, + {kAlloc, tracker.BufferAt(while_op, {0})}, + {kAlloc, tracker.BufferAt(while_op, {1})}, + + // Now the while body param is allocated and freed. + {kAlloc, tracker.BufferAt(body_param, {})}, + {kAlloc, tracker.BufferAt(body_param, {0})}, + {kAlloc, tracker.BufferAt(body_param, {1})}, + {kFree, tracker.BufferAt(body_param, {})}, + {kFree, tracker.BufferAt(body_param, {0})}, + {kFree, tracker.BufferAt(body_param, {1})}, + + // Now the while cond param is allocated. The GTE instructions just alias + // the param elements, so the param tuple can immediately be freed. + {kAlloc, tracker.BufferAt(cond_param, {})}, + {kAlloc, tracker.BufferAt(cond_param, {0})}, + {kAlloc, tracker.BufferAt(cond_param, {1})}, + {kFree, tracker.BufferAt(cond_param, {})}, + + // Now the final cond less-than buffer is allocated. + {kAlloc, tracker.BufferAt(cond_lt, {})}, + + // The order of the remaining Free calls is based on the LogicalBuffer.id, + // which is deterministic, but not obvious. + {kFree, tracker.BufferAt(param, {})}, + {kFree, tracker.BufferAt(param, {0})}, + {kFree, tracker.BufferAt(param, {1})}, + + {kFree, tracker.BufferAt(while_op, {})}, + {kFree, tracker.BufferAt(while_op, {0})}, + {kFree, tracker.BufferAt(while_op, {1})}, + + {kFree, tracker.BufferAt(cond_param, {0})}, + {kFree, tracker.BufferAt(cond_param, {1})}, + {kFree, tracker.BufferAt(cond_lt, {})}, + + {kFinish, nullptr}, + }); +} + // Base class for heap algorithm tests. class HeapAlgorithmTestBase : public ::testing::Test { protected: - HeapAlgorithmTestBase() { + HeapAlgorithmTestBase() : builder_("heap_simulator_test") { buffer_a_ = DummyLogicalBuffer(); buffer_b_ = DummyLogicalBuffer(); buffer_c_ = DummyLogicalBuffer(); @@ -385,15 +505,16 @@ class HeapAlgorithmTestBase : public ::testing::Test { const LogicalBuffer* buffer_i_; private: - // Create a dummy LogicalBuffer to pass to the heap algorithm. Since the - // algorithms only use the buffer as a handle, we don't need to fill in much - // other than the id. + // Create a dummy LogicalBuffer to pass to the heap algorithm. const LogicalBuffer* DummyLogicalBuffer() { const LogicalBuffer::Id id = buffers_.size(); - buffers_.emplace_back(MakeUnique(nullptr, ShapeIndex{}, id)); + auto const0 = builder_.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(1.0))); + buffers_.emplace_back(MakeUnique(const0, ShapeIndex{}, id)); return buffers_.back().get(); } + HloComputation::Builder builder_; std::vector> buffers_; }; diff --git a/tensorflow/compiler/xla/service/hlo.proto b/tensorflow/compiler/xla/service/hlo.proto new file mode 100644 index 0000000000000000000000000000000000000000..af853385d634b06d31cef94216fb4059dfcadc3d --- /dev/null +++ b/tensorflow/compiler/xla/service/hlo.proto @@ -0,0 +1,189 @@ +/* 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. +==============================================================================*/ + +// DO NOT USE THESE PROTO MESSAGES FOR ANYTHING OTHER THAN DEBUGGING. +// +// Don't use these protos in the real compilation or execution codepaths. The +// data format is meant for debugging only, and may change without notice. +// +// Many of the protos below are simple 1-to-1 serializations of the +// corresponding C++ classes. +// +// FIELD NAMES ARE IMPORTANT +// +// Unlike most protos, you can't safely change the names of fields, even if you +// keep the numeric ids the same. This is because we sometimes serialize these +// protos as JSON, which includes the field names in the serialization. + +syntax = "proto3"; + +package xla; +import "tensorflow/compiler/xla/xla_data.proto"; + +option cc_enable_arenas = true; + +// Serialization of HloInstruction. +message HloInstructionProto { + string name = 1; + string opcode = 2; + xla.Shape shape = 3; + repeated string operand_names = 4; + repeated string control_predecessor_names = 5; + repeated string called_computation_names = 6; + + xla.OpMetadata metadata = 7; + + // Literal, only present for kConstant. + xla.LiteralProto literal = 8; + + // Parameter info, only present for kParameter. + int64 parameter_number = 9; + string parameter_name = 10; + + // Fusion state, only present for kFusion. + string fusion_kind = 11; + HloComputationProto fused_instructions_computation = 12; + + // Index for kGetTupleElement. + int64 tuple_index = 13; +} + +// Serialization of HloComputation. +message HloComputationProto { + string name = 1; + + // The array of instructions is always in a valid dependency order, where + // operands appear before their users. + repeated HloInstructionProto instructions = 2; +} + +// Serialization of HloModule. +message HloModuleProto { + string name = 1; + string entry_computation_name = 2; + + // The array of computations is always in a valid dependency order, where + // callees appear before their callers. + repeated HloComputationProto computations = 3; +} + +// Serialization of HloOrdering. +message HloOrderingProto { + // NOTE: currently only sequential orderings are serialized. + message SequentialComputation { + string computation_name = 1; + repeated string instruction_names = 2; + } + repeated SequentialComputation sequential_computations = 1; +} + +// Serialization of LogicalBuffer. +message LogicalBufferProto { + // Location represents an instruction and its shape index, which uniquely + // identifies a point where a buffer is needed. + message Location { + // NOTE: module_name isn't necessary, since all LogicalBuffers are + // associated with a single HloModule. + string computation_name = 1; + string instruction_name = 2; + repeated int64 shape_index = 3; + } + + int64 id = 1; + int64 size = 2; + + // The location where the buffer is defined. + Location defined_at = 3; + + int64 color = 4; +} + +// Serialization of BufferAllocation. +message BufferAllocationProto { + // Assigned represents a single LogicalBuffer that is assigned to this + // BufferAllocation. + message Assigned { + int64 logical_buffer_id = 1; + int64 offset = 2; + int64 size = 3; + } + + int64 index = 1; + int64 size = 2; + bool is_thread_local = 3; + bool is_reusable = 4; + bool is_entry_computation_parameter = 5; + int64 parameter_number = 6; + bool maybe_live_out = 7; + int64 color = 8; + repeated Assigned assigned = 9; +} + +// A trace of a HeapSimulator run. +message HeapSimulatorTrace { + // The trace includes a list of events, where each event describes one action + // performed by the heap simulator. + message Event { + enum Kind { + ALLOC = 0; // A memory region was allocated for the buffer. + FREE = 1; // A memory region was freed for the buffer. + + // A buffer was shared with another (canonical) buffer. This is similar to + // ALLOC, except that instead of allocating a new region of memory, the + // memory region of the canonical buffer is directly re-used. Multiple + // buffers may share with the same canonical buffer. The lifetime of the + // canonical buffer is extended to the union of all lifetimes. + SHARE_WITH = 2; + } + Kind kind = 1; + + // The id of the LogicalBuffer that the event applies to. + int64 buffer_id = 2; + + // The HloInstruction that the simulation was processing that caused this + // event to occur, identified by its computation and instruction name. E.g. + // buffers defined by instruction A are allocated when processing A. + string computation_name = 3; + string instruction_name = 4; + + // The id of the canonical LogicalBuffer that the buffer shares with. Only + // set for SHARE_WITH events. + int64 share_with_canonical_id = 5; + } + repeated Event events = 1; + bool whole_module_simulation = 2; +} + +// Serialization of BufferAssignment. +message BufferAssignmentProto { + // Alias represents a source LogicalBuffer, and the buffer location that + // aliases it. + message BufferAlias { + int64 source_buffer_id = 1; + LogicalBufferProto.Location location = 2; + } + + repeated LogicalBufferProto logical_buffers = 1; + repeated BufferAlias buffer_aliases = 2; + repeated BufferAllocationProto buffer_allocations = 3; + repeated HeapSimulatorTrace heap_simulator_traces = 4; +} + +// Grouping message that contains all of the information above. +message HloProto { + HloModuleProto hlo_module = 1; + HloOrderingProto hlo_ordering = 2; + BufferAssignmentProto buffer_assignment = 3; +} diff --git a/tensorflow/compiler/xla/service/hlo_alias_analysis.cc b/tensorflow/compiler/xla/service/hlo_alias_analysis.cc new file mode 100644 index 0000000000000000000000000000000000000000..3dd8ac6dc5fa46b80328e080e6d1b4e8c402e8b0 --- /dev/null +++ b/tensorflow/compiler/xla/service/hlo_alias_analysis.cc @@ -0,0 +1,452 @@ +/* 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/hlo_alias_analysis.h" + +#include +#include +#include +#include + +#include "tensorflow/compiler/xla/map_util.h" +#include "tensorflow/compiler/xla/service/hlo_buffer.h" +#include "tensorflow/compiler/xla/service/hlo_instruction.h" +#include "tensorflow/compiler/xla/service/hlo_value.h" +#include "tensorflow/compiler/xla/shape_util.h" +#include "tensorflow/compiler/xla/types.h" +#include "tensorflow/compiler/xla/util.h" +#include "tensorflow/core/lib/core/errors.h" +#include "tensorflow/core/lib/strings/str_util.h" +#include "tensorflow/core/lib/strings/strcat.h" +#include "tensorflow/core/platform/logging.h" + +namespace xla { + +using ::tensorflow::strings::StrAppend; +using ::tensorflow::strings::StrCat; + +// Data structure used to construct the alias analysis. Thrown away after alias +// analysis is complete. This data structure keeps track of which sets of +// HloValues must be in the same HloBuffer. This is maintained as a map from a +// buffer identifier (BufferNumber) to set of HLoValues. +// +// Initially each value is its own buffer. In MergeAliasedBuffers, sets of +// values which must share the same buffer are merged together. The end result +// is a partitioning of all HloValues into sets where each set needs its own +// HloBuffer. By performing this analysis without constructing HloBuffers on the +// fly, we can after-the-fact construct a vector of contiguously numbered +// HloBuffers after the buffer requirement has been determined. +class BufferValueMap { + public: + // A unique identifier for a set of colocated values which must share the same + // buffer. This is not necessarily the same as the HloBuffer::Id which will + // ultimately contain the values. The reason is that HloBuffer::Id's are + // contiguous, while BufferNumbers may not be. BufferNumbers may not be + // dense because buffers may be created and destroyed during the analysis + // construction process. + using BufferNumber = int64; + + explicit BufferValueMap(const HloDataflowAnalysis& dataflow) + : dataflow_(dataflow) { + buffers_.reserve(dataflow_.values().size()); + value_to_buffer_number_.reserve(dataflow_.values().size()); + for (const HloValue* value : dataflow_.values()) { + BufferNumber buffer_number = next_buffer_number_++; + buffers_[buffer_number].insert(value); + value_to_buffer_number_[value] = buffer_number; + } + } + + // Merge together sets of HloValues which must be in the same HloBuffer + // because of aliasing rules (eg, in-place kWhile instruction). + void MergeAliasedBuffers() { + for (const HloValue* value : dataflow_.values()) { + VLOG(3) << "Merging colocated values, value: " << value->ToShortString(); + + // Gather the set of buffers with aliasing rules (eg, kWhile) which this + // value must be contained in. + std::vector aliased_buffers = ComputeAliasedBuffers(*value); + + BufferNumber current_buffer = value_to_buffer_number_.at(value); + if (aliased_buffers.empty()) { + // The buffer containing 'value' aliases no other buffers. If the buffer + // containing 'value' already only contains 'value', then no change is + // necessary. If the buffer containing 'value' does contain other + // values, then remove 'value' from the buffer and create a new buffer + // containing only 'value' + if (buffers_.at(current_buffer).size() == 1) { + CHECK_EQ(*buffers_.at(current_buffer).begin(), value); + } else { + MoveValueToNewBuffer(*value); + } + } else { + // If multiple buffers are aliased merge these buffers together into a + // single buffer (arbitrarily chosen as the first buffer in the vector). + if (aliased_buffers.size() > 1) { + for (int64 i = 1; i < aliased_buffers.size(); ++i) { + MergeBuffers(/*from=*/aliased_buffers[i], + /*to=*/aliased_buffers[0]); + } + } + BufferNumber new_buffer = aliased_buffers[0]; + if (current_buffer != new_buffer) { + MoveValueToBuffer(*value, new_buffer); + } + } + } + } + + // Compute and return a sorted vector of all BufferNumbers. Can be used to + // iterate through all buffers stabily. + std::vector ComputeSortedBufferNumbers() const { + std::vector buffer_numbers; + for (const auto& pair : buffers_) { + buffer_numbers.push_back(pair.first); + } + std::sort(buffer_numbers.begin(), buffer_numbers.end()); + return buffer_numbers; + } + + // Return a set of all the values in the given buffer. + const tensorflow::gtl::FlatSet& GetValuesInBuffer( + BufferNumber buffer_number) const { + return buffers_.at(buffer_number); + } + + private: + // Create a new buffer. + void NewBuffer(const HloValue& value) { + BufferNumber buffer_number = next_buffer_number_++; + buffers_[buffer_number].insert(&value); + value_to_buffer_number_[&value] = buffer_number; + } + + // Move the given value into a new buffer containing only the value. + void MoveValueToNewBuffer(const HloValue& value) { + BufferNumber new_buffer_number = next_buffer_number_++; + buffers_[new_buffer_number]; + MoveValueToBuffer(value, new_buffer_number); + } + + // Move the given value into the given buffer. + void MoveValueToBuffer(const HloValue& value, BufferNumber buffer_number) { + BufferNumber old_buffer_number = value_to_buffer_number_.at(&value); + buffers_.at(old_buffer_number).erase(&value); + if (buffers_.at(old_buffer_number).empty()) { + buffers_.erase(old_buffer_number); + } + + buffers_.at(buffer_number).insert(&value); + value_to_buffer_number_.at(&value) = buffer_number; + } + + // Merge the buffer 'from' into the buffer 'to'. + void MergeBuffers(BufferNumber from, BufferNumber to) { + auto& from_value_set = buffers_.at(from); + buffers_.at(to).insert(from_value_set.begin(), from_value_set.end()); + // NOTE: using a union-find algorithm to hold the colocated values might be + // faster. + for (const HloValue* value : from_value_set) { + value_to_buffer_number_.at(value) = to; + } + buffers_.erase(from); + } + + BufferNumber GetBufferForValue(const HloValue& value) { + return value_to_buffer_number_.at(&value); + } + + // Compute and return a vector of buffers that the given value must be + // contained in due to HLO aliasing rules. + std::vector ComputeAliasedBuffers(const HloValue& value) { + // Value is init of a while (use is while). + std::vector aliased_buffers; + for (const HloUse& use : value.uses()) { + VLOG(1) << "use of value " << value.ToShortString() << ": " << use; + if (use.instruction->opcode() == HloOpcode::kWhile) { + // Determine the while value that this shares a buffer with. + const HloValue& while_value = + dataflow_.GetUniqueValueAt(use.instruction, use.operand_index); + aliased_buffers.push_back(GetBufferForValue(while_value)); + VLOG(3) << " value is init value to a while; must share buffer with " + "while value " + << while_value.ToShortString(); + } + } + + // Value is a parameter of a while body/condition. + if (value.defining_instruction()->opcode() == HloOpcode::kParameter) { + const HloComputation* computation = + value.defining_instruction()->parent(); + const CallGraphNode& call_graph_node = + dataflow_.call_graph().GetNode(computation); + for (const CallSite& callsite : call_graph_node.caller_callsites()) { + if (callsite.instruction()->opcode() == HloOpcode::kWhile) { + // Call graph must have been flattened. + CHECK_EQ(call_graph_node.caller_callsites().size(), 1); + + const HloValue& while_value = dataflow_.GetUniqueValueAt( + callsite.instruction(), value.defining_index()); + VLOG(3) << " value is parameter value of the body or condition of a " + "while; must share buffer with while value " + << while_value.ToShortString(); + aliased_buffers.push_back(GetBufferForValue(while_value)); + } + } + } + + // Value is the root of a while body. + for (const HloPosition& position : value.positions()) { + const HloComputation* computation = position.instruction->parent(); + const CallGraphNode& call_graph_node = + dataflow_.call_graph().GetNode(computation); + if (position.instruction == computation->root_instruction()) { + for (const CallSite& callsite : call_graph_node.caller_callsites()) { + if (callsite.instruction()->opcode() == HloOpcode::kWhile && + callsite.instruction()->while_body() == computation) { + // Call graph must have been flattened. + CHECK_EQ(call_graph_node.caller_callsites().size(), 1); + + const HloValue& while_value = dataflow_.GetUniqueValueAt( + callsite.instruction(), position.index); + VLOG(3) << " value is root the body computation of a while; must " + "share buffer with while value " + << while_value.ToShortString(); + aliased_buffers.push_back(GetBufferForValue(while_value)); + } + } + } + } + + // Value is the output of the while instruction itself. + if (value.defining_instruction()->opcode() == HloOpcode::kWhile) { + VLOG(3) << " value is output of a while instruction"; + aliased_buffers.push_back(GetBufferForValue(value)); + } + + // Uniquify aliased buffers. + std::sort(aliased_buffers.begin(), aliased_buffers.end()); + aliased_buffers.erase( + std::unique(aliased_buffers.begin(), aliased_buffers.end()), + aliased_buffers.end()); + + return aliased_buffers; + } + + // Dataflow analysis used to construct the buffer map. + const HloDataflowAnalysis& dataflow_; + + // A map containing the set of values contained in each buffer. + tensorflow::gtl::FlatMap> + buffers_; + + // A map indicating which buffer each value is contained in. + tensorflow::gtl::FlatMap + value_to_buffer_number_; + + // The buffer number of the next buffer to be created. + BufferNumber next_buffer_number_ = 0; +}; + +HloAliasAnalysis::HloAliasAnalysis(HloModule* module) : module_(module) {} + +const HloBuffer& HloAliasAnalysis::GetUniqueBufferAt( + const HloInstruction* instruction, const ShapeIndex& index) const { + std::vector buffers = ComputeBuffersAt(instruction, index); + CHECK_EQ(buffers.size(), 1); + return *buffers[0]; +} + +HloBuffer& HloAliasAnalysis::GetUniqueBufferAt( + const HloInstruction* instruction, const ShapeIndex& index) { + return GetBuffer(static_cast(this) + ->GetUniqueBufferAt(instruction, index) + .id()); +} + +std::vector HloAliasAnalysis::ComputeBuffersAt( + const HloInstruction* instruction, const ShapeIndex& index) const { + std::vector buffers; + for (const HloValue* value : + dataflow_analysis_->GetValueSet(instruction, index).values()) { + buffers.push_back(&GetBufferContainingValue(*value)); + } + + // Sort and uniquify vector before returning. + std::sort(buffers.begin(), buffers.end(), HloBuffer::IdLessThan); + buffers.erase(std::unique(buffers.begin(), buffers.end()), buffers.end()); + + return buffers; +} + +bool HloAliasAnalysis::InstructionBuffersAreAmbiguous( + const HloInstruction* instruction) const { + for (const auto& pair : + dataflow_analysis_->GetInstructionValueSet(instruction)) { + const HloValueSet& value_set = pair.second; + const HloBuffer* buffer = nullptr; + for (const HloValue* value : value_set.values()) { + if (buffer == nullptr) { + buffer = &GetBufferContainingValue(*value); + } else if (buffer != &GetBufferContainingValue(*value)) { + return true; + } + } + } + return false; +} + +bool HloAliasAnalysis::InstructionBuffersAreDistinct( + const HloInstruction* instruction) const { + tensorflow::gtl::FlatSet buffers_seen; + for (const auto& pair : + dataflow_analysis_->GetInstructionValueSet(instruction)) { + const HloValueSet& value_set = pair.second; + if (value_set.values().size() == 1) { + if (!buffers_seen + .insert(&GetBufferContainingValue(value_set.GetUniqueValue())) + .second) { + return false; + } + } else { + // It's possible for multiple values at this index to have the same + // HloBuffer. This does not result in non-distictness. To account for + // this case, add all of the buffers at this index after checking + // whether each buffer exists at an earlier index. This is a corner + // case, however, as the number of values at an index is almost always + // one. + std::vector buffers_at_this_index; + for (const HloValue* value : value_set.values()) { + const HloBuffer* buffer = &GetBufferContainingValue(*value); + if (ContainsKey(buffers_seen, buffer)) { + return false; + } + buffers_at_this_index.push_back(buffer); + } + buffers_seen.insert(buffers_at_this_index.begin(), + buffers_at_this_index.end()); + } + } + return true; +} + +Status HloAliasAnalysis::Verify() const { + // Verify consistency between the value_to_buffer_ map and + // HloBuffer::values(). + 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()); + } + + for (HloBuffer::Id id = 0; id < buffers_.size(); ++id) { + const HloBuffer& buffer = buffers_[id]; + TF_RET_CHECK(buffer.id() == id); + + HloValue::Id last_value_id = -1; + for (const HloValue* value : buffer.values()) { + TF_RET_CHECK(GetBufferContainingValue(*value) == buffer); + + // Also verify the values in HloBuffer are unique and sorted by id. + TF_RET_CHECK(value->id() > last_value_id); + last_value_id = value->id(); + } + } + + return Status::OK(); +} + +string HloAliasAnalysis::ToString() const { + string out = StrCat("HloAliasAnalysis, module ", module_->name(), "\n"); + StrAppend(&out, " Buffers at each position:\n"); + for (const std::unique_ptr& computation : + module_->computations()) { + for (const std::unique_ptr& instruction : + computation->instructions()) { + StrAppend(&out, " ", instruction->name(), ":\n"); + if (ShapeUtil::IsTuple(instruction->shape())) { + ShapeUtil::ForEachSubshape( + instruction->shape(), + [&out, &instruction, this](const Shape&, const ShapeIndex& index) { + StrAppend(&out, " tuple index ", index.ToString(), ":\n"); + for (const HloBuffer* buffer : + ComputeBuffersAt(instruction.get(), index)) { + StrAppend(&out, " ", buffer->ToString(), "\n"); + } + }); + } else { + for (const HloBuffer* buffer : + ComputeBuffersAt(instruction.get(), /*index=*/{})) { + StrAppend(&out, " ", buffer->ToString(), "\n"); + } + } + } + } + + StrAppend(&out, " Buffers:\n"); + for (const HloBuffer& buffer : buffers()) { + StrAppend(&out, " ", buffer.ToString(), "\n"); + StrAppend(&out, " positions:\n"); + for (const HloPosition& position : buffer.ComputePositions()) { + StrAppend(&out, " ", position.ToString(), "\n"); + } + } + + return out; +} + +/* static */ +StatusOr> HloAliasAnalysis::Run( + HloModule* module) { + VLOG(1) << "HloAliasAnalysis::Run on module " << module->name(); + XLA_VLOG_LINES(2, module->ToString()); + + auto alias_analysis = WrapUnique(new HloAliasAnalysis(module)); + TF_ASSIGN_OR_RETURN( + alias_analysis->dataflow_analysis_, + HloDataflowAnalysis::Run(module, /*ssa_form=*/true, + /*bitcast_defines_value=*/false)); + + BufferValueMap buffer_map(alias_analysis->dataflow_analysis()); + buffer_map.MergeAliasedBuffers(); + + // Create a vector of HloBuffers, one for each set of values in the + // BufferValueMap. Create the HloBuffers as a vector of contiguously numbered + // buffers. + std::vector sorted_buffer_numbers = + buffer_map.ComputeSortedBufferNumbers(); + alias_analysis->buffers_.reserve(sorted_buffer_numbers.size()); + HloBuffer::Id next_id = 0; + for (BufferValueMap::BufferNumber buffer_number : sorted_buffer_numbers) { + 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); + alias_analysis->buffers_.emplace_back(next_id++, sorted_values); + for (const HloValue* value : sorted_values) { + alias_analysis->value_to_buffer_[value] = + &alias_analysis->buffers_.back(); + } + } + + TF_DCHECK_OK(alias_analysis->Verify()); + + XLA_VLOG_LINES(1, alias_analysis->ToString()); + return std::move(alias_analysis); +} + +} // namespace xla diff --git a/tensorflow/compiler/xla/service/hlo_alias_analysis.h b/tensorflow/compiler/xla/service/hlo_alias_analysis.h new file mode 100644 index 0000000000000000000000000000000000000000..39554e466488007bfca666b5453ebaa555f598bf --- /dev/null +++ b/tensorflow/compiler/xla/service/hlo_alias_analysis.h @@ -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. +==============================================================================*/ + +#ifndef TENSORFLOW_COMPILER_XLA_SERVICE_HLO_ALIAS_ANALYSIS_H_ +#define TENSORFLOW_COMPILER_XLA_SERVICE_HLO_ALIAS_ANALYSIS_H_ + +#include +#include +#include + +#include "tensorflow/compiler/xla/service/hlo_buffer.h" +#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/status.h" +#include "tensorflow/compiler/xla/statusor.h" +#include "tensorflow/compiler/xla/types.h" +#include "tensorflow/compiler/xla/xla_data.pb.h" +#include "tensorflow/core/lib/gtl/array_slice.h" +#include "tensorflow/core/platform/macros.h" + +namespace xla { + +// Analysis which allocates HloBuffers to HloValues. +class HloAliasAnalysis { + public: + // The callgraph of the given HloModule must be flattened + // (xla::FlattenCallGraph) prior to running the analysis. + static StatusOr> Run(HloModule* module); + + string ToString() const; + + // Return the buffer containing the given value. + const HloBuffer& GetBufferContainingValue(const HloValue& value) const { + return *value_to_buffer_.at(&value); + } + HloBuffer& GetBufferContainingValue(const HloValue& value) { + return *value_to_buffer_.at(&value); + } + + // Return the HloBuffer with the given ID. + const HloBuffer& GetBuffer(HloBuffer::Id buffer_id) const { + return buffers_.at(buffer_id); + } + HloBuffer& GetBuffer(HloBuffer::Id buffer_id) { + return buffers_.at(buffer_id); + } + + // Returns the unique buffer at the given position. CHECK fails if the buffer + // set at that position does not contain exactly one buffer. + const HloBuffer& GetUniqueBufferAt(const HloInstruction* instruction, + const ShapeIndex& index = {}) const; + HloBuffer& GetUniqueBufferAt(const HloInstruction* instruction, + const ShapeIndex& index = {}); + + // Compute the set of buffers at the given instruction and index and return as + // a vector. This set is exactly the union of the buffers containing the + // HloValues at this position. + std::vector ComputeBuffersAt( + const HloInstruction* instruction, const ShapeIndex& index = {}) const; + + // Return a vector of all HloBuffers stabily sorted by HloBuffer::Id. This + // vector is lazily computed. Mutating operations on HloAliasAnalysis may + // invalidate the underlying vector requiring recomputation. + const std::vector& buffers() const { return buffers_; } + + // Returns the underlying dataflow analysis used by this alias analysis. + const HloDataflowAnalysis& dataflow_analysis() const { + return *dataflow_analysis_; + } + + // Returns true if any index in the output of the given instruction has more + // than one buffer. That is, ComputeBuffersAt returns a vector with more than + // one element. + bool InstructionBuffersAreAmbiguous(const HloInstruction* instruction) const; + + // Returns true if no HloBuffer appears in more than one shape index in the + // output of the given instruction. + bool InstructionBuffersAreDistinct(const HloInstruction* instruction) const; + + // Compare the dataflow analysis against a clean recomputation of the + // analysis. Returns an error status if there is a mismatch. Useful for + // verifying the correctness after updates to the analysis. + Status VerifyAgainstReference() const; + + protected: + explicit HloAliasAnalysis(HloModule* module); + + // Verify various invariants of the alias analysis. + Status Verify() const; + + HloModule* module_; + + // The underlying dataflow analysis used by this alias analysis. + std::unique_ptr dataflow_analysis_; + + // A map indicating which buffer a value is contained in. + tensorflow::gtl::FlatMap value_to_buffer_; + + // A lazily constructed vector containing all HloBuffers sorted by + // HloBuffer::Id. + std::vector buffers_; +}; + +} // namespace xla + +#endif // TENSORFLOW_COMPILER_XLA_SERVICE_HLO_ALIAS_ANALYSIS_H_ diff --git a/tensorflow/compiler/xla/service/hlo_alias_analysis_test.cc b/tensorflow/compiler/xla/service/hlo_alias_analysis_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..6e311e25fb92f32ae8266bab0c3daad43d2349a3 --- /dev/null +++ b/tensorflow/compiler/xla/service/hlo_alias_analysis_test.cc @@ -0,0 +1,820 @@ +/* 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/hlo_alias_analysis.h" + +#include +#include + +#include "tensorflow/compiler/xla/literal_util.h" +#include "tensorflow/compiler/xla/service/flatten_call_graph.h" +#include "tensorflow/compiler/xla/service/hlo_matchers.h" +#include "tensorflow/compiler/xla/service/hlo_opcode.h" +#include "tensorflow/compiler/xla/service/instruction_fusion.h" +#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/logging.h" +#include "tensorflow/core/platform/test.h" + +namespace xla { +namespace { + +using ::testing::UnorderedElementsAre; + +class HloAliasAnalysisTest : public HloTestBase { + protected: + HloAliasAnalysisTest() : module_(CreateNewModule()) {} + + // Run alias analysis on the member module. For convenience returns a + // reference to the generated analysis stored in analysis_. + HloAliasAnalysis& RunAnalysis() { + analysis_ = HloAliasAnalysis::Run(module_.get()).ConsumeValueOrDie(); + return *analysis_; + } + + // Return a vector of the buffers in the buffer set at the current position + // sorted by buffer id. + std::vector GetBuffersAt(const HloInstruction* instruction, + const ShapeIndex& index = {}) const { + std::set buffer_ids; + for (const HloValue* value : analysis_->dataflow_analysis() + .GetValueSet(instruction, index) + .values()) { + buffer_ids.insert(analysis_->GetBufferContainingValue(*value).id()); + } + + std::vector buffers; + for (HloBuffer::Id id : buffer_ids) { + buffers.push_back(analysis_->GetBuffer(id)); + } + return buffers; + } + + // Return a vector containing all of the HloValues in the given buffer. + std::vector GetValuesInBuffer(const HloBuffer& buffer) { + std::vector values; + for (const HloValue* value : buffer.values()) { + values.push_back(*value); + } + return values; + } + + // Return the HloValue defined at the given position. + const HloValue& GetValueDefinedAt(const HloInstruction* instruction, + const ShapeIndex& index = {}) const { + return analysis_->dataflow_analysis().GetValueDefinedAt(instruction, index); + } + + // Returns true if any values held in the same buffer interfere. Generally, in + // the compiler pipeline copy-insertion will guarantee that this interference + // never occurs, but HLO graphs with interference can be explicitly + // constructed. + bool AnyValuesInSameBufferInterfere() { + DependencyHloOrdering ordering(module_.get()); + for (const HloBuffer& buffer : analysis_->buffers()) { + for (const HloValue* value_a : buffer.values()) { + for (const HloValue* value_b : buffer.values()) { + if (*value_a != *value_b && + ordering.MayInterfere(*value_a, *value_b)) { + VLOG(1) << *value_a << " interferes with " << *value_b + << " in buffer: " << buffer; + return true; + } + } + } + } + return false; + } + + std::unique_ptr module_; + std::unique_ptr analysis_; + + const Shape scalar_shape_ = ShapeUtil::MakeShape(F32, {}); +}; + +TEST_F(HloAliasAnalysisTest, BinaryOperation) { + // Test the analysis on a single binary operation (Add). + auto builder = HloComputation::Builder(TestName()); + auto constant1 = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(1.0))); + auto constant2 = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(2.0))); + auto add = builder.AddInstruction(HloInstruction::CreateBinary( + scalar_shape_, HloOpcode::kAdd, constant1, constant2)); + module_->AddEntryComputation(builder.Build()); + + const HloAliasAnalysis& analysis = RunAnalysis(); + + EXPECT_EQ(analysis.buffers().size(), 3); + + // All of the buffer sets should trivially contain a single buffer containing + // a single value. + for (const HloInstruction* instruction : {constant1, constant2, add}) { + EXPECT_EQ(analysis.GetUniqueBufferAt(instruction).GetUniqueValue(), + GetValueDefinedAt(instruction)); + } + + EXPECT_FALSE(analysis.InstructionBuffersAreAmbiguous(add)); + EXPECT_TRUE(analysis.InstructionBuffersAreDistinct(add)); + + EXPECT_FALSE(AnyValuesInSameBufferInterfere()); +} + +TEST_F(HloAliasAnalysisTest, TupleAndGtes) { + // Verify the analysis for a Tuple and GetTupleElement instructions. + auto builder = HloComputation::Builder(TestName()); + auto param0 = builder.AddInstruction( + HloInstruction::CreateParameter(0, scalar_shape_, "param0")); + auto param1 = builder.AddInstruction( + HloInstruction::CreateParameter(1, scalar_shape_, "param1")); + auto tuple = + builder.AddInstruction(HloInstruction::CreateTuple({param0, param1})); + auto gte0 = builder.AddInstruction( + HloInstruction::CreateGetTupleElement(scalar_shape_, tuple, 0)); + auto gte1 = builder.AddInstruction( + HloInstruction::CreateGetTupleElement(scalar_shape_, tuple, 1)); + builder.AddInstruction( + HloInstruction::CreateBinary(scalar_shape_, HloOpcode::kAdd, gte0, gte1)); + module_->AddEntryComputation(builder.Build()); + + const HloAliasAnalysis& analysis = RunAnalysis(); + + EXPECT_EQ(analysis.buffers().size(), 4); + + // Verify the expected aliasing of the tuple elements. + EXPECT_EQ(analysis.GetUniqueBufferAt(tuple, /*index=*/{}).GetUniqueValue(), + GetValueDefinedAt(tuple, /*index=*/{})); + EXPECT_EQ(analysis.GetUniqueBufferAt(tuple, /*index=*/{0}).GetUniqueValue(), + GetValueDefinedAt(param0)); + EXPECT_EQ(analysis.GetUniqueBufferAt(tuple, /*index=*/{1}).GetUniqueValue(), + GetValueDefinedAt(param1)); + + // The tuple operand, tuple element, and result of the GTE instruction should + // all be the same buffer. + EXPECT_EQ(analysis.GetUniqueBufferAt(param0), + analysis.GetUniqueBufferAt(tuple, /*index=*/{0})); + EXPECT_EQ(analysis.GetUniqueBufferAt(param0), + analysis.GetUniqueBufferAt(gte0)); + + // Verify the positions of an aliased buffer. + EXPECT_THAT( + analysis.GetUniqueBufferAt(param0).ComputePositions(), + UnorderedElementsAre(HloPosition{param0, {}}, HloPosition{tuple, {0}}, + HloPosition{gte0, {}})); + + EXPECT_FALSE(analysis.InstructionBuffersAreAmbiguous(tuple)); + EXPECT_TRUE(analysis.InstructionBuffersAreDistinct(tuple)); + + EXPECT_FALSE(AnyValuesInSameBufferInterfere()); +} + +TEST_F(HloAliasAnalysisTest, NondistinctTuple) { + // Test a expression with a non-distinct buffer set. + auto builder = HloComputation::Builder(TestName()); + auto param0 = builder.AddInstruction( + HloInstruction::CreateParameter(0, scalar_shape_, "param0")); + auto param1 = builder.AddInstruction( + HloInstruction::CreateParameter(1, scalar_shape_, "param1")); + // param0 is included twice in the tuple. + auto tuple = builder.AddInstruction( + HloInstruction::CreateTuple({param0, param1, param0})); + module_->AddEntryComputation(builder.Build()); + + const HloAliasAnalysis& analysis = RunAnalysis(); + + EXPECT_THAT( + analysis.GetUniqueBufferAt(param0).ComputePositions(), + UnorderedElementsAre(HloPosition{param0, {}}, HloPosition{tuple, {0}}, + HloPosition{tuple, {2}})); + + EXPECT_FALSE(analysis.InstructionBuffersAreAmbiguous(tuple)); + EXPECT_FALSE(analysis.InstructionBuffersAreDistinct(tuple)); + + EXPECT_FALSE(AnyValuesInSameBufferInterfere()); +} + +TEST_F(HloAliasAnalysisTest, SingleCall) { + // Test a single call of a subcomputation. The subcomputation adds its two + // array-shaped parameters. + auto subbuilder = HloComputation::Builder("Subcomputation"); + auto subparam0 = subbuilder.AddInstruction( + HloInstruction::CreateParameter(0, scalar_shape_, "param0")); + auto subparam1 = subbuilder.AddInstruction( + HloInstruction::CreateParameter(1, scalar_shape_, "param1")); + auto add = subbuilder.AddInstruction(HloInstruction::CreateBinary( + scalar_shape_, HloOpcode::kAdd, subparam0, subparam1)); + HloComputation* called_computation = + module_->AddEmbeddedComputation(subbuilder.Build()); + + auto builder = HloComputation::Builder(TestName()); + auto constant1 = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(1.0))); + auto constant2 = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(2.0))); + auto call = builder.AddInstruction(HloInstruction::CreateCall( + scalar_shape_, {constant1, constant2}, called_computation)); + module_->AddEntryComputation(builder.Build()); + + const HloAliasAnalysis& analysis = RunAnalysis(); + + // Verify aliasing of the kCall operands and the subcomputation parameters. + EXPECT_THAT(analysis.GetUniqueBufferAt(constant1).ComputePositions(), + UnorderedElementsAre(HloPosition{constant1, {}}, + HloPosition{subparam0, {}})); + EXPECT_THAT(analysis.GetUniqueBufferAt(constant2).ComputePositions(), + UnorderedElementsAre(HloPosition{constant2, {}}, + HloPosition{subparam1, {}})); + + // The subcomputation root and the kCall itself should alias. + EXPECT_THAT( + analysis.GetUniqueBufferAt(add).ComputePositions(), + UnorderedElementsAre(HloPosition{add, {}}, HloPosition{call, {}})); + + EXPECT_FALSE(AnyValuesInSameBufferInterfere()); +} + +TEST_F(HloAliasAnalysisTest, ComputationCalledTwice) { + // Test a subcomputation which is called twice with different argument values. + auto subbuilder = HloComputation::Builder("Subcomputation"); + auto subparam0 = subbuilder.AddInstruction( + HloInstruction::CreateParameter(0, scalar_shape_, "param0")); + auto subparam1 = subbuilder.AddInstruction( + HloInstruction::CreateParameter(1, scalar_shape_, "param1")); + auto add = subbuilder.AddInstruction(HloInstruction::CreateBinary( + scalar_shape_, HloOpcode::kAdd, subparam0, subparam1)); + HloComputation* called_computation = + module_->AddEmbeddedComputation(subbuilder.Build()); + + auto builder = HloComputation::Builder(TestName()); + auto constant1 = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(1.0))); + auto constant2 = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(2.0))); + auto call1 = builder.AddInstruction(HloInstruction::CreateCall( + scalar_shape_, {constant1, constant2}, called_computation)); + auto call2 = builder.AddInstruction(HloInstruction::CreateCall( + scalar_shape_, {call1, constant2}, called_computation)); + module_->AddEntryComputation(builder.Build()); + + const HloAliasAnalysis& analysis = RunAnalysis(); + + EXPECT_THAT(analysis.GetUniqueBufferAt(constant1).ComputePositions(), + UnorderedElementsAre(HloPosition{constant1, {}}, + HloPosition{subparam0, {}})); + EXPECT_THAT(analysis.GetUniqueBufferAt(constant2).ComputePositions(), + UnorderedElementsAre(HloPosition{constant2, {}}, + HloPosition{subparam1, {}})); + + // The 'add' (root of the subcomputation) aliases the two call instruction, + // and the first parameter of the subcomputation because 'call1' it is passed + // as an argument to the subcomputation in 'call2'. + EXPECT_THAT( + analysis.GetUniqueBufferAt(add).ComputePositions(), + UnorderedElementsAre(HloPosition{add, {}}, HloPosition{call1, {}}, + HloPosition{subparam0, {}}, HloPosition{call2, {}})); + + EXPECT_THAT(GetBuffersAt(subparam0), + UnorderedElementsAre(analysis.GetUniqueBufferAt(constant1), + analysis.GetUniqueBufferAt(add))); + EXPECT_THAT(GetBuffersAt(subparam1), + UnorderedElementsAre(analysis.GetUniqueBufferAt(constant2))); + + EXPECT_TRUE(analysis.InstructionBuffersAreAmbiguous(subparam0)); + EXPECT_FALSE(analysis.InstructionBuffersAreAmbiguous(subparam1)); + EXPECT_TRUE(analysis.InstructionBuffersAreDistinct(subparam0)); + EXPECT_TRUE(analysis.InstructionBuffersAreDistinct(subparam1)); + + EXPECT_FALSE(AnyValuesInSameBufferInterfere()); +} + +TEST_F(HloAliasAnalysisTest, SingleWhile) { + // Test a simple single while instruction. The while body includes a + // pass-through value. HLO: + // + // body((F32[], F32[]) %tuple_param): + // %add = Add(%tuple_param{0}, %tuple_param{1}) + // return Tuple(%tuple_param{0}, %add) + // + // condition((F32[], F32[]) %tuple_param): + // return Constant(false) + // + // entry: + // %constant1 = Constant(1.0) + // %constant2 = Constant(2.0) + // %tuple = Tuple(%constant1, %constant2) + // return While(%tuple, body, condition) + // + const Shape tuple_shape = + ShapeUtil::MakeTupleShape({scalar_shape_, scalar_shape_}); + + // Element 0 passes transparently through the body. + auto body_builder = HloComputation::Builder("body"); + auto body_param = body_builder.AddInstruction( + HloInstruction::CreateParameter(0, tuple_shape, "param")); + auto body_element_0 = body_builder.AddInstruction( + HloInstruction::CreateGetTupleElement(scalar_shape_, body_param, 0)); + auto body_element_1 = body_builder.AddInstruction( + HloInstruction::CreateGetTupleElement(scalar_shape_, body_param, 1)); + auto add = body_builder.AddInstruction(HloInstruction::CreateBinary( + scalar_shape_, HloOpcode::kAdd, body_element_0, body_element_1)); + auto body_tuple = body_builder.AddInstruction( + HloInstruction::CreateTuple({body_element_0, add})); + HloComputation* body = module_->AddEmbeddedComputation(body_builder.Build()); + + // Condition computation trivially returns a constant "false". + auto cond_builder = HloComputation::Builder("condition"); + auto cond_param = cond_builder.AddInstruction( + HloInstruction::CreateParameter(0, tuple_shape, "param")); + cond_builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(false))); + HloComputation* condition = + module_->AddEmbeddedComputation(cond_builder.Build()); + + auto builder = HloComputation::Builder(TestName()); + auto constant1 = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(1.0))); + auto constant2 = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(2.0))); + auto tuple = builder.AddInstruction( + HloInstruction::CreateTuple({constant1, constant2})); + auto xla_while = builder.AddInstruction( + HloInstruction::CreateWhile(tuple_shape, condition, body, tuple)); + module_->AddEntryComputation(builder.Build()); + + const HloAliasAnalysis& analysis = RunAnalysis(); + + // Verify the positions of the aliased while buffers. + EXPECT_THAT( + analysis.GetUniqueBufferAt(xla_while, /*index=*/{}).ComputePositions(), + UnorderedElementsAre(HloPosition{tuple, {}}, HloPosition{xla_while, {}}, + HloPosition{body_param, {}}, + HloPosition{body_tuple, {}}, + HloPosition{cond_param, {}})); + EXPECT_THAT( + analysis.GetUniqueBufferAt(xla_while, /*index=*/{0}).ComputePositions(), + UnorderedElementsAre( + HloPosition{constant1, {}}, HloPosition{tuple, {0}}, + HloPosition{xla_while, {0}}, HloPosition{body_param, {0}}, + HloPosition{body_element_0, {}}, HloPosition{body_tuple, {0}}, + HloPosition{cond_param, {0}})); + EXPECT_THAT( + analysis.GetUniqueBufferAt(xla_while, /*index=*/{1}).ComputePositions(), + UnorderedElementsAre( + HloPosition{constant2, {}}, HloPosition{tuple, {1}}, + HloPosition{xla_while, {1}}, HloPosition{body_param, {1}}, + HloPosition{body_element_1, {}}, HloPosition{add, {}}, + HloPosition{body_tuple, {1}}, HloPosition{cond_param, {1}})); + + EXPECT_THAT( + GetValuesInBuffer(analysis.GetUniqueBufferAt(xla_while, /*index=*/{0})), + UnorderedElementsAre(GetValueDefinedAt(constant1))); + EXPECT_THAT( + GetValuesInBuffer(analysis.GetUniqueBufferAt(xla_while, /*index=*/{1})), + UnorderedElementsAre(GetValueDefinedAt(constant2), + GetValueDefinedAt(xla_while, /*index=*/{1}), + GetValueDefinedAt(body_param, {1}), + GetValueDefinedAt(cond_param, {1}), + GetValueDefinedAt(add))); + + EXPECT_FALSE(AnyValuesInSameBufferInterfere()); +} + +TEST_F(HloAliasAnalysisTest, SequentialWhiles) { + // Test sequential while instructions. The while body includes a + // pass-through value. HLO: + // + // body((F32[], F32[]) %tuple_param): + // %add = Add(%tuple_param{0}, %tuple_param{1}) + // return Tuple(%tuple_param{0}, %add) + // + // condition((F32[], F32[]) %tuple_param): + // return Constant(false) + // + // entry: + // %constant1 = Constant(1.0) + // %constant2 = Constant(2.0) + // %tuple = Tuple(%constant1, %constant2) + // %while0 = While(%tuple, body, condition) + // %while1 = While(%while0, body, condition) + // return While(%while1, body, condition) + // + const Shape tuple_shape = + ShapeUtil::MakeTupleShape({scalar_shape_, scalar_shape_}); + + // Element 0 passes transparently through the body. + auto body_builder = HloComputation::Builder("body"); + auto body_param = body_builder.AddInstruction( + HloInstruction::CreateParameter(0, tuple_shape, "param")); + auto body_element_0 = body_builder.AddInstruction( + HloInstruction::CreateGetTupleElement(scalar_shape_, body_param, 0)); + auto body_element_1 = body_builder.AddInstruction( + HloInstruction::CreateGetTupleElement(scalar_shape_, body_param, 1)); + auto add = body_builder.AddInstruction(HloInstruction::CreateBinary( + scalar_shape_, HloOpcode::kAdd, body_element_0, body_element_1)); + body_builder.AddInstruction( + HloInstruction::CreateTuple({body_element_0, add})); + HloComputation* body = module_->AddEmbeddedComputation(body_builder.Build()); + + auto cond_builder = HloComputation::Builder("condition"); + cond_builder.AddInstruction( + HloInstruction::CreateParameter(0, tuple_shape, "param")); + cond_builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(false))); + HloComputation* condition = + module_->AddEmbeddedComputation(cond_builder.Build()); + + auto builder = HloComputation::Builder(TestName()); + auto constant1 = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(1.0))); + auto constant2 = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(2.0))); + auto tuple = builder.AddInstruction( + HloInstruction::CreateTuple({constant1, constant2})); + auto xla_while0 = builder.AddInstruction( + HloInstruction::CreateWhile(tuple_shape, condition, body, tuple)); + auto xla_while1 = builder.AddInstruction( + HloInstruction::CreateWhile(tuple_shape, condition, body, xla_while0)); + auto xla_while2 = builder.AddInstruction( + HloInstruction::CreateWhile(tuple_shape, condition, body, xla_while1)); + module_->AddEntryComputation(builder.Build()); + + FlattenCallGraph flattener; + TF_ASSERT_OK(flattener.Run(module_.get()).status()); + + const HloAliasAnalysis& analysis = RunAnalysis(); + + EXPECT_EQ(analysis.GetUniqueBufferAt(tuple, /*index=*/{}), + analysis.GetUniqueBufferAt(xla_while2, /*index=*/{})); + EXPECT_EQ(analysis.GetUniqueBufferAt(constant1), + analysis.GetUniqueBufferAt(xla_while2, /*index=*/{0})); + EXPECT_EQ(analysis.GetUniqueBufferAt(constant2), + analysis.GetUniqueBufferAt(xla_while2, /*index=*/{1})); +} + +TEST_F(HloAliasAnalysisTest, NestedWhiles) { + // Test nested while instructions. The inner body passes through element 0 of + // its parameter, and the outer body passes through element 1. HLO: + // + // inner_body((F32[], F32[]) %tuple_param): + // %add = Add(%tuple_param{0}, %tuple_param{1}) + // return Tuple(%tuple_param{0}, %add) + // + // outer_body((F32[], F32[]) %tuple_param): + // %negate = Negate(%tuple_param{0}) + // %tuple = Tuple(%negate, %tuple_param{1}) + // return While(%tuple, inner_body, condition) + // + // entry: + // %constant1 = Constant(1.0) + // %constant2 = Constant(2.0) + // %tuple = Tuple(%constant1, %constant2) + // return While(%tuple, outer_body, condition) + // + const Shape tuple_shape = + ShapeUtil::MakeTupleShape({scalar_shape_, scalar_shape_}); + + auto build_cond_computation = [&tuple_shape]() { + auto cond_builder = HloComputation::Builder("condition"); + cond_builder.AddInstruction( + HloInstruction::CreateParameter(0, tuple_shape, "param")); + cond_builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(false))); + return cond_builder.Build(); + }; + // Build separate condition computations so the call graph is flat. The + // callgraph is always flattened in the compiler pipeline, and the flattened + // callgraph enables representative interference analysis. + HloComputation* condition1 = + module_->AddEmbeddedComputation(build_cond_computation()); + HloComputation* condition2 = + module_->AddEmbeddedComputation(build_cond_computation()); + + // Element 0 passes transparently through the body. + auto inner_builder = HloComputation::Builder("inner_body"); + auto inner_param = inner_builder.AddInstruction( + HloInstruction::CreateParameter(0, tuple_shape, "param")); + auto inner_element_0 = inner_builder.AddInstruction( + HloInstruction::CreateGetTupleElement(scalar_shape_, inner_param, 0)); + auto inner_element_1 = inner_builder.AddInstruction( + HloInstruction::CreateGetTupleElement(scalar_shape_, inner_param, 1)); + auto add = inner_builder.AddInstruction(HloInstruction::CreateBinary( + scalar_shape_, HloOpcode::kAdd, inner_element_0, inner_element_1)); + inner_builder.AddInstruction( + HloInstruction::CreateTuple({inner_element_0, add})); + HloComputation* inner_body = + module_->AddEmbeddedComputation(inner_builder.Build()); + + // Element 1 passes transparently through the body. + auto outer_builder = HloComputation::Builder("outer_body"); + auto outer_param = outer_builder.AddInstruction( + HloInstruction::CreateParameter(0, tuple_shape, "param")); + auto outer_element_0 = outer_builder.AddInstruction( + HloInstruction::CreateGetTupleElement(scalar_shape_, outer_param, 0)); + auto negate = outer_builder.AddInstruction(HloInstruction::CreateUnary( + scalar_shape_, HloOpcode::kNegate, outer_element_0)); + auto outer_element_1 = outer_builder.AddInstruction( + HloInstruction::CreateGetTupleElement(scalar_shape_, outer_param, 1)); + auto outer_tuple = outer_builder.AddInstruction( + HloInstruction::CreateTuple({negate, outer_element_1})); + auto nested_while = outer_builder.AddInstruction(HloInstruction::CreateWhile( + tuple_shape, condition1, inner_body, outer_tuple)); + HloComputation* outer_body = + module_->AddEmbeddedComputation(outer_builder.Build()); + + auto builder = HloComputation::Builder(TestName()); + auto constant1 = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(1.0))); + auto constant2 = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(2.0))); + auto tuple = builder.AddInstruction( + HloInstruction::CreateTuple({constant1, constant2})); + auto entry_while = builder.AddInstruction( + HloInstruction::CreateWhile(tuple_shape, condition2, outer_body, tuple)); + module_->AddEntryComputation(builder.Build()); + + const HloAliasAnalysis& analysis = RunAnalysis(); + + EXPECT_EQ(analysis.GetUniqueBufferAt(constant1), + analysis.GetUniqueBufferAt(entry_while, /*index=*/{0})); + EXPECT_EQ(analysis.GetUniqueBufferAt(constant1), + analysis.GetUniqueBufferAt(nested_while, /*index=*/{0})); + EXPECT_EQ(analysis.GetUniqueBufferAt(constant1), + analysis.GetUniqueBufferAt(inner_element_0)); + + EXPECT_EQ(analysis.GetUniqueBufferAt(constant2), + analysis.GetUniqueBufferAt(entry_while, /*index=*/{1})); + EXPECT_EQ(analysis.GetUniqueBufferAt(constant2), + analysis.GetUniqueBufferAt(nested_while, /*index=*/{1})); + EXPECT_EQ(analysis.GetUniqueBufferAt(constant2), + analysis.GetUniqueBufferAt(inner_element_1)); + + EXPECT_FALSE(AnyValuesInSameBufferInterfere()); +} + +TEST_F(HloAliasAnalysisTest, SwizzlingWhile) { + // Test a while instruction with a body which permutes it's tuple parameter + // elements. HLO: + // + // body((F32[], F32[], F32[]) %tuple_param): + // return Tuple(%tuple_param{1}, %tuple_param{2}, %tuple_param{0}) + // + // condition((F32[], F32[]) %tuple_param): + // return Constant(false) + // + // entry: + // %constant1 = Constant(1.0) + // %constant2 = Constant(2.0) + // %constant3 = Constant(3.0) + // %tuple = Tuple(%constant1, %constant2, %constant3) + // return While(%tuple, body, condition) + // + const Shape tuple_shape = + ShapeUtil::MakeTupleShape({scalar_shape_, scalar_shape_, scalar_shape_}); + + auto body_builder = HloComputation::Builder("body"); + auto body_param = body_builder.AddInstruction( + HloInstruction::CreateParameter(0, tuple_shape, "param")); + auto body_element_0 = body_builder.AddInstruction( + HloInstruction::CreateGetTupleElement(scalar_shape_, body_param, 0)); + auto body_element_1 = body_builder.AddInstruction( + HloInstruction::CreateGetTupleElement(scalar_shape_, body_param, 1)); + auto body_element_2 = body_builder.AddInstruction( + HloInstruction::CreateGetTupleElement(scalar_shape_, body_param, 2)); + body_builder.AddInstruction(HloInstruction::CreateTuple( + {body_element_1, body_element_2, body_element_0})); + HloComputation* body = module_->AddEmbeddedComputation(body_builder.Build()); + + auto cond_builder = HloComputation::Builder("condition"); + cond_builder.AddInstruction( + HloInstruction::CreateParameter(0, tuple_shape, "param")); + auto cond_constant = cond_builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(false))); + HloComputation* condition = + module_->AddEmbeddedComputation(cond_builder.Build()); + + auto builder = HloComputation::Builder(TestName()); + auto constant1 = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(1.0))); + auto constant2 = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(2.0))); + auto constant3 = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(3.0))); + auto tuple = builder.AddInstruction( + HloInstruction::CreateTuple({constant1, constant2, constant3})); + auto xla_while = builder.AddInstruction( + HloInstruction::CreateWhile(tuple_shape, condition, body, tuple)); + module_->AddEntryComputation(builder.Build()); + + const HloAliasAnalysis& analysis = RunAnalysis(); + + // The swizzling while makes most positions in the module alias leaving only 3 + // HloBuffers. + EXPECT_THAT( + analysis.buffers(), + UnorderedElementsAre(analysis.GetUniqueBufferAt(constant1), + analysis.GetUniqueBufferAt(tuple, /*index=*/{}), + analysis.GetUniqueBufferAt(cond_constant))); + + // The tuple elements of the while and the three constant inputs should all be + // smooshed into the same buffer. + EXPECT_EQ(analysis.GetUniqueBufferAt(xla_while, /*index=*/{0}), + analysis.GetUniqueBufferAt(xla_while, /*index=*/{1})); + EXPECT_EQ(analysis.GetUniqueBufferAt(xla_while, /*index=*/{0}), + analysis.GetUniqueBufferAt(xla_while, /*index=*/{2})); + EXPECT_EQ(analysis.GetUniqueBufferAt(xla_while, /*index=*/{0}), + analysis.GetUniqueBufferAt(constant1)); + EXPECT_EQ(analysis.GetUniqueBufferAt(constant1), + analysis.GetUniqueBufferAt(constant2)); + EXPECT_EQ(analysis.GetUniqueBufferAt(constant1), + analysis.GetUniqueBufferAt(constant3)); + + // All elements in of the loop state tuple are forced into the same buffer + // resulting liveness interference. + EXPECT_TRUE(AnyValuesInSameBufferInterfere()); +} + +TEST_F(HloAliasAnalysisTest, TupleSelect) { + // Test a kSelect of a tuple value. Non-top-level element flow through the + // instruction. + auto builder = HloComputation::Builder(TestName()); + auto pred = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(false))); + auto constant1 = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(1.0))); + auto constant2 = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(2.0))); + auto constant3 = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(3.0))); + auto constant4 = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(4.0))); + auto tuple1 = + builder.AddInstruction(HloInstruction::CreateTuple({constant1})); + auto tuple2 = + builder.AddInstruction(HloInstruction::CreateTuple({constant2})); + auto tuple3 = + builder.AddInstruction(HloInstruction::CreateTuple({constant3})); + auto tuple4 = + builder.AddInstruction(HloInstruction::CreateTuple({constant4})); + const Shape tuple_shape = tuple1->shape(); + auto select11 = builder.AddInstruction(HloInstruction::CreateTernary( + tuple_shape, HloOpcode::kSelect, pred, tuple1, tuple1)); + auto select12 = builder.AddInstruction(HloInstruction::CreateTernary( + tuple_shape, HloOpcode::kSelect, pred, tuple1, tuple2)); + auto select34 = builder.AddInstruction(HloInstruction::CreateTernary( + tuple_shape, HloOpcode::kSelect, pred, tuple3, tuple4)); + auto select1234 = builder.AddInstruction(HloInstruction::CreateTernary( + tuple_shape, HloOpcode::kSelect, pred, select12, select34)); + + module_->AddEntryComputation(builder.Build()); + + const HloAliasAnalysis& analysis = RunAnalysis(); + + // Verify the buffer sets of each select. + EXPECT_THAT(GetBuffersAt(select11, /*index=*/{0}), + UnorderedElementsAre(analysis.GetUniqueBufferAt(constant1))); + EXPECT_THAT(GetBuffersAt(select12, /*index=*/{0}), + UnorderedElementsAre(analysis.GetUniqueBufferAt(constant1), + analysis.GetUniqueBufferAt(constant2))); + EXPECT_THAT(GetBuffersAt(select34, /*index=*/{0}), + UnorderedElementsAre(analysis.GetUniqueBufferAt(constant3), + analysis.GetUniqueBufferAt(constant4))); + EXPECT_THAT(GetBuffersAt(select1234, /*index=*/{0}), + UnorderedElementsAre(analysis.GetUniqueBufferAt(constant1), + analysis.GetUniqueBufferAt(constant2), + analysis.GetUniqueBufferAt(constant3), + analysis.GetUniqueBufferAt(constant4))); + + EXPECT_FALSE(analysis.InstructionBuffersAreAmbiguous(select11)); + EXPECT_TRUE(analysis.InstructionBuffersAreAmbiguous(select12)); + EXPECT_TRUE(analysis.InstructionBuffersAreAmbiguous(select34)); + EXPECT_TRUE(analysis.InstructionBuffersAreAmbiguous(select1234)); + + EXPECT_TRUE(analysis.InstructionBuffersAreDistinct(select11)); + EXPECT_TRUE(analysis.InstructionBuffersAreDistinct(select12)); + EXPECT_TRUE(analysis.InstructionBuffersAreDistinct(select34)); + EXPECT_TRUE(analysis.InstructionBuffersAreDistinct(select1234)); + + EXPECT_FALSE(AnyValuesInSameBufferInterfere()); +} + +TEST_F(HloAliasAnalysisTest, TupleSelectToWhile) { + // Test a tuple-shaped kSelect feeding a kWhile instruction. HLO: + // + // body((F32[], F32[]) %tuple_param): + // %negate = Negate(%tuple_param{0}) + // return Tuple(%negate) + // + // condition((F32[], F32[]) %tuple_param): + // return Constant(false) + // + // entry: + // %constant1 = Constant(1.0) + // %constant2 = Constant(2.0) + // %tuple1 = Tuple(%constant1) + // %tuple2 = Tuple(%constant2) + // %select = Select(%tuple1, %tuple2) + // return While(%select, body, condition) + // + auto builder = HloComputation::Builder(TestName()); + + const Shape tuple_shape = ShapeUtil::MakeTupleShape({scalar_shape_}); + + // Element 0 passes transparently through the body. + auto body_builder = HloComputation::Builder("body"); + auto body_param = body_builder.AddInstruction( + HloInstruction::CreateParameter(0, tuple_shape, "param")); + auto body_element = body_builder.AddInstruction( + HloInstruction::CreateGetTupleElement(scalar_shape_, body_param, 0)); + auto negate = body_builder.AddInstruction(HloInstruction::CreateUnary( + scalar_shape_, HloOpcode::kNegate, body_element)); + body_builder.AddInstruction(HloInstruction::CreateTuple({negate})); + HloComputation* body = module_->AddEmbeddedComputation(body_builder.Build()); + + auto cond_builder = HloComputation::Builder("condition"); + auto cond_param = cond_builder.AddInstruction( + HloInstruction::CreateParameter(0, tuple_shape, "param")); + cond_builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(false))); + HloComputation* condition = + module_->AddEmbeddedComputation(cond_builder.Build()); + + auto pred = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(false))); + auto constant1 = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(1.0))); + auto constant2 = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(2.0))); + auto tuple1 = + builder.AddInstruction(HloInstruction::CreateTuple({constant1})); + auto tuple2 = + builder.AddInstruction(HloInstruction::CreateTuple({constant2})); + auto select = builder.AddInstruction(HloInstruction::CreateTernary( + tuple_shape, HloOpcode::kSelect, pred, tuple1, tuple2)); + auto xla_while = builder.AddInstruction( + HloInstruction::CreateWhile(tuple_shape, condition, body, select)); + + module_->AddEntryComputation(builder.Build()); + + const HloAliasAnalysis& analysis = RunAnalysis(); + + // The while should flatten the ambiguous select buffer set so that the buffer + // set contents (constant1 and constant2) becomes a single buffer. + EXPECT_EQ(analysis.GetUniqueBufferAt(constant1), + analysis.GetUniqueBufferAt(constant2)); + EXPECT_EQ(analysis.GetUniqueBufferAt(constant1), + analysis.GetUniqueBufferAt(xla_while, /*index=*/{0})); + + EXPECT_THAT(GetValuesInBuffer(analysis.GetUniqueBufferAt(constant1)), + UnorderedElementsAre(GetValueDefinedAt(constant1), + GetValueDefinedAt(constant2), + GetValueDefinedAt(xla_while, /*index=*/{0}), + GetValueDefinedAt(body_param, /*index=*/{0}), + GetValueDefinedAt(cond_param, /*index=*/{0}), + GetValueDefinedAt(negate))); + EXPECT_FALSE(analysis.InstructionBuffersAreAmbiguous(select)); + EXPECT_FALSE(analysis.InstructionBuffersAreAmbiguous(xla_while)); + + EXPECT_TRUE(analysis.InstructionBuffersAreDistinct(select)); + EXPECT_TRUE(analysis.InstructionBuffersAreDistinct(xla_while)); + + // The two operands of the select get flattened into the same buffer resulting + // in liveness interference. + EXPECT_TRUE(AnyValuesInSameBufferInterfere()); +} + +TEST_F(HloAliasAnalysisTest, Bitcast) { + // Bitcasting a value should not produce a new buffer. + auto builder = HloComputation::Builder(TestName()); + auto constant = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(1.0))); + auto bitcast = builder.AddInstruction(HloInstruction::CreateUnary( + scalar_shape_, HloOpcode::kBitcast, constant)); + + module_->AddEntryComputation(builder.Build()); + + const HloAliasAnalysis& analysis = RunAnalysis(); + + EXPECT_EQ(analysis.buffers().size(), 1); + + EXPECT_EQ(analysis.GetUniqueBufferAt(constant), + analysis.GetUniqueBufferAt(bitcast)); +} + +} // namespace +} // namespace xla diff --git a/tensorflow/compiler/xla/service/hlo_buffer.cc b/tensorflow/compiler/xla/service/hlo_buffer.cc new file mode 100644 index 0000000000000000000000000000000000000000..e16413f361fb0216792b47c3c67ef3c1357c2221 --- /dev/null +++ b/tensorflow/compiler/xla/service/hlo_buffer.cc @@ -0,0 +1,73 @@ +/* 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/hlo_buffer.h" + +#include +#include +#include +#include + +#include "tensorflow/compiler/xla/map_util.h" +#include "tensorflow/compiler/xla/service/hlo_instruction.h" +#include "tensorflow/compiler/xla/shape_util.h" +#include "tensorflow/compiler/xla/types.h" +#include "tensorflow/compiler/xla/util.h" +#include "tensorflow/core/lib/core/errors.h" +#include "tensorflow/core/lib/gtl/flatset.h" +#include "tensorflow/core/lib/strings/str_util.h" +#include "tensorflow/core/lib/strings/strcat.h" +#include "tensorflow/core/platform/logging.h" + +namespace xla { + +using ::tensorflow::str_util::Join; +using ::tensorflow::strings::StrCat; + +bool HloBuffer::operator==(const HloBuffer& other) const { + bool equal = id() == other.id(); + if (equal) { + // DCHECK because these comparisons are expensive (linear time). + DCHECK(values_ == other.values_); + } + return equal; +} + +std::vector HloBuffer::ComputePositions() const { + std::vector positions; + for (const HloValue* value : values_) { + positions.insert(positions.end(), value->positions().begin(), + value->positions().end()); + } + // Remove duplicates and sort positions. + std::sort(positions.begin(), positions.end()); + positions.erase(std::unique(positions.begin(), positions.end()), + positions.end()); + return positions; +} + +string HloBuffer::ToString() const { + return StrCat("HloBuffer ", id_, ", values: ", + Join(values_, ", ", [](string* result, const HloValue* value) { + result->append(value->ToShortString()); + })); +} + +std::ostream& operator<<(std::ostream& out, const HloBuffer& buffer) { + out << buffer.ToString(); + return out; +} + +} // namespace xla diff --git a/tensorflow/compiler/xla/service/hlo_buffer.h b/tensorflow/compiler/xla/service/hlo_buffer.h new file mode 100644 index 0000000000000000000000000000000000000000..4873463b2ea4fee3ee39dff31fc3429a4998142f --- /dev/null +++ b/tensorflow/compiler/xla/service/hlo_buffer.h @@ -0,0 +1,123 @@ +/* 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_BUFFER_H_ +#define TENSORFLOW_COMPILER_XLA_SERVICE_HLO_BUFFER_H_ + +#include +#include +#include + +#include "tensorflow/compiler/xla/service/hlo_value.h" +#include "tensorflow/compiler/xla/shape_tree.h" +#include "tensorflow/compiler/xla/types.h" +#include "tensorflow/compiler/xla/xla_data.pb.h" +#include "tensorflow/core/platform/macros.h" + +namespace xla { + +// A container which can hold one or more HloValues. An HLO buffer abstractly +// represents the allocation which HLO instructions write into and read +// from. Generally there is a one-to-one correspondence between HloBuffers and +// HloValue where each HloValue in the module is held in a unique HloBuffer. An +// exception is the while instruction which updates the loop state in-place. In +// this case, we have a single HloBuffer for each HloPosition in the loop state, +// but multiple HloValues. For example: +// +// %init = ... +// %while = While(%init, body, condition) +// +// body: +// %body_param = Param(0) +// ... +// %body_root = ... +// +// condition: +// %cond_param = Param(0) +// ... +// +// For simplicity, assume that %while is array-shaped. In this case, we have a +// single HloBuffer which holds the following HloValues: HloValue{%init}, +// HloValue{%while}, HloValue{%body_param}, HloValue{%body_root}, and +// HloValue{%cond_param}. +// +// HloBuffers may appear at different HloPositions in the module mirroring the +// same propery of HloValues. For example: +// +// %sub = Sub(...) +// %add = Add(...) +// %tuple = Tuple(%add, %sub) +// %gte = GetTupleElement(%tuple, 0) +// +// In this case, the HloBuffer containing %add appears at the following +// positions: HloPosition{%add, {}}, HloPosition{%tuple, {0}}, and +// HloPosition{%gte, {}}. +// +// Different HloPositions which share the same HloBuffer indicate mandatory +// aliasing in the HLO module. These positions must share the same memory +// allocation for correctness (the backends rely on this property). This differs +// from incidental aliasing introduced by memory reuse in BufferAssignment where +// different instructions may happen to get the same allocation. +class HloBuffer { + public: + using Id = int64; + + // Predicate comparing HloBuffers by increasing id, useful for std::sort. + static bool IdLessThan(const HloBuffer* a, const HloBuffer* b) { + return a->id() < b->id(); + } + + // Predicate comparing HloBuffers by equal id, useful for std::unique. + static bool IdEqual(const HloBuffer* a, const HloBuffer* b) { + return a->id() == b->id(); + } + + HloBuffer(Id id, tensorflow::gtl::ArraySlice values) + : id_(id), values_(values.begin(), values.end()) {} + + // Return the unique identifier for this HloBuffer. + Id id() const { return id_; } + + // Return all values contained in this buffer. + const std::vector& values() const { return values_; } + + // Return the unique HLO value in the buffer. CHECK fails if the buffer does + // not contain exactly one value. + const HloValue& GetUniqueValue() const { + CHECK_EQ(values_.size(), 1); + return *values_[0]; + } + + std::vector ComputePositions() const; + + string ToString() const; + + bool operator==(const HloBuffer& other) const; + bool operator!=(const HloBuffer& other) const { return !(*this == other); } + + private: + // Unique identifier for this HloBuffer. + const Id id_; + + // The set of values contained in this buffer. Vector contains no duplicates + // and is sorted stably by HloValue::Id. + const std::vector values_; +}; + +std::ostream& operator<<(std::ostream& out, const HloBuffer& buffer); + +} // namespace xla + +#endif // TENSORFLOW_COMPILER_XLA_SERVICE_HLO_BUFFER_H_ diff --git a/tensorflow/compiler/xla/service/hlo_computation.cc b/tensorflow/compiler/xla/service/hlo_computation.cc index 35f8dcb7ca614f5660850c9022049eea908f323c..2d077846196bdaf5183f6ee43ab582ede4ef4f52 100644 --- a/tensorflow/compiler/xla/service/hlo_computation.cc +++ b/tensorflow/compiler/xla/service/hlo_computation.cc @@ -27,6 +27,7 @@ limitations under the License. #include "tensorflow/compiler/xla/map_util.h" #include "tensorflow/compiler/xla/ptr_util.h" #include "tensorflow/compiler/xla/service/dfs_hlo_visitor_with_default.h" +#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/status_macros.h" @@ -35,10 +36,14 @@ limitations under the License. #include "tensorflow/core/lib/core/errors.h" #include "tensorflow/core/lib/core/status.h" #include "tensorflow/core/lib/gtl/flatset.h" +#include "tensorflow/core/lib/strings/str_util.h" +#include "tensorflow/core/lib/strings/strcat.h" #include "tensorflow/core/platform/logging.h" namespace xla { +using ::tensorflow::strings::StrCat; + std::unique_ptr HloComputation::Builder::Build( HloInstruction* root_instruction) { int parameter_count = 0; @@ -52,31 +57,35 @@ std::unique_ptr HloComputation::Builder::Build( root_instruction ? root_instruction : last_added_instruction_; CHECK_NE(nullptr, root); - return WrapUnique( - new HloComputation(name_, parameter_count, &instructions_, root)); + return WrapUnique(new HloComputation(name_, parameter_count, &instructions_, + root, fusion_instruction_)); } HloComputation::HloComputation( const string& name, int parameter_count, std::vector>* instructions, - HloInstruction* root_instruction) + HloInstruction* root_instruction, HloInstruction* fusion_instruction) : name_(name), root_instruction_(root_instruction), - instruction_name_uniquer_(/*separator=*/".") { + fusion_instruction_(fusion_instruction) { param_instructions_.resize(parameter_count, nullptr); bool root_found = false; for (auto& instruction : *instructions) { if (instruction->opcode() == HloOpcode::kParameter) { int64 param_no = instruction->parameter_number(); - CHECK_GE(param_no, 0); - CHECK_LT(param_no, param_instructions_.size()); - CHECK_EQ(nullptr, param_instructions_[param_no]); + CHECK(param_no >= 0 && param_no < parameter_count) + << "\nERROR: invalid parameter number. Expected [0, " + << parameter_count << "), got " << param_no; + CHECK(param_instructions_[param_no] == nullptr) + << "\nERROR: parameter number " << param_no + << " already allocated in this computation"; param_instructions_[param_no] = instruction.get(); } root_found |= instruction.get() == root_instruction_; AddInstructionInternal(std::move(instruction)); } - CHECK(root_found); + CHECK(root_found) + << "\nERROR: root instruction is not present in computation."; } HloInstruction* HloComputation::AddInstruction( @@ -89,9 +98,10 @@ HloInstruction* HloComputation::AddInstruction( HloInstruction* HloComputation::AddInstructionInternal( std::unique_ptr instruction) { - // Generate a unique name for the instruction. - instruction->set_name( - instruction_name_uniquer_.GetUniqueName(instruction->name())); + if (parent() != nullptr) { + instruction->UniquifyName(&parent()->instruction_name_uniquer()); + instruction->SetUniqueId(parent()->NewUniqueInstructionId()); + } Reparent(instruction.get()); HloInstruction* pinst = instruction.get(); instruction_iterators_[pinst] = @@ -99,19 +109,88 @@ HloInstruction* HloComputation::AddInstructionInternal( return pinst; } -void HloComputation::Reparent(HloInstruction* instruction) { +HloInstruction* HloComputation::AddParameter( + std::unique_ptr instruction) { + CHECK(instruction->opcode() == HloOpcode::kParameter); + CHECK(IsFusionComputation()); + CHECK(fusion_instruction_->operand_count() == param_instructions_.size()); instruction->set_parent(this); - if (instruction->opcode() == HloOpcode::kFusion) { - for (auto& i : instruction->fused_instructions()) { - Reparent(i.get()); + param_instructions_.push_back(instruction.get()); + AddInstructionInternal(std::move(instruction)); + return instructions_.back().get(); +} + +Status HloComputation::RemoveParameter(int64 param_no) { + CHECK_GE(param_no, 0); + CHECK_LT(param_no, param_instructions_.size()); + CHECK(IsFusionComputation()); + HloInstruction* param_instruction = param_instructions_[param_no]; + auto param_instruction_iterator = param_instructions_.begin() + param_no; + param_instructions_.erase(param_instruction_iterator); + // Throw removed fused parameter instruction away. + TF_RETURN_IF_ERROR(RemoveInstruction(param_instruction)); + + while (param_no < param_instructions_.size()) { + param_instruction = param_instructions_[param_no]; + string param_name = param_instruction->parameter_name(); + // Fusion parameters are named foo.param_1, bar.param_2, etc. We are + // renumbering the parameters so replace the final number in the name with + // the updated value. + const string param_underscore = ".param_"; + size_t index = param_name.rfind(param_underscore); + if (index == string::npos) { + string after_param = name().substr(index + param_underscore.size()); + int64 numeric_suffix; + if (tensorflow::strings::safe_strto64(after_param, &numeric_suffix)) { + param_name = + StrCat(param_name.substr(0, index), param_underscore, param_no); + } } + + HloInstruction* new_instr = + AddInstructionInternal(HloInstruction::CreateParameter( + param_no, param_instruction->shape(), param_name)); + TF_RETURN_IF_ERROR(param_instruction->ReplaceAllUsesWith(new_instr)); + param_instructions_[param_no] = new_instr; + TF_RETURN_IF_ERROR(RemoveInstruction(param_instruction)); + param_no++; } + + return Status::OK(); } -/* static */ bool HloComputation::IsRemovable(const HloOpcode& opcode) { - return !(opcode == HloOpcode::kParameter || opcode == HloOpcode::kRecv || - opcode == HloOpcode::kSend || opcode == HloOpcode::kTrace || - opcode == HloOpcode::kOutfeed); +void HloComputation::Reparent(HloInstruction* instruction) { + instruction->set_parent(this); +} + +bool HloComputation::IsRemovable(const HloInstruction* instruction) { + // If the instruction has control predecessors or successors then we cannot + // remove the instruction without violating ordering constraints (added, for + // example, to avert interference due to buffer aliasing). + if (!instruction->control_predecessors().empty() || + !instruction->control_successors().empty()) { + return false; + } + + if (instruction->opcode() == HloOpcode::kParameter && + !IsFusionComputation()) { + return false; + } + + if (instruction->HasSideEffect()) { + return false; + } + + return true; +} + +bool HloComputation::HasSideEffect() const { + for (auto& instruction : instructions()) { + if (instruction->HasSideEffect()) { + return true; + } + } + return false; } Status HloComputation::RemoveInstructionAndUnusedOperands( @@ -119,7 +198,7 @@ Status HloComputation::RemoveInstructionAndUnusedOperands( TF_RET_CHECK(root_instruction() != instruction); TF_RET_CHECK(instruction->user_count() == 0); - TF_RET_CHECK(HloComputation::IsRemovable(instruction->opcode())); + TF_RET_CHECK(IsRemovable(instruction)); std::unordered_set removed; std::queue worklist; worklist.push(instruction); @@ -128,8 +207,7 @@ Status HloComputation::RemoveInstructionAndUnusedOperands( worklist.pop(); if (removed.count(item) != 0 || item->user_count() != 0 || - item == root_instruction() || - !HloComputation::IsRemovable(item->opcode())) { + item == root_instruction() || !IsRemovable(item)) { continue; } for (int i = 0; i < item->operand_count(); ++i) { @@ -145,7 +223,8 @@ Status HloComputation::RemoveInstructionAndUnusedOperands( Status HloComputation::RemoveInstruction(HloInstruction* instruction) { VLOG(2) << "Removing instruction " << instruction->name() << " from computation " << name(); - TF_RET_CHECK(IsRemovable(instruction->opcode())); + TF_RET_CHECK(IsRemovable(instruction)) + << "cannot remove instruction: " << instruction->ToString(); TF_RET_CHECK(root_instruction() != instruction) << "cannot remove root instruction " << instruction->name(); TF_RET_CHECK(instruction->user_count() == 0) @@ -177,12 +256,15 @@ Status HloComputation::ReplaceUsesOfInstruction( void HloComputation::set_root_instruction( HloInstruction* new_root_instruction) { - // The shape of the root (ignoring layout) is an invariant of the computation. - CHECK(ShapeUtil::Compatible(new_root_instruction->shape(), - root_instruction_->shape())) - << new_root_instruction->shape().ShortDebugString() - << " is incompatible with " - << root_instruction_->shape().ShortDebugString(); + // The shape of the root (ignoring layout) is an invariant of the computation + // for non-fusion cases. + if (!IsFusionComputation()) { + CHECK(ShapeUtil::Compatible(new_root_instruction->shape(), + root_instruction_->shape())) + << new_root_instruction->shape().ShortDebugString() + << " is incompatible with " + << root_instruction_->shape().ShortDebugString(); + } bool root_found = false; for (auto& instruction : instructions_) { if (new_root_instruction == instruction.get()) { @@ -250,7 +332,6 @@ void ComputeComputationPostOrder( visited->insert(computation); post_order->push_back(computation); - return; } } // namespace @@ -295,25 +376,41 @@ std::list HloComputation::MakeEmbeddedComputationsList() return post_order; } -string HloComputation::ToString() const { +string HloComputation::ToString(int nested_level) const { std::ostringstream s; + for (int i = 0; i < nested_level; i++) { + s << " "; + } s << name() << " " << ShapeUtil::HumanString(ComputeProgramShape()) << " { \n"; for (const HloInstruction* instruction : MakeInstructionPostOrder()) { + for (int i = 0; i < nested_level; i++) { + s << " "; + } s << " " << instruction->ToString() << "\n"; if (instruction->opcode() == HloOpcode::kFusion) { - tensorflow::gtl::FlatSet added_instructions; - auto fused_instructions = InstructionPostOrderer::GetOrder( - instruction->fused_expression_root(), &added_instructions); - for (const auto& fused_instruction : fused_instructions) { - s << " " << fused_instruction->ToString() << "\n"; - } + s << instruction->fused_instructions_computation()->ToString( + nested_level + 1) + << "\n"; } } + for (int i = 0; i < nested_level; i++) { + s << " "; + } s << "}"; return s.str(); } +HloComputationProto HloComputation::ToProto() const { + HloComputationProto proto; + proto.set_name(name_); + for (const HloInstruction* instruction : MakeInstructionPostOrder()) { + HloInstructionProto instruction_proto = instruction->ToProto(); + proto.add_instructions()->Swap(&instruction_proto); + } + return proto; +} + void HloComputation::FuseInstructionsInto( tensorflow::gtl::ArraySlice instructions_to_fuse, HloInstruction* fusion_instruction) { @@ -357,49 +454,65 @@ HloInstruction* HloComputation::CreateFusionInstructionForBackwardConvolution( return fusion_instruction; } -StatusOr HloComputation::DeepCopyTuple( - HloInstruction* instruction) { - TF_RET_CHECK(ShapeUtil::IsTuple(instruction->shape())); - std::vector element_copies; - for (int64 i = 0; i < ShapeUtil::TupleElementCount(instruction->shape()); - ++i) { - HloInstruction* gte = AddInstruction(HloInstruction::CreateGetTupleElement( - ShapeUtil::GetSubshape(instruction->shape(), {i}), instruction, i)); - // Recurse to copy tuple elements. For array elements, insert a kCopy - // because GetTupleElement forwards a pointer to the tuple element buffer. - HloInstruction* element_copy; - if (ShapeUtil::IsTuple(gte->shape())) { - TF_ASSIGN_OR_RETURN(element_copy, DeepCopyTuple(gte)); +StatusOr HloComputation::DeepCopyHelper( + HloInstruction* instruction, const ShapeTree* indices_to_copy, + ShapeTree* copies_added, ShapeIndex* index) { + if (ShapeUtil::IsArray(instruction->shape())) { + if (indices_to_copy == nullptr || indices_to_copy->element(*index)) { + // Use kCopy to copy array elements + HloInstruction* copy = AddInstruction(HloInstruction::CreateUnary( + instruction->shape(), HloOpcode::kCopy, instruction)); + if (copies_added != nullptr) { + *copies_added->mutable_element(*index) = copy; + } + return copy; } else { - element_copy = AddInstruction( - HloInstruction::CreateUnary(gte->shape(), HloOpcode::kCopy, gte)); + // Array elements which are not to be copied are passed through + // transparently. + return instruction; + } + } else if (ShapeUtil::IsTuple(instruction->shape())) { + std::vector elements; + for (int64 i = 0; i < ShapeUtil::TupleElementCount(instruction->shape()); + i++) { + HloInstruction* gte = + AddInstruction(HloInstruction::CreateGetTupleElement( + ShapeUtil::GetTupleElementShape(instruction->shape(), i), + instruction, i)); + + index->push_back(i); + TF_ASSIGN_OR_RETURN( + HloInstruction * element, + DeepCopyHelper(gte, indices_to_copy, copies_added, index)); + elements.push_back(element); + index->pop_back(); } - element_copies.push_back(element_copy); + return AddInstruction(HloInstruction::CreateTuple(elements)); + } else { + return FailedPrecondition( + "Can only copy array and tuple shaped instructions"); } - - // Gather element copies into a tuple with a new Tuple instruction. - return AddInstruction(HloInstruction::CreateTuple(element_copies)); } StatusOr HloComputation::DeepCopyInstruction( - HloInstruction* instruction) { + HloInstruction* instruction, const ShapeTree* indices_to_copy, + ShapeTree* copies_added) { if (instruction->parent() != this) { return FailedPrecondition( "Can't deep copy instruction %s: instruction is not in computation %s", instruction->name().c_str(), name().c_str()); } - // For tuple instructions, perform a deep copy. For array instructions, copy - // with a kCopy instruction. - if (ShapeUtil::IsTuple(instruction->shape())) { - return DeepCopyTuple(instruction); - } else if (ShapeUtil::IsArray(instruction->shape())) { - return AddInstruction(HloInstruction::CreateUnary( - instruction->shape(), HloOpcode::kCopy, instruction)); - } else { + if (indices_to_copy != nullptr && + !ShapeUtil::Compatible(instruction->shape(), indices_to_copy->shape())) { return FailedPrecondition( - "Can only copy array and tuple shaped instructions"); + "Can't deep copy instruction %s: given shape tree of indices to copy " + "has incompatible shape", + instruction->name().c_str()); } + + ShapeIndex index; + return DeepCopyHelper(instruction, indices_to_copy, copies_added, &index); } ProgramShape HloComputation::ComputeProgramShape() const { @@ -422,7 +535,9 @@ bool HloComputation::operator==(const HloComputation& other) const { // 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; + 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) { @@ -445,72 +560,60 @@ Status HloComputation::ReplaceInstruction(HloInstruction* old_instruction, new_instruction->shape())); VLOG(10) << "transformed " << old_instruction->ToString() << " to " << new_instruction->ToString(); + // Try to add metadata for HLO instructions that are created to replace + // existing HLO instructions (e.g. during optimizations). The assumption is + // that the old instruction and the new instruction would perform the same + // function, and that they would be correlated to the same TF op. This might + // not always be correct since HLO optimizations can cross TF op boundaries. + // But still this seems to be better than nothing. + if (new_instruction->metadata().op_name().empty()) { + new_instruction->set_metadata(old_instruction->metadata()); + } TF_RETURN_IF_ERROR( ReplaceUsesOfInstruction(old_instruction, new_instruction)); return RemoveInstructionAndUnusedOperands(old_instruction); } -HloComputation::ReachabilityMap::ReachabilityMap( - const std::list& all_instructions) { - const int n = all_instructions.size(); - int next_id = 0; - for (const auto* hlo : all_instructions) { - ids_[hlo] = next_id; - next_id++; +std::unique_ptr HloComputation::ComputeReachability() + const { + const std::list all = MakeInstructionPostOrder(); + auto result = MakeUnique(all); + + 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()); + result->SetReachabilityToUnion(inputs, hlo); } - DCHECK_EQ(n, ids_.size()); // instructions should be unique - matrix_.Reset(n * n); -} - -void HloComputation::ReachabilityMap::SetReachable(const HloInstruction* a, - const HloInstruction* b) { - const int id_a = FindOrDie(ids_, a); - const int id_b = FindOrDie(ids_, b); - matrix_.set(id_a * ids_.size() + id_b); -} - -bool HloComputation::ReachabilityMap::IsReachable( - const HloInstruction* a, const HloInstruction* b) const { - const int id_a = FindOrDie(ids_, a); - const int id_b = FindOrDie(ids_, b); - return matrix_.get(id_a * ids_.size() + id_b); + return result; } -bool HloComputation::ReachabilityMap::IsConnected( - const HloInstruction* a, const HloInstruction* b) const { - const int id_a = FindOrDie(ids_, a); - const int id_b = FindOrDie(ids_, b); - return matrix_.get(id_a * ids_.size() + id_b) || - matrix_.get(id_b * ids_.size() + id_a); -} +void HloComputation::UpdateReachabilityThroughInstruction( + const HloInstruction* instruction, HloReachabilityMap* reachability_map) { + std::queue worklist; + worklist.push(instruction); -void HloComputation::ReachabilityMap::SetReachableAndTransitiveClosure( - const HloInstruction* a, const HloInstruction* b) { - const int id_a = FindOrDie(ids_, a); - const int id_b = FindOrDie(ids_, b); - const int n = ids_.size(); - matrix_.set(id_a * n + id_b); + std::vector inputs; - // Copy transitive set for b into entries for a - for (int i = 0; i < n; i++) { - if (matrix_.get(id_b * n + i)) { - matrix_.set(id_a * n + i); - } - } -} + while (!worklist.empty()) { + const HloInstruction* item = worklist.front(); + worklist.pop(); -std::unique_ptr -HloComputation::ComputeTransitiveOperands() const { - const auto all = MakeInstructionPostOrder(); - auto result = MakeUnique(all); + inputs.assign(item->operands().begin(), item->operands().end()); + inputs.insert(inputs.end(), item->control_predecessors().begin(), + item->control_predecessors().end()); - // Fill in the dependency bit matrix - for (const auto* hlo : all) { - for (const HloInstruction* operand : hlo->operands()) { - result->SetReachableAndTransitiveClosure(hlo, operand); + 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); + } } } - return result; } std::vector HloComputation::CollectUnreachableRoots() const { @@ -522,6 +625,12 @@ std::vector HloComputation::CollectUnreachableRoots() const { unreachable_roots.push_back(instruction.get()); } } + VLOG(3) << "Unreachable roots:" + << tensorflow::str_util::Join( + unreachable_roots, "\n\t", + [](string* out, const HloInstruction* hlo) { + tensorflow::strings::StrAppend(out, hlo->ToString()); + }); return unreachable_roots; } @@ -530,6 +639,7 @@ Status HloComputation::Accept(DfsHloVisitor* visitor) const { // visited root, which would invalidate iterators if the unreachable roots // weren't computed ahead of time. for (HloInstruction* root : CollectUnreachableRoots()) { + VLOG(3) << "Traversing unreachable root: " << root->ToString(); // Call FinishVisit only at the end. TF_RETURN_IF_ERROR(root->Accept(visitor, /*call_finish_visit=*/false)); } @@ -556,9 +666,15 @@ Status HloComputation::AcceptWithOperandOrder( Status HloComputation::AcceptOrdered( DfsHloVisitor* visitor, const std::vector& 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(order.size() == instruction_count()); std::unordered_set visited; for (const HloInstruction* instruction : order) { + VLOG(3) << "Visiting ordered: " << instruction->ToString(); TF_RET_CHECK(instruction_iterators_.count(instruction) == 1) << "Instruction " << instruction->name() << " is not in computation " << name(); @@ -583,4 +699,44 @@ Status HloComputation::Accept( return this->Accept(&visitor); } +std::unique_ptr HloComputation::Clone(const string& suffix) { + VLOG(1) << "Cloning " << name() << " --> " << suffix << "\n"; + auto postorder = MakeInstructionPostOrder(); + std::unordered_map clone_map; + std::vector> instructions; + std::unique_ptr new_instr = nullptr; + for (auto instr : postorder) { + std::vector new_operands; + for (auto operand : instr->operands()) { + HloInstruction* new_operand = FindOrDie(clone_map, operand); + CHECK(new_operand != nullptr); + new_operands.push_back(new_operand); + } + + new_instr = instr->CloneWithNewOperands(instr->shape(), new_operands); + InsertOrDie(&clone_map, instr, new_instr.get()); + instructions.push_back(std::move(new_instr)); + } + Builder builder(name() + suffix); + for (auto& instr : instructions) { + builder.AddInstruction(std::move(instr)); + } + auto result = builder.Build( + /*root_instruction=*/FindOrDie(clone_map, root_instruction())); + + // Clone control dependencies. + for (auto instr : postorder) { + HloInstruction* new_instr = FindOrDie(clone_map, instr); + for (auto successor : instr->control_successors()) { + TF_CHECK_OK( + new_instr->AddControlDependencyTo(FindOrDie(clone_map, successor))); + } + } + return result; +} + +void HloComputation::UniquifyName(NameUniquer* name_uniquer) { + name_ = name_uniquer->GetUniqueName(name_); +} + } // namespace xla diff --git a/tensorflow/compiler/xla/service/hlo_computation.h b/tensorflow/compiler/xla/service/hlo_computation.h index dddcf51974932fb13eff73d256e029b53573ce36..576c44a9f344160fd6184bf2bd590044676a27d6 100644 --- a/tensorflow/compiler/xla/service/hlo_computation.h +++ b/tensorflow/compiler/xla/service/hlo_computation.h @@ -27,12 +27,14 @@ limitations under the License. #include "tensorflow/compiler/xla/map_util.h" #include "tensorflow/compiler/xla/service/dfs_hlo_visitor.h" #include "tensorflow/compiler/xla/service/dfs_hlo_visitor_with_default.h" +#include "tensorflow/compiler/xla/service/hlo.pb.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" #include "tensorflow/compiler/xla/types.h" #include "tensorflow/compiler/xla/xla_data.pb.h" -#include "tensorflow/core/lib/core/bitmap.h" #include "tensorflow/core/lib/core/status.h" #include "tensorflow/core/lib/gtl/array_slice.h" #include "tensorflow/core/lib/gtl/flatmap.h" @@ -54,8 +56,11 @@ class HloComputation { // Builder class for HloComputation. class Builder { public: - explicit Builder(const string& name) - : name_(name), last_added_instruction_(nullptr) {} + explicit Builder(const string& name, + HloInstruction* fusion_instruction = nullptr) + : name_(name), + last_added_instruction_(nullptr), + fusion_instruction_(fusion_instruction) {} // Build and return an HloComputation. The parameter root_instruction // specifies the already-added instruction to use as the root. If @@ -74,6 +79,7 @@ class HloComputation { private: const string name_; HloInstruction* last_added_instruction_; + HloInstruction* fusion_instruction_; std::vector> instructions_; }; @@ -81,6 +87,16 @@ class HloComputation { // the instruction. HloInstruction* AddInstruction(std::unique_ptr instruction); + // Remove the param_no'th parameter from the computation. + // Note this is only applicatable to the computation for the fusion + // instruction. + Status RemoveParameter(int64 param_no); + + // Add new parameter instruction to the computation. + // This should be a new parameter. Instruction will be appended to parameters + // and inserted to the instruction list. + HloInstruction* AddParameter(std::unique_ptr instruction); + // Remove an instruction from the computation. The instruction must have no // users. Instruction is deallocated with this call. Status RemoveInstruction(HloInstruction* instruction); @@ -98,8 +114,8 @@ class HloComputation { // Set the root of the computation to the given instruction. The instruction // must have already been added to the computation and have the same shape as - // the result of the computation. - void set_root_instruction(HloInstruction* instruction); + // the result of the computation for non fusion computations. + void set_root_instruction(HloInstruction* new_root_instruction); // Return the root instruction of the computation. The root instruction is the // instruction which produces the output of the computation. @@ -111,7 +127,8 @@ class HloComputation { // Returns the parameter instruction for the given parameter number. HloInstruction* parameter_instruction(int64 param_no) const { CHECK_GE(param_no, 0); - CHECK_LT(param_no, static_cast(param_instructions_.size())); + CHECK_LT(param_no, static_cast(param_instructions_.size())) + << "Computation " << name() << " has no parameter number " << param_no; return param_instructions_[param_no]; } @@ -120,10 +137,16 @@ class HloComputation { } const string& name() const { return name_; } - void set_name(const string& name) { name_ = name; } + + // Use the given NameUniquer to select a unique name for the computation based + // on the computation's existing name. + void UniquifyName(NameUniquer* name_uniquer); // Return a string representation of the computation. - string ToString() const; + string ToString(int nested_level = 0) const; + + // Returns a serialized representation of this computation. + HloComputationProto ToProto() const; const std::list>& instructions() const { return instructions_; @@ -133,9 +156,18 @@ class HloComputation { // this order, definitions of values always appear before their uses. std::list MakeInstructionPostOrder() const; - // Computes and returns the mapping from HLO to its transitive operands. - class ReachabilityMap; - std::unique_ptr ComputeTransitiveOperands() 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 instructions_.size(); } @@ -170,8 +202,16 @@ class HloComputation { // producing the copied result. All instructions performing the copy are added // to the computation. For array-shaped values, this method trivially returns // a kCopy instruction. For tuple-shaped instructions, the copy is performed - // with a series of kGetTupleElement and kTuple instructions. - StatusOr DeepCopyInstruction(HloInstruction* instruction); + // with a series of kGetTupleElement and kTuple instructions. If + // indices_to_copy is non-null then this ShapeTree indicates which elements + // (arrays) of the shape to copy. Non-copied elements are passed through + // transparently. If copies_added is non-null, then the added kCopy + // instructions will be inserted in the respective index in the given + // ShapeTree. + StatusOr DeepCopyInstruction( + HloInstruction* instruction, + const ShapeTree* indices_to_copy = nullptr, + ShapeTree* copies_added = nullptr); // Computes and returns the ProgramShape of this computation (shape of // parameters and result without layout). @@ -220,17 +260,33 @@ class HloComputation { // Same as Accept() above, but the visitor is given as a function. Status Accept(const FunctionVisitor::VisitorFunction& visitor_func) const; - // Returns true if instructions of the given opcode can be removed from the + // Returns a deep copy of this computation including all instructions. + std::unique_ptr Clone(const string& suffix = "clone"); + + // Returns true if the given instruction can be removed from the // computation. Instructions such as parameters and send/receive instructions // cannot be removed without violating invariants of the HLO computation or - // module. - static bool IsRemovable(const HloOpcode& opcode); + // module with the exception of fusion computation. A parameter instruction + // is removable for a fusion computation. + bool IsRemovable(const HloInstruction* instruction); + + // 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; + + // Returns if this computation is a fusion computation. + bool IsFusionComputation() const { return fusion_instruction_ != nullptr; } + + // Returns the owning fusion instruction, or nullptr if this is not a fusion + // computation. + HloInstruction* FusionInstruction() const { return fusion_instruction_; } private: explicit HloComputation( const string& name, int parameter_count, std::vector>* instructions, - HloInstruction* root_instruction); + HloInstruction* root_instruction, + HloInstruction* fusion_instruction = nullptr); // Internal helper for adding instructions. HloInstruction* AddInstructionInternal( @@ -238,10 +294,6 @@ class HloComputation { // Helper for setting the parent of instructions that are added to this // computation. - // - // Because we clone HLO instructions without knowing what computation they're - // destined to be added to, this is required to appropriate set the parent on - // fused instruction sequences. void Reparent(HloInstruction* instruction); // Fuses HLOs in instructions_to_fuse into fusion_instruction. @@ -251,9 +303,11 @@ class HloComputation { tensorflow::gtl::ArraySlice instructions_to_fuse, HloInstruction* fusion_instruction); - // Internal helper for copying a tuple value. Creates and returns a deep copy - // of the given instruction. The given instruction must be tuple-shaped. - StatusOr DeepCopyTuple(HloInstruction* instruction); + // Internal helper for recursive copying of an instruction. Creates and + // returns a deep copy of the given instruction. + StatusOr DeepCopyHelper( + HloInstruction* instruction, const ShapeTree* indices_to_copy, + ShapeTree* copies_added, ShapeIndex* index); // Internal helper to collect unreachable roots. std::vector CollectUnreachableRoots() const; @@ -261,6 +315,10 @@ class HloComputation { string name_; HloInstruction* root_instruction_; + // If this computation is a fusion computation, this field points to the + // corresponding fusion instruction. Otherwise, this is null. + HloInstruction* fusion_instruction_; + // Module containing this computation. HloModule* parent_ = nullptr; @@ -282,34 +340,6 @@ class HloComputation { TF_DISALLOW_COPY_AND_ASSIGN(HloComputation); }; -class HloComputation::ReachabilityMap { - public: - // Sets up an empty reachable matrix for the full set of - // instructions specified in "all_instructions" - explicit ReachabilityMap(const std::list& all_instructions); - // Sets entry so that IsReachable(a, b) will return true - void SetReachable(const HloInstruction* a, const HloInstruction* b); - - // Sets IsReachable(a_inst, b_inst) as well as IsReachable(a_inst, trans) - // for all "trans" s.t. "IsReachable(b_inst, trans)" is true - void SetReachableAndTransitiveClosure(const HloInstruction* a_inst, - const HloInstruction* b_inst); - - // Returns true if "b" is reachable from "a" - bool IsReachable(const HloInstruction* a, const HloInstruction* b) const; - - // Returns true if "b" is reachable from "a" or "a" is reachable from "b" - bool IsConnected(const HloInstruction* a, const HloInstruction* b) const; - - private: - friend class HloComputation; - - // dense id assignment from HloInstruction* to number - tensorflow::gtl::FlatMap ids_; - // matrix_(a,b) is true iff b is reachable from a - tensorflow::core::Bitmap matrix_; -}; - } // namespace xla #endif // TENSORFLOW_COMPILER_XLA_SERVICE_HLO_COMPUTATION_H_ diff --git a/tensorflow/compiler/xla/service/hlo_computation_test.cc b/tensorflow/compiler/xla/service/hlo_computation_test.cc index 12a568339627bea412dbbf478474df0f7e8190a6..abd99e64dfc6a353ffd663e3f0b41fb6908bda8a 100644 --- a/tensorflow/compiler/xla/service/hlo_computation_test.cc +++ b/tensorflow/compiler/xla/service/hlo_computation_test.cc @@ -20,15 +20,22 @@ limitations under the License. #include "tensorflow/compiler/xla/literal_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_matchers.h" #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/test_helpers.h" #include "tensorflow/compiler/xla/tests/hlo_test_base.h" +namespace op = xla::testing::opcode_matchers; + namespace xla { namespace { +using ::testing::ElementsAre; +using ::testing::UnorderedElementsAre; + class HloComputationTest : public HloTestBase { protected: HloComputationTest() {} @@ -58,55 +65,63 @@ class HloComputationTest : public HloTestBase { }; TEST_F(HloComputationTest, GetEmbeddedComputationsEmpty) { - auto negate_computation = CreateNegateComputation(); + auto module = CreateNewModule(); + auto negate_computation = + module->AddEntryComputation(CreateNegateComputation()); EXPECT_TRUE(negate_computation->MakeEmbeddedComputationsList().empty()); } TEST_F(HloComputationTest, GetEmbeddedComputationsOneComputation) { // Create computation which calls one other computation. - auto negate_computation = CreateNegateComputation(); - auto map_computation = CreateMapComputation(negate_computation.get()); + auto module = CreateNewModule(); + auto negate_computation = + module->AddEmbeddedComputation(CreateNegateComputation()); + auto map_computation = + module->AddEntryComputation(CreateMapComputation(negate_computation)); EXPECT_TRUE(negate_computation->MakeEmbeddedComputationsList().empty()); - EXPECT_EQ(map_computation->MakeEmbeddedComputationsList().front(), - negate_computation.get()); + EXPECT_THAT(map_computation->MakeEmbeddedComputationsList(), + ElementsAre(negate_computation)); } TEST_F(HloComputationTest, GetEmbeddedComputationsDiamond) { // Create computations with a diamond-shaped callgraph. - auto negate_computation = CreateNegateComputation(); - auto map1_computation = CreateMapComputation(negate_computation.get()); - auto map2_computation = CreateMapComputation(negate_computation.get()); + auto module = CreateNewModule(); + auto negate_computation = + module->AddEmbeddedComputation(CreateNegateComputation()); + auto map1_computation = + module->AddEmbeddedComputation(CreateMapComputation(negate_computation)); + auto map2_computation = + module->AddEmbeddedComputation(CreateMapComputation(negate_computation)); auto builder = HloComputation::Builder(TestName()); auto param = builder.AddInstruction( HloInstruction::CreateParameter(0, r0f32_, "param0")); auto map1 = builder.AddInstruction( - HloInstruction::CreateMap(r0f32_, {param}, map1_computation.get())); + HloInstruction::CreateMap(r0f32_, {param}, map1_computation)); auto map2 = builder.AddInstruction( - HloInstruction::CreateMap(r0f32_, {param}, map2_computation.get())); + HloInstruction::CreateMap(r0f32_, {param}, map2_computation)); builder.AddInstruction( HloInstruction::CreateBinary(r0f32_, HloOpcode::kAdd, map1, map2)); - auto computation = builder.Build(); + auto computation = module->AddEntryComputation(builder.Build()); auto embedded_computations = computation->MakeEmbeddedComputationsList(); EXPECT_EQ(3, embedded_computations.size()); // GetEmbeddedComputations returns a post order of the embedded computations, // so the negate computation must come first. - EXPECT_EQ(negate_computation.get(), *embedded_computations.begin()); - EXPECT_MATCH(testing::ListToVec(embedded_computations), - testing::UnorderedMatcher( - negate_computation.get(), map1_computation.get(), - map2_computation.get())); + EXPECT_EQ(negate_computation, *embedded_computations.begin()); + EXPECT_THAT(embedded_computations, + UnorderedElementsAre(negate_computation, map1_computation, + map2_computation)); } TEST_F(HloComputationTest, PostOrderSingleton) { // Test GetInstructionPostOrder for a computation with one instruction. auto builder = HloComputation::Builder(TestName()); auto constant = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(42.0f))); - auto computation = builder.Build(); - - EXPECT_EQ(computation->MakeInstructionPostOrder().front(), constant); + HloInstruction::CreateConstant(Literal::CreateR0(42.0f))); + auto module = CreateNewModule(); + auto computation = module->AddEntryComputation(builder.Build()); + EXPECT_THAT(computation->MakeInstructionPostOrder(), ElementsAre(constant)); } TEST_F(HloComputationTest, PostOrderSimple) { @@ -114,37 +129,33 @@ TEST_F(HloComputationTest, PostOrderSimple) { // instructions. auto builder = HloComputation::Builder(TestName()); auto constant = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(42.0f))); + HloInstruction::CreateConstant(Literal::CreateR0(42.0f))); auto negate1 = builder.AddInstruction( HloInstruction::CreateUnary(r0f32_, HloOpcode::kNegate, constant)); auto negate2 = builder.AddInstruction( HloInstruction::CreateUnary(r0f32_, HloOpcode::kNegate, negate1)); - auto computation = builder.Build(); - - EXPECT_MATCH( - testing::ListToVec( - computation->MakeInstructionPostOrder()), - testing::OrderedMatcher(constant, negate1, negate2)); + auto module = CreateNewModule(); + auto computation = module->AddEntryComputation(builder.Build()); + EXPECT_THAT(computation->MakeInstructionPostOrder(), + ElementsAre(constant, negate1, negate2)); } TEST_F(HloComputationTest, PostOrderTrace) { // Test GetInstructionPostOrder for a computation with a trace instruction. auto builder = HloComputation::Builder(TestName()); auto constant = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(42.0f))); + HloInstruction::CreateConstant(Literal::CreateR0(42.0f))); auto negate1 = builder.AddInstruction( HloInstruction::CreateUnary(r0f32_, HloOpcode::kNegate, constant)); auto trace = builder.AddInstruction(HloInstruction::CreateTrace("foobar", negate1)); auto negate2 = builder.AddInstruction( HloInstruction::CreateUnary(r0f32_, HloOpcode::kNegate, negate1)); - auto computation = builder.Build(); - + auto module = CreateNewModule(); + auto computation = module->AddEntryComputation(builder.Build()); // Trace instructions should be at the end of the sort. - EXPECT_MATCH(testing::ListToVec( - computation->MakeInstructionPostOrder()), - testing::OrderedMatcher(constant, negate1, - negate2, trace)); + EXPECT_THAT(computation->MakeInstructionPostOrder(), + ElementsAre(constant, negate1, negate2, trace)); } TEST_F(HloComputationTest, PostOrderDisconnectedInstructions) { @@ -152,19 +163,17 @@ TEST_F(HloComputationTest, PostOrderDisconnectedInstructions) { // which are not connected. auto builder = HloComputation::Builder(TestName()); auto constant1 = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(42.0f))); + HloInstruction::CreateConstant(Literal::CreateR0(42.0f))); auto constant2 = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(42.0f))); + HloInstruction::CreateConstant(Literal::CreateR0(42.0f))); auto constant3 = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(42.0f))); + HloInstruction::CreateConstant(Literal::CreateR0(42.0f))); auto constant4 = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(42.0f))); - auto computation = builder.Build(); - - EXPECT_MATCH(testing::ListToVec( - computation->MakeInstructionPostOrder()), - testing::UnorderedMatcher( - constant1, constant2, constant3, constant4)); + HloInstruction::CreateConstant(Literal::CreateR0(42.0f))); + auto module = CreateNewModule(); + auto computation = module->AddEntryComputation(builder.Build()); + EXPECT_THAT(computation->MakeInstructionPostOrder(), + UnorderedElementsAre(constant1, constant2, constant3, constant4)); } TEST_F(HloComputationTest, PostOrderWithMultipleRoots) { @@ -172,24 +181,23 @@ TEST_F(HloComputationTest, PostOrderWithMultipleRoots) { // which are not connected. auto builder = HloComputation::Builder(TestName()); auto constant1 = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(42.0f))); + HloInstruction::CreateConstant(Literal::CreateR0(42.0f))); auto constant2 = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(42.0f))); + HloInstruction::CreateConstant(Literal::CreateR0(42.0f))); auto constant3 = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(42.0f))); + HloInstruction::CreateConstant(Literal::CreateR0(42.0f))); auto add1 = builder.AddInstruction(HloInstruction::CreateBinary( r0f32_, HloOpcode::kAdd, constant1, constant2)); auto add2 = builder.AddInstruction(HloInstruction::CreateBinary( r0f32_, HloOpcode::kAdd, constant2, constant3)); auto add3 = builder.AddInstruction(HloInstruction::CreateBinary( r0f32_, HloOpcode::kAdd, constant1, constant3)); - auto computation = builder.Build(); - + auto module = CreateNewModule(); + auto computation = module->AddEntryComputation(builder.Build()); auto post_order = computation->MakeInstructionPostOrder(); EXPECT_EQ(6, post_order.size()); - EXPECT_MATCH(testing::ListToVec(post_order), - testing::UnorderedMatcher( - constant1, constant2, constant3, add1, add2, add3)); + EXPECT_THAT(post_order, UnorderedElementsAre(constant1, constant2, constant3, + add1, add2, add3)); } TEST_F(HloComputationTest, VisitWithMultipleRoots) { @@ -197,11 +205,11 @@ TEST_F(HloComputationTest, VisitWithMultipleRoots) { // computation has multiple roots (dead code). auto builder = HloComputation::Builder(TestName()); auto constant1 = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(42.0f))); + HloInstruction::CreateConstant(Literal::CreateR0(42.0f))); auto constant2 = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(42.0f))); + HloInstruction::CreateConstant(Literal::CreateR0(42.0f))); auto constant3 = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(42.0f))); + HloInstruction::CreateConstant(Literal::CreateR0(42.0f))); // Add three disconnected add expressions. builder.AddInstruction(HloInstruction::CreateBinary(r0f32_, HloOpcode::kAdd, constant1, constant2)); @@ -209,8 +217,8 @@ TEST_F(HloComputationTest, VisitWithMultipleRoots) { constant2, constant3)); builder.AddInstruction(HloInstruction::CreateBinary(r0f32_, HloOpcode::kAdd, constant1, constant3)); - auto computation = builder.Build(); - + auto module = CreateNewModule(); + auto computation = module->AddEntryComputation(builder.Build()); // Visitor which keeps track of which instructions have been visited. class TestVisitor : public DfsHloVisitorWithDefault { public: @@ -236,7 +244,7 @@ TEST_F(HloComputationTest, VisitWithMultipleRoots) { HloInstruction* last_visited_ = nullptr; }; - TestVisitor visitor(computation.get()); + TestVisitor visitor(computation); EXPECT_IS_OK(computation->Accept(&visitor)); EXPECT_EQ(6, visitor.visited_set_.size()); @@ -248,62 +256,140 @@ TEST_F(HloComputationTest, DeepCopyArray) { // Test that DeepCopyInstruction properly copies an array. auto builder = HloComputation::Builder(TestName()); auto constant = builder.AddInstruction(HloInstruction::CreateConstant( - LiteralUtil::CreateR1({1.0, 2.0, 3.0}))); - auto computation = builder.Build(); - + Literal::CreateR1({1.0, 2.0, 3.0}))); + auto module = CreateNewModule(); + auto computation = module->AddEntryComputation(builder.Build()); auto copy = computation->DeepCopyInstruction(constant).ValueOrDie(); - EXPECT_EQ(HloOpcode::kCopy, copy->opcode()); - EXPECT_EQ(constant, copy->operand(0)); + EXPECT_THAT(copy, op::Copy(constant)); } TEST_F(HloComputationTest, DeepCopyTuple) { // Test that DeepCopyInstruction properly copies a tuple. auto builder = HloComputation::Builder(TestName()); auto constant1 = builder.AddInstruction(HloInstruction::CreateConstant( - LiteralUtil::CreateR1({1.0, 2.0, 3.0}))); + Literal::CreateR1({1.0, 2.0, 3.0}))); auto constant2 = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(42.0))); + HloInstruction::CreateConstant(Literal::CreateR0(42.0))); auto tuple = builder.AddInstruction( HloInstruction::CreateTuple({constant1, constant2})); + auto module = CreateNewModule(); + 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_EQ(0, tuple_copy->operand(0)->operand(0)->tuple_index()); + EXPECT_EQ(1, tuple_copy->operand(1)->operand(0)->tuple_index()); +} + +TEST_F(HloComputationTest, DeepCopyArrayAtIndices) { + // Test that DeepCopyInstruction properly handles an array when the indices to + // copy are specified. + auto builder = HloComputation::Builder(TestName()); + auto constant = builder.AddInstruction(HloInstruction::CreateConstant( + Literal::CreateR1({1.0, 2.0, 3.0}))); auto computation = builder.Build(); - auto tuple_copy = computation->DeepCopyInstruction(tuple).ValueOrDie(); + { + // If the index is true, then a copy should be made. + ShapeTree indices_to_copy(constant->shape(), /*init_value=*/true); + EXPECT_THAT(computation->DeepCopyInstruction(constant, &indices_to_copy) + .ValueOrDie(), + op::Copy(constant)); + } - EXPECT_EQ(HloOpcode::kTuple, tuple_copy->opcode()); - EXPECT_EQ(HloOpcode::kCopy, tuple_copy->operand(0)->opcode()); - const HloInstruction* gte0 = tuple_copy->operand(0)->operand(0); - EXPECT_EQ(HloOpcode::kGetTupleElement, gte0->opcode()); - EXPECT_EQ(0, gte0->tuple_index()); - EXPECT_EQ(tuple, gte0->operand(0)); - - EXPECT_EQ(HloOpcode::kCopy, tuple_copy->operand(1)->opcode()); - const HloInstruction* gte1 = tuple_copy->operand(1)->operand(0); - EXPECT_EQ(HloOpcode::kGetTupleElement, gte1->opcode()); - EXPECT_EQ(1, gte1->tuple_index()); - EXPECT_EQ(tuple, gte1->operand(0)); + { + // If the index is false, then no copy should be made. + ShapeTree indices_to_copy(constant->shape(), /*init_value=*/false); + EXPECT_EQ(computation->DeepCopyInstruction(constant, &indices_to_copy) + .ValueOrDie(), + constant); + } +} + +TEST_F(HloComputationTest, DeepCopyTupleAtIndices) { + // Test that DeepCopyInstruction properly copies elements of a a tuple as + // specified by the given indices. + auto builder = HloComputation::Builder(TestName()); + auto constant1 = builder.AddInstruction(HloInstruction::CreateConstant( + Literal::CreateR1({1.0, 2.0, 3.0}))); + auto constant2 = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(42.0))); + auto tuple = builder.AddInstruction( + HloInstruction::CreateTuple({constant1, constant2})); + auto computation = builder.Build(); + + { + // All true values should copy all array elements. + ShapeTree indices_to_copy(tuple->shape(), /*init_value=*/true); + ShapeTree copies_added(tuple->shape(), + /*init_value=*/nullptr); + HloInstruction* deep_copy = + 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}))); + } + + { + // All false elements should copy no array elements, but the GTE and tuple + // instruction scaffolding should be built. + ShapeTree indices_to_copy(tuple->shape(), /*init_value=*/false); + ShapeTree copies_added(tuple->shape(), + /*init_value=*/nullptr); + HloInstruction* deep_copy = + computation->DeepCopyInstruction(tuple, &indices_to_copy, &copies_added) + .ValueOrDie(); + + EXPECT_THAT(deep_copy, op::Tuple(op::GetTupleElement(tuple), + op::GetTupleElement(tuple))); + EXPECT_TRUE(copies_added.element({}) == nullptr); + EXPECT_TRUE(copies_added.element({0}) == nullptr); + EXPECT_TRUE(copies_added.element({1}) == nullptr); + } + + { + // Verify one element copied, the other not. + ShapeTree indices_to_copy(tuple->shape(), /*init_value=*/false); + *indices_to_copy.mutable_element({0}) = true; + ShapeTree copies_added(tuple->shape(), + /*init_value=*/nullptr); + HloInstruction* deep_copy = + computation->DeepCopyInstruction(tuple, &indices_to_copy, &copies_added) + .ValueOrDie(); + + EXPECT_THAT(deep_copy, op::Tuple(op::Copy(op::GetTupleElement(tuple)), + op::GetTupleElement(tuple))); + EXPECT_TRUE(copies_added.element({}) == nullptr); + EXPECT_TRUE(copies_added.element({0}) != nullptr); + EXPECT_TRUE(copies_added.element({1}) == nullptr); + } } TEST_F(HloComputationTest, CycleDetection) { // Test whether the visitor can detect cycles in the graph. auto builder = HloComputation::Builder(TestName()); auto constant = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(42.0f))); + HloInstruction::CreateConstant(Literal::CreateR0(42.0f))); auto negate = builder.AddInstruction( HloInstruction::CreateUnary(r0f32_, HloOpcode::kNegate, constant)); auto add = builder.AddInstruction( HloInstruction::CreateBinary(r0f32_, HloOpcode::kAdd, negate, negate)); - auto computation = builder.Build(); - + auto module = CreateNewModule(); + auto computation = module->AddEntryComputation(builder.Build()); // Add a control dependency to create a cycle. ASSERT_IS_OK(add->AddControlDependencyTo(negate)); const auto visitor = [](HloInstruction* instruction) { return Status::OK(); }; auto visit_status = computation->Accept(visitor); ASSERT_FALSE(visit_status.ok()); - ASSERT_MATCH(visit_status.error_message(), - testing::ContainsRegex("cycle is detecte")); + ASSERT_THAT(visit_status.error_message(), + ::testing::ContainsRegex("cycle is detecte")); } TEST_F(HloComputationTest, RemoveInstructionWithDuplicateOperand) { @@ -312,24 +398,162 @@ TEST_F(HloComputationTest, RemoveInstructionWithDuplicateOperand) { // twice. auto builder = HloComputation::Builder(TestName()); auto constant = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(42.0f))); + HloInstruction::CreateConstant(Literal::CreateR0(42.0f))); auto dead_negate = builder.AddInstruction( HloInstruction::CreateUnary(r0f32_, HloOpcode::kNegate, constant)); auto dead_add = builder.AddInstruction(HloInstruction::CreateBinary( r0f32_, HloOpcode::kAdd, dead_negate, dead_negate)); auto negate = builder.AddInstruction( HloInstruction::CreateUnary(r0f32_, HloOpcode::kNegate, constant)); - auto computation = builder.Build(); - + auto module = CreateNewModule(); + auto computation = module->AddEntryComputation(builder.Build()); EXPECT_EQ(4, computation->instruction_count()); + EXPECT_THAT(computation->root_instruction(), op::Negate(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_EQ(negate, computation->root_instruction()); } +TEST_F(HloComputationTest, CloneWithControlDependency) { + auto builder = HloComputation::Builder(TestName()); + auto constant1 = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(1.0f))); + auto constant2 = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(2.0f))); + auto add = builder.AddInstruction(HloInstruction::CreateBinary( + r0f32_, HloOpcode::kAdd, constant1, constant2)); + + auto param = builder.AddInstruction( + HloInstruction::CreateParameter(0, r0f32_, "param0")); + auto negate = builder.AddInstruction( + HloInstruction::CreateUnary(r0f32_, HloOpcode::kNegate, param)); + auto module = CreateNewModule(); + auto computation = + module->AddEntryComputation(builder.Build(/*root_instruction=*/add)); + + TF_CHECK_OK(negate->AddControlDependencyTo(add)); + + auto clone = computation->Clone(); + + auto cloned_add = clone->root_instruction(); + EXPECT_EQ(cloned_add->opcode(), HloOpcode::kAdd); + + auto predecessors = cloned_add->control_predecessors(); + EXPECT_EQ(1, predecessors.size()); + EXPECT_EQ(HloOpcode::kNegate, predecessors[0]->opcode()); + auto successors = predecessors[0]->control_successors(); + 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'. + auto builder = HloComputation::Builder(TestName()); + auto constant1 = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(1.0f))); + auto constant2 = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::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(); + 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)); +} + } // namespace } // namespace xla + +int main(int argc, char** argv) { + return xla::ParseDebugOptionsFlagsAndRunTests(argc, argv); +} diff --git a/tensorflow/compiler/xla/service/hlo_constant_folding.cc b/tensorflow/compiler/xla/service/hlo_constant_folding.cc index 9a5345dc13d6db42553e9c343f7c81cd0e6c9d0e..58761cb4a487ef12b0cbeefd6820b415d724733c 100644 --- a/tensorflow/compiler/xla/service/hlo_constant_folding.cc +++ b/tensorflow/compiler/xla/service/hlo_constant_folding.cc @@ -15,70 +15,79 @@ limitations under the License. #include "tensorflow/compiler/xla/service/hlo_constant_folding.h" -#include -#include #include -#include #include #include #include #include "tensorflow/compiler/xla/layout_util.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_evaluator.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/shape_util.h" #include "tensorflow/compiler/xla/types.h" -#include "tensorflow/compiler/xla/xla_data.pb.h" #include "tensorflow/core/lib/core/errors.h" namespace xla { StatusOr HloConstantFolding::Run(HloModule* module) { + auto evaluator = MakeUnique(); + + XLA_VLOG_LINES(2, + "HloConstantFolding::Run(), before:\n" + module->ToString()); bool changed = false; + for (auto& computation : module->computations()) { + if (computation->IsFusionComputation()) { + continue; + } for (auto instruction : computation->MakeInstructionPostOrder()) { // Skip dead code. if (instruction->user_count() == 0 && computation->root_instruction() != instruction) { continue; } - // Depending on the opcode, choose how to handle constant operands. - // - // TODO(b/35975797): Fold constant computations for more than reshapes and - // transposes. - switch (instruction->opcode()) { - case HloOpcode::kReshape: { - if (instruction->operand(0)->opcode() == HloOpcode::kConstant) { - TF_ASSIGN_OR_RETURN( - auto reshaped_literal, - LiteralUtil::Reshape( - instruction->operand(0)->literal(), - AsInt64Slice(instruction->shape().dimensions()))); - TF_CHECK_OK(computation->ReplaceWithNewInstruction( - instruction, - HloInstruction::CreateConstant(std::move(reshaped_literal)))); - changed = true; - } - break; - } - case HloOpcode::kTranspose: { - if (instruction->operand(0)->opcode() == HloOpcode::kConstant) { - auto transposed_literal = LiteralUtil::Transpose( - instruction->operand(0)->literal(), instruction->dimensions()); - TF_CHECK_OK(computation->ReplaceWithNewInstruction( - instruction, - HloInstruction::CreateConstant(std::move(transposed_literal)))); - changed = true; - } - break; - } - default: - break; + // Skip Constant, Parameter, Reduce operation. + // TODO(b/35975797): Enable Reduce operation once arbitary computation are + // supported by the evaluator. + // TODO(b/64407269): Enable Tuple once the timeout issue is resolved. + if (instruction->opcode() == HloOpcode::kParameter || + instruction->opcode() == HloOpcode::kConstant || + instruction->opcode() == HloOpcode::kTuple || + instruction->opcode() == HloOpcode::kReduce) { + continue; } + // Skip instructions with non-constant operands. + if (!hlo_query::AllOperandsAreConstants(*instruction)) { + continue; + } + + // Broadcasts dramatically increase the size of constants with is often + // detrimental to performance and memory capacity so do not fold + // broadcasts. + if (instruction->opcode() == HloOpcode::kBroadcast) { + continue; + } + + std::unique_ptr result = evaluator->TryEvaluate(instruction); + // Currently we skip unimplemented operations. + // TODO(b/35975797): Fold constant computations for more operations. + if (result == nullptr) { + VLOG(2) << "Constant folding failed for instruction: " + << instruction->ToString(); + continue; + } + + TF_RETURN_IF_ERROR(computation->ReplaceWithNewInstruction( + instruction, HloInstruction::CreateConstant(std::move(result)))); + changed = true; } } + XLA_VLOG_LINES(2, "HloConstantFolding::Run(), after:\n" + module->ToString()); return changed; } diff --git a/tensorflow/compiler/xla/service/hlo_constant_folding.h b/tensorflow/compiler/xla/service/hlo_constant_folding.h index 514bb8164c1e1fa10a36ceeeac63dc946de2ab5a..331480bd029727fa15476cb9ced2e7b7afd170f3 100644 --- a/tensorflow/compiler/xla/service/hlo_constant_folding.h +++ b/tensorflow/compiler/xla/service/hlo_constant_folding.h @@ -21,16 +21,14 @@ limitations under the License. namespace xla { -// A pass which performs constant folding in order to avoid unecessary +// A pass which performs constant folding in order to avoid unnecessary // computation on constants. class HloConstantFolding : public HloPassInterface { public: - explicit HloConstantFolding() {} - ~HloConstantFolding() override {} tensorflow::StringPiece name() const override { return "constant_folding"; } - // Run ConstantFolding on the given module. Returns whether the module was - // changed (common subexpressions were found and eliminated). + // Run constant folding operations on the given module. Returns whether the + // module was changed (constant expressions folded). StatusOr Run(HloModule* module) override; }; diff --git a/tensorflow/compiler/xla/service/hlo_constant_folding_test.cc b/tensorflow/compiler/xla/service/hlo_constant_folding_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..3ae499d5e0c37532ae0a83a4a247cab85fd2c84e --- /dev/null +++ b/tensorflow/compiler/xla/service/hlo_constant_folding_test.cc @@ -0,0 +1,210 @@ +/* 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/hlo_constant_folding.h" + +#include +#include + +#include "tensorflow/compiler/xla/layout_util.h" +#include "tensorflow/compiler/xla/literal_util.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_pass_fix.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/literal_test_util.h" +#include "tensorflow/compiler/xla/types.h" + +namespace op = xla::testing::opcode_matchers; + +namespace xla { +namespace { + +using HloConstantFoldingTest = HloTestBase; + +TEST_F(HloConstantFoldingTest, ConvertF32ToS64) { + HloComputation::Builder builder(TestName()); + HloInstruction* input = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(42.0f))); + builder.AddInstruction( + HloInstruction::CreateConvert(ShapeUtil::MakeShape(S64, {}), input)); + + auto module = CreateNewModule(); + auto computation = module->AddEntryComputation(builder.Build()); + + EXPECT_THAT(computation->root_instruction(), op::Convert(input)); + + HloConstantFolding const_folder; + TF_ASSERT_OK_AND_ASSIGN(bool result, const_folder.Run(module.get())); + EXPECT_TRUE(result); + + EXPECT_THAT(computation->root_instruction(), op::Constant()); + EXPECT_EQ(computation->root_instruction()->literal().GetFirstElement(), + 42); +} + +TEST_F(HloConstantFoldingTest, ConvertS64ToF32) { + HloComputation::Builder builder(TestName()); + HloInstruction* input = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(42))); + builder.AddInstruction( + HloInstruction::CreateConvert(ShapeUtil::MakeShape(F32, {}), input)); + + auto module = CreateNewModule(); + auto computation = module->AddEntryComputation(builder.Build()); + + EXPECT_THAT(computation->root_instruction(), op::Convert(input)); + + HloConstantFolding const_folder; + TF_ASSERT_OK_AND_ASSIGN(bool result, const_folder.Run(module.get())); + EXPECT_TRUE(result); + + EXPECT_THAT(computation->root_instruction(), op::Constant()); + EXPECT_EQ(computation->root_instruction()->literal().GetFirstElement(), + 42.0f); +} + +TEST_F(HloConstantFoldingTest, ConvertF32ArrayToS64Array) { + HloComputation::Builder builder(TestName()); + HloInstruction* input = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR1({42.0f, 19.0f}))); + builder.AddInstruction( + HloInstruction::CreateConvert(ShapeUtil::MakeShape(S64, {2}), input)); + + auto module = CreateNewModule(); + auto computation = module->AddEntryComputation(builder.Build()); + + EXPECT_THAT(computation->root_instruction(), op::Convert(input)); + + HloConstantFolding const_folder; + TF_ASSERT_OK_AND_ASSIGN(bool result, const_folder.Run(module.get())); + EXPECT_TRUE(result); + + EXPECT_THAT(computation->root_instruction(), op::Constant()); + EXPECT_EQ(computation->root_instruction()->literal().Get({0}), 42); + EXPECT_EQ(computation->root_instruction()->literal().Get({1}), 19); +} + +TEST_F(HloConstantFoldingTest, Concatenate) { + const struct TestConfig { + int concat_dimension; + tensorflow::gtl::ArraySlice dimensions; + tensorflow::gtl::ArraySlice concat_sizes; + } test_configs[] = { + {1, {11, 0, 7, 5, 9}, {2, 5, 7, 11}}, + {3, {1, 4, 17, 0, 8}, {1, 3, 9, 12}}, + }; + + for (auto& test_config : test_configs) { + HloComputation::Builder builder(TestName()); + std::vector dimensions(test_config.dimensions.begin(), + test_config.dimensions.end()); + int64 concat_size = 0; + std::vector operands; + for (auto csize : test_config.concat_sizes) { + dimensions[test_config.concat_dimension] = csize; + concat_size += csize; + auto literal = Literal::CreateFromDimensions(F32, dimensions); + HloInstruction* insn = builder.AddInstruction( + HloInstruction::CreateConstant(std::move(literal))); + operands.push_back(insn); + } + dimensions[test_config.concat_dimension] = concat_size; + Shape shape = ShapeUtil::MakeShape(F32, dimensions); + builder.AddInstruction(HloInstruction::CreateConcatenate( + shape, operands, test_config.concat_dimension)); + auto module = CreateNewModule(); + auto computation = module->AddEntryComputation(builder.Build()); + + HloConstantFolding const_folder; + 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_TRUE(ShapeUtil::Equal(root->shape(), shape)); + } +} + +TEST_F(HloConstantFoldingTest, Slice) { + HloComputation::Builder builder(TestName()); + const int64 dimensions[] = {11, 8, 7, 5, 9}; + const int64 slice_start[] = {4, 2, 3, 1, 5}; + const int64 slice_limits[] = {10, 8, 6, 5, 9}; + const int64 slice_strides[] = {1, 1, 1, 1, 1}; + TF_ASSERT_OK_AND_ASSIGN(auto literal, + LiteralTestUtil::CreateRandomLiteral( + ShapeUtil::MakeShape(F32, dimensions), 0.0, 1.0)); + HloInstruction* literal_instruction = builder.AddInstruction( + HloInstruction::CreateConstant(std::move(literal))); + 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 computation = module->AddEntryComputation(builder.Build()); + + HloConstantFolding const_folder; + 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_TRUE(ShapeUtil::Equal(root->shape(), shape)); +} + +TEST_F(HloConstantFoldingTest, TransposeConstantFold) { + HloComputation::Builder builder(TestName()); + const int64 dimensions[] = {11, 8, 7, 5, 9}; + TF_ASSERT_OK_AND_ASSIGN(auto literal, + LiteralTestUtil::CreateRandomLiteral( + ShapeUtil::MakeShape(F32, dimensions), 0.0, 1.0)); + auto literal_clone = literal->Literal::CloneToUnique(); + HloInstruction* literal_instruction = builder.AddInstruction( + HloInstruction::CreateConstant(std::move(literal))); + Shape shape = ShapeUtil::MakeShape(F32, {8, 7, 11, 9, 5}); + const int64 permutation[] = {1, 2, 0, 4, 3}; + builder.AddInstruction( + HloInstruction::CreateTranspose(shape, literal_instruction, permutation)); + auto module = CreateNewModule(); + auto computation = module->AddEntryComputation(builder.Build()); + + HloConstantFolding const_folder; + 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_TRUE(ShapeUtil::Compatible(root->shape(), shape)); + + using NativeT = typename primitive_util::PrimitiveTypeToNative::type; + bool matched = true; + root->literal().EachCell( + [&](tensorflow::gtl::ArraySlice indices, NativeT value) { + std::vector rindexes = Permute(permutation, indices); + matched = matched && (value == literal_clone->Get(rindexes)); + }); + EXPECT_TRUE(matched); +} + +} // namespace +} // namespace xla + +int main(int argc, char** argv) { + return xla::ParseDebugOptionsFlagsAndRunTests(argc, argv); +} diff --git a/tensorflow/compiler/xla/service/hlo_cost_analysis.cc b/tensorflow/compiler/xla/service/hlo_cost_analysis.cc index 5fc7d6b22e924a701812bac3655f95e00799004d..f6b764732b495a1b60bd7dac114ee99bc70bd1b6 100644 --- a/tensorflow/compiler/xla/service/hlo_cost_analysis.cc +++ b/tensorflow/compiler/xla/service/hlo_cost_analysis.cc @@ -25,34 +25,57 @@ limitations under the License. namespace xla { +constexpr char HloCostAnalysis::kFlopsKey[]; +constexpr char HloCostAnalysis::kTranscendentalsKey[]; +constexpr char HloCostAnalysis::kBytesAccessedKey[]; +constexpr char HloCostAnalysis::kSecondsKey[]; + +HloCostAnalysis::HloCostAnalysis(const ShapeSizeFunction& shape_size) + : HloCostAnalysis(shape_size, {}) {} + +HloCostAnalysis::HloCostAnalysis(const ShapeSizeFunction& shape_size, + const Properties& per_second_rates) + : shape_size_(shape_size), per_second_rates_(per_second_rates) {} + Status HloCostAnalysis::Preprocess(HloInstruction* hlo) { // Set current instruction cost values to reasonable default values. Each - // handler can overwrite these values. In Postprocess, these value are + // handler can overwrite these values. In Postprocess, these values are // accumulated and written to the per-instruction maps. - current_flop_count_ = 0; - current_transcendental_count_ = 0; + current_properties_.clear(); + current_should_compute_bottleneck_time_ = true; - // The default element count for an instruction is the sum of elements in the - // operands and output. The default ShapeUtil::ByteSizeOf does not handle - // opaque types. - current_bytes_accessed_ = shape_size_(hlo->shape()); + // The default number of bytes accessed for an instruction is the sum of the + // sizes of the inputs and outputs. The default ShapeUtil::ByteSizeOf does not + // handle opaque types. + float bytes_accessed = shape_size_(hlo->shape()); for (const HloInstruction* operand : hlo->operands()) { - current_bytes_accessed_ += shape_size_(operand->shape()); + bytes_accessed += shape_size_(operand->shape()); } + current_properties_[kBytesAccessedKey] = bytes_accessed; return Status::OK(); } Status HloCostAnalysis::Postprocess(HloInstruction* hlo) { - // Accumulate cost values and write into per-instruction maps. - flop_count_ += current_flop_count_; - hlo_to_flop_count_[hlo] = current_flop_count_; - - transcendental_count_ += current_transcendental_count_; - hlo_to_transcendental_count_[hlo] = current_transcendental_count_; + if (current_should_compute_bottleneck_time_) { + // Compute the time as the time of the bottleneck, i.e. the slowest property + // given the per-second rate of each property. + float max_seconds = 0.0f; + for (const auto& property : current_properties_) { + if (property.first != kSecondsKey) { + max_seconds = std::max( + max_seconds, + property.second / + GetProperty(property.first, per_second_rates_, INFINITY)); + } + } + current_properties_[kSecondsKey] = max_seconds; + } - bytes_accessed_ += current_bytes_accessed_; - hlo_to_bytes_accessed_[hlo] = current_bytes_accessed_; + TF_RET_CHECK(hlo_properties_.emplace(hlo, current_properties_).second); + for (const auto& property : current_properties_) { + properties_sum_[property.first] += property.second; + } return Status::OK(); } @@ -63,27 +86,43 @@ Status HloCostAnalysis::HandleElementwiseOp(HloInstruction* hlo_instruction) { // number of elements in the output shape. auto computation_count = ShapeUtil::ElementsIn(shape); auto opcode = hlo_instruction->opcode(); - // We treat the two opcodes (kExp, kPower) as transcendental operations. - if (opcode == HloOpcode::kExp || opcode == HloOpcode::kPower) { - current_transcendental_count_ = computation_count; + // We treat transcendental operations separately since one transcendental + // operation can correspond to several floating point ops. + if (opcode == HloOpcode::kExp || opcode == HloOpcode::kPower || + opcode == HloOpcode::kTanh || opcode == HloOpcode::kSin || + opcode == HloOpcode::kCos) { + current_properties_[kTranscendentalsKey] = computation_count; } else { // Note: transcendental operations are considered a separate category from // FLOPs. - current_flop_count_ = computation_count; + current_properties_[kFlopsKey] = computation_count; } return Status::OK(); } -Status HloCostAnalysis::HandleElementwiseUnary(HloInstruction* hlo, - HloOpcode opcode, - HloInstruction* operand) { +/*static*/ float HloCostAnalysis::GetProperty(const string& key, + const Properties& properties, + const float default_value) { + auto key_value = properties.find(key); + return key_value == properties.end() ? default_value : key_value->second; +} + +/*static*/ float HloCostAnalysis::GetPropertyForHlo( + const HloInstruction& hlo, const string& key, + const HloToProperties& hlo_to_properties) { + auto it = hlo_to_properties.find(&hlo); + if (it == hlo_to_properties.end()) { + return 0.0f; + } else { + return GetProperty(key, it->second); + } +} + +Status HloCostAnalysis::HandleElementwiseUnary(HloInstruction* hlo) { return HandleElementwiseOp(hlo); } -Status HloCostAnalysis::HandleElementwiseBinary(HloInstruction* hlo, - HloOpcode opcode, - HloInstruction* lhs, - HloInstruction* rhs) { +Status HloCostAnalysis::HandleElementwiseBinary(HloInstruction* hlo) { return HandleElementwiseOp(hlo); } @@ -100,14 +139,18 @@ Status HloCostAnalysis::HandleClamp(HloInstruction* clamp, return HandleElementwiseOp(clamp); } +Status HloCostAnalysis::HandleReducePrecision(HloInstruction* hlo) { + return HandleElementwiseOp(hlo); +} + Status HloCostAnalysis::HandleParameter(HloInstruction* parameter) { - current_bytes_accessed_ = 0; + current_properties_[kBytesAccessedKey] = 0; return Status::OK(); } Status HloCostAnalysis::HandleConstant(HloInstruction* constant, const Literal& literal) { - current_bytes_accessed_ = 0; + current_properties_[kBytesAccessedKey] = 0; return Status::OK(); } @@ -115,7 +158,7 @@ Status HloCostAnalysis::HandleGetTupleElement(HloInstruction* get_tuple_element, HloInstruction* operand) { // GetTupleElement forwards a pointer and does not touch each element in the // output. - current_bytes_accessed_ = 0; + current_properties_[kBytesAccessedKey] = 0; return Status::OK(); } @@ -136,9 +179,9 @@ Status HloCostAnalysis::HandleSlice(HloInstruction* slice, return Status::OK(); } -Status HloCostAnalysis::HandleDynamicSlice( - HloInstruction* slice, - tensorflow::gtl::ArraySlice operands) { +Status HloCostAnalysis::HandleDynamicSlice(HloInstruction* dynamic_slice, + HloInstruction* operand, + HloInstruction* start_indices) { return Status::OK(); } @@ -153,8 +196,9 @@ Status HloCostAnalysis::HandleTuple( tensorflow::gtl::ArraySlice operands) { // The tuple instruction only gathers pointers from inputs (it doesn't iterate // through them). The memory touched is then only the size of the output - // buffer. - current_bytes_accessed_ = shape_size_(tuple->shape()); + // index table of the tuple. + + current_properties_[kBytesAccessedKey] = shape_size_(tuple->shape()); return Status::OK(); } @@ -164,13 +208,11 @@ Status HloCostAnalysis::HandleConcatenate( return Status::OK(); } -Status HloCostAnalysis::HandleConvert(HloInstruction* convert, - HloInstruction* operand) { +Status HloCostAnalysis::HandleConvert(HloInstruction* convert) { return HandleElementwiseOp(convert); } -Status HloCostAnalysis::HandleCopy(HloInstruction* copy, - HloInstruction* operand) { +Status HloCostAnalysis::HandleCopy(HloInstruction* copy) { return Status::OK(); } @@ -194,7 +236,7 @@ Status HloCostAnalysis::HandleDot(HloInstruction* dot, } // We count an FMA operation as 2 floating point operations. - current_flop_count_ = kFmaFlops * fma_count; + current_properties_[kFlopsKey] = kFmaFlops * fma_count; return Status::OK(); } @@ -210,16 +252,17 @@ Status HloCostAnalysis::HandleMap( HloInstruction* map, tensorflow::gtl::ArraySlice operands, HloComputation* function, tensorflow::gtl::ArraySlice /*static_operands*/) { - // Compute the cost of the user function. - HloInstruction* function_instruction = function->root_instruction(); - HloCostAnalysis visitor(shape_size_); - TF_RETURN_IF_ERROR(function_instruction->Accept(&visitor)); + // Compute properties of the mapped function. + TF_ASSIGN_OR_RETURN(const Properties sub_properties, + ProcessSubcomputation(function)); // Compute the cost of all elements for this Map operation. - int64 element_count = ShapeUtil::ElementsIn(map->shape()); - current_transcendental_count_ = - element_count * visitor.transcendental_count(); - current_flop_count_ = element_count * visitor.flop_count(); + const int64 element_count = ShapeUtil::ElementsIn(map->shape()); + for (const auto& property : sub_properties) { + if (property.first != kBytesAccessedKey) { + current_properties_[property.first] = property.second * element_count; + } + } return Status::OK(); } @@ -227,16 +270,17 @@ Status HloCostAnalysis::HandleReduce( HloInstruction* reduce, HloInstruction* arg, HloInstruction* init_value, tensorflow::gtl::ArraySlice dimensions, HloComputation* function) { // Compute the cost of the user function. - HloInstruction* function_instruction = function->root_instruction(); - HloCostAnalysis visitor(shape_size_); - TF_RETURN_IF_ERROR(function_instruction->Accept(&visitor)); + TF_ASSIGN_OR_RETURN(const Properties sub_properties, + ProcessSubcomputation(function)); // Compute the cost of all elements for this Reduce operation. int64 reduction_count = ShapeUtil::ElementsIn(arg->shape()) - ShapeUtil::ElementsIn(reduce->shape()); - current_flop_count_ = reduction_count * visitor.flop_count(); - current_transcendental_count_ = - reduction_count * visitor.transcendental_count(); + for (const auto& property : sub_properties) { + if (property.first != kBytesAccessedKey) { + current_properties_[property.first] = property.second * reduction_count; + } + } return Status::OK(); } @@ -244,55 +288,63 @@ Status HloCostAnalysis::HandleReduceWindow(HloInstruction* reduce_window, HloInstruction* operand, const Window& window, HloComputation* function) { - // Compute the cost of the user function. - HloInstruction* function_instruction = function->root_instruction(); - HloCostAnalysis visitor(shape_size_); - TF_RETURN_IF_ERROR(function_instruction->Accept(&visitor)); + // Compute the properties of the reduction function. + TF_ASSIGN_OR_RETURN(const Properties sub_properties, + ProcessSubcomputation(function)); // Compute the cost of all elements for this ReduceWindow operation. For each - // output element, (window_size - 1) number of user computations are applied. - auto output_size = ShapeUtil::ElementsIn(reduce_window->shape()); - int64 window_size = 1; + // output element there are window_size - 1 reductions to perform. + int64 window_element_count = 1; for (const auto& dimension : window.dimensions()) { - window_size *= dimension.size(); + window_element_count *= dimension.size(); + } + const int64 output_element_count = + ShapeUtil::ElementsIn(reduce_window->shape()); + const int64 reduction_count = + (window_element_count - 1) * output_element_count; + for (const auto& property : sub_properties) { + if (property.first != kBytesAccessedKey) { + current_properties_[property.first] = property.second * reduction_count; + } } - current_flop_count_ = output_size * (window_size - 1) * visitor.flop_count(); - current_transcendental_count_ = - output_size * (window_size - 1) * visitor.transcendental_count(); return Status::OK(); } Status HloCostAnalysis::HandleSelectAndScatter(HloInstruction* instruction) { - // Compute the cost of the select and scatter function. - HloInstruction* select = instruction->select()->root_instruction(); - HloCostAnalysis select_visitor(shape_size_); - TF_RETURN_IF_ERROR(select->Accept(&select_visitor)); - HloInstruction* scatter = instruction->scatter()->root_instruction(); - HloCostAnalysis scatter_visitor(shape_size_); - TF_RETURN_IF_ERROR(scatter->Accept(&scatter_visitor)); + // 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())); + TF_ASSIGN_OR_RETURN(const Properties scatter_properties, + ProcessSubcomputation(instruction->scatter())); // Compute the cost of all elements for this operation. For each scatter - // source element, (window_size - 1) number of select computations and 1 - // scatter computation are applied. + // source element there are window_size - 1 select computations to perform and + // 1 scatter computation to perform. const auto source = instruction->operand(1); const auto source_element_count = ShapeUtil::ElementsIn(source->shape()); - int64 window_size = 1; + int64 window_element_count = 1; for (const auto& dimension : instruction->window().dimensions()) { - window_size *= dimension.size(); + window_element_count *= dimension.size(); + } + const int64 select_count = source_element_count * (window_element_count - 1); + for (const auto& property : select_properties) { + if (property.first != kBytesAccessedKey) { + current_properties_[property.first] += property.second * select_count; + } + } + for (const auto& property : scatter_properties) { + if (property.first != kBytesAccessedKey) { + current_properties_[property.first] += + property.second * source_element_count; + } } - current_flop_count_ = - source_element_count * ((window_size - 1) * select_visitor.flop_count() + - scatter_visitor.flop_count()); - current_transcendental_count_ = - source_element_count * - ((window_size - 1) * select_visitor.transcendental_count() + - scatter_visitor.transcendental_count()); return Status::OK(); } Status HloCostAnalysis::HandleBitcast(HloInstruction* bitcast) { // A bitcast does no computation and touches no memory. - current_bytes_accessed_ = 0; + current_properties_[kBytesAccessedKey] = 0; return Status::OK(); } @@ -314,6 +366,23 @@ Status HloCostAnalysis::HandleReshape(HloInstruction* reshape) { return Status::OK(); } +Status HloCostAnalysis::HandleBatchNormTraining( + HloInstruction* batchNormTraining) { + // TODO(b/62294698): Implement cost analysis for batch-norm-training. + return Status::OK(); +} + +Status HloCostAnalysis::HandleBatchNormInference( + HloInstruction* batchNormInference) { + // TODO(b/62294698): Implement cost analysis for batch-norm-inference. + return Status::OK(); +} + +Status HloCostAnalysis::HandleBatchNormGrad(HloInstruction* batchNormGrad) { + // TODO(b/62294698): Implement cost analysis for batch-norm-grad. + return Status::OK(); +} + Status HloCostAnalysis::HandleTranspose(HloInstruction* transpose) { return Status::OK(); } @@ -326,12 +395,13 @@ Status HloCostAnalysis::HandleConvolution(HloInstruction* convolution, const int64 output_features = convolution->shape().dimensions(dnums.feature_dimension()); - // For each output element, we do one fma per element in the - // kernel at some given output feature index. + // For each output element, we do one fma per element in the kernel at some + // given output feature index. const int64 fmas_per_output_element = ShapeUtil::ElementsIn(rhs_instruction->shape()) / output_features; const int64 output_elements = ShapeUtil::ElementsIn(convolution->shape()); - current_flop_count_ = output_elements * fmas_per_output_element * kFmaFlops; + current_properties_[kFlopsKey] = + output_elements * fmas_per_output_element * kFmaFlops; return Status::OK(); } @@ -341,7 +411,7 @@ Status HloCostAnalysis::HandleCrossReplicaSum(HloInstruction* crs) { // // TODO(b/33004697): Compute correct cost here, taking the actual number of // replicas into account. - current_flop_count_ = ShapeUtil::ElementsIn(crs->shape()); + current_properties_[kFlopsKey] = ShapeUtil::ElementsIn(crs->shape()); return Status::OK(); } @@ -350,33 +420,43 @@ Status HloCostAnalysis::HandleRng(HloInstruction* random, // TODO(b/26346211): Implement better estimates for the RNG cost, since the // cost changes with the implementation and the distribution. For now, assume // the cost of each RNG is same as a transcendental operation. - current_transcendental_count_ = ShapeUtil::ElementsIn(random->shape()); + current_properties_[kTranscendentalsKey] = + ShapeUtil::ElementsIn(random->shape()); return Status::OK(); } Status HloCostAnalysis::HandleFusion(HloInstruction* fusion) { - // Compute the cost of the fused expression. - HloInstruction* fused_expression_root = fusion->fused_expression_root(); - // Don't compute sizes inside of fused ops. We don't use the size here and the - // operations inside might not have a layout. - HloCostAnalysis visitor([](const Shape&) { return 0; }); - TF_RETURN_IF_ERROR(fused_expression_root->Accept(&visitor)); + // Compute the properties of the fused expression and attribute them to the + // fusion node. Use a dummy shape_size to avoid any errors from trying to + // calculate the size of a shape that does not have a layout, since nodes + // inside fusion nodes do not necessarily have a layout assigned. + ShapeSizeFunction shape_size = [](const Shape& shape) { return 0; }; + TF_ASSIGN_OR_RETURN( + current_properties_, + ProcessSubcomputation(fusion->fused_instructions_computation(), + &shape_size)); + + // 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 + // prevent intermediate data from touching slow memory. + current_properties_[kBytesAccessedKey] = 0; + ShapeUtil::ForEachSubshape( + fusion->shape(), + [this](const Shape& subshape, const ShapeIndex& /*shape_index*/) { + current_properties_[kBytesAccessedKey] += shape_size_(subshape); + }); + + for (const HloInstruction* operand : fusion->operands()) { + current_properties_[kBytesAccessedKey] += shape_size_(operand->shape()); + } - // Attribute the cost of the fused expression to the fusion node. - current_transcendental_count_ = visitor.transcendental_count(); - current_flop_count_ = visitor.flop_count(); return Status::OK(); } -Status HloCostAnalysis::HandleCall( - HloInstruction* call, tensorflow::gtl::ArraySlice operands, - HloComputation* computation) { - HloCostAnalysis computation_visitor(shape_size_); - TF_RETURN_IF_ERROR(computation->Accept(&computation_visitor)); - - current_flop_count_ = computation_visitor.flop_count(); - current_transcendental_count_ = computation_visitor.transcendental_count(); - current_bytes_accessed_ = computation_visitor.bytes_accessed(); +Status HloCostAnalysis::HandleCall(HloInstruction* call) { + TF_ASSIGN_OR_RETURN(current_properties_, + ProcessSubcomputation(call->to_apply())); + current_should_compute_bottleneck_time_ = false; return Status::OK(); } @@ -384,37 +464,38 @@ Status HloCostAnalysis::HandleCustomCall( HloInstruction* custom_call, tensorflow::gtl::ArraySlice operands, tensorflow::StringPiece custom_call_target) { - return Unimplemented("custom-call"); + return Unimplemented("Custom-call is not implemented for HLO cost analysis."); } Status HloCostAnalysis::HandleSort(HloInstruction* sort, HloInstruction* operand_instruction) { - // The cost of sort is implementation dependent, so cannot determine at HLO - // level. Assume comparison based N*log(N) sorting. + // This assumes a comparison based N*log(N) algorithm. As for all ops, the + // actual properties of the op depend on the backend implementation. int64 elements = ShapeUtil::ElementsIn(operand_instruction->shape()); - current_flop_count_ = elements * tensorflow::Log2Ceiling(elements); + current_properties_[kFlopsKey] = elements * tensorflow::Log2Ceiling(elements); return Status::OK(); } -Status HloCostAnalysis::HandleWhile(HloInstruction* xla_while, - HloInstruction* init, - HloComputation* condition, - HloComputation* body) { - // Since the number of iterations of the while node is not statically - // determined, 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 node. - HloCostAnalysis body_visitor(shape_size_); - TF_RETURN_IF_ERROR(body->Accept(&body_visitor)); - HloCostAnalysis condition_visitor(shape_size_); - TF_RETURN_IF_ERROR(condition->Accept(&condition_visitor)); +Status HloCostAnalysis::HandleWhile(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())); + + TF_ASSIGN_OR_RETURN(const Properties condition_properties, + ProcessSubcomputation(xla_while->while_condition())); - current_flop_count_ = - body_visitor.flop_count() + condition_visitor.flop_count(); - current_transcendental_count_ = body_visitor.transcendental_count() + - condition_visitor.transcendental_count(); - current_bytes_accessed_ = - body_visitor.bytes_accessed() + condition_visitor.bytes_accessed(); + current_properties_.clear(); + for (const auto& property : body_properties) { + current_properties_[property.first] += property.second; + } + for (const auto& property : condition_properties) { + current_properties_[property.first] += property.second; + } + current_should_compute_bottleneck_time_ = false; return Status::OK(); } @@ -423,19 +504,46 @@ Status HloCostAnalysis::FinishVisit(HloInstruction* root) { return Status::OK(); } +float HloCostAnalysis::flop_count() const { + return GetProperty(kFlopsKey, properties_sum_); +} + +float HloCostAnalysis::transcendental_count() const { + return GetProperty(kTranscendentalsKey, properties_sum_); +} + +float HloCostAnalysis::bytes_accessed() const { + return GetProperty(kBytesAccessedKey, properties_sum_); +} + +float HloCostAnalysis::seconds() const { + return GetProperty(kSecondsKey, properties_sum_); +} + int64 HloCostAnalysis::flop_count(const HloInstruction& hlo) const { - auto it = hlo_to_flop_count_.find(&hlo); - return it == hlo_to_flop_count_.end() ? 0 : it->second; + return GetPropertyForHlo(hlo, kFlopsKey, hlo_properties_); } int64 HloCostAnalysis::transcendental_count(const HloInstruction& hlo) const { - auto it = hlo_to_transcendental_count_.find(&hlo); - return it == hlo_to_transcendental_count_.end() ? 0 : it->second; + return GetPropertyForHlo(hlo, kTranscendentalsKey, hlo_properties_); } int64 HloCostAnalysis::bytes_accessed(const HloInstruction& hlo) const { - auto it = hlo_to_bytes_accessed_.find(&hlo); - return it == hlo_to_bytes_accessed_.end() ? 0 : it->second; + return GetPropertyForHlo(hlo, kBytesAccessedKey, hlo_properties_); +} + +float HloCostAnalysis::seconds(const HloInstruction& hlo) const { + return GetPropertyForHlo(hlo, kSecondsKey, hlo_properties_); +} + +StatusOr HloCostAnalysis::ProcessSubcomputation( + HloComputation* computation, const ShapeSizeFunction* shape_size) { + if (shape_size == nullptr) { + shape_size = &shape_size_; + } + HloCostAnalysis visitor(*shape_size, per_second_rates_); + TF_RETURN_IF_ERROR(computation->Accept(&visitor)); + return visitor.properties(); } } // namespace xla diff --git a/tensorflow/compiler/xla/service/hlo_cost_analysis.h b/tensorflow/compiler/xla/service/hlo_cost_analysis.h index e6f059f53379df51c9f0b99e0e01f34f1aebb52a..eeb3d4edd1be3bb0204d37e3e6591058a687712e 100644 --- a/tensorflow/compiler/xla/service/hlo_cost_analysis.h +++ b/tensorflow/compiler/xla/service/hlo_cost_analysis.h @@ -36,17 +36,21 @@ namespace xla { // operations separately from transcendental operations. class HloCostAnalysis : public DfsHloVisitor { public: + // Each HLO is associated to a vector of properties with the indices given + // below. Sub-classes can add further properties. + typedef std::map Properties; + static constexpr char kFlopsKey[] = "flops"; + static constexpr char kTranscendentalsKey[] = "transcendentals"; + static constexpr char kBytesAccessedKey[] = "bytes accessed"; + static constexpr char kSecondsKey[] = "seconds"; + // shape_size is a function which returns the size in bytes of the top-level // buffer of a shape. using ShapeSizeFunction = std::function; - explicit HloCostAnalysis(const ShapeSizeFunction& shape_size) - : shape_size_(shape_size) {} - - Status HandleElementwiseUnary(HloInstruction* hlo, HloOpcode opcode, - HloInstruction* operand) override; - Status HandleElementwiseBinary(HloInstruction* hlo, HloOpcode opcode, - HloInstruction* lhs, - HloInstruction* rhs) override; + explicit HloCostAnalysis(const ShapeSizeFunction& shape_size); + + Status HandleElementwiseUnary(HloInstruction* hlo) override; + Status HandleElementwiseBinary(HloInstruction* hlo) override; Status HandleConstant(HloInstruction* constant, const Literal& literal) override; Status HandleGetTupleElement(HloInstruction* get_tuple_element, @@ -58,14 +62,14 @@ class HloCostAnalysis : public DfsHloVisitor { HloInstruction* lhs, HloInstruction* rhs) override; Status HandleClamp(HloInstruction* clamp, HloInstruction* min, HloInstruction* arg, HloInstruction* max) override; + Status HandleReducePrecision(HloInstruction* hlo) override; Status HandleConcatenate( HloInstruction* concatenate, tensorflow::gtl::ArraySlice operands) override; Status HandleSend(HloInstruction* send) override; Status HandleRecv(HloInstruction* recv) override; - Status HandleConvert(HloInstruction* convert, - HloInstruction* operand) override; - Status HandleCopy(HloInstruction* copy, HloInstruction* operand) override; + Status HandleConvert(HloInstruction* convert) override; + Status HandleCopy(HloInstruction* copy) override; Status HandleDot(HloInstruction* dot, HloInstruction* lhs, HloInstruction* rhs) override; Status HandleConvolution(HloInstruction* convolution, HloInstruction* lhs, @@ -83,17 +87,18 @@ class HloCostAnalysis : public DfsHloVisitor { HloInstruction* init_value, tensorflow::gtl::ArraySlice dimensions, HloComputation* function_handle) override; + Status HandleBatchNormTraining(HloInstruction* batchNormTraining) override; + Status HandleBatchNormInference(HloInstruction* batchNormInference) override; + Status HandleBatchNormGrad(HloInstruction* batchNormGrad) override; Status HandleFusion(HloInstruction* fusion) override; - Status HandleCall(HloInstruction* call, - tensorflow::gtl::ArraySlice operands, - HloComputation* computation) override; + Status HandleCall(HloInstruction* call) override; Status HandleCustomCall(HloInstruction* custom_call, tensorflow::gtl::ArraySlice operands, tensorflow::StringPiece custom_call_target) override; Status HandleSlice(HloInstruction* slice, HloInstruction* operand) override; - Status HandleDynamicSlice( - HloInstruction* slice, - tensorflow::gtl::ArraySlice operands) override; + Status HandleDynamicSlice(HloInstruction* dynamic_slice, + HloInstruction* operand, + HloInstruction* start_indices) override; Status HandleDynamicUpdateSlice(HloInstruction* dynamic_update_slice, HloInstruction* operand, HloInstruction* update, @@ -115,55 +120,96 @@ class HloCostAnalysis : public DfsHloVisitor { Status HandlePad(HloInstruction* pad) override; Status HandleReshape(HloInstruction* reshape) override; Status HandleTranspose(HloInstruction* transpose) override; - Status HandleWhile(HloInstruction* xla_while, HloInstruction* init, - HloComputation* condition, HloComputation* body) override; + Status HandleWhile(HloInstruction* xla_while) override; Status FinishVisit(HloInstruction* root) override; Status Preprocess(HloInstruction* hlo) override; Status Postprocess(HloInstruction* hlo) override; - // Returns the amount of computations in the graph. - int64 flop_count() const { return flop_count_; } - int64 transcendental_count() const { return transcendental_count_; } + // Set the rates used to calculate the time taken by the computation. These + // need to be set before visiting starts. + void set_flops_per_second(float value) { + per_second_rates_[kFlopsKey] = value; + } + void set_transcendentals_per_second(float value) { + per_second_rates_[kTranscendentalsKey] = value; + } + void set_bytes_per_second(float value) { + per_second_rates_[kBytesAccessedKey] = value; + } + + // Returns properties for the computation. + float flop_count() const; + float transcendental_count() const; + float bytes_accessed() const; + float seconds() const; // Returns the respective cost computed for a particular HLO instruction, or 0 // if the HLO was not found to have a cost in the analysis. int64 flop_count(const HloInstruction& hlo) const; int64 transcendental_count(const HloInstruction& hlo) const; - - // Returns the number of bytes read/written. int64 bytes_accessed(const HloInstruction& hlo) const; - int64 bytes_accessed() const { return bytes_accessed_; } + float seconds(const HloInstruction& hlo) const; + + const Properties& properties() const { return properties_sum_; } + const float property(const string& key) const { + return GetProperty(key, properties()); + } - private: - // An FMA counts as two floating point operations in these analyses. + protected: + typedef std::unordered_map HloToProperties; + + // An FMA counts as two floating point operations in these analyzes. static constexpr int64 kFmaFlops = 2; + HloCostAnalysis(const ShapeSizeFunction& shape_size, + const Properties& per_second_rates); + + // Returns the properties computed from visiting the computation rooted at the + // given hlo. Uses shape_size_ to calculate shape sizes if shape_size is null, + // otherwise uses shape_size_. + StatusOr ProcessSubcomputation( + HloComputation* computation, + const ShapeSizeFunction* shape_size = nullptr); + // Utility function to handle all element-wise operations. Status HandleElementwiseOp(HloInstruction* hlo_instruction); + // Returns the default value if the key is not present in the + // properties. Otherwise, returns the value that the key maps to from the + // properties parameter. + static float GetProperty(const string& key, const Properties& properties, + float default_value = 0.0f); + + // Returns 0.0f if the hlo is not present in hlo_to_properties or if the key + // is not present in hlo_to_properties[hlo]. Otherwise, returns the value that + // the key maps to in the properties of the given hlo. + static float GetPropertyForHlo(const HloInstruction& hlo, const string& key, + const HloToProperties& hlo_to_properties); + // Function which computes the size of the top-level of a given shape (not // including nested elements, if any). If null then bytes_accessed methods // return an error. const ShapeSizeFunction shape_size_; - // The total number of floating point operations, transcendental operations, - // and bytes accesses (read or written) in the computation. - int64 flop_count_ = 0; - int64 transcendental_count_ = 0; - int64 bytes_accessed_ = 0; - - // Cost counts of the current instruction. These should be set by each - // handlers if different from the default values computed in Preprocess. - int64 current_flop_count_; - int64 current_transcendental_count_; - int64 current_bytes_accessed_; - - // Mapping from HLO instructions to the cost we computed for them in the - // course of the graph analysis. - std::map hlo_to_flop_count_; - std::map hlo_to_transcendental_count_; - std::map hlo_to_bytes_accessed_; + HloToProperties hlo_properties_; + + // If true, the time taken will be computed from the rates for each property + // and the total time will be the maximum time, which is the time of the + // bottleneck. + bool current_should_compute_bottleneck_time_; + + // The properties of the currently visited instruction. A HandleFoo method can + // modify these to change the default values computed in Preprocess. + Properties current_properties_; + + // The sum of the properties of all HLOs in the computation. + Properties properties_sum_; + + // How much of each property can be processed per second. E.g. if the property + // is bytes accessed, this is the number of bytes that can be processed per + // second. Is empty if no rates have been set. + Properties per_second_rates_; TF_DISALLOW_COPY_AND_ASSIGN(HloCostAnalysis); }; diff --git a/tensorflow/compiler/xla/service/hlo_cost_analysis_test.cc b/tensorflow/compiler/xla/service/hlo_cost_analysis_test.cc index 22e782da27db25e9d59622b27d942426fec0f1dc..0a288a77ada840451915561b4b0865785b39ade7 100644 --- a/tensorflow/compiler/xla/service/hlo_cost_analysis_test.cc +++ b/tensorflow/compiler/xla/service/hlo_cost_analysis_test.cc @@ -31,6 +31,7 @@ limitations under the License. #include "tensorflow/compiler/xla/service/user_computation.h" #include "tensorflow/compiler/xla/service/versioned_computation_handle.h" #include "tensorflow/compiler/xla/shape_util.h" +#include "tensorflow/compiler/xla/tests/hlo_test_base.h" #include "tensorflow/core/platform/logging.h" #include "tensorflow/compiler/xla/statusor.h" @@ -54,7 +55,7 @@ class HloCostAnalysisTest : public ::testing::Test { HloCostAnalysisTest() : client_(ClientLibrary::LocalClientOrDie()), // Accessing service instance is required for the unit tests to enable - // whitebox acccesses to the user computation built from the client, + // whitebox accesses to the user computation built from the client, // as shown in the BuildHloGraph functions below. service_(static_cast(ClientLibrary::GetXlaService( static_cast(client_)->platform()))), @@ -127,7 +128,8 @@ class HloCostAnalysisTest : public ::testing::Test { VersionedComputationHandle versioned_handle = user_computation->GetVersionedHandle(); return std::move( - computation_tracker_.BuildHloModule(versioned_handle).ValueOrDie()); + computation_tracker_.BuildHloModule(versioned_handle, HloModuleConfig()) + .ValueOrDie()); } Client* client_; @@ -328,51 +330,67 @@ TEST_F(HloCostAnalysisTest, MatmulAndConvolutionCanBeTheSameComputation) { EXPECT_EQ(conv_analysis.flop_count(), matmul_analysis.flop_count()); } -using FusionCostAnalysis = ::testing::Test; +using FusionCostAnalysis = HloTestBase; TEST_F(FusionCostAnalysis, LoopFusion) { - Shape r2f32 = ShapeUtil::MakeShape(F32, {2, 2}); - - // Fuse all instructions in complicated expression: - // - // add = Add(C1, C2) - // clamp = Clamp(C2, add, add) - // exp = Exp(add) - // mul = Mul(exp, C3) - // sub = Sub(mul, clamp) - // tuple = Tuple({sub, sub, mul, C1}) - auto c1 = HloInstruction::CreateConstant(LiteralUtil::CreateR2F32Linspace( - /*from=*/0.0f, /*to=*/1.0f, /*rows=*/2, /*cols=*/2)); - auto c2 = HloInstruction::CreateConstant(LiteralUtil::CreateR2F32Linspace( - /*from=*/1.0f, /*to=*/2.0f, /*rows=*/2, /*cols=*/2)); - auto c3 = HloInstruction::CreateConstant(LiteralUtil::CreateR2F32Linspace( - /*from=*/2.0f, /*to=*/3.0f, /*rows=*/2, /*cols=*/2)); - - auto add = - HloInstruction::CreateBinary(r2f32, HloOpcode::kAdd, c1.get(), c2.get()); - auto clamp = HloInstruction::CreateTernary(r2f32, HloOpcode::kClamp, c2.get(), - add.get(), add.get()); - auto exp = HloInstruction::CreateUnary(r2f32, HloOpcode::kExp, add.get()); - auto mul = HloInstruction::CreateBinary(r2f32, HloOpcode::kMultiply, - exp.get(), c3.get()); - auto sub = HloInstruction::CreateBinary(r2f32, HloOpcode::kSubtract, - mul.get(), clamp.get()); - auto tuple = - HloInstruction::CreateTuple({sub.get(), sub.get(), mul.get(), c1.get()}); - - auto fusion = HloInstruction::CreateFusion( - r2f32, HloInstruction::FusionKind::kLoop, tuple.get()); - fusion->FuseInstruction(sub.get()); - fusion->FuseInstruction(mul.get()); - fusion->FuseInstruction(exp.get()); - fusion->FuseInstruction(clamp.get()); - fusion->FuseInstruction(add.get()); - - HloCostAnalysis fusion_analysis(ShapeSize); - ASSERT_IS_OK(fusion->Accept(&fusion_analysis)); - - EXPECT_EQ(fusion_analysis.flop_count(), 16); - EXPECT_EQ(fusion_analysis.transcendental_count(), 4); + // Do this 4 times with different per-second rates to test the computation of + // bottleneck time on fusion nodes. + for (int i = 0; i < 4; ++i) { + Shape r2f32 = ShapeUtil::MakeShape(F32, {2, 2}); + + // Fuse all instructions in complicated expression: + // + // add = Add(C1, C2) + // clamp = Clamp(C2, add, add) + // exp = Exp(add) + // mul = Mul(exp, C3) + // sub = Sub(mul, clamp) + // tuple = Tuple({sub, sub, mul, C1}) + HloComputation::Builder builder(TestName()); + auto c1 = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR2F32Linspace( + /*from=*/0.0f, /*to=*/1.0f, /*rows=*/2, /*cols=*/2))); + auto c2 = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR2F32Linspace( + /*from=*/1.0f, /*to=*/2.0f, /*rows=*/2, /*cols=*/2))); + auto c3 = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR2F32Linspace( + /*from=*/2.0f, /*to=*/3.0f, /*rows=*/2, /*cols=*/2))); + auto add = builder.AddInstruction( + HloInstruction::CreateBinary(r2f32, HloOpcode::kAdd, c1, c2)); + auto clamp = builder.AddInstruction( + HloInstruction::CreateTernary(r2f32, HloOpcode::kClamp, c2, add, add)); + auto exp = builder.AddInstruction( + HloInstruction::CreateUnary(r2f32, HloOpcode::kExp, add)); + auto mul = builder.AddInstruction( + HloInstruction::CreateBinary(r2f32, HloOpcode::kMultiply, exp, c3)); + auto sub = builder.AddInstruction( + HloInstruction::CreateBinary(r2f32, HloOpcode::kSubtract, mul, clamp)); + auto tuple = HloInstruction::CreateTuple({sub, sub, mul, c1}); + + HloModule module(TestName()); + auto* computation = module.AddEntryComputation(builder.Build()); + auto* fusion = computation->CreateFusionInstruction( + {sub, mul, exp, clamp, add}, HloInstruction::FusionKind::kLoop); + + // The time given these rates at i == 0 is exactly even among the properties + // at 1.0 seconds. For other values, one of the rates is slower so that it + // becomes the bottleneck. + HloCostAnalysis fusion_analysis(ShapeSize); + fusion_analysis.set_flops_per_second(16 * (i == 1 ? 1 / 2.0 : 1.0)); + fusion_analysis.set_transcendentals_per_second(4 * + (i == 2 ? 1 / 4.0 : 1.0)); + fusion_analysis.set_bytes_per_second(64 * (i == 3 ? 1 / 8.0 : 1.0)); + ASSERT_IS_OK(fusion->Accept(&fusion_analysis)); + + EXPECT_EQ(fusion_analysis.flop_count(), 16); + EXPECT_EQ(fusion_analysis.transcendental_count(), 4); + constexpr int64 bytes_accessed = sizeof(float) * 4 * 2 * 2; + static_assert(bytes_accessed == 64, ""); + EXPECT_EQ(fusion_analysis.bytes_accessed(), bytes_accessed); + + EXPECT_EQ(fusion_analysis.seconds(), 1 << i); + } } TEST_F(FusionCostAnalysis, NoLayout) { @@ -381,19 +399,21 @@ TEST_F(FusionCostAnalysis, NoLayout) { Shape shape_without_layout = shape_with_layout; shape_without_layout.clear_layout(); - auto c1 = HloInstruction::CreateConstant( - LiteralUtil::CreateR4FromArray4D(Array4D(2, 3, 4, 5))); - auto c2 = - HloInstruction::CreateConstant(LiteralUtil::CreateR1({1, 2, 3})); - - auto broadcast = - HloInstruction::CreateBroadcast(shape_without_layout, c2.get(), {1}); - auto add = HloInstruction::CreateBinary(shape_with_layout, HloOpcode::kAdd, - c1.get(), broadcast.get()); - - auto fusion = HloInstruction::CreateFusion( - shape_with_layout, HloInstruction::FusionKind::kLoop, add.get()); - fusion->FuseInstruction(broadcast.get()); + HloComputation::Builder builder(TestName()); + auto c1 = builder.AddInstruction(HloInstruction::CreateConstant( + Literal::CreateR4FromArray4D(Array4D(2, 3, 4, 5)))); + auto c2 = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR1({1, 2, 3}))); + + auto broadcast = builder.AddInstruction( + HloInstruction::CreateBroadcast(shape_without_layout, c2, {1})); + auto add = builder.AddInstruction(HloInstruction::CreateBinary( + shape_with_layout, HloOpcode::kAdd, c1, broadcast)); + + HloModule module(TestName()); + auto* computation = module.AddEntryComputation(builder.Build()); + auto* fusion = computation->CreateFusionInstruction( + {add, broadcast}, HloInstruction::FusionKind::kLoop); HloCostAnalysis fusion_analysis(ShapeSize); ASSERT_IS_OK(fusion->Accept(&fusion_analysis)); diff --git a/tensorflow/compiler/xla/service/hlo_cse.cc b/tensorflow/compiler/xla/service/hlo_cse.cc index 4c6af5c40fa563d1c656eb152819e454aae5fb69..690c084efb131e9b075ced17bfcd0b23a23218f1 100644 --- a/tensorflow/compiler/xla/service/hlo_cse.cc +++ b/tensorflow/compiler/xla/service/hlo_cse.cc @@ -68,7 +68,7 @@ bool CombineConstants(HloComputation* computation, bool is_layout_sensitive) { auto range = constants.equal_range(shape_string); HloInstruction* match = nullptr; for (auto it = range.first; it != range.second; ++it) { - if (LiteralUtil::Equal(instruction->literal(), it->second->literal())) { + if (instruction->literal().Equal(it->second->literal())) { match = it->second; break; } @@ -92,6 +92,9 @@ bool CombineConstants(HloComputation* computation, bool is_layout_sensitive) { StatusOr HloCSE::Run(HloModule* module) { bool changed = false; for (auto& computation : module->computations()) { + if (computation->IsFusionComputation()) { + continue; + } changed |= CombineConstants(computation.get(), is_layout_sensitive_); std::list post_order = diff --git a/tensorflow/compiler/xla/service/hlo_cse_test.cc b/tensorflow/compiler/xla/service/hlo_cse_test.cc index ec8161f55fd56c95bb088a0c539255aed2fe6993..8b0b9c8bbd0cf442149b32a4539277b2daeed90e 100644 --- a/tensorflow/compiler/xla/service/hlo_cse_test.cc +++ b/tensorflow/compiler/xla/service/hlo_cse_test.cc @@ -25,6 +25,7 @@ limitations under the License. #include "tensorflow/compiler/xla/ptr_util.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/shape_util.h" @@ -36,6 +37,8 @@ limitations under the License. #include "tensorflow/compiler/xla/types.h" #include "tensorflow/core/platform/types.h" +namespace op = xla::testing::opcode_matchers; + namespace xla { namespace { @@ -48,13 +51,13 @@ TEST_F(HloCseTest, CombineTwoConstants) { // Test that two identical constants are commoned. auto builder = HloComputation::Builder(TestName()); auto constant1 = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(42.0f))); + HloInstruction::CreateConstant(Literal::CreateR0(42.0f))); auto constant2 = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(42.0f))); + HloInstruction::CreateConstant(Literal::CreateR0(42.0f))); builder.AddInstruction(HloInstruction::CreateBinary( constant1->shape(), HloOpcode::kAdd, constant1, constant2)); - auto module = MakeUnique(TestName()); + auto module = CreateNewModule(); auto computation = module->AddEntryComputation(builder.Build()); EXPECT_EQ(3, computation->instruction_count()); @@ -64,10 +67,10 @@ TEST_F(HloCseTest, CombineTwoConstants) { EXPECT_EQ(2, computation->instruction_count()); HloInstruction* constant = computation->instructions().begin()->get(); - EXPECT_EQ(42.0f, LiteralUtil::Get(constant->literal(), {})); + EXPECT_EQ(42.0f, constant->literal().Get({})); auto result = ExecuteAndTransfer(std::move(module), {}); - auto expected = LiteralUtil::CreateR0(84.0); + auto expected = Literal::CreateR0(84.0); LiteralTestUtil::ExpectNear(*expected, *result, ErrorSpec(1e-4)); } @@ -84,20 +87,22 @@ TEST_F(HloCseTest, CombineTwoConstantsDifferentLayoutsAndInsensitive) { auto add = builder.AddInstruction(HloInstruction::CreateBinary( constant1->shape(), HloOpcode::kAdd, constant1, constant2)); - auto module = MakeUnique(TestName()); + auto module = CreateNewModule(); auto computation = module->AddEntryComputation(builder.Build()); EXPECT_EQ(3, computation->instruction_count()); - EXPECT_NE(add->operand(0), add->operand(1)); + EXPECT_THAT(add, op::Add(constant1, constant2)); HloCSE cse(/*is_layout_sensitive=*/false); EXPECT_TRUE(cse.Run(module.get()).ValueOrDie()); EXPECT_EQ(2, computation->instruction_count()); - EXPECT_EQ(add->operand(0), add->operand(1)); + auto first_operand = add->operand(0); + EXPECT_THAT(first_operand, ::testing::AnyOf(constant1, constant2)); + EXPECT_THAT(add, op::Add(first_operand, first_operand)); auto result = ExecuteAndTransfer(std::move(module), {}); - auto expected = LiteralUtil::CreateR2({{2.0, 4.0}, {6.0, 8.0}}); + auto expected = Literal::CreateR2({{2.0, 4.0}, {6.0, 8.0}}); LiteralTestUtil::ExpectNear(*expected, *result, ErrorSpec(1e-4)); } @@ -114,22 +119,20 @@ TEST_F(HloCseTest, CombineTwoConstantsDifferentLayoutsAndSensitive) { auto add = builder.AddInstruction(HloInstruction::CreateBinary( constant1->shape(), HloOpcode::kAdd, constant1, constant2)); - auto module = MakeUnique(TestName()); + auto module = CreateNewModule(); auto computation = module->AddEntryComputation(builder.Build()); EXPECT_EQ(3, computation->instruction_count()); - EXPECT_EQ(constant1, add->operand(0)); - EXPECT_EQ(constant2, add->operand(1)); + EXPECT_THAT(add, op::Add(constant1, constant2)); HloCSE cse(/*is_layout_sensitive=*/true); EXPECT_FALSE(cse.Run(module.get()).ValueOrDie()); EXPECT_EQ(3, computation->instruction_count()); - EXPECT_EQ(constant1, add->operand(0)); - EXPECT_EQ(constant2, add->operand(1)); + EXPECT_THAT(add, op::Add(constant1, constant2)); auto result = ExecuteAndTransfer(std::move(module), {}); - auto expected = LiteralUtil::CreateR2({{2.0, 4.0}, {6.0, 8.0}}); + auto expected = Literal::CreateR2({{2.0, 4.0}, {6.0, 8.0}}); LiteralTestUtil::ExpectNear(*expected, *result, ErrorSpec(1e-4)); } @@ -138,28 +141,28 @@ TEST_F(HloCseTest, ConstantsSameValueDifferentType) { // commoned. auto builder = HloComputation::Builder(TestName()); builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(42))); + HloInstruction::CreateConstant(Literal::CreateR0(42))); builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(42))); + HloInstruction::CreateConstant(Literal::CreateR0(42))); builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(42.0))); + HloInstruction::CreateConstant(Literal::CreateR0(42.0))); builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(42.0))); + HloInstruction::CreateConstant(Literal::CreateR0(42.0))); builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(42.0))); + HloInstruction::CreateConstant(Literal::CreateR0(42.0))); builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(42.0f))); + HloInstruction::CreateConstant(Literal::CreateR0(42.0f))); // Duplicate the float constant to verify something happens. builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(42.0f))); + HloInstruction::CreateConstant(Literal::CreateR0(42.0f))); - HloModule module(TestName()); - auto computation = module.AddEntryComputation(builder.Build()); + auto module = CreateNewModule(); + auto computation = module->AddEntryComputation(builder.Build()); EXPECT_EQ(7, computation->instruction_count()); HloCSE cse(/*is_layout_sensitive=*/false); - EXPECT_TRUE(cse.Run(&module).ValueOrDie()); + EXPECT_TRUE(cse.Run(module.get()).ValueOrDie()); EXPECT_EQ(6, computation->instruction_count()); } @@ -168,40 +171,42 @@ TEST_F(HloCseTest, NonscalarConstants) { // Test that identical nonscalar constants are merged. auto builder = HloComputation::Builder(TestName()); auto common_constant1 = builder.AddInstruction(HloInstruction::CreateConstant( - LiteralUtil::CreateR2({{1.0, 2.0}, {3.0, 4.0}}))); + Literal::CreateR2({{1.0, 2.0}, {3.0, 4.0}}))); auto common_constant2 = builder.AddInstruction(HloInstruction::CreateConstant( - LiteralUtil::CreateR2({{1.0, 2.0}, {3.0, 4.0}}))); + Literal::CreateR2({{1.0, 2.0}, {3.0, 4.0}}))); // Create a constant which has the same shape but a different value. auto uncommon_constant = builder.AddInstruction(HloInstruction::CreateConstant( - LiteralUtil::CreateR2({{2.0, 4.0}, {6.0, 8.0}}))); + Literal::CreateR2({{2.0, 4.0}, {6.0, 8.0}}))); // Tie the constants together with a tuple. This makes it easier to refer to // the constant instructions via their use. auto tuple = builder.AddInstruction(HloInstruction::CreateTuple( {common_constant1, common_constant2, uncommon_constant})); - HloModule module(TestName()); - auto computation = module.AddEntryComputation(builder.Build()); + auto module = CreateNewModule(); + auto computation = module->AddEntryComputation(builder.Build()); EXPECT_EQ(4, computation->instruction_count()); + EXPECT_THAT(tuple, + 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()); - - EXPECT_EQ(tuple->operand(0), tuple->operand(1)); - EXPECT_EQ(uncommon_constant, tuple->operand(2)); - EXPECT_TRUE(tuple->operand(0) == common_constant1 || - tuple->operand(0) == common_constant2); + auto first_operand = tuple->operand(0); + EXPECT_THAT(first_operand, + ::testing::AnyOf(common_constant1, common_constant2)); + EXPECT_THAT(tuple, + op::Tuple(first_operand, first_operand, uncommon_constant)); } TEST_F(HloCseTest, IdenticalInstructions) { // Test that three identical instructions are commoned. auto builder = HloComputation::Builder(TestName()); auto constant = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(42.0))); + HloInstruction::CreateConstant(Literal::CreateR0(42.0))); auto exp1 = builder.AddInstruction(HloInstruction::CreateUnary( constant->shape(), HloOpcode::kExp, constant)); auto exp2 = builder.AddInstruction(HloInstruction::CreateUnary( @@ -211,20 +216,19 @@ TEST_F(HloCseTest, IdenticalInstructions) { auto tuple = builder.AddInstruction(HloInstruction::CreateTuple({exp1, exp2, exp3})); - HloModule module(TestName()); - auto computation = module.AddEntryComputation(builder.Build()); + auto module = CreateNewModule(); + auto computation = module->AddEntryComputation(builder.Build()); EXPECT_EQ(5, computation->instruction_count()); - EXPECT_NE(tuple->operand(0), tuple->operand(1)); - EXPECT_NE(tuple->operand(1), tuple->operand(2)); - EXPECT_NE(tuple->operand(0), tuple->operand(2)); + EXPECT_THAT(tuple, op::Tuple(exp1, exp2, exp3)); 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()); - EXPECT_EQ(tuple->operand(0), tuple->operand(1)); - EXPECT_EQ(tuple->operand(1), tuple->operand(2)); + auto first_operand = tuple->operand(0); + EXPECT_THAT(first_operand, ::testing::AnyOf(exp1, exp2, exp3)); + EXPECT_THAT(tuple, op::Tuple(first_operand, first_operand, first_operand)); } TEST_F(HloCseTest, IdenticalInstructionsDifferentLayoutsSensitive) { @@ -232,7 +236,7 @@ TEST_F(HloCseTest, IdenticalInstructionsDifferentLayoutsSensitive) { // commoned if the pass is layout sensitive. auto builder = HloComputation::Builder(TestName()); auto constant = builder.AddInstruction(HloInstruction::CreateConstant( - LiteralUtil::CreateR2({{1.0, 2.0}, {3.0, 4.0}}))); + Literal::CreateR2({{1.0, 2.0}, {3.0, 4.0}}))); auto exp1 = builder.AddInstruction(HloInstruction::CreateUnary( constant->shape(), HloOpcode::kExp, constant)); @@ -245,17 +249,17 @@ TEST_F(HloCseTest, IdenticalInstructionsDifferentLayoutsSensitive) { auto tuple = builder.AddInstruction(HloInstruction::CreateTuple({exp1, exp2})); - HloModule module(TestName()); - auto computation = module.AddEntryComputation(builder.Build()); + auto module = CreateNewModule(); + auto computation = module->AddEntryComputation(builder.Build()); EXPECT_EQ(4, computation->instruction_count()); - EXPECT_NE(tuple->operand(0), tuple->operand(1)); + 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_NE(tuple->operand(0), tuple->operand(1)); + EXPECT_THAT(tuple, op::Tuple(exp1, exp2)); } TEST_F(HloCseTest, IdenticalInstructionsDifferentLayoutsInsensitive) { @@ -263,7 +267,7 @@ TEST_F(HloCseTest, IdenticalInstructionsDifferentLayoutsInsensitive) { // the pass is layout insensitive. auto builder = HloComputation::Builder(TestName()); auto constant = builder.AddInstruction(HloInstruction::CreateConstant( - LiteralUtil::CreateR2({{1.0, 2.0}, {3.0, 4.0}}))); + Literal::CreateR2({{1.0, 2.0}, {3.0, 4.0}}))); auto exp1 = builder.AddInstruction(HloInstruction::CreateUnary( constant->shape(), HloOpcode::kExp, constant)); @@ -276,17 +280,19 @@ TEST_F(HloCseTest, IdenticalInstructionsDifferentLayoutsInsensitive) { auto tuple = builder.AddInstruction(HloInstruction::CreateTuple({exp1, exp2})); - HloModule module(TestName()); - auto computation = module.AddEntryComputation(builder.Build()); + auto module = CreateNewModule(); + auto computation = module->AddEntryComputation(builder.Build()); EXPECT_EQ(4, computation->instruction_count()); - EXPECT_NE(tuple->operand(0), tuple->operand(1)); + 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()); - EXPECT_EQ(tuple->operand(0), tuple->operand(1)); + auto first_operand = tuple->operand(0); + EXPECT_THAT(first_operand, ::testing::AnyOf(exp1, exp2)); + EXPECT_THAT(tuple, op::Tuple(first_operand, first_operand)); } TEST_F(HloCseTest, IdenticalExpressions) { @@ -305,7 +311,7 @@ TEST_F(HloCseTest, IdenticalExpressions) { // The *1 instructions should be merged with the *2 instructions. auto builder = HloComputation::Builder(TestName()); auto constant = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(42.0))); + HloInstruction::CreateConstant(Literal::CreateR0(42.0))); auto negate1 = builder.AddInstruction(HloInstruction::CreateUnary( constant->shape(), HloOpcode::kNegate, constant)); @@ -324,39 +330,44 @@ TEST_F(HloCseTest, IdenticalExpressions) { auto tuple = builder.AddInstruction(HloInstruction::CreateTuple({add1, add2})); - HloModule module(TestName()); - auto computation = module.AddEntryComputation(builder.Build()); + auto module = CreateNewModule(); + auto computation = module->AddEntryComputation(builder.Build()); EXPECT_EQ(8, computation->instruction_count()); - EXPECT_NE(tuple->operand(0), tuple->operand(1)); + 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()); - EXPECT_EQ(tuple->operand(0), tuple->operand(1)); - EXPECT_EQ(HloOpcode::kAdd, tuple->operand(0)->opcode()); + auto operand = tuple->operand(0); + EXPECT_THAT(tuple, op::Tuple(operand, operand)); + EXPECT_THAT(operand, op::Add(op::Negate(), op::Exp())); } TEST_F(HloCseTest, DoNotCombineRng) { // Test that two RNG ops are not commoned. auto builder = HloComputation::Builder(TestName()); auto constant1 = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(0.0f))); + HloInstruction::CreateConstant(Literal::CreateR0(0.0f))); auto constant2 = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(1.0f))); + HloInstruction::CreateConstant(Literal::CreateR0(1.0f))); auto rng1 = builder.AddInstruction(HloInstruction::CreateRng( ShapeUtil::MakeShape(F32, {}), RandomDistribution::RNG_UNIFORM, {constant1, constant2})); auto rng2 = builder.AddInstruction(HloInstruction::CreateRng( ShapeUtil::MakeShape(F32, {}), RandomDistribution::RNG_UNIFORM, {constant1, constant2})); + builder.AddInstruction(HloInstruction::CreateBinary( constant1->shape(), HloOpcode::kAdd, rng1, rng2)); - auto module = MakeUnique(TestName()); + auto module = CreateNewModule(); auto computation = module->AddEntryComputation(builder.Build()); + HloInstruction* root = computation->root_instruction(); + EXPECT_THAT(root, op::Add(rng1, rng2)); + uint32 count_before = computation->instruction_count(); HloCSE cse(/*is_layout_sensitive=*/false); @@ -364,11 +375,8 @@ TEST_F(HloCseTest, DoNotCombineRng) { uint32 count_after = computation->instruction_count(); EXPECT_EQ(count_before, count_after); - HloInstruction* root = computation->root_instruction(); - EXPECT_EQ(root->opcode(), HloOpcode::kAdd); - EXPECT_EQ(root->operand(0)->opcode(), HloOpcode::kRng); - EXPECT_EQ(root->operand(1)->opcode(), HloOpcode::kRng); - EXPECT_NE(root->operand(0), root->operand(1)); + root = computation->root_instruction(); + EXPECT_THAT(root, op::Add(rng1, rng2)); } // TODO(b/28245743): Handle impure functions correctly in CSE. @@ -376,7 +384,7 @@ TEST_F(HloCseTest, DISABLED_DoNotCombineCallsToImpureFunctions) { // Test that two calls to an impure function are not commoned. RNG // is the source of the impurity. - auto module = MakeUnique(TestName()); + auto module = CreateNewModule(); // rng_function is an impure function because it does RNG. HloComputation* rng_function = nullptr; @@ -384,9 +392,9 @@ TEST_F(HloCseTest, DISABLED_DoNotCombineCallsToImpureFunctions) { Shape scalar_shape = ShapeUtil::MakeShape(F32, {}); auto builder = HloComputation::Builder(TestName() + "_rng_fun"); auto constant1 = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(0.0f))); + HloInstruction::CreateConstant(Literal::CreateR0(0.0f))); auto constant2 = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(1.0f))); + HloInstruction::CreateConstant(Literal::CreateR0(1.0f))); auto rng = builder.AddInstruction(HloInstruction::CreateRng( scalar_shape, RandomDistribution::RNG_UNIFORM, {constant1, constant2})); auto param = builder.AddInstruction(HloInstruction::CreateParameter( @@ -401,7 +409,7 @@ TEST_F(HloCseTest, DISABLED_DoNotCombineCallsToImpureFunctions) { { auto builder = HloComputation::Builder(TestName()); auto constant = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR1({5.0f}))); + HloInstruction::CreateConstant(Literal::CreateR1({5.0f}))); auto rng1 = builder.AddInstruction( HloInstruction::CreateMap(constant->shape(), {constant}, rng_function)); auto rng2 = builder.AddInstruction( @@ -412,17 +420,22 @@ TEST_F(HloCseTest, DISABLED_DoNotCombineCallsToImpureFunctions) { } EXPECT_EQ(4, computation->instruction_count()); + HloInstruction* root = computation->root_instruction(); + EXPECT_THAT(root, op::Add(op::Map(), op::Map())); HloCSE cse(/*is_layout_sensitive=*/false); EXPECT_TRUE(cse.Run(module.get()).ValueOrDie()); EXPECT_EQ(4, computation->instruction_count()); - HloInstruction* root = computation->root_instruction(); - EXPECT_EQ(root->opcode(), HloOpcode::kAdd); - EXPECT_EQ(root->operand(0)->opcode(), HloOpcode::kMap); - EXPECT_EQ(root->operand(1)->opcode(), HloOpcode::kMap); - EXPECT_NE(root->operand(0), root->operand(1)); + root = computation->root_instruction(); + auto operand = root->operand(0)->operand(0); + EXPECT_THAT(operand, op::Map()); + EXPECT_THAT(root, op::Add(operand, operand)); } } // namespace } // namespace xla + +int main(int argc, char** argv) { + return xla::ParseDebugOptionsFlagsAndRunTests(argc, argv); +} diff --git a/tensorflow/compiler/xla/service/hlo_dataflow_analysis.cc b/tensorflow/compiler/xla/service/hlo_dataflow_analysis.cc new file mode 100644 index 0000000000000000000000000000000000000000..2be1645f1b05dc5824faf7f485c3619716726d77 --- /dev/null +++ b/tensorflow/compiler/xla/service/hlo_dataflow_analysis.cc @@ -0,0 +1,692 @@ +/* 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/hlo_dataflow_analysis.h" + +#include +#include +#include + +#include "tensorflow/compiler/xla/map_util.h" +#include "tensorflow/compiler/xla/ptr_util.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/liveness_util.h" +#include "tensorflow/compiler/xla/shape_util.h" +#include "tensorflow/compiler/xla/status.h" +#include "tensorflow/compiler/xla/types.h" +#include "tensorflow/compiler/xla/util.h" +#include "tensorflow/core/lib/core/errors.h" +#include "tensorflow/core/lib/strings/str_util.h" +#include "tensorflow/core/lib/strings/strcat.h" +#include "tensorflow/core/platform/logging.h" + +namespace xla { + +using ::tensorflow::strings::StrAppend; +using ::tensorflow::strings::StrCat; + +HloDataflowAnalysis::HloDataflowAnalysis(HloModule* module, bool ssa_form, + bool bitcast_defines_value) + : module_(module), + ssa_form_(ssa_form), + bitcast_defines_value_(bitcast_defines_value), + call_graph_(CallGraph::Build(module)) {} + +bool HloDataflowAnalysis::ValueIsDefinedAt(const HloInstruction* instruction, + const ShapeIndex& index) const { + const HloValueSet& value_set = GetValueSet(instruction, index); + if (value_set.values().size() != 1) { + return false; + } + return value_set.GetUniqueValue().defining_instruction() == instruction; +} + +const HloValue& HloDataflowAnalysis::GetValueDefinedAt( + const HloInstruction* instruction, const ShapeIndex& index) const { + CHECK(ValueIsDefinedAt(instruction, index)); + return GetUniqueValueAt(instruction, index); +} + +HloValue& HloDataflowAnalysis::GetValueDefinedAt( + const HloInstruction* instruction, const ShapeIndex& index) { + CHECK(ValueIsDefinedAt(instruction, index)); + return GetUniqueValueAt(instruction, index); +} + +HloValue* HloDataflowAnalysis::NewHloValue(HloInstruction* instruction, + const ShapeIndex& index, + bool is_phi) { + const int64 value_id = next_value_id_++; + auto emplaced = values_.emplace( + std::piecewise_construct, std::forward_as_tuple(value_id), + std::forward_as_tuple(value_id, instruction, index, is_phi)); + CHECK(emplaced.second); + + return &emplaced.first->second; +} + +void HloDataflowAnalysis::DeleteHloValue(HloValue::Id value_id) { + values_.erase(value_id); +} + +string HloDataflowAnalysis::ToString() const { + string out = StrCat("HloDataflowAnalysis, module ", module_->name(), "\n"); + StrAppend(&out, " Instruction value sets:\n"); + for (const std::unique_ptr& computation : + module_->computations()) { + for (const std::unique_ptr& instruction : + computation->instructions()) { + StrAppend(&out, " ", instruction->name(), ":\n"); + if (ShapeUtil::IsTuple(instruction->shape())) { + GetInstructionValueSet(instruction.get()) + .ForEachElement([this, &instruction, &out]( + const ShapeIndex& index, + const HloValueSet& value_set) { + StrAppend(&out, " tuple index ", index.ToString(), ":\n"); + for (const HloValue* value : value_set.values()) { + StrAppend( + &out, " ", value->ToShortString(), + ValueIsDefinedAt(instruction.get(), index) ? " (def)" : "", + "\n"); + } + }); + } else { + const HloValueSet& top_level_value_set = + GetValueSet(instruction.get(), /*index=*/{}); + for (const HloValue* value : top_level_value_set.values()) { + StrAppend(&out, " ", value->ToShortString(), + ValueIsDefinedAt(instruction.get()) ? " (def)" : "", "\n"); + } + } + } + } + StrAppend(&out, " HloValues:\n"); + for (const HloValue* value : values()) { + StrAppend(&out, value->ToString(/*indent=*/4)); + } + return out; +} + +bool HloDataflowAnalysis::Phi( + HloInstruction* instruction, + tensorflow::gtl::ArraySlice inputs) { + CHECK(ssa_form_); + + for (const InstructionValueSet* input : inputs) { + DCHECK(ShapeUtil::Compatible(instruction->shape(), input->shape())); + } + + bool changed = false; + for (auto& pair : GetInstructionValueSet(instruction)) { + const ShapeIndex& index = pair.first; + HloValueSet& value_set = pair.second; + + // Positions with phi values should never have more than one value in the + // value set. + CHECK_LE(value_set.values().size(), 1); + const HloValue* current_value = + value_set.values().size() == 1 ? value_set.values()[0] : nullptr; + + // Construct a vector of unique value IDs of the inputs. + std::vector input_value_ids; + for (const InstructionValueSet* input : inputs) { + for (const HloValue* value : input->element(index).values()) { + input_value_ids.push_back(value->id()); + } + } + std::sort(input_value_ids.begin(), input_value_ids.end()); + input_value_ids.erase( + std::unique(input_value_ids.begin(), input_value_ids.end()), + input_value_ids.end()); + + // Remove the existing phi value (if it exists). The phi can be its own + // input, for example, in while body parameters where the body passes + // through the parameter value. + bool current_value_defined_here = + (current_value != nullptr && + current_value->defining_instruction() == instruction && + current_value->defining_index() == index); + if (current_value_defined_here) { + CHECK(current_value->is_phi()); + auto it = std::find(input_value_ids.begin(), input_value_ids.end(), + current_value->id()); + if (it != input_value_ids.end()) { + input_value_ids.erase(it); + } + } + + if (input_value_ids.empty()) { + // A value set which has at least one element should never have its value + // set reduced to zero elements. During dataflow value sets only can go + // from empty to non-empty, not the reverse. + CHECK_EQ(value_set.values().size(), 0) + << "Instruction " << instruction->name() << " at index " << index + << " previously had non-empty value set. Value set: " << value_set; + } else if (input_value_ids.size() == 1) { + // Only a single value reaches this point. There should be no phi, and + // this value set should contain this single value. + const HloValue& new_value = GetValue(input_value_ids[0]); + if (current_value == nullptr) { + value_set.Clear(); + value_set.AddValue(&new_value); + changed = true; + } else if (current_value != &new_value) { + if (current_value_defined_here) { + // Remove the existing phi. + DeleteHloValue(current_value->id()); + } + value_set.Clear(); + value_set.AddValue(&new_value); + changed = true; + } + } else { + // Multiple distinct values reach this point. A phi value is + // necessary. + CHECK_GT(input_value_ids.size(), 1); + if (current_value == nullptr || !current_value->is_phi()) { + value_set.Clear(); + value_set.AddValue(NewHloValue(instruction, index, /*is_phi=*/true)); + changed = true; + } + } + } + return changed; +} + +const HloValue& HloDataflowAnalysis::GetValue(HloValue::Id value_id) const { + return values_.at(value_id); +} + +HloValue& HloDataflowAnalysis::GetValue(HloValue::Id value_id) { + return values_.at(value_id); +} + +const HloValueSet& HloDataflowAnalysis::GetValueSet( + const HloInstruction* instruction, const ShapeIndex& index) const { + return GetInstructionValueSet(instruction).element(index); +} + +HloValueSet& HloDataflowAnalysis::GetValueSet(const HloInstruction* instruction, + const ShapeIndex& index) { + return *GetInstructionValueSet(instruction).mutable_element(index); +} + +const HloValueSet& HloDataflowAnalysis::GetValueSet( + const HloPosition& position) const { + return GetValueSet(position.instruction, position.index); +} + +HloValueSet& HloDataflowAnalysis::GetValueSet(const HloPosition& position) { + return GetValueSet(position.instruction, position.index); +} + +bool HloDataflowAnalysis::UpdateBitcastValueSet(HloInstruction* bitcast) { + CHECK_EQ(bitcast->opcode(), HloOpcode::kBitcast); + const InstructionValueSet& operand_set = + GetInstructionValueSet(bitcast->operand(0)); + InstructionValueSet& bitcast_set = GetInstructionValueSet(bitcast); + if (!bitcast_defines_value_ && operand_set != bitcast_set) { + bitcast_set = operand_set; + return true; + } + return false; +} + +bool HloDataflowAnalysis::UpdateCallValueSet(HloInstruction* call) { + CHECK_EQ(call->opcode(), HloOpcode::kCall); + InstructionValueSet& value_set = GetInstructionValueSet(call); + InstructionValueSet& root_value_set = + GetInstructionValueSet(call->to_apply()->root_instruction()); + if (value_set != root_value_set) { + value_set = root_value_set; + return true; + } + return false; +} + +bool HloDataflowAnalysis::UpdateCopyValueSet(HloInstruction* copy) { + CHECK_EQ(copy->opcode(), HloOpcode::kCopy); + bool changed = false; + for (auto& pair : GetInstructionValueSet(copy)) { + const ShapeIndex& index = pair.first; + if (index.empty()) { + // kCopy shallow copies and thus defines the top-level value so nothing to + // update. + continue; + } + + HloValueSet& value_set = pair.second; + HloValueSet& operand_value_set = GetValueSet(copy->operand(0), index); + if (value_set != operand_value_set) { + value_set = operand_value_set; + changed = true; + } + } + return changed; +} + +bool HloDataflowAnalysis::UpdateGetTupleElementValueSet(HloInstruction* gte) { + CHECK_EQ(gte->opcode(), HloOpcode::kGetTupleElement); + bool changed = false; + // The GetTupleElement instruction forwards the values from the specified + // tuple element. + for (auto& pair : GetInstructionValueSet(gte)) { + const ShapeIndex& index = pair.first; + HloValueSet& value_set = pair.second; + + // The corresponding ShapeIndex of the operand is simply the GTE ShapeIndex + // with the tuple element number prefixed. + ShapeIndex operand_index = {gte->tuple_index()}; + for (int64 i : index) { + operand_index.push_back(i); + } + + HloValueSet& operand_value_set = + GetValueSet(gte->operand(0), operand_index); + if (value_set != operand_value_set) { + value_set = operand_value_set; + changed = true; + } + } + return changed; +} + +bool HloDataflowAnalysis::UpdateParameterValueSet(HloInstruction* parameter) { + CHECK_EQ(parameter->opcode(), HloOpcode::kParameter); + const CallGraphNode& call_graph_node = + call_graph_->GetNode(parameter->parent()); + + // Subcomputations called in a parallel context (eg, map) do not have dataflow + // from the caller operands. + if (call_graph_node.context() == CallContext::kParallel || + call_graph_node.caller_callsites().empty()) { + return false; + } + CHECK_EQ(call_graph_node.context(), CallContext::kSequential); + + std::vector inputs; + bool called_from_while = false; + for (const CallSite& callsite : call_graph_node.caller_callsites()) { + if (callsite.instruction()->opcode() == HloOpcode::kCall) { + // The operand values of a call instruction are forwarded to the + // respective parameter instruction of the subcomputation. + inputs.push_back(&GetInstructionValueSet( + callsite.instruction()->operand(parameter->parameter_number()))); + } else if (callsite.instruction()->opcode() == HloOpcode::kWhile) { + // In a while instruction, the while operand (ie, the init value) and the + // backedge are dataflow inputs to the parameter instruction. This is the + // case for parameters of both the body and condition computations. + CHECK_EQ(parameter->parameter_number(), 0); + inputs.push_back( + &GetInstructionValueSet(callsite.instruction()->operand(0))); + // If the parameter *is* the root, then don't consider it's current state + // (InstructionValueSet) as we are recomputing its current + // state. Otherwise, the parameter state would never be updated. + if (parameter != + callsite.instruction()->while_body()->root_instruction()) { + inputs.push_back(&GetInstructionValueSet( + callsite.instruction()->while_body()->root_instruction())); + } + called_from_while = true; + } else { + LOG(FATAL) << "CallContext::kSequential computations should only be " + "called from call or while instructions"; + } + } + + if (ssa_form_ && called_from_while) { + return Phi(parameter, inputs); + } else { + return GetInstructionValueSet(parameter).AssignUnionOf(inputs); + } +} + +bool HloDataflowAnalysis::UpdateSelectValueSet(HloInstruction* select) { + CHECK_EQ(select->opcode(), HloOpcode::kSelect); + // A phi value is not defined at a kSelect instruction because kSelect does + // not create a new value. Rather it forwards a value from its operands. This + // contrasts with kWhile instruction (which does define a phi value) which has + // in-place update semantics. + bool changed = false; + for (auto& pair : GetInstructionValueSet(select)) { + const ShapeIndex& index = pair.first; + if (index.empty()) { + // kSelect copies (not forwards) the top-level value. + continue; + } + HloValueSet& value_set = pair.second; + changed |= + value_set.AssignUnionOf({&GetValueSet(select->operand(1), index), + &GetValueSet(select->operand(2), index)}); + } + return changed; +} + +bool HloDataflowAnalysis::UpdateTupleValueSet(HloInstruction* tuple) { + CHECK_EQ(tuple->opcode(), HloOpcode::kTuple); + bool changed = false; + for (int64 i = 0; i < tuple->operands().size(); ++i) { + // Copy the value set(s) of each operand into the respective position in the + // kTuple instruction's value sets. + for (auto& pair : GetInstructionValueSet(tuple->operand(i))) { + const ShapeIndex& operand_index = pair.first; + HloValueSet& operand_value_set = pair.second; + + ShapeIndex index = {i}; + for (int64 op_index : operand_index) { + index.push_back(op_index); + } + HloValueSet& value_set = GetValueSet(tuple, index); + + if (value_set != operand_value_set) { + value_set = operand_value_set; + changed = true; + } + } + } + return changed; +} + +bool HloDataflowAnalysis::UpdateWhileValueSet(HloInstruction* xla_while) { + CHECK_EQ(xla_while->opcode(), HloOpcode::kWhile); + std::vector inputs = { + &GetInstructionValueSet(xla_while->while_body()->root_instruction()), + &GetInstructionValueSet(xla_while->operand(0))}; + if (ssa_form_) { + return Phi(xla_while, inputs); + } else { + return GetInstructionValueSet(xla_while).AssignUnionOf(inputs); + } +} + +bool HloDataflowAnalysis::UpdateInstructionValueSet( + HloInstruction* instruction) { + // Recompute from operands. + switch (instruction->opcode()) { + case HloOpcode::kBitcast: + return UpdateBitcastValueSet(instruction); + case HloOpcode::kCopy: + return UpdateCopyValueSet(instruction); + case HloOpcode::kGetTupleElement: + return UpdateGetTupleElementValueSet(instruction); + case HloOpcode::kSelect: + return UpdateSelectValueSet(instruction); + case HloOpcode::kTuple: + return UpdateTupleValueSet(instruction); + case HloOpcode::kParameter: + return UpdateParameterValueSet(instruction); + case HloOpcode::kCall: + return UpdateCallValueSet(instruction); + case HloOpcode::kWhile: + return UpdateWhileValueSet(instruction); + default: + // Instruction does not forward HloValues (it defines all values in its + // output). No update is necessary. + return false; + } +} + +void HloDataflowAnalysis::UpdateInstructionsAndPropagate( + tensorflow::gtl::ArraySlice instructions) { + std::queue worklist; + for (HloInstruction* instruction : instructions) { + worklist.push(instruction); + } + + while (!worklist.empty()) { + HloInstruction* instruction = worklist.front(); + worklist.pop(); + + VLOG(3) << "Worklist top: " << instruction->name(); + VLOG(3) << ToString(); + + if (!UpdateInstructionValueSet(instruction)) { + // No change to the instruction's value set. + VLOG(4) << "No change."; + continue; + } + + VLOG(4) << "New value set for " << instruction->name() << ": " + << GetInstructionValueSet(instruction); + + // Instruction value was updated. Add users to work list. + for (HloInstruction* user : instruction->users()) { + worklist.push(user); + + // If user sequentially calls a computation, then the respective + // parameter(s) of the computation need to be updated. + for (HloComputation* called_computation : user->called_computations()) { + const CallGraphNode& call_graph_node = + call_graph_->GetNode(called_computation); + if (call_graph_node.context() == CallContext::kSequential) { + for (int64 operand_number : user->OperandIndices(instruction)) { + worklist.push( + called_computation->parameter_instruction(operand_number)); + } + } + } + } + + // If instruction is a root instruction, then propagate out to any calling + // instruction and across any while backedge. + if (instruction == instruction->parent()->root_instruction()) { + const CallGraphNode& call_graph_node = + call_graph_->GetNode(instruction->parent()); + for (const CallSite& callsite : call_graph_node.caller_callsites()) { + if (callsite.instruction()->opcode() == HloOpcode::kCall) { + worklist.push(callsite.instruction()); + } else if (callsite.instruction()->opcode() == HloOpcode::kWhile) { + // Add the while itself, and the body and condition parameters. + worklist.push(callsite.instruction()); + worklist.push( + callsite.instruction()->while_body()->parameter_instruction(0)); + worklist.push( + callsite.instruction()->while_condition()->parameter_instruction( + 0)); + } + } + } + } +} + +const InstructionValueSet& HloDataflowAnalysis::GetInstructionValueSet( + const HloInstruction* instruction) const { + return value_sets_.at(instruction); +} + +InstructionValueSet& HloDataflowAnalysis::GetInstructionValueSet( + const HloInstruction* instruction) { + return value_sets_.at(instruction); +} + +Status HloDataflowAnalysis::InitializeInstructionValueSets() { + for (const std::unique_ptr& computation : + module_->computations()) { + const CallGraphNode& call_graph_node = + call_graph_->GetNode(computation.get()); + + for (const std::unique_ptr& instruction : + computation->instructions()) { + // Create an empty shape tree. + value_sets_.emplace(std::piecewise_construct, + std::forward_as_tuple(instruction.get()), + std::forward_as_tuple(instruction->shape())); + + // Lambda to set the value set to define all values in the output of the + // instruction. + auto define_all_values = [this, &instruction](bool is_phi = false) { + for (auto& pair : GetInstructionValueSet(instruction.get())) { + const ShapeIndex& index = pair.first; + HloValue* value = + NewHloValue(instruction.get(), index, /*is_phi=*/false); + GetValueSet(instruction.get(), index).AddValue(value); + } + }; + + // Lambda to set the value set to define only the top-level buffer in the + // output of the instruction. Any other values flow from the operands of + // the instruction (or from cross-computation dataflow). + auto define_top_level_only = [this, &instruction]() { + HloValue* value = + NewHloValue(instruction.get(), /*index=*/{}, /*is_phi=*/false); + GetValueSet(instruction.get(), /*index=*/{}).AddValue(value); + }; + + switch (instruction->opcode()) { + case HloOpcode::kBitcast: + if (bitcast_defines_value_) { + define_all_values(); + } + break; + case HloOpcode::kWhile: + case HloOpcode::kCall: + case HloOpcode::kGetTupleElement: + // These instructions define no values. The values in their output + // flow from their operands or from cross computation dataflow. + break; + case HloOpcode::kParameter: + if (call_graph_node.context() == CallContext::kBoth) { + // We do not support a subcomputation that is called from both a + // parallel and sequential context. In this case, the parameter + // would both define a value and propagate a value from its + // caller. This limitation is not really a problem because the call + // graph is typically flattened. + return Unimplemented( + "Computation %s is called in both a parallel (eg, kMap) and " + "sequential (eg, kCall) context", + computation->name().c_str()); + } + if (call_graph_node.caller_callsites().empty() || + call_graph_node.context() == CallContext::kParallel) { + // Parameters of computations called in a parallel context (eg, map + // and reduce) as well as parameters of dead computations define all + // values in their output. Otherwise the values of the parameter + // come from the caller (eg, operands to the kCall instruction). + define_all_values(); + } + break; + case HloOpcode::kCopy: + case HloOpcode::kSelect: + case HloOpcode::kTuple: + // These instructions only define their top-level values. Any other + // values flow from their operands. + define_top_level_only(); + break; + default: + define_all_values(); + break; + } + } + } + + return Status::OK(); +} + +/* static */ +StatusOr> HloDataflowAnalysis::Run( + HloModule* module, bool ssa_form, bool bitcast_defines_value) { + VLOG(1) << "HloDataflowAnalysis::Run on module " << module->name(); + XLA_VLOG_LINES(2, module->ToString()); + + auto dataflow_analysis = WrapUnique( + new HloDataflowAnalysis(module, ssa_form, bitcast_defines_value)); + + TF_RETURN_IF_ERROR(dataflow_analysis->InitializeInstructionValueSets()); + + // Construct list of all instructions to initialize the worklist to propagate + // the data flow. For efficiency sort the instruction in post order so + // producers appear before consumers. + std::vector all_instructions; + for (const HloComputation* computation : module->MakeComputationPostOrder()) { + for (HloInstruction* instruction : + computation->MakeInstructionPostOrder()) { + all_instructions.push_back(instruction); + } + } + dataflow_analysis->UpdateInstructionsAndPropagate(all_instructions); + + // Add in positions to all values. + for (const std::unique_ptr& computation : + module->computations()) { + for (const std::unique_ptr& instruction : + computation->instructions()) { + for (const auto& pair : + dataflow_analysis->GetInstructionValueSet(instruction.get())) { + const ShapeIndex& index = pair.first; + const HloValueSet& value_set = pair.second; + for (const HloValue* value : value_set.values()) { + if (value->defining_instruction() != instruction.get()) { + dataflow_analysis->GetValue(value->id()) + .AddPosition(instruction.get(), index); + } + } + } + } + } + + // Construct vector of values. + dataflow_analysis->values_vector_.reserve(dataflow_analysis->values_.size()); + 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); + + TF_DCHECK_OK(dataflow_analysis->Verify()); + + XLA_VLOG_LINES(1, dataflow_analysis->ToString()); + + return std::move(dataflow_analysis); +} + +Status HloDataflowAnalysis::Verify() const { + // Verify each HloValue appears in the value sets that the value's positions() + // indicate. + 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()) + << "Value set at position " << position << " does not contain value " + << value->ToShortString(); + } + } + + // For each value in each value set, verify that the value set's position + // appears in the value's positions(). + for (const auto& computation : module_->computations()) { + for (const auto& instruction : computation->instructions()) { + for (const auto& pair : GetInstructionValueSet(instruction.get())) { + const ShapeIndex& index = pair.first; + const HloValueSet& value_set = pair.second; + const HloPosition position{instruction.get(), index}; + for (const HloValue* value : value_set.values()) { + TF_RET_CHECK(std::find(value->positions().begin(), + value->positions().end(), + position) != value->positions().end()) + << "Value set at position " << position + << " unexpectedly contains value " << value->ToShortString(); + } + } + } + } + + return Status::OK(); +} + +} // namespace xla diff --git a/tensorflow/compiler/xla/service/hlo_dataflow_analysis.h b/tensorflow/compiler/xla/service/hlo_dataflow_analysis.h new file mode 100644 index 0000000000000000000000000000000000000000..aae257dd09e8ee37e040b8c7b673059355615ed4 --- /dev/null +++ b/tensorflow/compiler/xla/service/hlo_dataflow_analysis.h @@ -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. +==============================================================================*/ + +// Analysis for determining the possible set of values for all positions +// (instructions and ShapeIndexes) in the HLO module. Analysis is module-scoped +// tracking values across computation boundaries. + +#ifndef TENSORFLOW_COMPILER_XLA_SERVICE_HLO_DATAFLOW_ANALYSIS_H_ +#define TENSORFLOW_COMPILER_XLA_SERVICE_HLO_DATAFLOW_ANALYSIS_H_ + +#include +#include +#include +#include + +#include "tensorflow/compiler/xla/service/call_graph.h" +#include "tensorflow/compiler/xla/service/hlo_instruction.h" +#include "tensorflow/compiler/xla/service/hlo_module.h" +#include "tensorflow/compiler/xla/service/hlo_ordering.h" +#include "tensorflow/compiler/xla/service/hlo_value.h" +#include "tensorflow/compiler/xla/shape_util.h" +#include "tensorflow/compiler/xla/status.h" +#include "tensorflow/compiler/xla/statusor.h" +#include "tensorflow/compiler/xla/types.h" +#include "tensorflow/compiler/xla/xla_data.pb.h" +#include "tensorflow/core/lib/gtl/array_slice.h" +#include "tensorflow/core/platform/macros.h" + +namespace xla { + +// Analysis which identifies all HLO values and their uses in an HLO module. +class HloDataflowAnalysis { + public: + // Run dataflow analysis on the given module. Parameters: + // + // ssa_form : If true then new values are defined at the merge points of + // kWhile instructions. Abusing nomenclature somewhat, we call these "phi + // values". The merge is formed by the init value and loop backedge. The + // SSA form is minimal in that a new phi value is defined only if the + // merge point is reachable by multiple different values. The SSA form is + // also in loop-closed form in that no values defined inside of a loop + // (while body) is used outside of the loop. + // + // If ssa_form is false, then merge points do not define new + // values. Rather, the HloValueSet for the merge point contains the union + // of the merged HloValues. + // + // bitcast_defines_value : If true then the Bitcast HLO instruction defines + // a new HLO value in the analysis. If false then Bitcast forwards the + // value of its operand. + static StatusOr> Run( + HloModule* module, bool ssa_form = false, + bool bitcast_defines_value = false); + + // Returns true if 'instruction' defines an HLO value at the given shape index + // of its output. + bool ValueIsDefinedAt(const HloInstruction* instruction, + const ShapeIndex& index = {}) const; + + // Return the HloValue defined by 'instruction' at the given shape index of + // its output. + // + // Precondition: ValueIsDefinedAt is true for this instruction and index. + const HloValue& GetValueDefinedAt(const HloInstruction* instruction, + const ShapeIndex& index = {}) const; + HloValue& GetValueDefinedAt(const HloInstruction* instruction, + const ShapeIndex& index = {}); + + // Return the InstructionValueSet for the given instruction. + const InstructionValueSet& GetInstructionValueSet( + const HloInstruction* instruction) const; + InstructionValueSet& GetInstructionValueSet( + const HloInstruction* instruction); + + // Return the HloValueSet for the given instruction at the given index or the + // given position. + const HloValueSet& GetValueSet(const HloInstruction* instruction, + const ShapeIndex& index = {}) const; + const HloValueSet& GetValueSet(const HloPosition& position) const; + HloValueSet& GetValueSet(const HloPosition& position); + HloValueSet& GetValueSet(const HloInstruction* instruction, + const ShapeIndex& index = {}); + + // Return the unique value in the HloValueSet at the given instruction and + // shape index. CHECKs if the value set does not contain a exactly one value. + const HloValue& GetUniqueValueAt(const HloInstruction* instruction, + const ShapeIndex& index = {}) const { + return GetValueSet(instruction, index).GetUniqueValue(); + } + HloValue& GetUniqueValueAt(const HloInstruction* instruction, + const ShapeIndex& index = {}) { + return GetValue(GetValueSet(instruction, index).GetUniqueValue().id()); + } + + // Return the HloValue with the given Id. + const HloValue& GetValue(HloValue::Id value_id) const; + HloValue& GetValue(HloValue::Id value_id); + + // Return the total number of HloValues. + int64 value_count() const { return values_.size(); } + + // Return a vector of all HloValues stabily sorted by HloValue::Id. + const std::vector& values() const { return values_vector_; } + + // Return the call graph used for computing the dataflow. + const CallGraph& call_graph() const { return *call_graph_; } + + string ToString() const; + + protected: + HloDataflowAnalysis(HloModule* module, bool ssa_form, + bool bitcast_defines_value = false); + + // Returns a new HloValue defined at the given instruction and shape index. + HloValue* NewHloValue(HloInstruction* instruction, const ShapeIndex& index, + bool is_phi = false); + + // Delete the HloValue with the given ID. + void DeleteHloValue(HloValue::Id value_id); + + // Constructs and initializes the InstructionValueSets of all instructions to + // contain exactly the HloValues defined by each instruction. These values can + // then propagated throughout the HLO graph by calling + // UpdateInstructionsAndPropagate. + Status InitializeInstructionValueSets(); + + // Updates the value set of the given instruction based on the values flowing + // into the instruction (operands and cross-computation dataflow). + bool UpdateInstructionValueSet(HloInstruction* instruction); + + // Updates the value set for a particular instruction type. Returns whether + // the instruction value set changed. + bool UpdateBitcastValueSet(HloInstruction* bitcast); + bool UpdateCallValueSet(HloInstruction* call); + bool UpdateCopyValueSet(HloInstruction* copy); + bool UpdateGetTupleElementValueSet(HloInstruction* gte); + bool UpdateParameterValueSet(HloInstruction* parameter); + bool UpdateSelectValueSet(HloInstruction* select); + bool UpdateTupleValueSet(HloInstruction* tuple); + bool UpdateWhileValueSet(HloInstruction* xla_while); + + // Update the value sets of the given instructions and propagate the + // changes to fixed point. + void UpdateInstructionsAndPropagate( + tensorflow::gtl::ArraySlice instructions); + + // Return the result of the SSA Phi function applied to the given inputs at + // the given instruction. If skip_top_level is true, then the top level of the + // value set of 'instruction' is not modified. + bool Phi(HloInstruction* instruction, + tensorflow::gtl::ArraySlice inputs); + + // Updates the positions of the HloValues in the output of the given + // instruction. This should be called after the instruction value set of + // 'instruction' has been changed. 'prev_value_set' must point to the previous + // state of the value set prior to the change. 'prev_value_set' may be null if + // this is the first time positions are being computed. The previous state is + // necessary to efficiently remove positions which have been eliminated due to + // changes in the instructions' InstructionValueSet. + void UpdatePositionsOfValuesAt( + HloInstruction* instruction, const InstructionValueSet& new_value_set, + const InstructionValueSet* prev_value_set = nullptr); + + // Verify various invariants of the dataflow analysis. + Status Verify() const; + + HloModule* const module_; + const bool ssa_form_; + const bool bitcast_defines_value_; + + std::unique_ptr call_graph_; + + // The map of all HloValues in the module. We pass around pointers to the + // mapped HloValues, so the underlying container must keep them valid despite + // mutations touching other map entries. + std::unordered_map values_; + + // A map from instruction to InstructionValueSet. + std::unordered_map value_sets_; + + // A vector containing all HloValues sorted by HloValue::Id. + std::vector values_vector_; + + // The Id to use for the next HloValue. + HloValue::Id next_value_id_ = 0; +}; + +} // namespace xla + +#endif // TENSORFLOW_COMPILER_XLA_SERVICE_HLO_DATAFLOW_ANALYSIS_H_ diff --git a/tensorflow/compiler/xla/service/hlo_dataflow_analysis_test.cc b/tensorflow/compiler/xla/service/hlo_dataflow_analysis_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..ef0fa1d745ae38a7f899fe92ee2c5f77e270ec2f --- /dev/null +++ b/tensorflow/compiler/xla/service/hlo_dataflow_analysis_test.cc @@ -0,0 +1,1490 @@ +/* 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/hlo_dataflow_analysis.h" + +#include "tensorflow/compiler/xla/literal_util.h" +#include "tensorflow/compiler/xla/service/hlo_computation.h" +#include "tensorflow/compiler/xla/service/hlo_matchers.h" +#include "tensorflow/compiler/xla/service/hlo_opcode.h" +#include "tensorflow/compiler/xla/service/instruction_fusion.h" +#include "tensorflow/compiler/xla/shape_util.h" +#include "tensorflow/compiler/xla/status_macros.h" +#include "tensorflow/compiler/xla/test.h" +#include "tensorflow/compiler/xla/test_helpers.h" +#include "tensorflow/compiler/xla/tests/hlo_test_base.h" +#include "tensorflow/compiler/xla/xla_data.pb.h" +#include "tensorflow/core/platform/logging.h" +#include "tensorflow/core/platform/test.h" + +namespace xla { +namespace { + +using ::testing::UnorderedElementsAre; + +// Test is parameterized on a bool which is whether the dataflow analysis is +// performed with SSA form. +class HloDataflowAnalysisTest : public HloTestBase, + public ::testing::WithParamInterface { + protected: + HloDataflowAnalysisTest() : module_(CreateNewModule()) {} + + // Run dataflow analysis on the member module. For convenience returns a + // reference to the generated analysis stored in analysis_. + const HloDataflowAnalysis& RunAnalysis(bool ssa_form, + bool bitcast_defines_value = false) { + analysis_ = + HloDataflowAnalysis::Run(module_.get(), ssa_form, bitcast_defines_value) + .ConsumeValueOrDie(); + return *analysis_; + } + + // Return a vector of the HloValues at the given program position. + std::vector HloValuesAt(const HloInstruction* instruction, + const ShapeIndex& index = {}) { + CHECK(analysis_ != nullptr); + std::vector values; + for (const HloValue* value : + analysis_->GetValueSet(instruction, index).values()) { + values.push_back(*value); + } + return values; + } + + // Returns true if the top-level values for instructions 'a' and 'b' may + // interfere. Precondition: 'a' and 'b' define array-shaped values. + bool InstructionsMayInterfere(const HloOrdering& ordering, + const HloInstruction* a, + const HloInstruction* b) { + EXPECT_FALSE(ShapeUtil::IsTuple(a->shape())); + EXPECT_FALSE(ShapeUtil::IsTuple(b->shape())); + return ordering.MayInterfere(analysis_->GetValueDefinedAt(a), + analysis_->GetValueDefinedAt(b)); + } + + std::unique_ptr module_; + std::unique_ptr analysis_; + + const Shape scalar_shape_ = ShapeUtil::MakeShape(F32, {}); + const Shape vector_shape_ = ShapeUtil::MakeShape(F32, {42}); +}; + +TEST_P(HloDataflowAnalysisTest, BinaryOperation) { + // Test the dataflow for a simple binary operation (Add). + auto builder = HloComputation::Builder(TestName()); + auto constant1 = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(1.0))); + auto constant2 = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(2.0))); + auto add = builder.AddInstruction(HloInstruction::CreateBinary( + scalar_shape_, HloOpcode::kAdd, constant1, constant2)); + module_->AddEntryComputation(builder.Build()); + + bool ssa_form = GetParam(); + const HloDataflowAnalysis& analysis = RunAnalysis(ssa_form); + + // Each instruction should define a single value. + EXPECT_EQ(analysis.values().size(), 3); + EXPECT_TRUE(analysis.ValueIsDefinedAt(constant1)); + EXPECT_TRUE(analysis.ValueIsDefinedAt(constant2)); + EXPECT_TRUE(analysis.ValueIsDefinedAt(add)); + + // Verify the positions of the values. These positions are all trivial because + // there are no instructions which forward values. + EXPECT_THAT(analysis.GetValueDefinedAt(constant1).positions(), + UnorderedElementsAre(HloPosition{constant1, {}})); + EXPECT_THAT(analysis.GetValueDefinedAt(constant2).positions(), + UnorderedElementsAre(HloPosition{constant2, {}})); + EXPECT_THAT(analysis.GetValueDefinedAt(add).positions(), + UnorderedElementsAre(HloPosition{add, {}})); + + // Verify the uses of the values. + EXPECT_THAT(analysis.GetValueDefinedAt(constant1).uses(), + UnorderedElementsAre(HloUse{add, 0, {}})); + EXPECT_THAT(analysis.GetValueDefinedAt(constant2).uses(), + UnorderedElementsAre(HloUse{add, 1, {}})); + EXPECT_TRUE(analysis.GetValueDefinedAt(add).uses().empty()); + + // Verify liveout values from the module. + EXPECT_FALSE(analysis.GetValueDefinedAt(constant1).live_out_of_module()); + EXPECT_FALSE(analysis.GetValueDefinedAt(constant2).live_out_of_module()); + EXPECT_TRUE(analysis.GetValueDefinedAt(add).live_out_of_module()); +} + +TEST_P(HloDataflowAnalysisTest, TupleAndGtes) { + // Verify the dataflow through a Tuple and GetTupleElement instructions. + auto builder = HloComputation::Builder(TestName()); + auto param0 = builder.AddInstruction( + HloInstruction::CreateParameter(0, scalar_shape_, "param0")); + auto param1 = builder.AddInstruction( + HloInstruction::CreateParameter(1, scalar_shape_, "param1")); + auto tuple = + builder.AddInstruction(HloInstruction::CreateTuple({param0, param1})); + auto gte0 = builder.AddInstruction( + HloInstruction::CreateGetTupleElement(scalar_shape_, tuple, 0)); + auto gte1 = builder.AddInstruction( + HloInstruction::CreateGetTupleElement(scalar_shape_, tuple, 1)); + auto add = builder.AddInstruction( + HloInstruction::CreateBinary(scalar_shape_, HloOpcode::kAdd, gte0, gte1)); + module_->AddEntryComputation(builder.Build()); + + bool ssa_form = GetParam(); + const HloDataflowAnalysis& analysis = RunAnalysis(ssa_form); + + // The two params, tuple, and add should each define one value. + EXPECT_EQ(analysis.values().size(), 4); + + EXPECT_TRUE(analysis.ValueIsDefinedAt(param0)); + EXPECT_TRUE(analysis.ValueIsDefinedAt(param1)); + EXPECT_TRUE(analysis.ValueIsDefinedAt(tuple, /*index=*/{})); + EXPECT_FALSE(analysis.ValueIsDefinedAt(tuple, /*index=*/{0})); + EXPECT_FALSE(analysis.ValueIsDefinedAt(tuple, /*index=*/{1})); + EXPECT_FALSE(analysis.ValueIsDefinedAt(gte0)); + EXPECT_FALSE(analysis.ValueIsDefinedAt(gte1)); + EXPECT_TRUE(analysis.ValueIsDefinedAt(add)); + + // Verify the positions of the values. + EXPECT_THAT( + analysis.GetValueDefinedAt(param0).positions(), + UnorderedElementsAre(HloPosition{param0, {}}, HloPosition{tuple, {0}}, + HloPosition{gte0, {}})); + EXPECT_THAT( + analysis.GetValueDefinedAt(param1).positions(), + UnorderedElementsAre(HloPosition{param1, {}}, HloPosition{tuple, {1}}, + HloPosition{gte1, {}})); + EXPECT_THAT(analysis.GetValueDefinedAt(tuple).positions(), + UnorderedElementsAre(HloPosition{tuple, {}})); + + // Verify uses. Of interest is that a GetTupleElement instruction is only a + // use of the top-level value in the tuple operand. + EXPECT_THAT(analysis.GetValueDefinedAt(param0).uses(), + UnorderedElementsAre(HloUse{add, 0, {}})); + EXPECT_THAT(analysis.GetValueDefinedAt(param1).uses(), + UnorderedElementsAre(HloUse{add, 1, {}})); + EXPECT_THAT(analysis.GetValueDefinedAt(tuple, /*index=*/{}).uses(), + UnorderedElementsAre(HloUse{gte0, 0, {}}, HloUse{gte1, 0, {}})); + EXPECT_TRUE(analysis.GetValueDefinedAt(add).live_out_of_module()); +} + +TEST_P(HloDataflowAnalysisTest, NestedTuple) { + // Verify the dataflow through a nested tuple. + auto builder = HloComputation::Builder(TestName()); + auto constant1 = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(1.0))); + auto constant2 = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(2.0))); + auto tuple = builder.AddInstruction( + HloInstruction::CreateTuple({constant1, constant2})); + auto nested_tuple = builder.AddInstruction( + HloInstruction::CreateTuple({tuple, tuple, constant1})); + auto gte_tuple = builder.AddInstruction( + HloInstruction::CreateGetTupleElement(tuple->shape(), nested_tuple, 1)); + auto gte_out = builder.AddInstruction( + HloInstruction::CreateGetTupleElement(scalar_shape_, gte_tuple, 0)); + module_->AddEntryComputation(builder.Build()); + + bool ssa_form = GetParam(); + const HloDataflowAnalysis& analysis = RunAnalysis(ssa_form); + + EXPECT_EQ(analysis.values().size(), 4); + + // Verify positions and uses. + EXPECT_THAT( + analysis.GetValueDefinedAt(constant1).positions(), + UnorderedElementsAre( + HloPosition{constant1, {}}, HloPosition{tuple, {0}}, + HloPosition{nested_tuple, {0, 0}}, HloPosition{nested_tuple, {1, 0}}, + HloPosition{nested_tuple, {2}}, HloPosition{gte_tuple, {0}}, + HloPosition{gte_out, {}})); + // Constant values should have no uses though one is live out. The positions + // where they appear as operands are on instructions which do not use the + // values (eg, Tuple). + EXPECT_TRUE(analysis.GetValueDefinedAt(constant1).uses().empty()); + EXPECT_TRUE(analysis.GetValueDefinedAt(constant2).uses().empty()); + + // The top-level tuple values are used in GTE instructions. + EXPECT_THAT(analysis.GetValueDefinedAt(tuple, /*index=*/{}).uses(), + UnorderedElementsAre(HloUse{gte_out, 0, {}})); + EXPECT_THAT(analysis.GetValueDefinedAt(nested_tuple, /*index=*/{}).uses(), + UnorderedElementsAre(HloUse{gte_tuple, 0, {}})); + + EXPECT_TRUE(analysis.GetValueDefinedAt(constant1).live_out_of_module()); + EXPECT_FALSE(analysis.GetValueDefinedAt(constant2).live_out_of_module()); + EXPECT_FALSE( + analysis.GetValueDefinedAt(tuple, /*index=*/{}).live_out_of_module()); + EXPECT_FALSE(analysis.GetValueDefinedAt(nested_tuple, /*index=*/{}) + .live_out_of_module()); +} + +TEST_P(HloDataflowAnalysisTest, SingleCall) { + // Test a single call of a subcomputation. The subcomputation adds its two + // array-shaped parameters. + auto subbuilder = HloComputation::Builder("Subcomputation"); + auto subparam0 = subbuilder.AddInstruction( + HloInstruction::CreateParameter(0, scalar_shape_, "param0")); + auto subparam1 = subbuilder.AddInstruction( + HloInstruction::CreateParameter(1, scalar_shape_, "param1")); + auto add = subbuilder.AddInstruction(HloInstruction::CreateBinary( + scalar_shape_, HloOpcode::kAdd, subparam0, subparam1)); + HloComputation* called_computation = + module_->AddEmbeddedComputation(subbuilder.Build()); + + auto builder = HloComputation::Builder(TestName()); + auto constant1 = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(1.0))); + auto constant2 = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(2.0))); + auto call = builder.AddInstruction(HloInstruction::CreateCall( + scalar_shape_, {constant1, constant2}, called_computation)); + module_->AddEntryComputation(builder.Build()); + + bool ssa_form = GetParam(); + const HloDataflowAnalysis& analysis = RunAnalysis(ssa_form); + + EXPECT_EQ(analysis.values().size(), 3); + + // The parameters of the subcomputation and the call instruction itself should + // not define values. Their values flow from elsewhere. + EXPECT_TRUE(analysis.ValueIsDefinedAt(constant1)); + EXPECT_TRUE(analysis.ValueIsDefinedAt(constant2)); + EXPECT_FALSE(analysis.ValueIsDefinedAt(subparam0)); + EXPECT_FALSE(analysis.ValueIsDefinedAt(subparam1)); + EXPECT_TRUE(analysis.ValueIsDefinedAt(add)); + EXPECT_FALSE(analysis.ValueIsDefinedAt(call)); + + EXPECT_EQ(analysis.GetUniqueValueAt(subparam0), + analysis.GetValueDefinedAt(constant1)); + EXPECT_EQ(analysis.GetUniqueValueAt(subparam1), + analysis.GetValueDefinedAt(constant2)); + EXPECT_EQ(analysis.GetUniqueValueAt(call), analysis.GetValueDefinedAt(add)); + + EXPECT_THAT(analysis.GetValueDefinedAt(constant1).uses(), + UnorderedElementsAre(HloUse{add, 0, {}})); + EXPECT_THAT(analysis.GetValueDefinedAt(constant2).uses(), + UnorderedElementsAre(HloUse{add, 1, {}})); + + EXPECT_TRUE(analysis.GetValueDefinedAt(add).live_out_of_module()); + EXPECT_TRUE(analysis.GetValueDefinedAt(add).live_out_of_computation()); +} + +TEST_P(HloDataflowAnalysisTest, ComputationCalledTwiceWithSameArguments) { + // Test a subcomputation which is called twice with identical values. + auto subbuilder = HloComputation::Builder("Subcomputation"); + auto subparam0 = subbuilder.AddInstruction( + HloInstruction::CreateParameter(0, scalar_shape_, "param0")); + auto subparam1 = subbuilder.AddInstruction( + HloInstruction::CreateParameter(1, scalar_shape_, "param1")); + auto add = subbuilder.AddInstruction(HloInstruction::CreateBinary( + scalar_shape_, HloOpcode::kAdd, subparam0, subparam1)); + HloComputation* called_computation = + module_->AddEmbeddedComputation(subbuilder.Build()); + + auto builder = HloComputation::Builder(TestName()); + auto constant1 = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(1.0))); + auto constant2 = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(2.0))); + auto call1 = builder.AddInstruction(HloInstruction::CreateCall( + scalar_shape_, {constant1, constant2}, called_computation)); + auto call2 = builder.AddInstruction(HloInstruction::CreateCall( + scalar_shape_, {constant1, constant2}, called_computation)); + auto sub = builder.AddInstruction(HloInstruction::CreateBinary( + scalar_shape_, HloOpcode::kSubtract, call1, call2)); + module_->AddEntryComputation(builder.Build()); + + bool ssa_form = GetParam(); + const HloDataflowAnalysis& analysis = RunAnalysis(ssa_form); + + EXPECT_EQ(analysis.values().size(), 4); + + // Definitions should be identical to the single callsite case. + EXPECT_TRUE(analysis.ValueIsDefinedAt(constant1)); + EXPECT_TRUE(analysis.ValueIsDefinedAt(constant2)); + EXPECT_FALSE(analysis.ValueIsDefinedAt(subparam0)); + EXPECT_FALSE(analysis.ValueIsDefinedAt(subparam1)); + EXPECT_TRUE(analysis.ValueIsDefinedAt(add)); + EXPECT_FALSE(analysis.ValueIsDefinedAt(call1)); + EXPECT_FALSE(analysis.ValueIsDefinedAt(call2)); + EXPECT_TRUE(analysis.ValueIsDefinedAt(sub)); + + EXPECT_THAT(analysis.GetValueDefinedAt(constant1).uses(), + UnorderedElementsAre(HloUse{add, 0, {}})); + EXPECT_THAT(analysis.GetValueDefinedAt(constant2).uses(), + UnorderedElementsAre(HloUse{add, 1, {}})); + // The Add from the subcomputation is used as both operands of the Subtract. + EXPECT_THAT(analysis.GetValueDefinedAt(add).uses(), + UnorderedElementsAre(HloUse{sub, 0, {}}, HloUse{sub, 1, {}})); + + EXPECT_FALSE(analysis.GetValueDefinedAt(add).live_out_of_module()); + EXPECT_TRUE(analysis.GetValueDefinedAt(add).live_out_of_computation()); + + EXPECT_TRUE(analysis.GetValueDefinedAt(sub).live_out_of_module()); + EXPECT_TRUE(analysis.GetValueDefinedAt(sub).live_out_of_computation()); +} + +TEST_P(HloDataflowAnalysisTest, ComputationCalledTwiceWithDifferentArguments) { + // Test a subcomputation which is called twice with different argument values. + auto subbuilder = HloComputation::Builder("Subcomputation"); + auto subparam0 = subbuilder.AddInstruction( + HloInstruction::CreateParameter(0, scalar_shape_, "param0")); + auto subparam1 = subbuilder.AddInstruction( + HloInstruction::CreateParameter(1, scalar_shape_, "param1")); + auto add = subbuilder.AddInstruction(HloInstruction::CreateBinary( + scalar_shape_, HloOpcode::kAdd, subparam0, subparam1)); + HloComputation* called_computation = + module_->AddEmbeddedComputation(subbuilder.Build()); + + auto builder = HloComputation::Builder(TestName()); + auto constant1 = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(1.0))); + auto constant2 = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(2.0))); + auto call1 = builder.AddInstruction(HloInstruction::CreateCall( + scalar_shape_, {constant1, constant2}, called_computation)); + auto call2 = builder.AddInstruction(HloInstruction::CreateCall( + scalar_shape_, {call1, constant2}, called_computation)); + module_->AddEntryComputation(builder.Build()); + + bool ssa_form = GetParam(); + const HloDataflowAnalysis& analysis = RunAnalysis(ssa_form); + + EXPECT_FALSE(analysis.ValueIsDefinedAt(call1)); + EXPECT_FALSE(analysis.ValueIsDefinedAt(call2)); + + EXPECT_FALSE(analysis.ValueIsDefinedAt(subparam0)); + + EXPECT_THAT(HloValuesAt(subparam0), + UnorderedElementsAre(analysis.GetValueDefinedAt(constant1), + analysis.GetValueDefinedAt(add))); + EXPECT_THAT(HloValuesAt(subparam1), + UnorderedElementsAre(analysis.GetValueDefinedAt(constant2))); + + EXPECT_TRUE(analysis.GetValueDefinedAt(add).live_out_of_module()); +} + +TEST_P(HloDataflowAnalysisTest, NestedCalls) { + // Test a module with nested computations. HLO is: + // + // F32[] inner_computation(F32[] %param0, F32[] %param1): + // %add = Add(%param0, %param1) + // + // F32[] outer_computation((F32[] %param0, F32[] %param1): + // ;; Note that parameters are interchanged in the call. + // %nested_call = Call(inner_computation, {%param1, %param0}) + // + // F32[] entry: + // %constant1 = Constant(1.0) + // %constant2 = Constant(2.0) + // %call = Call(outer_computation, {%constant1, %constant2}) + // + auto inner_builder = HloComputation::Builder("InnerComputation"); + auto inner_param0 = inner_builder.AddInstruction( + HloInstruction::CreateParameter(0, scalar_shape_, "param0")); + auto inner_param1 = inner_builder.AddInstruction( + HloInstruction::CreateParameter(1, scalar_shape_, "param1")); + auto add = inner_builder.AddInstruction(HloInstruction::CreateBinary( + scalar_shape_, HloOpcode::kAdd, inner_param0, inner_param1)); + HloComputation* inner_computation = + module_->AddEmbeddedComputation(inner_builder.Build()); + + auto outer_builder = HloComputation::Builder("OuterComputation"); + auto outer_param0 = outer_builder.AddInstruction( + HloInstruction::CreateParameter(0, scalar_shape_, "param0")); + auto outer_param1 = outer_builder.AddInstruction( + HloInstruction::CreateParameter(1, scalar_shape_, "param1")); + // Swizzle parameters. + outer_builder.AddInstruction(HloInstruction::CreateCall( + scalar_shape_, {outer_param1, outer_param0}, inner_computation)); + HloComputation* outer_computation = + module_->AddEmbeddedComputation(outer_builder.Build()); + + auto builder = HloComputation::Builder(TestName()); + auto constant1 = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(1.0))); + auto constant2 = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(2.0))); + builder.AddInstruction(HloInstruction::CreateCall( + scalar_shape_, {constant1, constant2}, outer_computation)); + module_->AddEntryComputation(builder.Build()); + + bool ssa_form = GetParam(); + const HloDataflowAnalysis& analysis = RunAnalysis(ssa_form); + + // Only three values should be defined. Most instructions just pass through + // their operand values. + EXPECT_EQ(analysis.values().size(), 3); + + // Verify that the uses of the constants are properly swizzled by parameter + // permutation in nested_call. + EXPECT_THAT(analysis.GetValueDefinedAt(constant1).uses(), + UnorderedElementsAre(HloUse{add, 1, {}})); + EXPECT_THAT(analysis.GetValueDefinedAt(constant2).uses(), + UnorderedElementsAre(HloUse{add, 0, {}})); + + EXPECT_TRUE(analysis.GetValueDefinedAt(add).live_out_of_module()); +} + +TEST_P(HloDataflowAnalysisTest, SingleWhile) { + // Test a simple single while instruction. The while body includes a + // pass-through value. HLO: + // + // body((F32[], F32[]) %tuple_param): + // %add = Add(%tuple_param{0}, %tuple_param{1}) + // return Tuple(%tuple_param{0}, %add) + // + // condition((F32[], F32[]) %tuple_param): + // return Constant(false) + // + // entry: + // %constant1 = Constant(1.0) + // %constant2 = Constant(2.0) + // %tuple = Tuple(%constant1, %constant2) + // return While(%tuple, body, condition) + // + const Shape tuple_shape = + ShapeUtil::MakeTupleShape({scalar_shape_, scalar_shape_}); + + // Element 0 passes transparently through the body. + auto body_builder = HloComputation::Builder("body"); + auto body_param = body_builder.AddInstruction( + HloInstruction::CreateParameter(0, tuple_shape, "param")); + auto body_element_0 = body_builder.AddInstruction( + HloInstruction::CreateGetTupleElement(scalar_shape_, body_param, 0)); + auto body_element_1 = body_builder.AddInstruction( + HloInstruction::CreateGetTupleElement(scalar_shape_, body_param, 1)); + auto add = body_builder.AddInstruction(HloInstruction::CreateBinary( + scalar_shape_, HloOpcode::kAdd, body_element_0, body_element_1)); + body_builder.AddInstruction( + HloInstruction::CreateTuple({body_element_0, add})); + HloComputation* body = module_->AddEmbeddedComputation(body_builder.Build()); + + // Condition computation trivially returns a constant "false". + auto cond_builder = HloComputation::Builder("condition"); + auto cond_param = cond_builder.AddInstruction( + HloInstruction::CreateParameter(0, tuple_shape, "param")); + auto cond_constant = cond_builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(false))); + HloComputation* condition = + module_->AddEmbeddedComputation(cond_builder.Build()); + + auto builder = HloComputation::Builder(TestName()); + auto constant1 = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(1.0))); + auto constant2 = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(2.0))); + auto tuple = builder.AddInstruction( + HloInstruction::CreateTuple({constant1, constant2})); + auto xla_while = builder.AddInstruction( + HloInstruction::CreateWhile(tuple_shape, condition, body, tuple)); + module_->AddEntryComputation(builder.Build()); + + bool ssa_form = GetParam(); + const HloDataflowAnalysis& analysis = RunAnalysis(ssa_form); + + EXPECT_TRUE( + analysis.GetValueDefinedAt(cond_constant).live_out_of_computation()); + EXPECT_FALSE(analysis.GetValueDefinedAt(cond_constant).live_out_of_module()); + + if (ssa_form) { + // Element 0 of the tuple passed through the body so no phi value is + // defined. + EXPECT_FALSE(analysis.ValueIsDefinedAt(xla_while, /*index=*/{0})); + EXPECT_FALSE(analysis.ValueIsDefinedAt(body_param, /*index=*/{0})); + EXPECT_FALSE(analysis.ValueIsDefinedAt(cond_param, /*index=*/{0})); + + // Element 1 of the tuple should be a phi value. + EXPECT_TRUE(analysis.ValueIsDefinedAt(xla_while, /*index=*/{1})); + EXPECT_TRUE(analysis.GetValueDefinedAt(xla_while, /*index=*/{1}).is_phi()); + EXPECT_TRUE(analysis.ValueIsDefinedAt(body_param, /*index=*/{1})); + EXPECT_TRUE(analysis.GetValueDefinedAt(body_param, /*index=*/{1}).is_phi()); + EXPECT_TRUE(analysis.ValueIsDefinedAt(cond_param, /*index=*/{1})); + EXPECT_TRUE(analysis.GetValueDefinedAt(cond_param, /*index=*/{1}).is_phi()); + + EXPECT_THAT( + analysis.GetValueDefinedAt(constant1).uses(), + UnorderedElementsAre(HloUse{add, 0, {}}, HloUse{xla_while, 0, {0}})); + + // Constant1 passes through the body and out of the module. + EXPECT_TRUE(analysis.GetValueDefinedAt(constant1).live_out_of_module()); + EXPECT_TRUE(analysis.GetValueDefinedAt(xla_while, /*index=*/{1}) + .live_out_of_module()); + + EXPECT_TRUE(analysis.GetValueDefinedAt(add).live_out_of_computation()); + EXPECT_FALSE(analysis.GetValueDefinedAt(add).live_out_of_module()); + } else { + // While instruction and subcomputation parameters should not define values + // in non-ssa form. + EXPECT_FALSE(analysis.ValueIsDefinedAt(xla_while, /*index=*/{0})); + EXPECT_FALSE(analysis.ValueIsDefinedAt(xla_while, /*index=*/{1})); + EXPECT_FALSE(analysis.ValueIsDefinedAt(body_param, /*index=*/{0})); + EXPECT_FALSE(analysis.ValueIsDefinedAt(body_param, /*index=*/{1})); + EXPECT_FALSE(analysis.ValueIsDefinedAt(cond_param, /*index=*/{0})); + EXPECT_FALSE(analysis.ValueIsDefinedAt(cond_param, /*index=*/{1})); + + EXPECT_TRUE(analysis.GetValueDefinedAt(constant1).live_out_of_module()); + EXPECT_TRUE(analysis.GetValueDefinedAt(add).live_out_of_module()); + EXPECT_TRUE(analysis.GetValueDefinedAt(add).live_out_of_computation()); + } +} + +TEST_P(HloDataflowAnalysisTest, SequentialWhiles) { + // Test sequential while instructions. The while body includes a + // pass-through value. HLO: + // + // body((F32[], F32[]) %tuple_param): + // %add = Add(%tuple_param{0}, %tuple_param{1}) + // return Tuple(%tuple_param{0}, %add) + // + // condition((F32[], F32[]) %tuple_param): + // return Constant(false) + // + // entry: + // %constant1 = Constant(1.0) + // %constant2 = Constant(2.0) + // %tuple = Tuple(%constant1, %constant2) + // %while0 = While(%tuple, body, condition) + // %while1 = While(%while0, body, condition) + // return While(%while1, body, condition) + // + const Shape tuple_shape = + ShapeUtil::MakeTupleShape({scalar_shape_, scalar_shape_}); + + // Element 0 passes transparently through the body. + auto body_builder = HloComputation::Builder("body"); + auto body_param = body_builder.AddInstruction( + HloInstruction::CreateParameter(0, tuple_shape, "param")); + auto body_element_0 = body_builder.AddInstruction( + HloInstruction::CreateGetTupleElement(scalar_shape_, body_param, 0)); + auto body_element_1 = body_builder.AddInstruction( + HloInstruction::CreateGetTupleElement(scalar_shape_, body_param, 1)); + auto add = body_builder.AddInstruction(HloInstruction::CreateBinary( + scalar_shape_, HloOpcode::kAdd, body_element_0, body_element_1)); + body_builder.AddInstruction( + HloInstruction::CreateTuple({body_element_0, add})); + HloComputation* body = module_->AddEmbeddedComputation(body_builder.Build()); + + auto cond_builder = HloComputation::Builder("condition"); + cond_builder.AddInstruction( + HloInstruction::CreateParameter(0, tuple_shape, "param")); + cond_builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(false))); + HloComputation* condition = + module_->AddEmbeddedComputation(cond_builder.Build()); + + auto builder = HloComputation::Builder(TestName()); + auto constant1 = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(1.0))); + auto constant2 = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(2.0))); + auto tuple = builder.AddInstruction( + HloInstruction::CreateTuple({constant1, constant2})); + auto xla_while0 = builder.AddInstruction( + HloInstruction::CreateWhile(tuple_shape, condition, body, tuple)); + auto xla_while1 = builder.AddInstruction( + HloInstruction::CreateWhile(tuple_shape, condition, body, xla_while0)); + auto xla_while2 = builder.AddInstruction( + HloInstruction::CreateWhile(tuple_shape, condition, body, xla_while1)); + module_->AddEntryComputation(builder.Build()); + + bool ssa_form = GetParam(); + const HloDataflowAnalysis& analysis = RunAnalysis(ssa_form); + + // Element 0 is passed through all the while instructions and out of the + // module.. + EXPECT_EQ(analysis.GetUniqueValueAt(xla_while0, /*index=*/{0}), + analysis.GetValueDefinedAt(constant1)); + EXPECT_EQ(analysis.GetUniqueValueAt(xla_while1, /*index=*/{0}), + analysis.GetValueDefinedAt(constant1)); + EXPECT_EQ(analysis.GetUniqueValueAt(xla_while2, /*index=*/{0}), + analysis.GetValueDefinedAt(constant1)); + EXPECT_TRUE(analysis.GetValueDefinedAt(constant1).live_out_of_module()); +} + +TEST_P(HloDataflowAnalysisTest, NestedWhiles) { + // Test nested while instructions. The inner body passes through element 0 of + // its parameter, and the outer body passes through element 1. HLO: + // + // inner_body((F32[], F32[]) %tuple_param): + // %add = Add(%tuple_param{0}, %tuple_param{1}) + // return Tuple(%tuple_param{0}, %add) + // + // outer_body((F32[], F32[]) %tuple_param): + // %negate = Negate(%tuple_param{0}) + // %tuple = Tuple(%negate, %tuple_param{1}) + // return While(%tuple, inner_body, condition) + // + // entry: + // %constant1 = Constant(1.0) + // %constant2 = Constant(2.0) + // %tuple = Tuple(%constant1, %constant2) + // return While(%tuple, outer_body, condition) + // + const Shape tuple_shape = + ShapeUtil::MakeTupleShape({scalar_shape_, scalar_shape_}); + + auto cond_builder = HloComputation::Builder("condition"); + cond_builder.AddInstruction( + HloInstruction::CreateParameter(0, tuple_shape, "param")); + cond_builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(false))); + HloComputation* condition = + module_->AddEmbeddedComputation(cond_builder.Build()); + + // Element 0 passes transparently through the body. + auto inner_builder = HloComputation::Builder("inner_body"); + auto inner_param = inner_builder.AddInstruction( + HloInstruction::CreateParameter(0, tuple_shape, "param")); + auto inner_element_0 = inner_builder.AddInstruction( + HloInstruction::CreateGetTupleElement(scalar_shape_, inner_param, 0)); + auto inner_element_1 = inner_builder.AddInstruction( + HloInstruction::CreateGetTupleElement(scalar_shape_, inner_param, 1)); + auto add = inner_builder.AddInstruction(HloInstruction::CreateBinary( + scalar_shape_, HloOpcode::kAdd, inner_element_0, inner_element_1)); + inner_builder.AddInstruction( + HloInstruction::CreateTuple({inner_element_0, add})); + HloComputation* inner_body = + module_->AddEmbeddedComputation(inner_builder.Build()); + + // Element 1 passes transparently through the body. + auto outer_builder = HloComputation::Builder("outer_body"); + auto outer_param = outer_builder.AddInstruction( + HloInstruction::CreateParameter(0, tuple_shape, "param")); + auto outer_element_0 = outer_builder.AddInstruction( + HloInstruction::CreateGetTupleElement(scalar_shape_, outer_param, 0)); + auto negate = outer_builder.AddInstruction(HloInstruction::CreateUnary( + scalar_shape_, HloOpcode::kNegate, outer_element_0)); + auto outer_element_1 = outer_builder.AddInstruction( + HloInstruction::CreateGetTupleElement(scalar_shape_, outer_param, 1)); + auto outer_tuple = outer_builder.AddInstruction( + HloInstruction::CreateTuple({negate, outer_element_1})); + auto nested_while = outer_builder.AddInstruction(HloInstruction::CreateWhile( + tuple_shape, condition, inner_body, outer_tuple)); + HloComputation* outer_body = + module_->AddEmbeddedComputation(outer_builder.Build()); + + auto builder = HloComputation::Builder(TestName()); + auto constant1 = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(1.0))); + auto constant2 = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(2.0))); + auto tuple = builder.AddInstruction( + HloInstruction::CreateTuple({constant1, constant2})); + auto entry_while = builder.AddInstruction( + HloInstruction::CreateWhile(tuple_shape, condition, outer_body, tuple)); + module_->AddEntryComputation(builder.Build()); + + bool ssa_form = GetParam(); + const HloDataflowAnalysis& analysis = RunAnalysis(ssa_form); + + EXPECT_THAT(HloValuesAt(inner_param, /*index=*/{0}), + UnorderedElementsAre(analysis.GetValueDefinedAt(negate))); + if (ssa_form) { + EXPECT_TRUE(analysis.ValueIsDefinedAt(inner_param, /*index=*/{1})); + EXPECT_TRUE( + analysis.GetValueDefinedAt(inner_param, /*index=*/{1}).is_phi()); + + // Element 0 of the nested while is %negate. + EXPECT_FALSE(analysis.ValueIsDefinedAt(nested_while, /*index=*/{0})); + EXPECT_THAT(HloValuesAt(inner_param, /*index=*/{0}), + UnorderedElementsAre(analysis.GetValueDefinedAt(negate))); + // Element 1 is a phi value (join of %add and %constant2). + EXPECT_TRUE(analysis.ValueIsDefinedAt(nested_while, /*index=*/{1})); + EXPECT_TRUE( + analysis.GetValueDefinedAt(nested_while, /*index=*/{1}).is_phi()); + + EXPECT_TRUE(analysis.ValueIsDefinedAt(entry_while, /*index=*/{0})); + EXPECT_TRUE( + analysis.GetValueDefinedAt(entry_while, /*index=*/{0}).is_phi()); + + EXPECT_TRUE(analysis.ValueIsDefinedAt(entry_while, /*index=*/{1})); + EXPECT_TRUE( + analysis.GetValueDefinedAt(entry_while, /*index=*/{1}).is_phi()); + } else { + EXPECT_THAT(HloValuesAt(inner_param, /*index=*/{1}), + UnorderedElementsAre(analysis.GetValueDefinedAt(add), + analysis.GetValueDefinedAt(constant2))); + + EXPECT_THAT(HloValuesAt(nested_while, /*index=*/{0}), + UnorderedElementsAre(analysis.GetValueDefinedAt(negate))); + EXPECT_THAT(HloValuesAt(nested_while, /*index=*/{1}), + UnorderedElementsAre(analysis.GetValueDefinedAt(add), + analysis.GetValueDefinedAt(constant2))); + + EXPECT_THAT(HloValuesAt(entry_while, /*index=*/{0}), + UnorderedElementsAre(analysis.GetValueDefinedAt(negate), + analysis.GetValueDefinedAt(constant1))); + EXPECT_THAT(HloValuesAt(entry_while, /*index=*/{1}), + UnorderedElementsAre(analysis.GetValueDefinedAt(add), + analysis.GetValueDefinedAt(constant2))); + } +} + +TEST_P(HloDataflowAnalysisTest, SwizzlingWhile) { + // Test a while instruction with a body which permutes it's tuple parameter + // elements. HLO: + // + // body((F32[], F32[]) %tuple_param): + // return Tuple(%tuple_param{1}, %tuple_param{0}) + // + // condition((F32[], F32[]) %tuple_param): + // return Constant(false) + // + // entry: + // %constant1 = Constant(1.0) + // %constant2 = Constant(2.0) + // %tuple = Tuple(%constant1, %constant2) + // return While(%tuple, body, condition) + // + const Shape tuple_shape = + ShapeUtil::MakeTupleShape({scalar_shape_, scalar_shape_}); + + auto body_builder = HloComputation::Builder("body"); + auto body_param = body_builder.AddInstruction( + HloInstruction::CreateParameter(0, tuple_shape, "param")); + auto body_element_0 = body_builder.AddInstruction( + HloInstruction::CreateGetTupleElement(scalar_shape_, body_param, 0)); + auto body_element_1 = body_builder.AddInstruction( + HloInstruction::CreateGetTupleElement(scalar_shape_, body_param, 1)); + body_builder.AddInstruction( + HloInstruction::CreateTuple({body_element_1, body_element_0})); + HloComputation* body = module_->AddEmbeddedComputation(body_builder.Build()); + + auto cond_builder = HloComputation::Builder("condition"); + auto cond_param = cond_builder.AddInstruction( + HloInstruction::CreateParameter(0, tuple_shape, "param")); + cond_builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(false))); + HloComputation* condition = + module_->AddEmbeddedComputation(cond_builder.Build()); + + auto builder = HloComputation::Builder(TestName()); + auto constant1 = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(1.0))); + auto constant2 = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(2.0))); + auto tuple = builder.AddInstruction( + HloInstruction::CreateTuple({constant1, constant2})); + auto xla_while = builder.AddInstruction( + HloInstruction::CreateWhile(tuple_shape, condition, body, tuple)); + module_->AddEntryComputation(builder.Build()); + + bool ssa_form = GetParam(); + const HloDataflowAnalysis& analysis = RunAnalysis(ssa_form); + + if (ssa_form) { + // Element 0 and 1 in the while should both be phi values. + EXPECT_TRUE(analysis.ValueIsDefinedAt(body_param, /*index=*/{0})); + EXPECT_TRUE(analysis.GetValueDefinedAt(body_param, /*index=*/{0}).is_phi()); + EXPECT_TRUE(analysis.ValueIsDefinedAt(body_param, /*index=*/{1})); + EXPECT_TRUE(analysis.GetValueDefinedAt(body_param, /*index=*/{1}).is_phi()); + + EXPECT_TRUE(analysis.ValueIsDefinedAt(xla_while, /*index=*/{0})); + EXPECT_TRUE(analysis.GetValueDefinedAt(xla_while, /*index=*/{0}).is_phi()); + EXPECT_TRUE(analysis.ValueIsDefinedAt(xla_while, /*index=*/{1})); + EXPECT_TRUE(analysis.GetValueDefinedAt(xla_while, /*index=*/{1}).is_phi()); + + EXPECT_TRUE(analysis.ValueIsDefinedAt(cond_param, /*index=*/{0})); + EXPECT_TRUE(analysis.GetValueDefinedAt(cond_param, /*index=*/{0}).is_phi()); + EXPECT_TRUE(analysis.ValueIsDefinedAt(cond_param, /*index=*/{1})); + EXPECT_TRUE(analysis.GetValueDefinedAt(cond_param, /*index=*/{1}).is_phi()); + + EXPECT_FALSE(analysis.GetValueDefinedAt(constant1).live_out_of_module()); + EXPECT_FALSE(analysis.GetValueDefinedAt(constant2).live_out_of_module()); + EXPECT_TRUE(analysis.GetValueDefinedAt(xla_while, /*index=*/{}) + .live_out_of_module()); + EXPECT_TRUE(analysis.GetValueDefinedAt(xla_while, /*index=*/{0}) + .live_out_of_module()); + EXPECT_TRUE(analysis.GetValueDefinedAt(xla_while, /*index=*/{1}) + .live_out_of_module()); + } else { + // Elements 0 and 1 have both constants as reaching definitions. + EXPECT_THAT(HloValuesAt(xla_while, /*index=*/{0}), + UnorderedElementsAre(analysis.GetValueDefinedAt(constant1), + analysis.GetValueDefinedAt(constant2))); + EXPECT_THAT(HloValuesAt(xla_while, /*index=*/{1}), + UnorderedElementsAre(analysis.GetValueDefinedAt(constant1), + analysis.GetValueDefinedAt(constant2))); + EXPECT_TRUE(analysis.GetValueDefinedAt(constant1).live_out_of_module()); + EXPECT_TRUE(analysis.GetValueDefinedAt(constant2).live_out_of_module()); + } +} + +TEST_P(HloDataflowAnalysisTest, ArraySelect) { + // Test a kSelect of an array value. + auto builder = HloComputation::Builder(TestName()); + auto pred = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(false))); + auto constant1 = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(1.0))); + auto constant2 = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(2.0))); + auto select = builder.AddInstruction(HloInstruction::CreateTernary( + scalar_shape_, HloOpcode::kSelect, pred, constant1, constant2)); + + module_->AddEntryComputation(builder.Build()); + + bool ssa_form = GetParam(); + const HloDataflowAnalysis& analysis = RunAnalysis(ssa_form); + + EXPECT_TRUE(analysis.ValueIsDefinedAt(select)); + EXPECT_FALSE(analysis.GetValueDefinedAt(constant1).live_out_of_module()); + EXPECT_FALSE(analysis.GetValueDefinedAt(constant2).live_out_of_module()); + EXPECT_TRUE(analysis.GetValueDefinedAt(select).live_out_of_module()); +} + +TEST_P(HloDataflowAnalysisTest, TupleSelect) { + // Test a kSelect of a tuple value. Non-top-level element flow through the + // instruction. + auto builder = HloComputation::Builder(TestName()); + auto pred = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(false))); + auto constant1 = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(1.0))); + auto constant2 = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(2.0))); + auto constant3 = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(3.0))); + auto constant4 = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(4.0))); + auto tuple1 = + builder.AddInstruction(HloInstruction::CreateTuple({constant1})); + auto tuple2 = + builder.AddInstruction(HloInstruction::CreateTuple({constant2})); + auto tuple3 = + builder.AddInstruction(HloInstruction::CreateTuple({constant3})); + auto tuple4 = + builder.AddInstruction(HloInstruction::CreateTuple({constant4})); + const Shape tuple_shape = tuple1->shape(); + auto select11 = builder.AddInstruction(HloInstruction::CreateTernary( + tuple_shape, HloOpcode::kSelect, pred, tuple1, tuple1)); + auto select12 = builder.AddInstruction(HloInstruction::CreateTernary( + tuple_shape, HloOpcode::kSelect, pred, tuple1, tuple2)); + auto select34 = builder.AddInstruction(HloInstruction::CreateTernary( + tuple_shape, HloOpcode::kSelect, pred, tuple3, tuple4)); + auto select1234 = builder.AddInstruction(HloInstruction::CreateTernary( + tuple_shape, HloOpcode::kSelect, pred, select12, select34)); + + module_->AddEntryComputation(builder.Build()); + + bool ssa_form = GetParam(); + const HloDataflowAnalysis& analysis = RunAnalysis(ssa_form); + + // Top-level value is always defined by a kSelect. + EXPECT_TRUE(analysis.ValueIsDefinedAt(select11)); + EXPECT_TRUE(analysis.ValueIsDefinedAt(select12)); + EXPECT_TRUE(analysis.ValueIsDefinedAt(select34)); + EXPECT_TRUE(analysis.ValueIsDefinedAt(select1234)); + + EXPECT_FALSE(analysis.ValueIsDefinedAt(select11, /*index=*/{0})); + EXPECT_FALSE(analysis.ValueIsDefinedAt(select12, /*index=*/{0})); + EXPECT_FALSE(analysis.ValueIsDefinedAt(select34, /*index=*/{0})); + EXPECT_FALSE(analysis.ValueIsDefinedAt(select1234, /*index=*/{0})); + + EXPECT_THAT(HloValuesAt(select11, /*index=*/{0}), + UnorderedElementsAre(analysis.GetValueDefinedAt(constant1))); + EXPECT_THAT(HloValuesAt(select12, /*index=*/{0}), + UnorderedElementsAre(analysis.GetValueDefinedAt(constant1), + analysis.GetValueDefinedAt(constant2))); + EXPECT_THAT(HloValuesAt(select34, /*index=*/{0}), + UnorderedElementsAre(analysis.GetValueDefinedAt(constant3), + analysis.GetValueDefinedAt(constant4))); + EXPECT_THAT(HloValuesAt(select1234, /*index=*/{0}), + UnorderedElementsAre(analysis.GetValueDefinedAt(constant1), + analysis.GetValueDefinedAt(constant2), + analysis.GetValueDefinedAt(constant3), + analysis.GetValueDefinedAt(constant4))); + + EXPECT_THAT( + analysis.GetValueDefinedAt(tuple1, /*index=*/{}).uses(), + UnorderedElementsAre(HloUse{select11, 1, {}}, HloUse{select11, 2, {}}, + HloUse{select12, 1, {}})); + + // The two constant values just pass through the Selects and are not + // used. They are live out however. + EXPECT_TRUE(analysis.GetValueDefinedAt(constant1).uses().empty()); + EXPECT_TRUE(analysis.GetValueDefinedAt(constant2).uses().empty()); + EXPECT_TRUE(analysis.GetValueDefinedAt(constant1).live_out_of_module()); + EXPECT_TRUE(analysis.GetValueDefinedAt(constant2).live_out_of_module()); +} + +TEST_P(HloDataflowAnalysisTest, NestedTupleSelect) { + // Test kSelect of a nested tuple. + auto builder = HloComputation::Builder(TestName()); + auto pred = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(false))); + auto constant1 = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(1.0))); + auto constant2 = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(2.0))); + auto constant3 = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(3.0))); + auto constant4 = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(4.0))); + auto constant5 = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(5.0))); + auto inner_tuple1 = builder.AddInstruction( + HloInstruction::CreateTuple({constant2, constant3})); + auto tuple1 = builder.AddInstruction( + HloInstruction::CreateTuple({constant1, inner_tuple1})); + auto inner_tuple2 = builder.AddInstruction( + HloInstruction::CreateTuple({constant5, constant3})); + auto tuple2 = builder.AddInstruction( + HloInstruction::CreateTuple({constant4, inner_tuple2})); + auto select = builder.AddInstruction(HloInstruction::CreateTernary( + tuple1->shape(), HloOpcode::kSelect, pred, tuple1, tuple2)); + + module_->AddEntryComputation(builder.Build()); + + bool ssa_form = GetParam(); + const HloDataflowAnalysis& analysis = RunAnalysis(ssa_form); + + EXPECT_TRUE(analysis.ValueIsDefinedAt(select)); + + EXPECT_THAT(HloValuesAt(select, /*index=*/{0}), + UnorderedElementsAre(analysis.GetValueDefinedAt(constant1), + analysis.GetValueDefinedAt(constant4))); + EXPECT_THAT(HloValuesAt(select, /*index=*/{1}), + UnorderedElementsAre(analysis.GetValueDefinedAt(inner_tuple1), + analysis.GetValueDefinedAt(inner_tuple2))); + EXPECT_THAT(HloValuesAt(select, /*index=*/{1, 0}), + UnorderedElementsAre(analysis.GetValueDefinedAt(constant2), + analysis.GetValueDefinedAt(constant5))); + EXPECT_THAT(HloValuesAt(select, /*index=*/{1, 1}), + UnorderedElementsAre(analysis.GetValueDefinedAt(constant3))); +} + +TEST_P(HloDataflowAnalysisTest, TupleSelectToWhile) { + // Test a tuple-shaped kSelect feeding a kWhile instruction. HLO: + // + // body((F32[], F32[]) %tuple_param): + // %add = Add(%tuple_param{0}, %tuple_param{1}) + // return Tuple(%tuple_param{0}, %add) + // + // condition((F32[], F32[]) %tuple_param): + // return Constant(false) + // + // entry: + // %constant1 = Constant(1.0) + // %constant2 = Constant(2.0) + // %constant3 = Constant(3.0) + // %tuple1 = Tuple(%constant1) + // %tuple2 = Tuple(%constant2) + // %select = Select(%tuple1, %tuple2) + // %gte = GetTupleElement(%select, 0) + // %tuple = Tuple(%gte, %constant3) + // return While(%tuple, body, condition) + // + auto builder = HloComputation::Builder(TestName()); + + const Shape tuple_shape = + ShapeUtil::MakeTupleShape({scalar_shape_, scalar_shape_}); + + // Element 0 passes transparently through the body. + auto body_builder = HloComputation::Builder("body"); + auto body_param = body_builder.AddInstruction( + HloInstruction::CreateParameter(0, tuple_shape, "param")); + auto body_element_0 = body_builder.AddInstruction( + HloInstruction::CreateGetTupleElement(scalar_shape_, body_param, 0)); + auto body_element_1 = body_builder.AddInstruction( + HloInstruction::CreateGetTupleElement(scalar_shape_, body_param, 1)); + auto add = body_builder.AddInstruction(HloInstruction::CreateBinary( + scalar_shape_, HloOpcode::kAdd, body_element_0, body_element_1)); + body_builder.AddInstruction( + HloInstruction::CreateTuple({body_element_0, add})); + HloComputation* body = module_->AddEmbeddedComputation(body_builder.Build()); + + auto cond_builder = HloComputation::Builder("condition"); + cond_builder.AddInstruction( + HloInstruction::CreateParameter(0, tuple_shape, "param")); + cond_builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(false))); + HloComputation* condition = + module_->AddEmbeddedComputation(cond_builder.Build()); + + auto pred = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(false))); + auto constant1 = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(1.0))); + auto constant2 = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(2.0))); + auto constant3 = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(3.0))); + auto tuple1 = + builder.AddInstruction(HloInstruction::CreateTuple({constant1})); + auto tuple2 = + builder.AddInstruction(HloInstruction::CreateTuple({constant2})); + auto select = builder.AddInstruction(HloInstruction::CreateTernary( + tuple1->shape(), HloOpcode::kSelect, pred, tuple1, tuple2)); + auto gte = builder.AddInstruction( + HloInstruction::CreateGetTupleElement(scalar_shape_, select, 0)); + auto tuple = + builder.AddInstruction(HloInstruction::CreateTuple({gte, constant3})); + auto xla_while = builder.AddInstruction( + HloInstruction::CreateWhile(tuple->shape(), condition, body, tuple)); + + module_->AddEntryComputation(builder.Build()); + + bool ssa_form = GetParam(); + const HloDataflowAnalysis& analysis = RunAnalysis(ssa_form); + + if (ssa_form) { + EXPECT_TRUE(analysis.ValueIsDefinedAt(xla_while, /*index=*/{0})); + EXPECT_TRUE(analysis.GetValueDefinedAt(xla_while, /*index=*/{0}).is_phi()); + EXPECT_TRUE(analysis.ValueIsDefinedAt(xla_while, /*index=*/{1})); + EXPECT_TRUE(analysis.GetValueDefinedAt(xla_while, /*index=*/{1}).is_phi()); + + EXPECT_FALSE(analysis.ValueIsDefinedAt(select, /*index=*/{0})); + + EXPECT_FALSE(analysis.GetValueDefinedAt(constant1).live_out_of_module()); + EXPECT_FALSE(analysis.GetValueDefinedAt(constant2).live_out_of_module()); + EXPECT_FALSE(analysis.GetValueDefinedAt(constant3).live_out_of_module()); + EXPECT_TRUE(analysis.GetValueDefinedAt(xla_while, /*index=*/{1}) + .live_out_of_module()); + } else { + EXPECT_THAT(HloValuesAt(gte), + UnorderedElementsAre(analysis.GetValueDefinedAt(constant1), + analysis.GetValueDefinedAt(constant2))); + EXPECT_THAT(HloValuesAt(xla_while, /*index=*/{0}), + UnorderedElementsAre(analysis.GetValueDefinedAt(constant1), + analysis.GetValueDefinedAt(constant2))); + EXPECT_THAT(HloValuesAt(xla_while, /*index=*/{1}), + UnorderedElementsAre(analysis.GetValueDefinedAt(add), + analysis.GetValueDefinedAt(constant3))); + EXPECT_TRUE(analysis.GetValueDefinedAt(constant1).live_out_of_module()); + EXPECT_TRUE(analysis.GetValueDefinedAt(constant2).live_out_of_module()); + EXPECT_TRUE(analysis.GetValueDefinedAt(constant3).live_out_of_module()); + } +} + +TEST_P(HloDataflowAnalysisTest, BitcastDefinesValue) { + // Test the bitcast_defines_value flag to the dataflow analysis. + auto builder = HloComputation::Builder(TestName()); + auto constant = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(1.0))); + auto bitcast = builder.AddInstruction(HloInstruction::CreateUnary( + scalar_shape_, HloOpcode::kBitcast, constant)); + + module_->AddEntryComputation(builder.Build()); + + bool ssa_form = GetParam(); + { + const HloDataflowAnalysis& analysis = + RunAnalysis(ssa_form, /*bitcast_defines_value=*/true); + + EXPECT_EQ(analysis.values().size(), 2); + + EXPECT_TRUE(analysis.ValueIsDefinedAt(constant)); + EXPECT_TRUE(analysis.ValueIsDefinedAt(bitcast)); + EXPECT_FALSE(analysis.GetValueDefinedAt(constant).live_out_of_module()); + EXPECT_TRUE(analysis.GetValueDefinedAt(bitcast).live_out_of_module()); + } + { + const HloDataflowAnalysis& analysis = + RunAnalysis(ssa_form, /*bitcast_defines_value=*/false); + EXPECT_EQ(analysis.values().size(), 1); + + EXPECT_TRUE(analysis.ValueIsDefinedAt(constant)); + EXPECT_FALSE(analysis.ValueIsDefinedAt(bitcast)); + EXPECT_TRUE(analysis.GetValueDefinedAt(constant).live_out_of_module()); + } +} + +TEST_P(HloDataflowAnalysisTest, TupleCopy) { + // Test that a tuple-shaped copy only copies (defines) the top-level value. + auto builder = HloComputation::Builder(TestName()); + auto param0 = builder.AddInstruction( + HloInstruction::CreateParameter(0, scalar_shape_, "param0")); + auto param1 = builder.AddInstruction( + HloInstruction::CreateParameter(1, scalar_shape_, "param1")); + auto tuple = + builder.AddInstruction(HloInstruction::CreateTuple({param0, param1})); + auto copy = builder.AddInstruction( + HloInstruction::CreateUnary(tuple->shape(), HloOpcode::kCopy, tuple)); + module_->AddEntryComputation(builder.Build()); + + bool ssa_form = GetParam(); + const HloDataflowAnalysis& analysis = RunAnalysis(ssa_form); + + EXPECT_EQ(analysis.values().size(), 4); + + EXPECT_TRUE(analysis.ValueIsDefinedAt(param0)); + EXPECT_TRUE(analysis.ValueIsDefinedAt(param1)); + EXPECT_TRUE(analysis.ValueIsDefinedAt(tuple, /*index=*/{})); + EXPECT_FALSE(analysis.ValueIsDefinedAt(tuple, /*index=*/{0})); + EXPECT_FALSE(analysis.ValueIsDefinedAt(tuple, /*index=*/{1})); + EXPECT_TRUE(analysis.ValueIsDefinedAt(copy, /*index=*/{})); + EXPECT_FALSE(analysis.ValueIsDefinedAt(copy, /*index=*/{0})); + EXPECT_FALSE(analysis.ValueIsDefinedAt(copy, /*index=*/{1})); + + EXPECT_THAT(HloValuesAt(copy, /*index=*/{0}), + UnorderedElementsAre(analysis.GetValueDefinedAt(param0))); + EXPECT_THAT(HloValuesAt(copy, /*index=*/{1}), + UnorderedElementsAre(analysis.GetValueDefinedAt(param1))); + EXPECT_TRUE( + analysis.GetValueDefinedAt(copy, /*index=*/{}).live_out_of_module()); +} + +TEST_P(HloDataflowAnalysisTest, ElementwiseChainInterference) { + // A simple chain of elementwise operations. No values should interfere. + // + // param --> negate -> exp -> log + // + auto builder = HloComputation::Builder(TestName()); + auto param = builder.AddInstruction( + HloInstruction::CreateParameter(0, vector_shape_, "param")); + auto negate = builder.AddInstruction( + HloInstruction::CreateUnary(vector_shape_, HloOpcode::kNegate, param)); + auto exp = builder.AddInstruction( + HloInstruction::CreateUnary(vector_shape_, HloOpcode::kExp, negate)); + auto log = builder.AddInstruction( + HloInstruction::CreateUnary(vector_shape_, HloOpcode::kLog, exp)); + + module_->AddEntryComputation(builder.Build()); + RunAnalysis(GetParam()); + + DependencyHloOrdering ordering(module_.get()); + + // No values should interfere. + EXPECT_FALSE(InstructionsMayInterfere(ordering, param, negate)); + EXPECT_FALSE(InstructionsMayInterfere(ordering, param, exp)); + EXPECT_FALSE(InstructionsMayInterfere(ordering, param, log)); + EXPECT_FALSE(InstructionsMayInterfere(ordering, negate, exp)); + EXPECT_FALSE(InstructionsMayInterfere(ordering, negate, log)); + EXPECT_FALSE(InstructionsMayInterfere(ordering, exp, negate)); + EXPECT_FALSE(InstructionsMayInterfere(ordering, exp, log)); + EXPECT_FALSE(InstructionsMayInterfere(ordering, log, negate)); + EXPECT_FALSE(InstructionsMayInterfere(ordering, log, exp)); + + // Values should interfere with itself. + EXPECT_TRUE(InstructionsMayInterfere(ordering, exp, exp)); +} + +TEST_P(HloDataflowAnalysisTest, MultipleEntryParameters_Sequential) { + // Two entry params, which interfere with each other. + // + // param0 --> negate ---------------\ + // param1 --> exp --> add + auto builder = HloComputation::Builder(TestName()); + auto param0 = builder.AddInstruction( + HloInstruction::CreateParameter(0, vector_shape_, "param0")); + auto param1 = builder.AddInstruction( + HloInstruction::CreateParameter(1, vector_shape_, "param1")); + auto negate = builder.AddInstruction( + HloInstruction::CreateUnary(vector_shape_, HloOpcode::kNegate, param0)); + auto exp = builder.AddInstruction( + HloInstruction::CreateUnary(vector_shape_, HloOpcode::kExp, param1)); + auto add = builder.AddInstruction(HloInstruction::CreateBinary( + vector_shape_, HloOpcode::kAdd, negate, exp)); + + auto entry = module_->AddEntryComputation(builder.Build()); + RunAnalysis(GetParam()); + + SequentialHloOrdering::HloModuleSequence sequence; + sequence.insert({entry, {param0, negate, param1, exp, add}}); + SequentialHloOrdering ordering(module_.get(), sequence); + + // Entry parameters interfere as if they are defined simultaneously at + // the very beginning. + EXPECT_TRUE(InstructionsMayInterfere(ordering, param0, param1)); + EXPECT_FALSE(InstructionsMayInterfere(ordering, param0, negate)); + EXPECT_FALSE(InstructionsMayInterfere(ordering, param0, exp)); + EXPECT_FALSE(InstructionsMayInterfere(ordering, param0, add)); + EXPECT_TRUE(InstructionsMayInterfere(ordering, param1, param0)); + EXPECT_TRUE(InstructionsMayInterfere(ordering, param1, negate)); + EXPECT_FALSE(InstructionsMayInterfere(ordering, param1, exp)); + EXPECT_FALSE(InstructionsMayInterfere(ordering, param1, add)); + + // Negate and exp still interfere. + EXPECT_TRUE(InstructionsMayInterfere(ordering, negate, exp)); + EXPECT_TRUE(InstructionsMayInterfere(ordering, exp, negate)); + + // But {negate, add} and {exp, add} don't interfere. + EXPECT_FALSE(InstructionsMayInterfere(ordering, negate, add)); + EXPECT_FALSE(InstructionsMayInterfere(ordering, add, negate)); + EXPECT_FALSE(InstructionsMayInterfere(ordering, exp, add)); + EXPECT_FALSE(InstructionsMayInterfere(ordering, add, exp)); +} + +TEST_P(HloDataflowAnalysisTest, WhileParameters_Sequential) { + // Similar to MultipleEntryParameters_Sequential, but the parameter is of + // while body computation. Body computation in the sequential order: + // + // %constant = Constant(...) + // %exp = Exp(%constant) + // %param = Param(0) + // %add = Add(%param, %exp) ;; Root of body + // %dead_constant = Constant(...) + // %dead_negate = Negate(%dead_constant) + // + // %constant and its only use %exp are ordered before 'param'. However, the + // %constant and %param values still interfere because the parameter is + // considered live into the while body. + // + // Similarly, %dead_constant and %dead_negate are ordered after the root of + // the body computation %add. However, %add is liveout of the computation so + // %dead_constant and %add interfere. + auto body_builder = HloComputation::Builder(TestName()); + auto body_param = body_builder.AddInstruction( + HloInstruction::CreateParameter(0, scalar_shape_, "body_param")); + auto constant = body_builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(1.0))); + auto exp = body_builder.AddInstruction( + HloInstruction::CreateUnary(scalar_shape_, HloOpcode::kExp, constant)); + auto add = body_builder.AddInstruction(HloInstruction::CreateBinary( + scalar_shape_, HloOpcode::kAdd, exp, body_param)); + auto dead_constant = body_builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(1.0))); + auto dead_negate = body_builder.AddInstruction(HloInstruction::CreateUnary( + scalar_shape_, HloOpcode::kNegate, dead_constant)); + HloComputation* body = module_->AddEmbeddedComputation( + body_builder.Build(/*root_instruction=*/add)); + + auto cond_builder = HloComputation::Builder("condition"); + auto cond_param = cond_builder.AddInstruction( + HloInstruction::CreateParameter(0, scalar_shape_, "cond_param")); + auto cond_constant = cond_builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(false))); + HloComputation* condition = + module_->AddEmbeddedComputation(cond_builder.Build()); + + auto builder = HloComputation::Builder(TestName()); + auto param = builder.AddInstruction( + HloInstruction::CreateParameter(0, scalar_shape_, "param")); + auto xla_while = builder.AddInstruction( + HloInstruction::CreateWhile(scalar_shape_, condition, body, param)); + + auto entry = module_->AddEntryComputation(builder.Build()); + bool ssa_form = GetParam(); + const HloDataflowAnalysis& analysis = RunAnalysis(ssa_form); + + SequentialHloOrdering::HloModuleSequence sequence; + sequence.insert({entry, {param, xla_while}}); + sequence.insert({condition, {cond_param, cond_constant}}); + // Construct the order such that 'constant' and its use 'exp' are before + // body_param. + sequence.insert({body, {constant, exp, body_param, add}}); + + SequentialHloOrdering ordering(module_.get(), sequence); + + // 'add' is the body root even though later instructions follow in the order + // like 'dead_negate'. Only 'add' should be live out of the computation. + EXPECT_TRUE(analysis.GetValueDefinedAt(add).live_out_of_computation()); + EXPECT_FALSE( + analysis.GetValueDefinedAt(dead_negate).live_out_of_computation()); + + // 'add' is live out of the body and will interfere with an later instructions + // such as 'dead_constant' and 'dead_negate'. + EXPECT_TRUE(InstructionsMayInterfere(ordering, add, dead_constant)); + EXPECT_TRUE(InstructionsMayInterfere(ordering, add, dead_negate)); + + // The remaining checks test phi values defined by body and condition + // parameters which only occur in the SSA form of the analysis. + if (ssa_form) { + // Though the ordering suggests 'constant' and 'param' should not interfere, + // 'param' is live in and thus interferes with any earlier instruction of + // the computation in the order (eg 'constant')' + EXPECT_TRUE(InstructionsMayInterfere(ordering, body_param, constant)); + EXPECT_TRUE(InstructionsMayInterfere(ordering, body_param, exp)); + EXPECT_FALSE(InstructionsMayInterfere(ordering, body_param, add)); + + // The following values end up in the same buffer: + // (1) the init value: 'param' + // (2) the body parameter: 'body_param' + // (3) the condition parameter: 'cond_param' + // (4) the root value of the while body: 'add' + // (5) the while value: 'xla_while' + // None should interfere. + EXPECT_FALSE(InstructionsMayInterfere(ordering, param, body_param)); + EXPECT_FALSE(InstructionsMayInterfere(ordering, param, cond_param)); + EXPECT_FALSE(InstructionsMayInterfere(ordering, param, add)); + EXPECT_FALSE(InstructionsMayInterfere(ordering, param, xla_while)); + + EXPECT_FALSE(InstructionsMayInterfere(ordering, body_param, cond_param)); + EXPECT_FALSE(InstructionsMayInterfere(ordering, body_param, add)); + EXPECT_FALSE(InstructionsMayInterfere(ordering, body_param, xla_while)); + + EXPECT_FALSE(InstructionsMayInterfere(ordering, cond_param, add)); + EXPECT_FALSE(InstructionsMayInterfere(ordering, cond_param, xla_while)); + + EXPECT_FALSE(InstructionsMayInterfere(ordering, add, xla_while)); + } +} + +TEST_P(HloDataflowAnalysisTest, NonElementwiseOperand) { + // A chain of operations with two elementwise and one non-elementwise. The + // elementwise op should not interfere with its operand, while the + // non-elementwise op should interfere. Entry params always interfere. + // + // param --> exp -> negate -> reverse + // + auto builder = HloComputation::Builder(TestName()); + auto param = builder.AddInstruction( + HloInstruction::CreateParameter(0, vector_shape_, "param")); + auto exp = builder.AddInstruction( + HloInstruction::CreateUnary(vector_shape_, HloOpcode::kExp, param)); + auto negate = builder.AddInstruction( + HloInstruction::CreateUnary(vector_shape_, HloOpcode::kNegate, exp)); + auto reverse = builder.AddInstruction( + HloInstruction::CreateReverse(vector_shape_, negate, {0})); + + module_->AddEntryComputation(builder.Build()); + RunAnalysis(GetParam()); + + DependencyHloOrdering ordering(module_.get()); + + EXPECT_FALSE(InstructionsMayInterfere(ordering, param, exp)); + EXPECT_FALSE(InstructionsMayInterfere(ordering, param, negate)); + EXPECT_FALSE(InstructionsMayInterfere(ordering, param, reverse)); + + // Negate is elementwise, so doesn't interfere with its operand. + // Reverse is non-elementwise, so does interfere with its operand. + EXPECT_FALSE(InstructionsMayInterfere(ordering, exp, negate)); + EXPECT_TRUE(InstructionsMayInterfere(ordering, negate, reverse)); +} + +TEST_P(HloDataflowAnalysisTest, OverlappedValues) { + // Verify simultaneously live values interfere (exp and negate). + // + // param --> negate -> add + // \---> exp -----/ + // + auto builder = HloComputation::Builder(TestName()); + auto param = builder.AddInstruction( + HloInstruction::CreateParameter(0, vector_shape_, "param")); + auto negate = builder.AddInstruction( + HloInstruction::CreateUnary(vector_shape_, HloOpcode::kNegate, param)); + auto exp = builder.AddInstruction( + HloInstruction::CreateUnary(vector_shape_, HloOpcode::kExp, param)); + auto add = builder.AddInstruction(HloInstruction::CreateBinary( + vector_shape_, HloOpcode::kAdd, negate, exp)); + + module_->AddEntryComputation(builder.Build()); + RunAnalysis(GetParam()); + + DependencyHloOrdering ordering(module_.get()); + + EXPECT_TRUE(InstructionsMayInterfere(ordering, param, negate)); + EXPECT_TRUE(InstructionsMayInterfere(ordering, param, exp)); + EXPECT_FALSE(InstructionsMayInterfere(ordering, param, add)); + + // Negate and exp interfere with each other, but not with add. + EXPECT_TRUE(InstructionsMayInterfere(ordering, negate, exp)); + EXPECT_TRUE(InstructionsMayInterfere(ordering, exp, negate)); + EXPECT_FALSE(InstructionsMayInterfere(ordering, negate, add)); + EXPECT_FALSE(InstructionsMayInterfere(ordering, add, negate)); + EXPECT_FALSE(InstructionsMayInterfere(ordering, exp, add)); + EXPECT_FALSE(InstructionsMayInterfere(ordering, add, exp)); +} + +TEST_P(HloDataflowAnalysisTest, OverlappedValuesSequentialOrder) { + // Identical to the test OverlappedValue but using a sequential ordering of + // HLO instructions. + // + // param --> negate -> add + // \---> exp -----/ + // + // Sequential order: + // param, negate, exp, add + // + // Liveness is identical to the DependencyHloOrdering. + auto builder = HloComputation::Builder(TestName()); + auto param = builder.AddInstruction( + HloInstruction::CreateParameter(0, vector_shape_, "param")); + auto negate = builder.AddInstruction( + HloInstruction::CreateUnary(vector_shape_, HloOpcode::kNegate, param)); + auto exp = builder.AddInstruction( + HloInstruction::CreateUnary(vector_shape_, HloOpcode::kExp, param)); + auto add = builder.AddInstruction(HloInstruction::CreateBinary( + vector_shape_, HloOpcode::kAdd, negate, exp)); + + auto entry = module_->AddEntryComputation(builder.Build()); + RunAnalysis(GetParam()); + + SequentialHloOrdering::HloModuleSequence sequence; + std::vector order = {param, negate, exp, add}; + sequence.emplace(entry, order); + + SequentialHloOrdering ordering(module_.get(), sequence); + + EXPECT_TRUE(InstructionsMayInterfere(ordering, param, negate)); + EXPECT_FALSE(InstructionsMayInterfere(ordering, param, exp)); + EXPECT_FALSE(InstructionsMayInterfere(ordering, param, add)); + + // Negate and exp interfere with each other, but not with add. + EXPECT_TRUE(InstructionsMayInterfere(ordering, negate, exp)); + EXPECT_TRUE(InstructionsMayInterfere(ordering, exp, negate)); + EXPECT_FALSE(InstructionsMayInterfere(ordering, negate, add)); + EXPECT_FALSE(InstructionsMayInterfere(ordering, add, negate)); + EXPECT_FALSE(InstructionsMayInterfere(ordering, exp, add)); + EXPECT_FALSE(InstructionsMayInterfere(ordering, add, exp)); +} + +TEST_P(HloDataflowAnalysisTest, EmbeddedComputationInterference) { + // Test MayInterfere() for embedded computation, specifically the interference + // of values in different computations. + // + // embedded_computation: + // %embedded_param = Param(0) + // %embedded_log = Log(%embedded_param) + // + // entry computation: + // %param = Param(0) + // %negate = Negate(%param) + // %exp = Negate(%exp) + // %call = Call(embedded_computation, {%exp}) + // %add = Add(%negate, %call) + // + // Note %negate is live across the call and should interfere with all values + // in the embedded computation. + auto embedded_builder = HloComputation::Builder(TestName() + "_embedded"); + auto embedded_param = embedded_builder.AddInstruction( + HloInstruction::CreateParameter(0, vector_shape_, "embedded_param")); + auto embedded_log = + embedded_builder.AddInstruction(HloInstruction::CreateUnary( + vector_shape_, HloOpcode::kLog, embedded_param)); + auto embedded_computation = + module_->AddEmbeddedComputation(embedded_builder.Build()); + + auto builder = HloComputation::Builder(TestName()); + auto param = builder.AddInstruction( + HloInstruction::CreateParameter(0, vector_shape_, "param")); + auto negate = builder.AddInstruction( + HloInstruction::CreateUnary(vector_shape_, HloOpcode::kNegate, param)); + auto exp = builder.AddInstruction( + HloInstruction::CreateUnary(vector_shape_, HloOpcode::kExp, param)); + auto call = builder.AddInstruction( + HloInstruction::CreateCall(vector_shape_, {exp}, embedded_computation)); + builder.AddInstruction(HloInstruction::CreateBinary( + vector_shape_, HloOpcode::kAdd, negate, call)); + module_->AddEntryComputation(builder.Build()); + RunAnalysis(GetParam()); + + DependencyHloOrdering ordering(module_.get()); + + // Exp only use is the call so it should not interfere with values inside the + // embedded computation. + EXPECT_FALSE(InstructionsMayInterfere(ordering, exp, embedded_log)); + + // Negate is live across the call and should interfere with values in the + // embedded computation + EXPECT_TRUE(InstructionsMayInterfere(ordering, negate, embedded_log)); +} + +INSTANTIATE_TEST_CASE_P(HloDataflowAnalysisInstantiation, + HloDataflowAnalysisTest, + ::testing::Values(false, true)); + +} // namespace +} // namespace xla diff --git a/tensorflow/compiler/xla/service/hlo_dce.cc b/tensorflow/compiler/xla/service/hlo_dce.cc index fdfbbf8baf65884fcb1eed846e6ce3eda07bc45d..5b2c57da4ff3a1f887f777c3304893d950b3d3a9 100644 --- a/tensorflow/compiler/xla/service/hlo_dce.cc +++ b/tensorflow/compiler/xla/service/hlo_dce.cc @@ -38,6 +38,9 @@ StatusOr HloDCE::Run(HloModule* module) { bool changed = false; for (auto& computation : module->computations()) { + if (computation->IsFusionComputation()) { + continue; + } std::unordered_set live_instructions; TF_RETURN_IF_ERROR(computation->root_instruction()->Accept( [&live_instructions](HloInstruction* instruction) { @@ -52,7 +55,7 @@ StatusOr HloDCE::Run(HloModule* module) { for (auto& instruction : computation->instructions()) { if (instruction->user_count() == 0 && live_instructions.count(instruction.get()) == 0 && - HloComputation::IsRemovable(instruction->opcode())) { + computation->IsRemovable(instruction.get())) { dead_roots.push_back(instruction.get()); } } diff --git a/tensorflow/compiler/xla/service/hlo_dce_test.cc b/tensorflow/compiler/xla/service/hlo_dce_test.cc index dcd9e00c56c76046e6c1de75558637b7e941e57e..7354ef4043721f1de69454bd8970377c5af7609f 100644 --- a/tensorflow/compiler/xla/service/hlo_dce_test.cc +++ b/tensorflow/compiler/xla/service/hlo_dce_test.cc @@ -30,6 +30,7 @@ limitations under the License. #include "tensorflow/compiler/xla/tests/test_utils.h" #include "tensorflow/compiler/xla/types.h" #include "tensorflow/compiler/xla/xla_data.pb.h" +#include "tensorflow/core/lib/core/status_test_util.h" #include "tensorflow/core/platform/types.h" namespace xla { @@ -38,19 +39,30 @@ namespace { class HloDceTest : public HloTestBase { protected: HloDceTest() {} + + // Returns whether the given instruction exists in the given computation. + bool HasInstruction(const HloComputation& computation, + const HloInstruction* instruction) { + for (auto& inst : computation.instructions()) { + if (inst.get() == instruction) { + return true; + } + } + return false; + } }; TEST_F(HloDceTest, NoDeadCode) { // Verify that no dead code is removed from a computation with no dead code. auto builder = HloComputation::Builder(TestName()); auto constant1 = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(42.0f))); + HloInstruction::CreateConstant(Literal::CreateR0(42.0f))); auto constant2 = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(123.0f))); + HloInstruction::CreateConstant(Literal::CreateR0(123.0f))); builder.AddInstruction(HloInstruction::CreateBinary( constant1->shape(), HloOpcode::kAdd, constant1, constant2)); - auto module = MakeUnique(TestName()); + auto module = CreateNewModule(); auto computation = module->AddEntryComputation(builder.Build()); EXPECT_EQ(3, computation->instruction_count()); @@ -80,7 +92,7 @@ TEST_F(HloDceTest, DeadParameters) { builder.AddInstruction(HloInstruction::CreateUnary( live_param->shape(), HloOpcode::kNegate, live_param)); - auto module = MakeUnique(TestName()); + auto module = CreateNewModule(); auto computation = module->AddEntryComputation(builder.Build()); EXPECT_EQ(5, computation->instruction_count()); @@ -93,5 +105,206 @@ TEST_F(HloDceTest, DeadParameters) { EXPECT_EQ(0, dead_param1->user_count()); } +TEST_F(HloDceTest, ControlDependencies) { + // Verify that instructions with control dependencies are not removed. + auto builder = HloComputation::Builder(TestName()); + auto constant1 = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(42.0f))); + auto constant2 = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(123.0f))); + + // Create two dead instructions: a negate and an add. + auto dead_negate = builder.AddInstruction(HloInstruction::CreateUnary( + constant1->shape(), HloOpcode::kNegate, constant1)); + auto dead_add = builder.AddInstruction(HloInstruction::CreateBinary( + constant1->shape(), HloOpcode::kAdd, constant1, constant2)); + + // Create the same two instructions again, but these will have a control + // dependency added. + auto dead_negate_with_control_dep = + builder.AddInstruction(HloInstruction::CreateUnary( + constant1->shape(), HloOpcode::kNegate, constant1)); + auto dead_add_with_control_dep = + builder.AddInstruction(HloInstruction::CreateBinary( + constant1->shape(), HloOpcode::kAdd, constant1, constant2)); + + // Create a root so the previously added instruction is dead. + builder.AddInstruction(HloInstruction::CreateBinary( + constant1->shape(), HloOpcode::kAdd, constant1, constant2)); + + auto module = CreateNewModule(); + auto computation = module->AddEntryComputation(builder.Build()); + + // Add a control dependency between two instructions. + TF_ASSERT_OK(dead_negate_with_control_dep->AddControlDependencyTo( + dead_add_with_control_dep)); + + EXPECT_EQ(7, computation->instruction_count()); + EXPECT_TRUE(HasInstruction(*computation, dead_negate)); + EXPECT_TRUE(HasInstruction(*computation, dead_add)); + EXPECT_TRUE(HasInstruction(*computation, dead_negate_with_control_dep)); + EXPECT_TRUE(HasInstruction(*computation, dead_add_with_control_dep)); + + HloDCE dce; + EXPECT_TRUE(dce.Run(module.get()).ValueOrDie()); + + EXPECT_EQ(5, computation->instruction_count()); + EXPECT_FALSE(HasInstruction(*computation, dead_negate)); + EXPECT_FALSE(HasInstruction(*computation, dead_add)); + EXPECT_TRUE(HasInstruction(*computation, dead_negate_with_control_dep)); + EXPECT_TRUE(HasInstruction(*computation, dead_add_with_control_dep)); +} + +// Tests that a dead call instruction is removed. +TEST_F(HloDceTest, DeadInstructionWithCalledComputation) { + auto module = CreateNewModule(); + Shape shape = ShapeUtil::MakeShape(F32, {}); + + // Called computation for the call instruction. + auto callee_builder = HloComputation::Builder(TestName() + "-callee"); + { + auto param = callee_builder.AddInstruction( + HloInstruction::CreateParameter(0, shape, "param")); + callee_builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, param)); + } + auto called_computation = + module->AddEmbeddedComputation(callee_builder.Build()); + + // Entry computation with a call instruction. + auto builder = HloComputation::Builder(TestName()); + auto param = builder.AddInstruction( + HloInstruction::CreateParameter(0, shape, "param")); + auto dead_call = builder.AddInstruction( + HloInstruction::CreateCall(shape, {param}, called_computation)); + builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, param)); + auto computation = module->AddEntryComputation(builder.Build()); + + EXPECT_EQ(3, computation->instruction_count()); + EXPECT_EQ(2, param->user_count()); + EXPECT_EQ(0, dead_call->user_count()); + EXPECT_TRUE(HasInstruction(*computation, dead_call)); + + HloDCE dce; + EXPECT_TRUE(dce.Run(module.get()).ValueOrDie()); + + EXPECT_EQ(2, computation->instruction_count()); + EXPECT_EQ(1, param->user_count()); + EXPECT_FALSE(HasInstruction(*computation, dead_call)); +} + +// 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(); + Shape shape = ShapeUtil::MakeShape(F32, {}); + + // Condition computation of a while instruction. + auto cond_builder = HloComputation::Builder(TestName() + "-cond"); + { + auto param = cond_builder.AddInstruction( + HloInstruction::CreateParameter(0, shape, "cond_param")); + auto constant = cond_builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(42.0f))); + cond_builder.AddInstruction(HloInstruction::CreateBinary( + ShapeUtil::MakeShape(PRED, {}), HloOpcode::kLt, param, constant)); + } + auto cond_computation = module->AddEmbeddedComputation(cond_builder.Build()); + + // Body computation of a while instruction. + auto body_builder = HloComputation::Builder(TestName() + "-body"); + { + auto param = body_builder.AddInstruction( + HloInstruction::CreateParameter(0, shape, "param")); + + auto infeed = + body_builder.AddInstruction(HloInstruction::CreateInfeed(shape, "")); + body_builder.AddInstruction( + HloInstruction::CreateBinary(shape, HloOpcode::kAdd, param, infeed)); + } + auto body_computation = module->AddEmbeddedComputation(body_builder.Build()); + + // Entry computation with a while instruction and a negate on the parameter. + auto builder = HloComputation::Builder(TestName()); + auto param = builder.AddInstruction( + HloInstruction::CreateParameter(0, shape, "param")); + auto live_while = builder.AddInstruction(HloInstruction::CreateWhile( + shape, cond_computation, body_computation, param)); + builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, param)); + auto computation = module->AddEntryComputation(builder.Build()); + + // Check the while instruction is not removed even if its user count is 0. + EXPECT_EQ(3, computation->instruction_count()); + EXPECT_EQ(2, param->user_count()); + EXPECT_EQ(0, live_while->user_count()); + EXPECT_TRUE(HasInstruction(*computation, live_while)); + + HloDCE dce; + EXPECT_FALSE(dce.Run(module.get()).ValueOrDie()); + + EXPECT_EQ(3, computation->instruction_count()); + EXPECT_EQ(2, param->user_count()); + EXPECT_EQ(0, live_while->user_count()); + EXPECT_TRUE(HasInstruction(*computation, live_while)); +} + +// Tests that a nested call instruction with a side effect is not removed. +TEST_F(HloDceTest, CalledComputationWithNestedSideEffect) { + auto module = CreateNewModule(); + Shape shape = ShapeUtil::MakeShape(F32, {}); + + // Nested called computation with a side effect. + auto nested_callee_builder = + HloComputation::Builder(TestName() + "-nested_callee"); + { + auto param = nested_callee_builder.AddInstruction( + HloInstruction::CreateParameter(0, shape, "param")); + nested_callee_builder.AddInstruction( + HloInstruction::CreateOutfeed(shape, param, "")); + } + auto nested_called_computation = + module->AddEmbeddedComputation(nested_callee_builder.Build()); + + // Outer called computation that calls the nested computation. + auto callee_builder = HloComputation::Builder(TestName() + "-callee"); + { + auto param = callee_builder.AddInstruction( + HloInstruction::CreateParameter(0, shape, "param")); + callee_builder.AddInstruction( + HloInstruction::CreateCall(shape, {param}, nested_called_computation)); + } + auto called_computation = + module->AddEmbeddedComputation(callee_builder.Build()); + + // Entry computation with a call instruction. + auto builder = HloComputation::Builder(TestName()); + auto param = builder.AddInstruction( + HloInstruction::CreateParameter(0, shape, "param")); + auto live_call = builder.AddInstruction( + HloInstruction::CreateCall(shape, {param}, called_computation)); + builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, param)); + auto computation = module->AddEntryComputation(builder.Build()); + + EXPECT_EQ(3, computation->instruction_count()); + EXPECT_EQ(2, param->user_count()); + EXPECT_EQ(0, live_call->user_count()); + EXPECT_TRUE(HasInstruction(*computation, live_call)); + + HloDCE dce; + EXPECT_FALSE(dce.Run(module.get()).ValueOrDie()); + + EXPECT_EQ(3, computation->instruction_count()); + EXPECT_EQ(2, param->user_count()); + EXPECT_EQ(0, live_call->user_count()); + EXPECT_TRUE(HasInstruction(*computation, live_call)); +} + } // namespace } // namespace xla + +int main(int argc, char** argv) { + return xla::ParseDebugOptionsFlagsAndRunTests(argc, argv); +} diff --git a/tensorflow/compiler/xla/service/hlo_evaluator.cc b/tensorflow/compiler/xla/service/hlo_evaluator.cc new file mode 100644 index 0000000000000000000000000000000000000000..e09c9d3beb2d33008adade37e30ee2828df721df --- /dev/null +++ b/tensorflow/compiler/xla/service/hlo_evaluator.cc @@ -0,0 +1,1507 @@ +/* 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/hlo_evaluator.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "tensorflow/compiler/xla/index_util.h" +#include "tensorflow/compiler/xla/layout_util.h" +#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/ptr_util.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/status.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/core/lib/core/bitmap.h" +#include "tensorflow/core/lib/core/errors.h" +#include "tensorflow/core/lib/core/status.h" +#include "tensorflow/core/lib/core/stringpiece.h" +#include "tensorflow/core/platform/logging.h" +#include "tensorflow/core/platform/protobuf.h" +#include "tensorflow/core/platform/types.h" + +namespace xla { + +namespace { + +template +StatusOr> Compare(const Shape& shape, HloOpcode opcode, + const Literal& lhs_literal, + const Literal& rhs_literal) { + std::function compare_op; + switch (opcode) { + case HloOpcode::kEq: + compare_op = [](OperandT lhs_el, OperandT rhs_el) { + return lhs_el == rhs_el; + }; + break; + case HloOpcode::kNe: + compare_op = [](OperandT lhs_el, OperandT rhs_el) { + return lhs_el != rhs_el; + }; + break; + case HloOpcode::kGe: + compare_op = [](OperandT lhs_el, OperandT rhs_el) { + return lhs_el >= rhs_el; + }; + break; + case HloOpcode::kGt: + compare_op = [](OperandT lhs_el, OperandT rhs_el) { + return lhs_el > rhs_el; + }; + break; + case HloOpcode::kLe: + compare_op = [](OperandT lhs_el, OperandT rhs_el) { + return lhs_el <= rhs_el; + }; + break; + case HloOpcode::kLt: + compare_op = [](OperandT lhs_el, OperandT rhs_el) { + return lhs_el < rhs_el; + }; + break; + default: + LOG(FATAL) << "unhandled HLO opcode for conversion to Comparison: " + << HloOpcodeString(opcode); + } + + auto result = Literal::CreateFromShape(shape); + TF_RETURN_IF_ERROR(result->Populate( + [&](tensorflow::gtl::ArraySlice multi_index) { + return compare_op(lhs_literal.Get(multi_index), + rhs_literal.Get(multi_index)); + })); + + return std::move(result); +} + +template +StatusOr> ElementWiseUnaryOpImpl( + HloInstruction* instruction, + const std::function& unary_op, + 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).c_str(), + ShapeUtil::HumanString(operand->shape()).c_str()); + } + + auto result = Literal::CreateFromShape(shape); + + TF_RETURN_IF_ERROR(result->Populate( + [&](tensorflow::gtl::ArraySlice multi_index) { + return unary_op(operand_literal.Get(multi_index)); + })); + return std::move(result); +} + +} // namespace + +template +class HloEvaluator::TypedVisitor : public DfsHloVisitorWithDefault { + public: + explicit TypedVisitor(HloEvaluator* p) : parent_(p) {} + + Status DefaultAction(HloInstruction* hlo_instruction) override { + return Unimplemented("unhandled HLO ops for HloEvaluator: %s.", + HloOpcodeString(hlo_instruction->opcode()).c_str()); + }; + + // TODO(b/35950897): many of the stl functions used in the handlers are not + // overloaded for every XLA primitive types. + + template ::value>::type* = + nullptr> + Status HandleAbs(HloInstruction* abs, HloInstruction* operand) { + TF_ASSIGN_OR_RETURN(parent_->evaluated_[abs], + ElementWiseUnaryOp(abs, [](NativeT elem_operand) { + return elem_operand; + })); + return Status::OK(); + } + + template < + typename NativeT, + typename std::enable_if::value>::type* = nullptr> + Status HandleAbs(HloInstruction* abs, HloInstruction* operand) { + TF_ASSIGN_OR_RETURN(parent_->evaluated_[abs], + ElementWiseUnaryOp(abs, [](NativeT elem_operand) { + return std::abs(elem_operand); + })); + return Status::OK(); + } + + Status HandleAbs(HloInstruction* abs, HloInstruction* operand) override { + return HandleAbs(abs, operand); + }; + + Status HandleBroadcast(HloInstruction* broadcast) override { + parent_->evaluated_[broadcast] = + Literal::CreateFromShape(broadcast->shape()); + auto output = parent_->evaluated_[broadcast].get(); + auto operand_to_broadcast = + parent_->GetEvaluatedLiteralFor(broadcast->operand(0)); + std::vector broadcast_indices( + ShapeUtil::Rank(broadcast->operand(0)->shape()), 0); + + // Special case for broadcasting scalars: ignore broadcast dimension and + // broadcast to whatever the output dimension is. + // TODO(b/64533549): Remove the need of this once this bug is resolved. + if (ShapeUtil::IsScalar(operand_to_broadcast.shape())) { + return output->Populate( + [&](tensorflow::gtl::ArraySlice multi_index) { + return operand_to_broadcast.Get({}); + }); + } + + TF_RET_CHECK(broadcast->dimensions().size() == + ShapeUtil::Rank(operand_to_broadcast.shape())) + << "broadcast dimensions is of size: " << broadcast->dimensions().size() + << " and rank of operand_to_broadcast is: " + << ShapeUtil::Rank(operand_to_broadcast.shape()); + // 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) { + TF_RET_CHECK(broadcast->shape().dimensions(broadcast->dimensions(i)) == + operand_to_broadcast.shape().dimensions(i)); + } + + return output->Populate( + [&](tensorflow::gtl::ArraySlice multi_index) { + for (int64 i = 0; i < broadcast->dimensions().size(); ++i) { + broadcast_indices[i] = multi_index[broadcast->dimensions(i)]; + } + return operand_to_broadcast.Get(broadcast_indices); + }); + }; + + Status HandleCeil(HloInstruction* ceil, HloInstruction* operand) override { + TF_ASSIGN_OR_RETURN(parent_->evaluated_[ceil], + ElementWiseUnaryOp(ceil, [](ReturnT elem_operand) { + return std::ceil(elem_operand); + })); + return Status::OK(); + }; + + Status HandleConvert(HloInstruction* convert) override { + const HloInstruction* operand = convert->operand(0); + TF_RET_CHECK(ShapeUtil::SameDimensions(operand->shape(), convert->shape())); + TF_ASSIGN_OR_RETURN(std::unique_ptr result, + parent_->GetEvaluatedLiteralFor(operand).Convert( + convert->shape().element_type())); + + if (LayoutUtil::LayoutsInShapesEqual(result->shape(), convert->shape())) { + parent_->evaluated_[convert] = std::move(result); + } else { + parent_->evaluated_[convert] = + result->Relayout(convert->shape().layout()); + } + return Status::OK(); + } + + Status HandleExp(HloInstruction* exp, HloInstruction* operand) override { + TF_ASSIGN_OR_RETURN(parent_->evaluated_[exp], + ElementWiseUnaryOp(exp, [](ReturnT elem_operand) { + return std::exp(elem_operand); + })); + return Status::OK(); + }; + + Status HandleFloor(HloInstruction* floor, HloInstruction* operand) override { + TF_ASSIGN_OR_RETURN(parent_->evaluated_[floor], + ElementWiseUnaryOp(floor, [](ReturnT elem_operand) { + return std::floor(elem_operand); + })); + return Status::OK(); + }; + + Status HandleLog(HloInstruction* log, HloInstruction* operand) override { + TF_ASSIGN_OR_RETURN(parent_->evaluated_[log], + ElementWiseUnaryOp(log, [](ReturnT elem_operand) { + return std::log(elem_operand); + })); + return Status::OK(); + }; + + Status HandleLogicalNot(HloInstruction* logical_not, + HloInstruction* operand) override { + TF_ASSIGN_OR_RETURN( + parent_->evaluated_[logical_not], + ElementWiseUnaryOp(logical_not, + [](ReturnT elem_operand) { return !elem_operand; })); + return Status::OK(); + }; + + Status HandleNegate(HloInstruction* negate, + HloInstruction* operand) override { + TF_ASSIGN_OR_RETURN(parent_->evaluated_[negate], + ElementWiseUnaryOp(negate, [](ReturnT elem_operand) { + return -elem_operand; + })); + return Status::OK(); + }; + + Status HandleSign(HloInstruction* sign, HloInstruction* operand) override { + TF_ASSIGN_OR_RETURN(parent_->evaluated_[sign], + ElementWiseUnaryOp(sign, [](ReturnT elem_operand) { + return (ReturnT(0) < elem_operand) - + (elem_operand < ReturnT(0)); + })); + return Status::OK(); + }; + + Status HandleTanh(HloInstruction* tanh, HloInstruction* operand) override { + TF_ASSIGN_OR_RETURN(parent_->evaluated_[tanh], + ElementWiseUnaryOp(tanh, [](ReturnT elem_operand) { + return std::tanh(elem_operand); + })); + return Status::OK(); + }; + + Status HandleMultiply(HloInstruction* multiply, HloInstruction* lhs, + HloInstruction* rhs) override { + TF_ASSIGN_OR_RETURN( + parent_->evaluated_[multiply], + ElementWiseBinaryOp(multiply, [](ReturnT lhs_elem, ReturnT rhs_elem) { + return lhs_elem * rhs_elem; + })); + return Status::OK(); + }; + + Status HandleSubtract(HloInstruction* subtract, HloInstruction* lhs, + HloInstruction* rhs) override { + TF_ASSIGN_OR_RETURN( + parent_->evaluated_[subtract], + ElementWiseBinaryOp(subtract, [](ReturnT lhs_elem, ReturnT rhs_elem) { + return lhs_elem - rhs_elem; + })); + return Status::OK(); + }; + + Status HandleAdd(HloInstruction* add, HloInstruction* lhs, + HloInstruction* rhs) override { + TF_ASSIGN_OR_RETURN( + parent_->evaluated_[add], + ElementWiseBinaryOp(add, [](ReturnT lhs_elem, ReturnT rhs_elem) { + return lhs_elem + rhs_elem; + })); + return Status::OK(); + }; + + Status HandleDivide(HloInstruction* divide, HloInstruction* lhs, + HloInstruction* rhs) override { + TF_ASSIGN_OR_RETURN( + parent_->evaluated_[divide], + ElementWiseBinaryOp(divide, [](ReturnT lhs_elem, ReturnT rhs_elem) { + return lhs_elem / rhs_elem; + })); + return Status::OK(); + }; + + Status HandleMaximum(HloInstruction* maximum) override { + TF_ASSIGN_OR_RETURN( + parent_->evaluated_[maximum], + ElementWiseBinaryOp(maximum, [](ReturnT lhs, ReturnT rhs) { + return std::fmax(lhs, rhs); + })); + return Status::OK(); + }; + + Status HandleMinimum(HloInstruction* minimum) override { + TF_ASSIGN_OR_RETURN( + parent_->evaluated_[minimum], + ElementWiseBinaryOp(minimum, [](ReturnT lhs_el, ReturnT rhs_el) { + return std::fmin(lhs_el, rhs_el); + })); + return Status::OK(); + }; + + Status HandlePower(HloInstruction* power, HloInstruction* lhs, + HloInstruction* rhs) override { + TF_ASSIGN_OR_RETURN( + parent_->evaluated_[power], + ElementWiseBinaryOp(power, [](ReturnT lhs_el, ReturnT rhs_el) { + return std::pow(lhs_el, rhs_el); + })); + return Status::OK(); + }; + + Status HandleRemainder(HloInstruction* remainder, HloInstruction* lhs, + HloInstruction* rhs) override { + TF_ASSIGN_OR_RETURN( + parent_->evaluated_[remainder], + ElementWiseBinaryOp(remainder, [](ReturnT lhs_el, ReturnT rhs_el) { + return std::fmod(lhs_el, rhs_el); + })); + return Status::OK(); + }; + + Status HandleLogicalAnd(HloInstruction* logical_and, HloInstruction* lhs, + HloInstruction* rhs) override { + TF_ASSIGN_OR_RETURN( + parent_->evaluated_[logical_and], + ElementWiseBinaryOp(logical_and, [](ReturnT lhs_el, ReturnT rhs_el) { + return lhs_el && rhs_el; + })); + return Status::OK(); + }; + + Status HandleLogicalOr(HloInstruction* logical_or, HloInstruction* lhs, + HloInstruction* rhs) override { + TF_ASSIGN_OR_RETURN( + parent_->evaluated_[logical_or], + ElementWiseBinaryOp(logical_or, [](ReturnT lhs_el, ReturnT rhs_el) { + return lhs_el || rhs_el; + })); + return Status::OK(); + }; + + Status HandleClamp(HloInstruction* clamp, HloInstruction* min, + HloInstruction* arg, HloInstruction* max) override { + std::function clamp_op = + [](ReturnT low, ReturnT high, ReturnT value) { + return std::fmax(low, std::fmin(value, high)); + }; + TF_ASSIGN_OR_RETURN(parent_->evaluated_[clamp], + ElementWiseTernaryOp(clamp, std::move(clamp_op))); + return Status::OK(); + }; + + Status HandleSelect(HloInstruction* select, HloInstruction* pred, + HloInstruction* on_true, + HloInstruction* on_false) override { + CHECK(!ShapeUtil::IsTuple(select->shape())); + std::function select_op = + [](bool pred, ReturnT on_true, ReturnT on_false) { + if (pred) { + return on_true; + } + return on_false; + }; + TF_ASSIGN_OR_RETURN(parent_->evaluated_[select], + ElementWiseTernaryOp(select, std::move(select_op))); + return Status::OK(); + }; + + Status HandleReverse(HloInstruction* reverse, + HloInstruction* operand) override { + const auto result_shape = reverse->shape(); + const auto reverse_dimensions = reverse->dimensions(); + + TF_ASSIGN_OR_RETURN(auto inferred_return_shape, + ShapeInference::InferReverseShape(operand->shape(), + reverse_dimensions)); + + TF_RET_CHECK(ShapeUtil::Compatible(result_shape, inferred_return_shape)) + << "return shape set to: " << ShapeUtil::HumanString(result_shape) + << " but is inferred to be: " + << ShapeUtil::HumanString(inferred_return_shape); + + auto operand_literal = parent_->GetEvaluatedLiteralFor(operand); + auto result = Literal::CreateFromShape(result_shape); + + TF_RETURN_IF_ERROR(result->Populate( + [&](tensorflow::gtl::ArraySlice out_index) { + std::vector from_index(out_index.begin(), out_index.end()); + for (const int64 dim : reverse_dimensions) { + from_index[dim] = result_shape.dimensions(dim) - 1 - out_index[dim]; + } + return operand_literal.Get(from_index); + })); + + parent_->evaluated_[reverse] = std::move(result); + return Status::OK(); + }; + + Status HandleConvolution(HloInstruction* conv, HloInstruction* lhs, + HloInstruction* rhs, const Window& window) override { + const Shape& result_shape = conv->shape(); + const Shape& lhs_shape = lhs->shape(); + const Shape& rhs_shape = rhs->shape(); + + 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(ShapeUtil::SameElementType(lhs_shape, rhs_shape)); + CHECK(ShapeUtil::SameElementType(lhs_shape, result_shape)); + + const auto& dnums = conv->convolution_dimension_numbers(); + const int64 num_spatial_dims = dnums.spatial_dimensions_size(); + CHECK_EQ(num_spatial_dims, dnums.kernel_spatial_dimensions_size()); + CHECK_GE(num_spatial_dims, 1); + CHECK_EQ(window.dimensions_size(), num_spatial_dims); + + const auto lhs_rank = ShapeUtil::Rank(lhs_shape); + const auto rhs_rank = ShapeUtil::Rank(rhs_shape); + + 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, + window, dnums)); + CHECK(ShapeUtil::Compatible(result_shape, inferred_return_shape)) + << "return shape set to: " << ShapeUtil::HumanString(result_shape) + << " but is inferred to be: " + << ShapeUtil::HumanString(inferred_return_shape); + + const Literal& lhs_literal = parent_->GetEvaluatedLiteralFor(lhs); + const Literal& rhs_literal = parent_->GetEvaluatedLiteralFor(rhs); + + // Dimension number applicable for both input (lhs), and output. + const int64 batch_dim = dnums.batch_dimension(); + const int64 z_dim = dnums.feature_dimension(); + // Dimension number applicable for kernel (rhs). + const int64 kernel_input_z_dim = dnums.kernel_input_feature_dimension(); + const int64 kernel_output_z_dim = dnums.kernel_output_feature_dimension(); + + const int64 z_size = ShapeUtil::GetDimension(lhs_shape, z_dim); + + std::vector window_dimension_sizes; + for (auto i : dnums.kernel_spatial_dimensions()) { + window_dimension_sizes.push_back(ShapeUtil::GetDimension(rhs_shape, i)); + } + + const Shape& window_shape = + ShapeUtil::MakeShape(rhs_shape.element_type(), window_dimension_sizes); + + DimensionVector lhs_index(lhs_rank); + DimensionVector rhs_index(rhs_rank); + DimensionVector rhs_spatial_index(dnums.kernel_spatial_dimensions_size()); + + auto func = [&](tensorflow::gtl::ArraySlice out_index) { + ReturnT result_val = static_cast(0); + + std::fill(lhs_index.begin(), lhs_index.end(), 0); + std::fill(rhs_index.begin(), rhs_index.end(), 0); + std::fill(rhs_spatial_index.begin(), rhs_spatial_index.end(), 0); + + lhs_index[batch_dim] = out_index[batch_dim]; + rhs_index[kernel_output_z_dim] = out_index[z_dim]; + + // Convolve input feature with kernel. + do { + for (int64 iz = 0; iz < z_size; ++iz) { + lhs_index[z_dim] = iz; + rhs_index[kernel_input_z_dim] = iz; + + // Find corresponding spatial dimension index for input (lhs). + for (int64 ki = 0; ki < rhs_spatial_index.size(); ++ki) { + // Spatial dimension number for input (lhs) and output. + const int64 spatial_dim = dnums.spatial_dimensions(ki); + + // Calculate lhs (input) index without taking base dilation into + // account. + const auto& window_dim = window.dimensions(ki); + const int64 undilated_index = + out_index[spatial_dim] * window_dim.stride() - + window_dim.padding_low() + + rhs_spatial_index[ki] * window_dim.window_dilation(); + // Skip if the lhs (input) index is to be dilated. + if (undilated_index % window_dim.base_dilation() != 0) { + goto cnt; + } + + // Calculate the actual lhs (input) index after dilation. + lhs_index[spatial_dim] = + undilated_index / window_dim.base_dilation(); + + // Skip if input index is not in bound. + if (!(lhs_index[spatial_dim] >= 0 && + lhs_index[spatial_dim] < lhs_shape.dimensions(spatial_dim))) { + goto cnt; + } + + rhs_index[dnums.kernel_spatial_dimensions(ki)] = + rhs_spatial_index[ki]; + } + + result_val += lhs_literal.Get(lhs_index) * + rhs_literal.Get(rhs_index); + } + cnt:; + } while (IndexUtil::BumpIndices(window_shape, &rhs_spatial_index)); + + return result_val; + }; + + auto result = Literal::CreateFromShape(result_shape); + TF_RETURN_IF_ERROR(result->Populate(func)); + + parent_->evaluated_[conv] = std::move(result); + return Status::OK(); + }; + + Status HandleDot(HloInstruction* dot, HloInstruction* lhs, + HloInstruction* rhs) override { + CHECK(ShapeUtil::IsArray(dot->shape())); + CHECK(ShapeUtil::IsArray(lhs->shape())); + CHECK(ShapeUtil::IsArray(rhs->shape())); + + // Dot only supports operands of rank 1 and 2. + const auto dot_rank = ShapeUtil::Rank(dot->shape()); + const auto lhs_rank = ShapeUtil::Rank(lhs->shape()); + const auto rhs_rank = ShapeUtil::Rank(rhs->shape()); + CHECK(lhs_rank > 0 && lhs_rank <= 2); + CHECK(rhs_rank > 0 && rhs_rank <= 2); + CHECK_EQ(dot_rank, lhs_rank + rhs_rank - 2); + + CHECK(ShapeUtil::SameElementType(lhs->shape(), rhs->shape())); + CHECK(ShapeUtil::SameElementType(lhs->shape(), dot->shape())); + + // Check contracted dimensions are the same. + // + // Determine the index of the contracted dimensions for input tensors. + // dimensions -1 of lhs and dimension 0 of rhs are contracted. + const int64 lhs_contracted_dimension = + ShapeUtil::GetDimensionNumber(lhs->shape(), -1); + const int64 rhs_contracted_dimension = 0; + CHECK_EQ(lhs->shape().dimensions(lhs_contracted_dimension), + rhs->shape().dimensions(rhs_contracted_dimension)) + << "lhs contracted dimension: " + << lhs->shape().dimensions(lhs_contracted_dimension) + << " rhs contracted dimension: " + << rhs->shape().dimensions(rhs_contracted_dimension); + const int64 contracted_dimension_size = + lhs->shape().dimensions(lhs_contracted_dimension); + + const Literal& lhs_literal = parent_->GetEvaluatedLiteralFor(lhs); + const Literal& rhs_literal = parent_->GetEvaluatedLiteralFor(rhs); + + auto result = Literal::CreateFromShape(dot->shape()); + TF_RETURN_IF_ERROR(result->Populate( + [&](tensorflow::gtl::ArraySlice multi_index) { + ReturnT result_val = static_cast(0); + + std::vector lhs_index(lhs_rank, 0); + std::vector rhs_index(rhs_rank, 0); + // Set index for non-contracted dimension for lhs and rhs. + if (lhs_rank > 1) { + lhs_index[0] = multi_index[0]; + } + if (rhs_rank > 1) { + rhs_index[1] = multi_index[multi_index.size() - 1]; + } + + // Accumulates resulting product along the contracted dimension. + for (int64 i = 0; i < contracted_dimension_size; ++i) { + lhs_index[lhs_contracted_dimension] = i; + rhs_index[rhs_contracted_dimension] = i; + + result_val += lhs_literal.Get(lhs_index) * + rhs_literal.Get(rhs_index); + } + + return result_val; + })); + + parent_->evaluated_[dot] = std::move(result); + return Status::OK(); + }; + + Status HandlePad(HloInstruction* pad) override { + CHECK(!ShapeUtil::IsTuple(pad->operand(0)->shape())); + // Padding value must be scalar. + CHECK(ShapeUtil::IsScalar(pad->operand(1)->shape())); + CHECK_EQ(ShapeUtil::Rank(pad->operand(0)->shape()), + pad->padding_config().dimensions_size()); + + TF_ASSIGN_OR_RETURN(auto inferred_return_shape, + ShapeInference::InferPadShape( + /*operand_shape=*/pad->operand(0)->shape(), + /*padding_value_shape=*/pad->operand(1)->shape(), + /*padding_config=*/pad->padding_config())); + CHECK(ShapeUtil::Compatible(pad->shape(), inferred_return_shape)) + << "return shape is set to: " << ShapeUtil::HumanString(pad->shape()) + << "but is inferred to be: " + << ShapeUtil::HumanString(inferred_return_shape); + + // Create new HLO of padded shape with padding value. + ReturnT scalar = + parent_->GetEvaluatedLiteralFor(pad->operand(1)).Get({}); + auto result = Literal::CreateFromShape(pad->shape()); + TF_RETURN_IF_ERROR(result->Populate( + [&scalar](tensorflow::gtl::ArraySlice multi_index) { + return scalar; + })); + + auto 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); + + // Loop through each element of the operand, assign them to the + // corresponding index of the resulting padded literal. + const PaddingConfig& pad_config = pad->padding_config(); + + auto func = [&](const std::vector& input_index) { + for (auto i = 0; i < input_index.size(); ++i) { + // Interior padding occurs logically before edge padding, so in the case + // of negative edge padding elements are removed from the + // interior-padded operand. + target_index[i] = + pad_config.dimensions(i).edge_padding_low() + + input_index[i] * (pad_config.dimensions(i).interior_padding() + 1); + + // Account for negative low and high padding: skip assignment if the + // any target index is out of range. + if (!(target_index[i] >= 0 && + target_index[i] < pad->shape().dimensions(i))) { + return true; + } + } + result->Set(target_index, + evaluated_operand.Get(input_index)); + return true; + }; + + std::vector zero_base(evaluated_operand.shape().dimensions_size(), + 0); + std::vector step(evaluated_operand.shape().dimensions_size(), 1); + + ShapeUtil::ForEachIndex( + evaluated_operand.shape(), zero_base, + AsInt64Slice(evaluated_operand.shape().dimensions()), step, func); + + parent_->evaluated_[pad] = std::move(result); + return Status::OK(); + }; + + Status HandleDynamicSlice(HloInstruction* dynamic_slice, + HloInstruction* operand, + HloInstruction* start_indices) override { + 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_RET_CHECK(ShapeUtil::Compatible(result_shape, inferred_return_shape)) + << "return shape is set to: " << ShapeUtil::HumanString(result_shape) + << "but is inferred to be: " + << ShapeUtil::HumanString(inferred_return_shape); + TF_RET_CHECK( + 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)); + } break; + case S64: { + TF_ASSIGN_OR_RETURN( + parent_->evaluated_[dynamic_slice], + DynamicSlice(operand_literal, start_indices_literal, + result_shape)); + } break; + case U32: { + TF_ASSIGN_OR_RETURN( + parent_->evaluated_[dynamic_slice], + DynamicSlice(operand_literal, start_indices_literal, + result_shape)); + } break; + case U64: { + TF_ASSIGN_OR_RETURN( + parent_->evaluated_[dynamic_slice], + DynamicSlice(operand_literal, start_indices_literal, + result_shape)); + } break; + default: + LOG(FATAL) << "HandleDynamicSlice: unhandled primitive type for " + "start_indices: " + << PrimitiveType_Name(start_indices->shape().element_type()); + } + + return Status::OK(); + }; + + Status HandleDynamicUpdateSlice(HloInstruction* dynamic_update_slice, + HloInstruction* operand, + HloInstruction* update, + HloInstruction* start_indices) override { + auto result_shape = dynamic_update_slice->shape(); + TF_ASSIGN_OR_RETURN( + auto inferred_return_shape, + ShapeInference::InferDynamicUpdateSliceShape( + operand->shape(), update->shape(), start_indices->shape())); + TF_RET_CHECK(ShapeUtil::Compatible(result_shape, inferred_return_shape)) + << "return shape is set to: " << ShapeUtil::HumanString(result_shape) + << "but is inferred to be: " + << ShapeUtil::HumanString(inferred_return_shape); + TF_RET_CHECK( + primitive_util::IsIntegralType(start_indices->shape().element_type())); + TF_RET_CHECK(ShapeUtil::Compatible(result_shape, operand->shape())); + + 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)); + } break; + case S64: { + TF_ASSIGN_OR_RETURN( + parent_->evaluated_[dynamic_update_slice], + DynamicUpdateSlice(operand_literal, update_literal, + start_indices_literal)); + } break; + case U32: { + TF_ASSIGN_OR_RETURN( + parent_->evaluated_[dynamic_update_slice], + DynamicUpdateSlice(operand_literal, update_literal, + start_indices_literal)); + } break; + case U64: { + TF_ASSIGN_OR_RETURN( + parent_->evaluated_[dynamic_update_slice], + DynamicUpdateSlice(operand_literal, update_literal, + start_indices_literal)); + } break; + default: + LOG(FATAL) << "HandleDynamicUpdateSlice: unhandled primitive type for " + "start_indices: " + << PrimitiveType_Name(start_indices->shape().element_type()); + } + + return Status::OK(); + }; + + Status HandleReduce(HloInstruction* reduce, HloInstruction* arg, + HloInstruction* init_value, + tensorflow::gtl::ArraySlice dimensions, + HloComputation* function) override { + TF_RET_CHECK(ShapeUtil::Rank(reduce->shape()) == + ShapeUtil::Rank(arg->shape()) - dimensions.size()); + TF_ASSIGN_OR_RETURN(auto inferred_return_shape, + ShapeInference::InferReduceShape( + /*arg=*/arg->shape(), + /*init_value=*/init_value->shape(), + /*dimensions_to_reduce=*/dimensions, + /*to_apply=*/function->ComputeProgramShape())); + TF_RET_CHECK(ShapeUtil::Compatible(reduce->shape(), inferred_return_shape)) + << "return shape is set to: " << ShapeUtil::HumanString(reduce->shape()) + << "but is inferred to be: " + << ShapeUtil::HumanString(inferred_return_shape); + + const Literal& arg_literal = parent_->GetEvaluatedLiteralFor(arg); + VLOG(3) << "HandleReduce arg_literal: " << arg_literal.ToString(); + const Literal& init_literal = parent_->GetEvaluatedLiteralFor(init_value); + VLOG(3) << "HandleReduce init_literal: " << init_literal.ToString(); + TF_RET_CHECK(ShapeUtil::IsScalar(init_literal.shape())); + auto init_scalar = init_literal.Get({}); + + auto result = Literal::CreateFromShape(reduce->shape()); + + const auto arg_dimensions = AsInt64Slice(arg_literal.shape().dimensions()); + std::vector arg_dim_steps(arg_dimensions.size()); + std::vector arg_dim_counts(arg_dimensions.size()); + for (const int64 dim : dimensions) { + arg_dim_steps[dim] = 1; + arg_dim_counts[dim] = arg_dimensions[dim]; + } + + // Create mapping from result index to arg index. + const int64 result_rank = ShapeUtil::Rank(result->shape()); + int64 result_dim = 0; + std::vector result_to_arg_index(result_rank); + for (int64 i = 0; i < arg_dimensions.size(); ++i) { + if (arg_dim_steps[i] == 0) { + result_to_arg_index[result_dim] = i; + ++result_dim; + } + } + + // For each resulting dimension, calculate and assign computed value. + TF_RETURN_IF_ERROR(result->Populate( + [&](tensorflow::gtl::ArraySlice multi_index) { + ReturnT result_val = init_scalar; + + std::vector base(arg_dimensions.size()); + for (int64 i = 0; i < multi_index.size(); ++i) { + base[result_to_arg_index[i]] = multi_index[i]; + } + + auto func = [&](const std::vector& input_index) { + auto curr_val = arg_literal.Get(input_index); + + // Evaluate computation with specified literal operands. + auto curr_val_literal = Literal::CreateR0(curr_val); + auto result_val_literal = Literal::CreateR0(result_val); + std::vector args = {curr_val_literal.get(), + result_val_literal.get()}; + + // We need a new visitor for each evaluation, so that the same + // computation can be visited more than once (with different + // inputs). + HloEvaluator embedded_evaluator; + std::unique_ptr computed_result = + embedded_evaluator.Evaluate(*function, args) + .ConsumeValueOrDie(); + + // Assign computed result to result_val. + result_val = computed_result->Get({}); + + return true; + }; + + ShapeUtil::ForEachIndex(arg_literal.shape(), base, arg_dim_counts, + arg_dim_steps, func); + + return result_val; + })); + + parent_->evaluated_[reduce] = std::move(result); + return Status::OK(); + }; + + Status HandleReduceWindow(HloInstruction* reduce_window, + HloInstruction* operand, const Window& window, + HloComputation* function) override { + TF_ASSIGN_OR_RETURN( + auto inferred_return_shape, + ShapeInference::InferReduceWindowShape( + /*operand_shape=*/reduce_window->operand(0)->shape(), + /*init_value=*/reduce_window->operand(1)->shape(), window, + /*to_apply_shape=*/function->ComputeProgramShape())); + TF_RET_CHECK( + ShapeUtil::Compatible(reduce_window->shape(), inferred_return_shape)) + << "return shape is set to: " + << ShapeUtil::HumanStringWithLayout(reduce_window->shape()) + << "but is inferred to be: " + << ShapeUtil::HumanStringWithLayout(inferred_return_shape); + + const Literal& operand_literal = + parent_->GetEvaluatedLiteralFor(reduce_window->operand(0)); + VLOG(3) << "HandleReduceWindow arg_literal: " << operand_literal.ToString(); + const Literal& init_literal = + parent_->GetEvaluatedLiteralFor(reduce_window->operand(1)); + VLOG(3) << "HandleReduceWindow init_literal: " << init_literal.ToString(); + TF_RET_CHECK(ShapeUtil::IsScalar(init_literal.shape())); + auto init_scalar = init_literal.Get({}); + + auto result = Literal::CreateFromShape(reduce_window->shape()); + + // Creates a Shape object from window, for iteration below. + std::vector window_dimension_sizes; + for (const auto& window_dimension : window.dimensions()) { + window_dimension_sizes.push_back(window_dimension.size()); + } + const Shape window_shape = ShapeUtil::MakeShape( + operand->shape().element_type(), window_dimension_sizes); + + DimensionVector window_index(window.dimensions_size()); + DimensionVector operand_index(ShapeUtil::Rank(operand_literal.shape())); + + // For each resulting dimension, calculate and assign computed value. + TF_RETURN_IF_ERROR(result->Populate( + [&](tensorflow::gtl::ArraySlice output_index) { + ReturnT result_val = init_scalar; + + std::fill(window_index.begin(), window_index.end(), 0); + std::fill(operand_index.begin(), operand_index.end(), 0); + + do { + // Set curr_val to 0 if out of bound (padded). + ReturnT curr_val = static_cast(0); + bool out_of_bound = false; + for (int i = 0; i < operand_index.size(); ++i) { + operand_index[i] = + output_index[i] * window.dimensions(i).stride() + + window_index[i] - window.dimensions(i).padding_low(); + if (operand_index[i] < 0 || + operand_index[i] >= operand_literal.shape().dimensions(i)) { + out_of_bound = true; + break; + } + } + if (!out_of_bound) { + curr_val = operand_literal.Get(operand_index); + } + // Evaluate computation with specified literal operands. + const auto curr_val_literal = Literal::CreateR0(curr_val); + const auto result_val_literal = + Literal::CreateR0(result_val); + const std::vector args = {curr_val_literal.get(), + result_val_literal.get()}; + // We need a new visitor for each evaluation, so that the same + // computation can be visited more than once (with different + // inputs). + HloEvaluator embedded_evaluator; + std::unique_ptr computed_result = + embedded_evaluator.Evaluate(*function, args) + .ConsumeValueOrDie(); + + result_val = computed_result->Get({}); + } while (IndexUtil::BumpIndices(window_shape, &window_index)); + + return result_val; + })); + + parent_->evaluated_[reduce_window] = std::move(result); + return Status::OK(); + }; + + Status HandleSlice(HloInstruction* slice, HloInstruction* operand) override { + const Shape& shape = slice->shape(); + TF_ASSIGN_OR_RETURN(auto inferred_return_shape, + ShapeInference::InferSliceShape( + operand->shape(), slice->slice_starts(), + slice->slice_limits(), slice->slice_strides())); + TF_RET_CHECK(ShapeUtil::Compatible(shape, inferred_return_shape)) + << "return shape set to: " << ShapeUtil::HumanString(shape) + << " but is inferred to be: " + << ShapeUtil::HumanString(inferred_return_shape); + + const int64 rank = ShapeUtil::Rank(operand->shape()); + auto operand_literal = parent_->GetEvaluatedLiteralFor(operand); + auto func = [&](tensorflow::gtl::ArraySlice out_index) { + DimensionVector operand_index(rank); + for (int64 i = 0; i < rank; ++i) { + operand_index[i] = + slice->slice_starts(i) + out_index[i] * slice->slice_strides(i); + } + return operand_literal.Get(operand_index); + }; + + auto result = Literal::CreateFromDimensions( + shape.element_type(), AsInt64Slice(shape.dimensions())); + TF_RETURN_IF_ERROR(result->Populate(func)); + parent_->evaluated_[slice] = std::move(result); + return Status::OK(); + }; + + private: + template + StatusOr> DynamicSlice( + const Literal& operand_literal, const Literal& start_indices_literal, + const Shape& result_shape) { + const auto& start_indices_typed = + start_indices_literal.GetArraySlice(); + std::vector start(start_indices_typed.begin(), + start_indices_typed.end()); + + std::vector operand_indices(start.size()); + + auto result = Literal::CreateFromShape(result_shape); + TF_RETURN_IF_ERROR(result->Populate( + [&](tensorflow::gtl::ArraySlice multi_index) { + for (int64 i = 0; i < operand_indices.size(); ++i) { + CHECK_GE(multi_index[i] + start[i], 0); + // Mod is only used here to be consistent with the existing + // backends' behavior. + operand_indices[i] = (multi_index[i] + start[i]) % + operand_literal.shape().dimensions(i); + } + + auto result = operand_literal.Get(operand_indices); + return result; + })); + + return std::move(result); + } + + template + StatusOr> DynamicUpdateSlice( + const Literal& operand_literal, const Literal& update_literal, + const Literal& start_indices_literal) { + const auto& start_indices_typed = + start_indices_literal.GetArraySlice(); + const std::vector start(start_indices_typed.begin(), + start_indices_typed.end()); + + auto result = MakeUnique(operand_literal); + std::vector result_index(ShapeUtil::Rank(result->shape()), 0); + + auto func = [&](const std::vector& update_index) { + std::transform(update_index.begin(), update_index.end(), start.begin(), + result_index.begin(), std::plus()); + + result->Set(result_index, + update_literal.Get(update_index)); + return true; + }; + + std::vector base(update_literal.shape().dimensions_size(), 0); + std::vector step(update_literal.shape().dimensions_size(), 1); + ShapeUtil::ForEachIndex(update_literal.shape(), base, + AsInt64Slice(update_literal.shape().dimensions()), + step, func); + + return std::move(result); + } + + StatusOr> ElementWiseUnaryOp( + HloInstruction* instruction, + const std::function& unary_op) { + const Literal& operand_literal = + parent_->GetEvaluatedLiteralFor(instruction->operand(0)); + return ElementWiseUnaryOpImpl(instruction, unary_op, + operand_literal); + } + + StatusOr> ElementWiseBinaryOp( + HloInstruction* instruction, + const std::function& binary_op) { + 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).c_str(), + ShapeUtil::HumanString(lhs->shape()).c_str(), + ShapeUtil::HumanString(rhs->shape()).c_str()); + } + + const Literal& lhs_literal = parent_->GetEvaluatedLiteralFor(lhs); + const Literal& rhs_literal = parent_->GetEvaluatedLiteralFor(rhs); + + auto result = Literal::CreateFromShape(shape); + + TF_RETURN_IF_ERROR(result->Populate( + [&](tensorflow::gtl::ArraySlice multi_index) { + return binary_op(lhs_literal.Get(multi_index), + rhs_literal.Get(multi_index)); + })); + return std::move(result); + } + + template + StatusOr> ElementWiseTernaryOp( + HloInstruction* instruction, + const std::function& ternary_op) { + const auto shape = instruction->shape(); + 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).c_str(), + ShapeUtil::HumanString(lhs->shape()).c_str(), + ShapeUtil::HumanString(rhs->shape()).c_str(), + ShapeUtil::HumanString(ehs->shape()).c_str()); + } + + const Literal& lhs_literal = parent_->GetEvaluatedLiteralFor(lhs); + const Literal& rhs_literal = parent_->GetEvaluatedLiteralFor(rhs); + const Literal& ehs_literal = parent_->GetEvaluatedLiteralFor(ehs); + + auto result = Literal::CreateFromShape(shape); + + TF_RETURN_IF_ERROR(result->Populate( + [&](tensorflow::gtl::ArraySlice multi_index) { + return ternary_op(lhs_literal.Get(multi_index), + rhs_literal.Get(multi_index), + ehs_literal.Get(multi_index)); + })); + + return std::move(result); + } + + HloEvaluator* parent_; +}; // namespace xla + +HloEvaluator::HloEvaluator() { + typed_visitors_[PRED] = MakeUnique>(this); + typed_visitors_[U8] = MakeUnique>(this); + typed_visitors_[U16] = MakeUnique([](HloInstruction*) { + return Unimplemented("unhandled primitive type: U16."); + }); + typed_visitors_[U32] = MakeUnique>(this); + typed_visitors_[U64] = MakeUnique>(this); + typed_visitors_[S8] = MakeUnique>(this); + typed_visitors_[S16] = MakeUnique([](HloInstruction*) { + return Unimplemented("unhandled primitive type: S16."); + }); + typed_visitors_[S32] = MakeUnique>(this); + typed_visitors_[S64] = MakeUnique>(this); + typed_visitors_[F16] = MakeUnique([](HloInstruction*) { + return Unimplemented("unhandled primitive type: F16."); + }); + typed_visitors_[F32] = MakeUnique>(this); + typed_visitors_[F64] = MakeUnique>(this); + typed_visitors_[TUPLE] = MakeUnique([](HloInstruction*) { + return Unimplemented("unhandled primitive type: TUPLE."); + }); + typed_visitors_[OPAQUE] = MakeUnique([](HloInstruction*) { + return Unimplemented("unhandled primitive type: OPAQUE."); + }); +} + +StatusOr> HloEvaluator::Evaluate( + const HloModule& module, + tensorflow::gtl::ArraySlice arg_literals) { + XLA_VLOG_LINES(2, "HloEvaluator::Evaluate module:\n" + module.ToString()); + + arg_literals_ = arg_literals; + evaluated_.clear(); + + TF_RETURN_IF_ERROR(module.entry_computation()->Accept(this)); + + return MakeUnique( + GetEvaluatedLiteralFor(module.entry_computation()->root_instruction())); +} + +StatusOr> HloEvaluator::Evaluate( + const HloComputation& computation, + tensorflow::gtl::ArraySlice arg_literals) { + XLA_VLOG_LINES( + 2, "HloEvaluator::Evaluate computation:\n" + computation.ToString()); + arg_literals_ = arg_literals; + evaluated_.clear(); + + TF_RETURN_IF_ERROR(computation.Accept(this)); + return MakeUnique( + GetEvaluatedLiteralFor(computation.root_instruction())); +} + +StatusOr> HloEvaluator::Evaluate( + HloInstruction* instruction, + tensorflow::gtl::ArraySlice operands) { + TF_RET_CHECK(hlo_query::AllOperandsAreParametersOrConstants(*instruction)); + TF_RETURN_IF_ERROR(ShapeUtil::ValidateShape(instruction->shape())); + + arg_literals_ = operands; + evaluated_.clear(); + + // 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] = MakeUnique(*input_literal); + } + } + + TF_RETURN_IF_ERROR(Preprocess(instruction)); + TF_RETURN_IF_ERROR(instruction->Visit(this)); + TF_RETURN_IF_ERROR(Postprocess(instruction)); + return MakeUnique(GetEvaluatedLiteralFor(instruction)); +} + +StatusOr> HloEvaluator::Evaluate( + HloInstruction* instruction) { + TF_RET_CHECK(hlo_query::AllOperandsAreConstants(*instruction)); + TF_RET_CHECK(instruction->opcode() != HloOpcode::kParameter); + TF_RETURN_IF_ERROR(ShapeUtil::ValidateShape(instruction->shape())); + + arg_literals_.clear(); + evaluated_.clear(); + + TF_RETURN_IF_ERROR(Preprocess(instruction)); + TF_RETURN_IF_ERROR(instruction->Visit(this)); + TF_RETURN_IF_ERROR(Postprocess(instruction)); + return MakeUnique(GetEvaluatedLiteralFor(instruction)); +} + +std::unique_ptr HloEvaluator::TryEvaluate( + HloInstruction* instruction) { + auto result_or = Evaluate(instruction); + if (!result_or.ok()) { + VLOG(1) << "TryEvaluate failed:" << result_or.status(); + return nullptr; + } + + return result_or.ConsumeValueOrDie(); +} + +Status HloEvaluator::HandleParameter(HloInstruction* parameter) { + 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())); + + evaluated_[parameter] = MakeUnique(*input_literal); + return Status::OK(); +} + +Status HloEvaluator::HandleConstant(HloInstruction* constant, + const Literal& literal) { + return Status::OK(); +} + +Status HloEvaluator::HandleReshape(HloInstruction* reshape) { + TF_ASSIGN_OR_RETURN( + evaluated_[reshape], + GetEvaluatedLiteralFor(reshape->operand(0)) + .Reshape(AsInt64Slice(reshape->shape().dimensions()))); + return Status::OK(); +} + +Status HloEvaluator::HandleTranspose(HloInstruction* transpose) { + evaluated_[transpose] = GetEvaluatedLiteralFor(transpose->operand(0)) + .Transpose(transpose->dimensions()); + return Status::OK(); +} + +Status HloEvaluator::HandleConcatenate( + HloInstruction* concatenate, + tensorflow::gtl::ArraySlice operands) { + // 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::IsTuple(reference_shape)); + const int64 rank = ShapeUtil::Rank(reference_shape); + const int64 concat_dim = concatenate->dimensions()[0]; + CHECK_GE(concat_dim, 0); + CHECK_LT(concat_dim, rank); + + DimensionVector concat_dimensions(reference_shape.dimensions().begin(), + reference_shape.dimensions().end()); + + for (int64 i = 1; i < operands.size(); ++i) { + const Shape& operand_shape = operands[i]->shape(); + CHECK(!ShapeUtil::IsTuple(operand_shape)); + // Accumulate the concat dimension from all tensors taking part to the + // operation. + concat_dimensions[concat_dim] += + ShapeUtil::GetDimension(operand_shape, concat_dim); + } + + auto result_literal = Literal::CreateFromDimensions( + reference_shape.element_type(), concat_dimensions); + DimensionVector source_indices(rank, 0); + DimensionVector dest_indices(concat_dimensions.size(), 0); + + for (auto operand : operands) { + const Shape& operand_shape = operand->shape(); + TF_RETURN_IF_ERROR(result_literal->Copy( + GetEvaluatedLiteralFor(operand), source_indices, dest_indices, + AsInt64Slice(operand_shape.dimensions()))); + dest_indices[concat_dim] += + ShapeUtil::GetDimension(operand_shape, concat_dim); + } + + evaluated_[concatenate] = std::move(result_literal); + return Status::OK(); +} + +Status HloEvaluator::HandleIsFinite(HloInstruction* is_finite, + HloInstruction* operand) { + if (!ShapeUtil::ElementIsFloating(operand->shape())) { + return InvalidArgument( + "expected element type in shape to be float for IsFinite op, got: %s", + PrimitiveType_Name(operand->shape().element_type()).c_str()); + } + + switch (operand->shape().element_type()) { + case F16: + return Unimplemented("unhandled primitive type: F16."); + case F32: { + auto result_or = ElementWiseUnaryOpImpl( + is_finite, + [](float elem_operand) { return std::isfinite(elem_operand); }, + GetEvaluatedLiteralFor(operand)); + TF_ASSIGN_OR_RETURN(evaluated_[is_finite], std::move(result_or)); + break; + } + case F64: { + auto result_or = ElementWiseUnaryOpImpl( + is_finite, + [](double elem_operand) { return std::isfinite(elem_operand); }, + GetEvaluatedLiteralFor(operand)); + TF_ASSIGN_OR_RETURN(evaluated_[is_finite], std::move(result_or)); + break; + } + default: + LOG(FATAL) << "HandleIsFinite: unknown/unhandled primitive type: " + << PrimitiveType_Name(operand->shape().element_type()); + } + + return Status::OK(); +} + +Status HloEvaluator::HandleCompare(HloInstruction* compare, HloOpcode opcode, + HloInstruction* lhs, HloInstruction* rhs) { + // TODO(b/35950897, b/27796129): add DCHECK back once implicit broadcast is + // removed. + if (!(ShapeUtil::SameDimensions(compare->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(compare->shape()).c_str(), + ShapeUtil::HumanString(lhs->shape()).c_str(), + ShapeUtil::HumanString(rhs->shape()).c_str()); + } + + TF_RET_CHECK(lhs->shape().element_type() == rhs->shape().element_type()); + + const Literal& lhs_literal = GetEvaluatedLiteralFor(lhs); + const Literal& rhs_literal = GetEvaluatedLiteralFor(rhs); + + // Note here we switch on the operand's type. + switch (lhs->shape().element_type()) { + case PRED: { + TF_ASSIGN_OR_RETURN( + evaluated_[compare], + Compare(compare->shape(), opcode, lhs_literal, rhs_literal)); + } break; + case U8: { + TF_ASSIGN_OR_RETURN( + evaluated_[compare], + Compare(compare->shape(), opcode, lhs_literal, rhs_literal)); + } break; + case U16: + return Unimplemented("unhandled primitive type: U16."); + case U32: { + TF_ASSIGN_OR_RETURN( + evaluated_[compare], + Compare(compare->shape(), opcode, lhs_literal, rhs_literal)); + } break; + case U64: { + TF_ASSIGN_OR_RETURN( + evaluated_[compare], + Compare(compare->shape(), opcode, lhs_literal, rhs_literal)); + } break; + case S8: { + TF_ASSIGN_OR_RETURN( + evaluated_[compare], + Compare(compare->shape(), opcode, lhs_literal, rhs_literal)); + } break; + case S16: + return Unimplemented("unhandled primitive type: S16."); + case S32: { + TF_ASSIGN_OR_RETURN( + evaluated_[compare], + Compare(compare->shape(), opcode, lhs_literal, rhs_literal)); + } break; + case S64: { + TF_ASSIGN_OR_RETURN( + evaluated_[compare], + Compare(compare->shape(), opcode, lhs_literal, rhs_literal)); + } break; + case F16: + return Unimplemented("unhandled primitive type: F16."); + case F32: { + TF_ASSIGN_OR_RETURN( + evaluated_[compare], + Compare(compare->shape(), opcode, lhs_literal, rhs_literal)); + } break; + case F64: { + 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()); + } + + return Status::OK(); +} + +Status HloEvaluator::HandleTuple( + HloInstruction* tuple, + tensorflow::gtl::ArraySlice operands) { + std::vector operand_literals; + for (auto operand : operands) { + operand_literals.push_back(&GetEvaluatedLiteralFor(operand)); + } + + evaluated_[tuple] = Literal::MakeTuple(operand_literals); + return Status::OK(); +} + +Status HloEvaluator::HandleGetTupleElement(HloInstruction* get_tuple_element, + HloInstruction* operand) { + const auto result_shape = get_tuple_element->shape(); + const int64 index = get_tuple_element->tuple_index(); + + TF_ASSIGN_OR_RETURN( + auto inferred_return_shape, + ShapeInference::InferGetTupleElementShape(operand->shape(), index)); + TF_RET_CHECK(ShapeUtil::Compatible(result_shape, inferred_return_shape)) + << "return shape set to: " << ShapeUtil::HumanString(result_shape) + << " but is inferred to be: " + << ShapeUtil::HumanString(inferred_return_shape); + + const Literal& operand_tuple_literal = GetEvaluatedLiteralFor(operand); + + evaluated_[get_tuple_element] = + MakeUnique(operand_tuple_literal.tuple_literals(index)); + + return Status::OK(); +} + +Status HloEvaluator::HandleCopy(HloInstruction* copy) { + TF_RET_CHECK(ShapeUtil::Compatible(copy->shape(), copy->operand(0)->shape())); + + auto result = MakeUnique(GetEvaluatedLiteralFor(copy->operand(0))); + evaluated_[copy] = std::move(result); + return Status::OK(); +} + +Status HloEvaluator::Preprocess(HloInstruction* hlo) { + VLOG(2) << "About to visit HLO: " << hlo->ToString(); + return Status::OK(); +} + +Status HloEvaluator::Postprocess(HloInstruction* hlo) { + VLOG(2) << "Finished visiting " << hlo->ToString() + << "; evaluated value is: " << GetEvaluatedLiteralFor(hlo).ToString(); + return Status::OK(); +} + +} // namespace xla diff --git a/tensorflow/compiler/xla/service/hlo_evaluator.h b/tensorflow/compiler/xla/service/hlo_evaluator.h new file mode 100644 index 0000000000000000000000000000000000000000..66a53e1fa5a219a60665198a03026ad36cc4c117 --- /dev/null +++ b/tensorflow/compiler/xla/service/hlo_evaluator.h @@ -0,0 +1,178 @@ +/* Copyright 2017 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#ifndef THIRD_PARTY_TENSORFLOW_COMPILER_XLA_SERVICE_HLO_EVALUATOR_H_ +#define THIRD_PARTY_TENSORFLOW_COMPILER_XLA_SERVICE_HLO_EVALUATOR_H_ + +#include + +#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" +#include "tensorflow/compiler/xla/statusor.h" +#include "tensorflow/compiler/xla/util.h" +#include "tensorflow/compiler/xla/xla_data.pb.h" +#include "tensorflow/core/lib/gtl/array_slice.h" +#include "tensorflow/core/lib/gtl/flatmap.h" +#include "tensorflow/core/platform/macros.h" + +namespace xla { + +// Responsible for evaluating HLO and obtain literal as the evaluation results. +// +// This class is not thread-safe. +class HloEvaluator : public DfsHloVisitorWithDefault { + public: + HloEvaluator(); + // Evaluates an HLO module and an array of pointers to literals. + // Returns the evaluated result as a literal if successful. + // Precondition: argument literals correspond to each input computation's + // parameters in their post-ordering. See comment below for example. + StatusOr> Evaluate( + const HloModule& module, + tensorflow::gtl::ArraySlice arg_literals); + + // Evaluates an HLO computation and an array of pointers to literals. + // Returns the evaluated result as a literal if successful. + // Precondition: argument literals correspond to the input computation's + // parameters in their post-ordering. For e.g., consider the following graph: + // + // * + // / \ + // + Parameter1 + // / \ + // / \ + // Parameter0 Constant + // + // The input literals array will have its first literal map to Parameter0 and + // the second map to Parameter1. + StatusOr> Evaluate( + const HloComputation& computation, + tensorflow::gtl::ArraySlice 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. + // TODO(b/35950897): implement more ops other than element-wise ops. + StatusOr> Evaluate( + HloInstruction* instruction, + tensorflow::gtl::ArraySlice 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. + StatusOr> Evaluate(HloInstruction* instruction); + + // Same as Evaluate, except returning nullptr on error. + std::unique_ptr TryEvaluate(HloInstruction* instruction); + + protected: + // Templated DfsHloVisitor. Typically ReturnT here indicates the resulting + // literal type of each evaluated Handle* method of a TypedVisitor. + // There are however a few notable exceptions to this is rule, notably: + // - HandleCompare and HandleIsFinite: where the resulting literal type is + // always boolean. + // These operations are handled outside of the parent HloEvaluator handlers + // instead of from within TypedVisitor. + template + class TypedVisitor; + + // Wraps around instruction handling to infer types before dispatching to + // the corresponding typed Visitor. + Status DefaultAction(HloInstruction* hlo) override { + return hlo->Visit(typed_visitors_.at(hlo->shape().element_type()).get()); + } + + Status Preprocess(HloInstruction* hlo) override; + + Status Postprocess(HloInstruction* hlo) override; + + // Operations that are type-agnostic or always return a specific type, such as + // HandleIsFinite where boolean is always returned. + // + Status HandleParameter(HloInstruction* parameter) override; + + Status HandleConstant(HloInstruction* constant, + const Literal& literal) override; + + Status HandleConcatenate( + HloInstruction* concatenate, + tensorflow::gtl::ArraySlice operands) override; + + Status HandleReshape(HloInstruction* reshape) override; + + Status HandleTranspose(HloInstruction* transpose) override; + + Status HandleIsFinite(HloInstruction* is_finite, + HloInstruction* operand) override; + + Status HandleCompare(HloInstruction* compare, HloOpcode opcode, + HloInstruction* lhs, HloInstruction* rhs) override; + Status HandleTuple( + HloInstruction* tuple, + tensorflow::gtl::ArraySlice operands) override; + + Status HandleGetTupleElement(HloInstruction* get_tuple_element, + HloInstruction* operand) override; + + Status HandleCopy(HloInstruction* copy) override; + + private: + // 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. + // Crash with log if the given instruction has not been evaluated previously. + const Literal& GetEvaluatedLiteralFor(const HloInstruction* hlo) { + if (hlo->IsConstant()) { + return hlo->literal(); + } + auto it = evaluated_.find(hlo); + CHECK(it != evaluated_.end()) + << "could not find evaluated value for: " << hlo->ToString(); + return *(it->second); + } + + // Map from a primitive type to its associated (templated) DfsHloVisitor. + // Note: the hash function here is only needed because current gcc std::hash + // does not specialize for enum types. This should however be fixed in the + // future: https://gcc.gnu.org/bugzilla/show_bug.cgi?id=60970#c5 + tensorflow::gtl::FlatMap, + std::hash> + typed_visitors_; + + // Tracks the HLO instruction and its evaluated literal result. + // 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. + tensorflow::gtl::FlatMap> + evaluated_; + + // Stores input literals, assuming they are in post-order. Literals are not + // owned by this class, and they must outlive the lifetime of the instance of + // this class. + tensorflow::gtl::ArraySlice arg_literals_; + + TF_DISALLOW_COPY_AND_ASSIGN(HloEvaluator); +}; + +} // namespace xla + +#endif // THIRD_PARTY_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 new file mode 100644 index 0000000000000000000000000000000000000000..a8265483492d5e4eaa0eb75599218437b3dc2f28 --- /dev/null +++ b/tensorflow/compiler/xla/service/hlo_evaluator_test.cc @@ -0,0 +1,1553 @@ +/* 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/hlo_evaluator.h" + +#include +#include +#include +#include +#include +#include + +#include "tensorflow/compiler/xla/client/computation_builder.h" +#include "tensorflow/compiler/xla/literal_util.h" +#include "tensorflow/compiler/xla/reference_util.h" +#include "tensorflow/compiler/xla/service/hlo_computation.h" +#include "tensorflow/compiler/xla/service/hlo_instruction.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/test.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" +#include "tensorflow/compiler/xla/util.h" +#include "tensorflow/compiler/xla/xla_data.pb.h" +#include "tensorflow/core/platform/types.h" + +namespace xla { +namespace { + +class HloEvaluatorTest : public HloTestBase { + protected: + HloEvaluatorTest() { evaluator_ = MakeUnique(); } + + std::unique_ptr evaluator_; +}; + +// Verifies that HloEvaluator evaluates a HLO instruction that performs clamp +// with 3 operands. +TEST_F(HloEvaluatorTest, DoesClamp) { + auto low = Literal::CreateR2({{0.f, 2.f}, {2.f, 4.f}}); + auto high = Literal::CreateR2({{2.f, 4.f}, {4.f, 4.f}}); + auto value = Literal::CreateR2({{0.f, 5.f}, {0.f, 4.f}}); + + Shape shape = low->shape(); + HloComputation::Builder b(TestName()); + auto c1 = b.AddInstruction(HloInstruction::CreateConstant(std::move(low))); + auto c2 = b.AddInstruction(HloInstruction::CreateConstant(std::move(high))); + auto c3 = b.AddInstruction(HloInstruction::CreateConstant(std::move(value))); + auto instruction = b.AddInstruction( + HloInstruction::CreateTernary(shape, HloOpcode::kClamp, c1, c2, c3)); + HloModule module(TestName()); + module.AddEntryComputation(b.Build()); + + std::unique_ptr result = + evaluator_->Evaluate(instruction, {}).ConsumeValueOrDie(); + + auto expected = Literal::CreateR2({{0, 4}, {2, 4}}); + + LiteralTestUtil::ExpectEqual(*expected, *result); +} + +// Verifies that HloEvaluator evaluates a HLO instruction that performs select +// with 3 operands. +TEST_F(HloEvaluatorTest, DoesSelect) { + auto pred = Literal::CreateR2({{true, false}, {false, true}}); + auto on_true = Literal::CreateR2({{2.f, 4.f}, {4.f, 4.f}}); + auto on_false = Literal::CreateR2({{0.f, 5.f}, {0.f, 4.f}}); + + Shape shape = on_true->shape(); + HloComputation::Builder b(TestName()); + auto c1 = b.AddInstruction(HloInstruction::CreateConstant(std::move(pred))); + auto c2 = + b.AddInstruction(HloInstruction::CreateConstant(std::move(on_true))); + auto c3 = + b.AddInstruction(HloInstruction::CreateConstant(std::move(on_false))); + auto instruction = b.AddInstruction( + HloInstruction::CreateTernary(shape, HloOpcode::kSelect, c1, c2, c3)); + HloModule module(TestName()); + module.AddEntryComputation(b.Build()); + + std::unique_ptr result = + evaluator_->Evaluate(instruction, {}).ConsumeValueOrDie(); + + auto expected = Literal::CreateR2({{2, 5}, {0, 4}}); + + LiteralTestUtil::ExpectEqual(*expected, *result); +} + +// Verifies that HloEvaluator evaluates a HLO instruction that performs +// element-wise addition with 2 operands. +TEST_F(HloEvaluatorTest, DoesAdd) { + auto lhs = Literal::CreateR2({{1, 0}, {-100, 4}}); + auto rhs = Literal::CreateR2({{2, 4}, {4, 4}}); + + Shape shape = ShapeUtil::MakeShape(S64, {2, 2}); + HloComputation::Builder b(TestName()); + auto c1 = b.AddInstruction(HloInstruction::CreateConstant(std::move(lhs))); + auto c2 = b.AddInstruction(HloInstruction::CreateConstant(std::move(rhs))); + auto instruction = b.AddInstruction( + HloInstruction::CreateBinary(shape, HloOpcode::kAdd, c1, c2)); + HloModule module(TestName()); + module.AddEntryComputation(b.Build()); + + std::unique_ptr result = + evaluator_->Evaluate(instruction, {}).ConsumeValueOrDie(); + + auto expected = Literal::CreateR2({{3, 4}, {-96, 8}}); + + LiteralTestUtil::ExpectEqual(*expected, *result); +} + +// Verifies that HloEvaluator evaluates a HLO instruction that performs +// element-wise divide with 2 operands. +TEST_F(HloEvaluatorTest, DoesDivide) { + { + auto lhs_s64 = Literal::CreateR2({{1, 0}, {-100, 4}}); + auto rhs_s64 = Literal::CreateR2({{2, 4}, {4, 4}}); + + Shape shape_s64 = ShapeUtil::MakeShape(S64, {2, 2}); + HloComputation::Builder b(TestName()); + auto c1_s64 = + b.AddInstruction(HloInstruction::CreateConstant(std::move(lhs_s64))); + auto c2_s64 = + b.AddInstruction(HloInstruction::CreateConstant(std::move(rhs_s64))); + auto instruction = b.AddInstruction(HloInstruction::CreateBinary( + shape_s64, HloOpcode::kDivide, c1_s64, c2_s64)); + HloModule module(TestName()); + module.AddEntryComputation(b.Build()); + + std::unique_ptr result = + evaluator_->Evaluate(instruction, {}).ConsumeValueOrDie(); + + auto expected = Literal::CreateR2({{0, 0}, {-25, 1}}); + + LiteralTestUtil::ExpectEqual(*expected, *result); + } + { + auto lhs_f64 = Literal::CreateR2({{1.0, 0.0}, {-100.0, 4.0}}); + auto rhs_f64 = Literal::CreateR2({{2.2, 4.0}, {4.0, 4.0}}); + + Shape shape_f64 = ShapeUtil::MakeShape(F64, {2, 2}); + HloComputation::Builder b(TestName()); + auto c1_f64 = + b.AddInstruction(HloInstruction::CreateConstant(std::move(lhs_f64))); + auto c2_f64 = + b.AddInstruction(HloInstruction::CreateConstant(std::move(rhs_f64))); + auto instruction = b.AddInstruction(HloInstruction::CreateBinary( + shape_f64, HloOpcode::kDivide, c1_f64, c2_f64)); + HloModule module(TestName()); + module.AddEntryComputation(b.Build()); + + auto result = evaluator_->Evaluate(instruction, {}).ConsumeValueOrDie(); + + auto expected = + Literal::CreateR2({{0.45454545454545453, 0}, {-25, 1}}); + + LiteralTestUtil::ExpectEqual(*expected, *result); + } +} + +// Verifies that HloEvaluator evaluates a HLO instruction that performs +// element-wise abs op with 1 operand. +TEST_F(HloEvaluatorTest, DoesAbs) { + { + auto operand = Literal::CreateR2({{1, -20}, {-100, 4}}); + const Shape& shape = ShapeUtil::MakeShape(S64, {2, 2}); + HloComputation::Builder b(TestName()); + auto c1 = + b.AddInstruction(HloInstruction::CreateConstant(std::move(operand))); + auto instruction = b.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kAbs, c1)); + HloModule module(TestName()); + module.AddEntryComputation(b.Build()); + + std::unique_ptr result = + evaluator_->Evaluate(instruction, {}).ConsumeValueOrDie(); + + auto expected = Literal::CreateR2({{1, 20}, {100, 4}}); + + LiteralTestUtil::ExpectEqual(*expected, *result); + } + + // For R0 literal. + { + const Shape& r0 = ShapeUtil::MakeShape(F32, {}); + auto operand = Literal::CreateR0(-1.0f); + HloComputation::Builder b(TestName()); + auto c1 = + b.AddInstruction(HloInstruction::CreateConstant(std::move(operand))); + auto instruction = + b.AddInstruction(HloInstruction::CreateUnary(r0, HloOpcode::kAbs, c1)); + HloModule module(TestName()); + module.AddEntryComputation(b.Build()); + + auto result = evaluator_->Evaluate(instruction).ConsumeValueOrDie(); + auto expected = Literal::CreateR0(1.0f); + + LiteralTestUtil::ExpectEqual(*expected, *result); + } + + // For R1 literal with dimension of size 0. + { + Shape empty_r1 = ShapeUtil::MakeShape(F32, {0}); + auto operand = Literal::CreateR1({}); + HloComputation::Builder b(TestName()); + auto c1 = + b.AddInstruction(HloInstruction::CreateConstant(std::move(operand))); + auto instruction = b.AddInstruction( + HloInstruction::CreateUnary(empty_r1, HloOpcode::kAbs, c1)); + HloModule module(TestName()); + module.AddEntryComputation(b.Build()); + + auto result = evaluator_->Evaluate(instruction).ConsumeValueOrDie(); + auto expected = Literal::CreateR1({}); + + LiteralTestUtil::ExpectEqual(*expected, *result); + } +} // namespace + +// Verifies that HloEvaluator evaluates a HLO Computation with non-parameter nor +// constant operands. +TEST_F(HloEvaluatorTest, DoesTraverseInstructions) { + auto lhs = Literal::CreateR2({{1, 0}, {-100, 4}}); + auto rhs = Literal::CreateR2({{2, 4}, {4, 4}}); + auto rhs2 = Literal::CreateR2({{1, -20}, {-100, 4}}); + std::vector args = {lhs.get(), rhs.get(), rhs2.get()}; + + Shape shape = ShapeUtil::MakeShape(S64, {2, 2}); + + HloComputation::Builder b(TestName()); + auto param_lhs = + b.AddInstruction(HloInstruction::CreateParameter(0, shape, "lhs")); + auto param_rhs = + b.AddInstruction(HloInstruction::CreateParameter(1, shape, "rhs")); + auto lhs_instruction = b.AddInstruction(HloInstruction::CreateBinary( + shape, HloOpcode::kAdd, param_lhs, param_rhs)); + + auto param_rhs2 = + b.AddInstruction(HloInstruction::CreateParameter(2, shape, "rhs2")); + b.AddInstruction(HloInstruction::CreateBinary(shape, HloOpcode::kAdd, + lhs_instruction, param_rhs2)); + HloModule module(TestName()); + auto computation = module.AddEntryComputation(b.Build()); + + std::unique_ptr result = + evaluator_->Evaluate(*computation, args).ConsumeValueOrDie(); + + auto expected = Literal::CreateR2({{4, -16}, {-196, 12}}); + + LiteralTestUtil::ExpectEqual(*expected, *result); +} + +// Verifies Reshape operation is correctly evaluated. +TEST_F(HloEvaluatorTest, DoesReshape) { + HloComputation::Builder b(TestName()); + const int64 dimensions[] = {11, 8, 7, 5, 9}; + TF_ASSERT_OK_AND_ASSIGN(auto literal, + LiteralTestUtil::CreateRandomLiteral( + ShapeUtil::MakeShape(F32, dimensions), 0.0, 1.0)); + auto literal_clone = literal->CloneToUnique(); + HloInstruction* literal_instruction = + b.AddInstruction(HloInstruction::CreateConstant(std::move(literal))); + + Shape shape = ShapeUtil::MakeShape(F32, {8, 7, 11, 9, 5}); + const int64 permutation[] = {1, 2, 0, 4, 3}; + b.AddInstruction( + HloInstruction::CreateTranspose(shape, literal_instruction, permutation)); + HloModule module(TestName()); + auto computation = module.AddEntryComputation(b.Build()); + + std::unique_ptr result = + evaluator_->Evaluate(*computation, {}).ConsumeValueOrDie(); + + using NativeT = typename primitive_util::PrimitiveTypeToNative::type; + result->EachCell( + [&](tensorflow::gtl::ArraySlice indices, NativeT value) { + std::vector rindexes = Permute(permutation, indices); + EXPECT_TRUE(value == literal_clone->Get(rindexes)); + }); +} + +// Verifies Broadcast operation is correctly evaluated. +TEST_F(HloEvaluatorTest, DoesBroadcast) { + HloComputation::Builder b(TestName()); + auto input_literal = Literal::CreateR2({{1, 2}, {3, 4}, {5, 6}}); + auto output_literal = Literal::CreateR3( + {{{1, 2}, {3, 4}, {5, 6}}, {{1, 2}, {3, 4}, {5, 6}}}); + HloInstruction* literal_instruction = b.AddInstruction( + HloInstruction::CreateConstant(std::move(input_literal))); + b.AddInstruction(HloInstruction::CreateBroadcast( + output_literal->shape(), literal_instruction, {1, 2})); + HloModule module(TestName()); + auto computation = module.AddEntryComputation(b.Build()); + + std::unique_ptr result = + evaluator_->Evaluate(*computation, {}).ConsumeValueOrDie(); + + LiteralTestUtil::ExpectEqual(*result, *output_literal); +} + +TEST_F(HloEvaluatorTest, DoesBroadcastScalar) { + HloComputation::Builder b(TestName()); + auto input_literal = Literal::CreateR0(111); + auto output_literal = Literal::CreateR2( + {{111, 111}, {111, 111}, {111, 111}, {111, 111}, {111, 111}, {111, 111}}); + + HloInstruction* literal_instruction = b.AddInstruction( + HloInstruction::CreateConstant(std::move(input_literal))); + // Broadcast dimension is ignored in the case of scalars. + b.AddInstruction(HloInstruction::CreateBroadcast( + output_literal->shape(), literal_instruction, + /*broadcast_dimensions=*/{1})); + HloModule module(TestName()); + auto computation = module.AddEntryComputation(b.Build()); + + std::unique_ptr result = + evaluator_->Evaluate(*computation, {}).ConsumeValueOrDie(); + + LiteralTestUtil::ExpectEqual(*result, *output_literal); +} + +TEST_F(HloEvaluatorTest, ConvertWithSameLayout) { + HloComputation::Builder b(TestName()); + + auto input_literal = Literal::CreateR2({{1, 2}, {3, 4}, {5, 6}}); + auto expected = + Literal::CreateR2({{1.0, 2.0}, {3.0, 4.0}, {5.0, 6.0}}); + ASSERT_TRUE(LayoutUtil::LayoutsInShapesEqual(input_literal->shape(), + expected->shape())); + + HloInstruction* constant = b.AddInstruction( + HloInstruction::CreateConstant(std::move(input_literal))); + b.AddInstruction(HloInstruction::CreateConvert(expected->shape(), constant)); + HloModule module(TestName()); + auto computation = module.AddEntryComputation(b.Build()); + + std::unique_ptr result = + evaluator_->Evaluate(*computation, {}).ConsumeValueOrDie(); + + LiteralTestUtil::ExpectEqual(*result, *expected); +} + +TEST_F(HloEvaluatorTest, ConvertWithDifferentLayout) { + HloComputation::Builder b(TestName()); + + auto input_literal = Literal::CreateR2WithLayout( + {{1, 2}, {3, 4}, {5, 6}}, LayoutUtil::MakeLayout({0, 1})); + auto expected = Literal::CreateR2WithLayout( + {{1.0, 2.0}, {3.0, 4.0}, {5.0, 6.0}}, LayoutUtil::MakeLayout({1, 0})); + ASSERT_FALSE(LayoutUtil::LayoutsInShapesEqual(input_literal->shape(), + expected->shape())); + + HloInstruction* constant = b.AddInstruction( + HloInstruction::CreateConstant(std::move(input_literal))); + b.AddInstruction(HloInstruction::CreateConvert(expected->shape(), constant)); + HloModule module(TestName()); + auto computation = module.AddEntryComputation(b.Build()); + + std::unique_ptr result = + evaluator_->Evaluate(*computation, {}).ConsumeValueOrDie(); + + LiteralTestUtil::ExpectEqual(*result, *expected); +} + +PaddingConfig CreatePaddingConfig( + std::initializer_list> padding_dimensions) { + PaddingConfig padding_config; + + for (auto& paddings_per_dim : padding_dimensions) { + auto dimension = padding_config.add_dimensions(); + dimension->set_edge_padding_low(paddings_per_dim[0]); + dimension->set_edge_padding_high(paddings_per_dim[1]); + dimension->set_interior_padding(paddings_per_dim[2]); + } + return padding_config; +} + +TEST_F(HloEvaluatorTest, Pad2DIntegerArrayWithZeroDimension) { + auto operand = Literal::CreateR2({{}, {}}); + HloComputation::Builder b(TestName()); + auto operand_instruction = + b.AddInstruction(HloInstruction::CreateConstant(std::move(operand))); + + constexpr int32 kPadValue = 10; + auto pad_value = Literal::CreateR0(kPadValue); + auto padding_value_instruction = + b.AddInstruction(HloInstruction::CreateConstant(std::move(pad_value))); + + auto padding_config = CreatePaddingConfig({{{1, 0, 2}}, {{0, 2, 1}}}); + Shape shape = ShapeUtil::MakeShape(S32, {5, 2}); + auto pad_instruction = b.AddInstruction(HloInstruction::CreatePad( + shape, operand_instruction, padding_value_instruction, padding_config)); + HloModule module(TestName()); + module.AddEntryComputation(b.Build()); + + auto result = evaluator_->Evaluate(pad_instruction).ConsumeValueOrDie(); + + auto expected = Literal::CreateR2( + {{10, 10}, {10, 10}, {10, 10}, {10, 10}, {10, 10}}); + + LiteralTestUtil::ExpectEqual(*expected, *result); +} + +TEST_F(HloEvaluatorTest, Pad4DFloatArrayWithInteriorPadding) { + HloComputation::Builder b(TestName()); + + Array4D input_array(3, 2, 1, 1, {1, 2, 3, 4, 5, 6}); + auto input = Literal::CreateR4FromArray4D(input_array); + HloInstruction* input_instruction = + b.AddInstruction(HloInstruction::CreateConstant(std::move(input))); + constexpr float kPadValue = 1.5; + auto pad_value = Literal::CreateR0(kPadValue); + HloInstruction* pad_instruction = + b.AddInstruction(HloInstruction::CreateConstant(std::move(pad_value))); + + Shape shape = ShapeUtil::MakeShape(F32, {8, 5, 1, 1}); + auto r4_padding_on_dim0_dim1 = + 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)); + HloModule module(TestName()); + auto computation = module.AddEntryComputation(b.Build()); + + std::unique_ptr result = + evaluator_->Evaluate(*computation, {}).ConsumeValueOrDie(); + + auto expected_array = MakeUnique>(8, 5, 1, 1); + expected_array->Fill(kPadValue); + (*expected_array)(1, 0, 0, 0) = 1.0f; + (*expected_array)(1, 2, 0, 0) = 2.0f; + (*expected_array)(4, 0, 0, 0) = 3.0f; + (*expected_array)(4, 2, 0, 0) = 4.0f; + (*expected_array)(7, 0, 0, 0) = 5.0f; + (*expected_array)(7, 2, 0, 0) = 6.0f; + + auto expected = Literal::CreateR4FromArray4D(*expected_array); + + LiteralTestUtil::ExpectEqual(*expected, *result); +} + +TEST_F(HloEvaluatorTest, NegativePadding2D) { + HloComputation::Builder b(TestName()); + + // input_array: + // f32[4,3] { + // { 1, 2, 3 }, + // { 5, 6, 7 }, + // { 9, 10, 11 }, + // { 13, 14, 15 }, + // } + auto input_array = MakeUnique>(4, 3); + input_array->FillUnique(1.0f); + auto input = Literal::CreateR2FromArray2D(*input_array); + HloInstruction* input_instruction = + b.AddInstruction(HloInstruction::CreateConstant(std::move(input))); + + auto pad_value_instruction = b.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(2.718f))); + + auto r2_padding_on_dim0_dim1 = + CreatePaddingConfig({{{-1, -2, 0}}, {{-2, 4, 0}}}); + Shape shape = ShapeUtil::MakeShape(F32, {1, 5}); + b.AddInstruction(HloInstruction::CreatePad(shape, input_instruction, + pad_value_instruction, + r2_padding_on_dim0_dim1)); + + HloModule module(TestName()); + auto computation = module.AddEntryComputation(b.Build()); + + std::unique_ptr result = + evaluator_->Evaluate(*computation, {}).ConsumeValueOrDie(); + + // f32[1,5] { 7.0, 2.718, 2.718, 2.718, 2.718 } + auto expected_array = MakeUnique>(1, 5); + (*expected_array)(0, 0) = 7.0f; + (*expected_array)(0, 1) = 2.718f; + (*expected_array)(0, 2) = 2.718f; + (*expected_array)(0, 3) = 2.718f; + (*expected_array)(0, 4) = 2.718f; + auto expected = Literal::CreateR2FromArray2D(*expected_array); + + LiteralTestUtil::ExpectEqual(*expected, *result); +} + +TEST_F(HloEvaluatorTest, NegativeAndInteriorPadding2D) { + HloComputation::Builder b(TestName()); + + // f32[4,3] { + // { 1, 2, 3 }, + // { 5, 6, 7 }, + // { 9, 10, 11 }, + // { 13, 14, 15 }, + // } + auto input_array = MakeUnique>(4, 3); + input_array->FillUnique(1.0f); + auto input = Literal::CreateR2FromArray2D(*input_array); + HloInstruction* input_instruction = + b.AddInstruction(HloInstruction::CreateConstant(std::move(input))); + + auto pad_value_instruction = b.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(2.718f))); + + PaddingConfig padding_config = MakeNoPaddingConfig(2); + + // Negative padding that results in zero dimensions. + auto r2_padding_on_dim0_dim1 = + CreatePaddingConfig({{{-2, -5, 1}}, {{-2, 4, 2}}}); + + Shape shape = ShapeUtil::MakeShape(F32, {0, 9}); + b.AddInstruction(HloInstruction::CreatePad(shape, input_instruction, + pad_value_instruction, + r2_padding_on_dim0_dim1)); + + HloModule module(TestName()); + auto computation = module.AddEntryComputation(b.Build()); + + std::unique_ptr result = + evaluator_->Evaluate(*computation, {}).ConsumeValueOrDie(); + + auto expected_array = MakeUnique>(0, 9); + auto expected = Literal::CreateR2FromArray2D(*expected_array); + + LiteralTestUtil::ExpectEqual(*expected, *result); +} + +TEST_F(HloEvaluatorTest, DotRank2AndRank1) { + HloComputation::Builder b(TestName()); + + // lhs: + // f32[4,1] { + // { 1 }, + // { 2 }, + // { 3 }, + // { 4 }, + // } + auto lhs_array = MakeUnique>(4, 1); + lhs_array->FillUnique(1.0f); + auto lhs_literal = Literal::CreateR2FromArray2D(*lhs_array); + HloInstruction* lhs_instruction = + b.AddInstruction(HloInstruction::CreateConstant(std::move(lhs_literal))); + + // rhs: + // f32[2] { 1, 2 }, + auto rhs_literal = Literal::CreateR2({{1, 2}}); + HloInstruction* rhs_instruction = + b.AddInstruction(HloInstruction::CreateConstant(std::move(rhs_literal))); + + Shape shape = ShapeUtil::MakeShape(F32, {4, 2}); + b.AddInstruction(HloInstruction::CreateBinary( + shape, HloOpcode::kDot, lhs_instruction, rhs_instruction)); + HloModule module(TestName()); + auto computation = module.AddEntryComputation(b.Build()); + + std::unique_ptr result = + evaluator_->Evaluate(*computation, {}).ConsumeValueOrDie(); + + // clang-format off + auto expected_array = Array2D({ + {1.f, 2.f}, + {2.f, 4.f}, + {3.f, 6.f}, + {4.f, 8.f}, + }); + // clang-format on + auto expected = Literal::CreateR2FromArray2D(expected_array); + + LiteralTestUtil::ExpectEqual(*expected, *result); +} + +TEST_F(HloEvaluatorTest, DotRank1AndRank2) { + HloComputation::Builder b(TestName()); + + // lhs: + // f32[3] + // { 1, 2, 3 }, + auto lhs_literal = Literal::CreateR1({1, 2, 3}); + HloInstruction* lhs_instruction = + b.AddInstruction(HloInstruction::CreateConstant(std::move(lhs_literal))); + + // rhs: + // f32[3,2] { + // { 1, 2 }, + // { 3, 4 }, + // { 5, 6 }, + // } + auto rhs_array = MakeUnique>(3, 2); + rhs_array->FillUnique(1.0f); + auto rhs_literal = Literal::CreateR2FromArray2D(*rhs_array); + HloInstruction* rhs_instruction = + b.AddInstruction(HloInstruction::CreateConstant(std::move(rhs_literal))); + + Shape shape = ShapeUtil::MakeShape(F32, {2}); + b.AddInstruction(HloInstruction::CreateBinary( + shape, HloOpcode::kDot, lhs_instruction, rhs_instruction)); + HloModule module(TestName()); + auto computation = module.AddEntryComputation(b.Build()); + + std::unique_ptr result = + evaluator_->Evaluate(*computation, {}).ConsumeValueOrDie(); + + auto expected = Literal::CreateR1({22.f, 28.f}); + + LiteralTestUtil::ExpectEqual(*expected, *result); +} + +TEST_F(HloEvaluatorTest, DotRank2AndRank2) { + HloComputation::Builder b(TestName()); + + // lhs: + // f32[4,3] { + // { 1, 2, 3 }, + // { 5, 6, 7 }, + // { 9, 10, 11 }, + // { 13, 14, 15 }, + // } + auto lhs_array = MakeUnique>(4, 3); + lhs_array->FillUnique(1.0f); + auto lhs_literal = Literal::CreateR2FromArray2D(*lhs_array); + HloInstruction* lhs_instruction = + b.AddInstruction(HloInstruction::CreateConstant(std::move(lhs_literal))); + + // rhs: + // f32[3,2] { + // { 1, 2 }, + // { 3, 4 }, + // { 5, 6 }, + // } + auto rhs_array = MakeUnique>(3, 2); + rhs_array->FillUnique(1.0f); + auto rhs_literal = Literal::CreateR2FromArray2D(*rhs_array); + HloInstruction* rhs_instruction = + b.AddInstruction(HloInstruction::CreateConstant(std::move(rhs_literal))); + + Shape shape = ShapeUtil::MakeShape(F32, {4, 2}); + b.AddInstruction(HloInstruction::CreateBinary( + shape, HloOpcode::kDot, lhs_instruction, rhs_instruction)); + HloModule module(TestName()); + auto computation = module.AddEntryComputation(b.Build()); + + std::unique_ptr result = + evaluator_->Evaluate(*computation, {}).ConsumeValueOrDie(); + + auto expected_array = Array2D({ + {22.f, 28.f}, + {58.f, 76.f}, + {94.f, 124.f}, + {130.f, 172.f}, + }); + auto expected = Literal::CreateR2FromArray2D(expected_array); + + LiteralTestUtil::ExpectEqual(*expected, *result); +} + +TEST_F(HloEvaluatorTest, SimpleConv1D) { + HloComputation::Builder b(TestName()); + + Array3D lhs_array = {{{1, 2, 3}}}; + auto lhs_literal = Literal::CreateR3FromArray3D(lhs_array); + HloInstruction* lhs_instruction = + b.AddInstruction(HloInstruction::CreateConstant(std::move(lhs_literal))); + + Array3D rhs_array = {{{3.f, 4.f}}}; + auto rhs_literal = Literal::CreateR3FromArray3D(rhs_array); + HloInstruction* rhs_instruction = + b.AddInstruction(HloInstruction::CreateConstant(std::move(rhs_literal))); + + Window window; + WindowDimension dim; + dim.set_size(2); + dim.set_stride(1); + dim.set_padding_low(0); + dim.set_padding_high(1); + dim.set_window_dilation(1); + dim.set_base_dilation(1); + *window.add_dimensions() = dim; + + ConvolutionDimensionNumbers dnums; + dnums.set_batch_dimension(0); + dnums.set_feature_dimension(1); + dnums.add_spatial_dimensions(2); + + dnums.set_kernel_output_feature_dimension(0); + dnums.set_kernel_input_feature_dimension(1); + dnums.add_kernel_spatial_dimensions(2); + + const Shape& shape = ShapeUtil::MakeShape(F32, {1, 1, 3}); + b.AddInstruction(HloInstruction::CreateConvolve( + shape, lhs_instruction, rhs_instruction, window, dnums)); + HloModule module(TestName()); + auto computation = module.AddEntryComputation(b.Build()); + + std::unique_ptr result = + evaluator_->Evaluate(*computation, {}).ConsumeValueOrDie(); + + Array3D expected_array = {{{11.f, 18.f, 9.f}}}; + auto expected = Literal::CreateR3FromArray3D(expected_array); + + LiteralTestUtil::ExpectEqual(*expected, *result); +} + +TEST_F(HloEvaluatorTest, Simple4x4Conv2DWith2x2Kernel) { + HloComputation::Builder b(TestName()); + + Array4D lhs_array(1, 1, 4, 4); + // clang-format off + lhs_array.FillWithYX(Array2D({ + {1, 2, 3, 4 }, + {5, 6, 7, 8 }, + {9, 10, 11, 12}, + {13, 14, 15, 16}, + })); + // clang-format on + auto lhs_literal = Literal::CreateR4FromArray4D(lhs_array); + HloInstruction* lhs_instruction = + b.AddInstruction(HloInstruction::CreateConstant(std::move(lhs_literal))); + + Array4D rhs_array(1, 1, 2, 2); + // clang-format off + rhs_array.FillWithYX(Array2D({ + {5, 6}, + {7, 8}, + })); + // clang-format on + auto rhs_literal = Literal::CreateR4FromArray4D(rhs_array); + HloInstruction* rhs_instruction = + b.AddInstruction(HloInstruction::CreateConstant(std::move(rhs_literal))); + + Window window; + WindowDimension dim; + dim.set_size(2); + dim.set_stride(1); + dim.set_padding_low(0); + dim.set_padding_high(1); + dim.set_window_dilation(1); + dim.set_base_dilation(1); + *window.add_dimensions() = dim; + *window.add_dimensions() = dim; + + ConvolutionDimensionNumbers dnums = + ComputationBuilder::CreateDefaultConvDimensionNumbers(2); + + const Shape& shape = ShapeUtil::MakeShape(F32, {1, 1, 4, 4}); + b.AddInstruction(HloInstruction::CreateConvolve( + shape, lhs_instruction, rhs_instruction, window, dnums)); + HloModule module(TestName()); + auto computation = module.AddEntryComputation(b.Build()); + + std::unique_ptr result = + evaluator_->Evaluate(*computation, {}).ConsumeValueOrDie(); + + Array4D expected_array(1, 1, 4, 4); + // clang-format off + expected_array.FillWithYX(Array2D({ + {100, 126, 152, 76}, + {204, 230, 256, 124}, + {308, 334, 360, 172}, + {149, 160, 171, 80}, + })); + // clang-format on + auto expected = Literal::CreateR4FromArray4D(expected_array); + + LiteralTestUtil::ExpectEqual(*expected, *result); +} + +TEST_F(HloEvaluatorTest, Conv2DGeneralDimensions) { + HloComputation::Builder b(TestName()); + + // clang-format off + // Input dimensions: [feature=2, height=3, batch=1, width=4] + Array4D 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}}} + }); + // Weight dimensions: + // [kernel_output_feature=1, width=3, kernel_input_feature=2, height=3] + Array4D weight({{ + {{1, 7, 13}, + {4, 10, 16}}, + {{2, 8, 14}, + {5, 11, 17}}, + {{3, 9, 15}, + {6, 12, 18}} + }}); + // clang-format on + + auto lhs_literal = Literal::CreateR4FromArray4D(input); + HloInstruction* lhs_instruction = + b.AddInstruction(HloInstruction::CreateConstant(std::move(lhs_literal))); + + auto rhs_literal = Literal::CreateR4FromArray4D(weight); + HloInstruction* rhs_instruction = + b.AddInstruction(HloInstruction::CreateConstant(std::move(rhs_literal))); + + Window window; + WindowDimension dim; + dim.set_size(3); + dim.set_stride(1); + dim.set_padding_low(0); + dim.set_padding_high(0); + dim.set_window_dilation(1); + dim.set_base_dilation(1); + *window.add_dimensions() = dim; + *window.add_dimensions() = dim; + + ConvolutionDimensionNumbers dnums; + dnums.set_batch_dimension(2); + dnums.set_feature_dimension(0); + dnums.add_spatial_dimensions(1); + dnums.add_spatial_dimensions(3); + + dnums.set_kernel_output_feature_dimension(0); + dnums.set_kernel_input_feature_dimension(2); + dnums.add_kernel_spatial_dimensions(3); + dnums.add_kernel_spatial_dimensions(1); + + const Shape& shape = ShapeUtil::MakeShape(F32, {1, 1, 1, 2}); + b.AddInstruction(HloInstruction::CreateConvolve( + shape, lhs_instruction, rhs_instruction, window, dnums)); + HloModule module(TestName()); + auto computation = module.AddEntryComputation(b.Build()); + + std::unique_ptr result = + evaluator_->Evaluate(*computation, {}).ConsumeValueOrDie(); + + // clang-format off + // Result dimensions: [feature=1, height=1, batch=1, width=2] + Array4D expected_array({{{{2514, 2685}}}}); + // clang-format on + auto expected = Literal::CreateR4FromArray4D(expected_array); + + LiteralTestUtil::ExpectEqual(*expected, *result); +} + +TEST_F(HloEvaluatorTest, DilatedBaseConv2DWithHighPadding) { + HloComputation::Builder b(TestName()); + + Array4D lhs_array(1, 1, 4, 4); + // clang-format off + lhs_array.FillWithYX(Array2D({ + {1, 2, 3, 4 }, + {5, 6, 7, 8 }, + {9, 10, 11, 12}, + {13, 14, 15, 16}, + })); + // clang-format on + auto lhs_literal = Literal::CreateR4FromArray4D(lhs_array); + HloInstruction* lhs_instruction = + b.AddInstruction(HloInstruction::CreateConstant(std::move(lhs_literal))); + + Array4D rhs_array(1, 1, 2, 2); + // clang-format off + rhs_array.FillWithYX(Array2D({ + {5, 6}, + {7, 8}, + })); + // clang-format on + auto rhs_literal = Literal::CreateR4FromArray4D(rhs_array); + HloInstruction* rhs_instruction = + b.AddInstruction(HloInstruction::CreateConstant(std::move(rhs_literal))); + + Window window; + WindowDimension dim; + dim.set_size(2); + dim.set_stride(1); + dim.set_padding_low(0); + dim.set_padding_high(1); + dim.set_window_dilation(1); + dim.set_base_dilation(2); + *window.add_dimensions() = dim; + *window.add_dimensions() = dim; + + ConvolutionDimensionNumbers dnums = + ComputationBuilder::CreateDefaultConvDimensionNumbers(2); + + const Shape& shape = ShapeUtil::MakeShape(F32, {1, 1, 7, 7}); + b.AddInstruction(HloInstruction::CreateConvolve( + shape, lhs_instruction, rhs_instruction, window, dnums)); + HloModule module(TestName()); + auto computation = module.AddEntryComputation(b.Build()); + + std::unique_ptr result = + evaluator_->Evaluate(*computation, {}).ConsumeValueOrDie(); + + Array4D expected_array(1, 1, 7, 7); + expected_array.FillWithYX(Array2D({ + {5, 12, 10, 18, 15, 24, 20}, + {35, 48, 42, 56, 49, 64, 56}, + {25, 36, 30, 42, 35, 48, 40}, + {63, 80, 70, 88, 77, 96, 84}, + {45, 60, 50, 66, 55, 72, 60}, + {91, 112, 98, 120, 105, 128, 112}, + {65, 84, 70, 90, 75, 96, 80}, + })); + auto expected = Literal::CreateR4FromArray4D(expected_array); + + LiteralTestUtil::ExpectEqual(*expected, *result); +} + +TEST_F(HloEvaluatorTest, DilatedBaseConv2DWithLowAndHighPadding) { + HloComputation::Builder b(TestName()); + + Array4D lhs_array(1, 1, 4, 4); + // clang-format off + lhs_array.FillWithYX(Array2D({ + {1, 2, 3, 4 }, + {5, 6, 7, 8 }, + {9, 10, 11, 12}, + {13, 14, 15, 16}, + })); + // clang-format on + auto lhs_literal = Literal::CreateR4FromArray4D(lhs_array); + HloInstruction* lhs_instruction = + b.AddInstruction(HloInstruction::CreateConstant(std::move(lhs_literal))); + + Array4D rhs_array(1, 1, 2, 2); + // clang-format off + rhs_array.FillWithYX(Array2D({ + {5, 6}, + {7, 8}, + })); + // clang-format on + auto rhs_literal = Literal::CreateR4FromArray4D(rhs_array); + HloInstruction* rhs_instruction = + b.AddInstruction(HloInstruction::CreateConstant(std::move(rhs_literal))); + + Window window; + WindowDimension dim; + dim.set_size(2); + dim.set_stride(1); + dim.set_padding_low(1); + dim.set_padding_high(1); + dim.set_window_dilation(1); + dim.set_base_dilation(2); + *window.add_dimensions() = dim; + *window.add_dimensions() = dim; + + ConvolutionDimensionNumbers dnums = + ComputationBuilder::CreateDefaultConvDimensionNumbers(2); + + const Shape& shape = ShapeUtil::MakeShape(F32, {1, 1, 8, 8}); + b.AddInstruction(HloInstruction::CreateConvolve( + shape, lhs_instruction, rhs_instruction, window, dnums)); + HloModule module(TestName()); + auto computation = module.AddEntryComputation(b.Build()); + + std::unique_ptr result = + evaluator_->Evaluate(*computation, {}).ConsumeValueOrDie(); + + Array4D expected_array(1, 1, 8, 8); + expected_array.FillWithYX(Array2D({ + {8, 7, 16, 14, 24, 21, 32, 28}, + {6, 5, 12, 10, 18, 15, 24, 20}, + {40, 35, 48, 42, 56, 49, 64, 56}, + {30, 25, 36, 30, 42, 35, 48, 40}, + {72, 63, 80, 70, 88, 77, 96, 84}, + {54, 45, 60, 50, 66, 55, 72, 60}, + {104, 91, 112, 98, 120, 105, 128, 112}, + {78, 65, 84, 70, 90, 75, 96, 80}, + })); + auto expected = Literal::CreateR4FromArray4D(expected_array); + + LiteralTestUtil::ExpectEqual(*expected, *result); +} + +TEST_F(HloEvaluatorTest, + DilatedWindowAndBaseConv2DWithDifferentLowAndHighPaddingAndStrides) { + HloComputation::Builder b(TestName()); + + Array4D lhs_array(1, 1, 4, 4); + // clang-format off + lhs_array.FillWithYX(Array2D({ + {1, 2, 3, 4 }, + {5, 6, 7, 8 }, + {9, 10, 11, 12}, + {13, 14, 15, 16}, + })); + // clang-format on + auto lhs_literal = Literal::CreateR4FromArray4D(lhs_array); + HloInstruction* lhs_instruction = + b.AddInstruction(HloInstruction::CreateConstant(std::move(lhs_literal))); + + Array4D rhs_array(1, 1, 2, 3); + // clang-format off + rhs_array.FillWithYX(Array2D({ + {5, 6, 7}, + {8, 9, 10}, + })); + // clang-format on + auto rhs_literal = Literal::CreateR4FromArray4D(rhs_array); + HloInstruction* rhs_instruction = + b.AddInstruction(HloInstruction::CreateConstant(std::move(rhs_literal))); + + Window window; + WindowDimension dim; + dim.set_size(2); + dim.set_stride(1); + dim.set_padding_low(2); + dim.set_padding_high(2); + dim.set_window_dilation(2); + dim.set_base_dilation(2); + *window.add_dimensions() = dim; + dim.set_size(3); + dim.set_stride(3); + dim.set_padding_low(2); + dim.set_padding_high(-1); + dim.set_window_dilation(1); + dim.set_base_dilation(3); + *window.add_dimensions() = dim; + + ConvolutionDimensionNumbers dnums = + ComputationBuilder::CreateDefaultConvDimensionNumbers(2); + + const Shape& shape = ShapeUtil::MakeShape(F32, {1, 1, 9, 3}); + b.AddInstruction(HloInstruction::CreateConvolve( + shape, lhs_instruction, rhs_instruction, window, dnums)); + HloModule module(TestName()); + auto computation = module.AddEntryComputation(b.Build()); + + std::unique_ptr result = + evaluator_->Evaluate(*computation, {}).ConsumeValueOrDie(); + + Array4D expected_array(1, 1, 9, 3); + expected_array.FillWithYX(Array2D({ + {10, 20, 30}, + {0, 0, 0}, + {57, 74, 91}, + {0, 0, 0}, + {125, 142, 159}, + {0, 0, 0}, + {193, 210, 227}, + {0, 0, 0}, + {91, 98, 105}, + })); + auto expected = Literal::CreateR4FromArray4D(expected_array); + + LiteralTestUtil::ExpectEqual(*expected, *result); +} + +TEST_F(HloEvaluatorTest, ReduceAdd) { + HloComputation::Builder b(TestName()); + + // arg: + // f32[2,3] { + // { 1, 2, 3 }, + // { 5, 6, 7 }, + // } + auto arg_array = MakeUnique>(2, 3); + arg_array->FillUnique(1.0f); + auto arg_literal = Literal::CreateR2FromArray2D(*arg_array); + + HloInstruction* arg_instruction = + b.AddInstruction(HloInstruction::CreateConstant(std::move(arg_literal))); + + auto init_value = b.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(0.f))); + + HloComputation::Builder add_computation("add"); + Shape scalar_shape = ShapeUtil::MakeShape(F32, {}); + auto param_lhs = add_computation.AddInstruction( + HloInstruction::CreateParameter(0, scalar_shape, "lhs")); + auto param_rhs = add_computation.AddInstruction( + HloInstruction::CreateParameter(1, scalar_shape, "rhs")); + add_computation.AddInstruction(HloInstruction::CreateBinary( + scalar_shape, HloOpcode::kAdd, param_lhs, param_rhs)); + HloModule module(TestName()); + auto add_func = module.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)); + + auto computation = module.AddEntryComputation(b.Build()); + + std::unique_ptr result = + evaluator_->Evaluate(*computation, {}).ConsumeValueOrDie(); + + auto expected = Literal::CreateR1({6, 18}); + + LiteralTestUtil::ExpectEqual(*expected, *result); +} + +TEST_F(HloEvaluatorTest, ReduceWindowMax) { + HloComputation::Builder b(TestName()); + + // arg: + // f32[2,3] { + // { 1, 2, 3 }, + // { 5, 6, 7 }, + // } + auto arg_array = MakeUnique>(2, 3); + arg_array->FillUnique(1.0f); + auto arg_literal = Literal::CreateR2FromArray2D(*arg_array); + + HloInstruction* arg_instruction = + b.AddInstruction(HloInstruction::CreateConstant(std::move(arg_literal))); + + auto init_value = b.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(0.f))); + + HloComputation::Builder max_computation("max"); + Shape scalar_shape = ShapeUtil::MakeShape(F32, {}); + auto param_lhs = max_computation.AddInstruction( + HloInstruction::CreateParameter(0, scalar_shape, "lhs")); + auto param_rhs = max_computation.AddInstruction( + HloInstruction::CreateParameter(1, scalar_shape, "rhs")); + max_computation.AddInstruction(HloInstruction::CreateBinary( + scalar_shape, HloOpcode::kMaximum, param_lhs, param_rhs)); + HloModule module(TestName()); + auto max_func = module.AddEmbeddedComputation(max_computation.Build()); + + Window window; + WindowDimension dim; + dim.set_size(2); + dim.set_stride(1); + dim.set_padding_low(0); + dim.set_padding_high(0); + dim.set_window_dilation(1); + dim.set_base_dilation(1); + *window.add_dimensions() = dim; + *window.add_dimensions() = dim; + + Shape shape = ShapeUtil::MakeShape(F32, {1, 2}); + b.AddInstruction(HloInstruction::CreateReduceWindow( + shape, arg_instruction, init_value, window, max_func)); + + auto computation = module.AddEntryComputation(b.Build()); + std::unique_ptr result = + evaluator_->Evaluate(*computation, {}).ConsumeValueOrDie(); + + auto expected = Literal::CreateR2({{6, 7}}); + LiteralTestUtil::ExpectEqual(*expected, *result); +} + +TEST_F(HloEvaluatorTest, ReduceWindowAdd) { + HloComputation::Builder b(TestName()); + + // arg: + // f32[2,3] { + // { 1, 2, 3 }, + // { 5, 6, 7 }, + // } + auto arg_array = MakeUnique>(2, 3); + arg_array->FillUnique(1.0f); + auto arg_literal = Literal::CreateR2FromArray2D(*arg_array); + + HloInstruction* arg_instruction = + b.AddInstruction(HloInstruction::CreateConstant(std::move(arg_literal))); + + auto init_value = b.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(0.f))); + + HloComputation::Builder add_computation("add"); + Shape scalar_shape = ShapeUtil::MakeShape(F32, {}); + auto param_lhs = add_computation.AddInstruction( + HloInstruction::CreateParameter(0, scalar_shape, "lhs")); + auto param_rhs = add_computation.AddInstruction( + HloInstruction::CreateParameter(1, scalar_shape, "rhs")); + add_computation.AddInstruction(HloInstruction::CreateBinary( + scalar_shape, HloOpcode::kAdd, param_lhs, param_rhs)); + HloModule module(TestName()); + auto add_func = module.AddEmbeddedComputation(add_computation.Build()); + + Window window; + WindowDimension dim; + dim.set_size(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); + *window.add_dimensions() = dim; + dim.set_size(2); + dim.set_stride(1); + dim.set_padding_low(1); + dim.set_padding_high(0); + dim.set_window_dilation(1); + dim.set_base_dilation(1); + *window.add_dimensions() = dim; + + Shape shape = ShapeUtil::MakeShape(F32, {2, 3}); + b.AddInstruction(HloInstruction::CreateReduceWindow( + shape, arg_instruction, init_value, window, add_func)); + + auto computation = module.AddEntryComputation(b.Build()); + std::unique_ptr result = + evaluator_->Evaluate(*computation, {}).ConsumeValueOrDie(); + + auto expected = Literal::CreateR2({{1, 3, 5}, {5, 11, 13}}); + LiteralTestUtil::ExpectEqual(*expected, *result); +} + +TEST_F(HloEvaluatorTest, ReduceWindowAdd6D) { + HloComputation::Builder b(TestName()); + + // arg: f32[4,4,4,4,4,4] full of ones. Using small dims to limit run-time. + std::vector input_dims(6, 4); + std::unique_ptr arg_literal = + Literal::CreateFullWithMonotonicDim0MajorLayout(input_dims, 1.0f); + + HloInstruction* arg_instruction = + b.AddInstruction(HloInstruction::CreateConstant(std::move(arg_literal))); + + auto init_value = b.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(0.f))); + + HloComputation::Builder add_computation("add"); + Shape scalar_shape = ShapeUtil::MakeShape(F32, {}); + auto param_lhs = add_computation.AddInstruction( + HloInstruction::CreateParameter(0, scalar_shape, "lhs")); + auto param_rhs = add_computation.AddInstruction( + HloInstruction::CreateParameter(1, scalar_shape, "rhs")); + add_computation.AddInstruction(HloInstruction::CreateBinary( + scalar_shape, HloOpcode::kAdd, param_lhs, param_rhs)); + HloModule module(TestName()); + auto add_func = module.AddEmbeddedComputation(add_computation.Build()); + + Window window; + + WindowDimension trivial_dim; + trivial_dim.set_size(1); + trivial_dim.set_stride(1); + trivial_dim.set_padding_low(0); + trivial_dim.set_padding_high(0); + trivial_dim.set_window_dilation(1); + trivial_dim.set_base_dilation(1); + + WindowDimension active_dim; + active_dim.set_size(2); + active_dim.set_stride(1); + active_dim.set_padding_low(0); + active_dim.set_padding_high(0); + active_dim.set_window_dilation(1); + active_dim.set_base_dilation(1); + + *window.add_dimensions() = trivial_dim; + *window.add_dimensions() = active_dim; + *window.add_dimensions() = active_dim; + *window.add_dimensions() = active_dim; + *window.add_dimensions() = trivial_dim; + *window.add_dimensions() = trivial_dim; + + Shape shape = ShapeUtil::MakeShape(F32, {4, 3, 3, 3, 4, 4}); + b.AddInstruction(HloInstruction::CreateReduceWindow( + shape, arg_instruction, init_value, window, add_func)); + + auto computation = module.AddEntryComputation(b.Build()); + std::unique_ptr result = + evaluator_->Evaluate(*computation, {}).ConsumeValueOrDie(); + + std::vector output_dims = {4, 3, 3, 3, 4, 4}; + std::unique_ptr result_literal = + Literal::CreateFullWithMonotonicDim0MajorLayout(output_dims, 8.0f); + LiteralTestUtil::ExpectEqual(*result_literal, *result); +} + +TEST_F(HloEvaluatorTest, StridedSlice) { + HloComputation::Builder b(TestName()); + + // arg: + // f32[3,5] { + // { 1, 2, 3, 4, 5 }, + // { 9, 10, 11, 12, 13 }, + // { 17, 18, 19, 20, 21 }, + // } + auto operand_array = MakeUnique>(3, 5); + operand_array->FillUnique(1.0f); + auto operand_literal = Literal::CreateR2FromArray2D(*operand_array); + + HloInstruction* operand = b.AddInstruction( + HloInstruction::CreateConstant(std::move(operand_literal))); + + Shape shape = ShapeUtil::MakeShape(F32, {2, 1}); + b.AddInstruction(HloInstruction::CreateSlice(shape, operand, + /*start_indices=*/{0, 2}, + /*limit_indices=*/{3, 5}, + /*strides=*/{2, 3})); + HloModule module(TestName()); + auto computation = module.AddEntryComputation(b.Build()); + + std::unique_ptr result = + evaluator_->Evaluate(*computation, {}).ConsumeValueOrDie(); + + auto expected = Literal::CreateR2({ + {3}, + {19}, + }); + + LiteralTestUtil::ExpectEqual(*expected, *result); +} + +TEST_F(HloEvaluatorTest, DynamicSlice) { + HloComputation::Builder b(TestName()); + + // arg: + // f32[2,4] { + // { 1, 2, 3, 4 }, + // { 5, 6, 7, 8 }, + // } + auto operand_array = MakeUnique>(2, 4); + operand_array->FillUnique(1.0f); + auto operand_literal = Literal::CreateR2FromArray2D(*operand_array); + + HloInstruction* operand = b.AddInstruction( + HloInstruction::CreateConstant(std::move(operand_literal))); + + auto start_indices = b.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR1({0, 1}))); + + Shape shape = ShapeUtil::MakeShape(F32, {2, 3}); + b.AddInstruction(HloInstruction::CreateDynamicSlice(shape, operand, + start_indices, {2, 3})); + HloModule module(TestName()); + auto computation = module.AddEntryComputation(b.Build()); + + std::unique_ptr result = + evaluator_->Evaluate(*computation, {}).ConsumeValueOrDie(); + + auto expected = Literal::CreateR2({ + {2, 3, 4}, + {6, 7, 8}, + }); + + LiteralTestUtil::ExpectEqual(*expected, *result); +} + +// Verifies that the HloEvaluator's implementation goes along with existing +// backends' behavior, although this is not required by the spec. +TEST_F(HloEvaluatorTest, DynamicSliceModSlice) { + HloComputation::Builder b(TestName()); + + // arg: + // f32[2,4] { + // { 1, 2, 3, 4 }, + // { 5, 6, 7, 8 }, + // } + auto operand_array = MakeUnique>(2, 4); + operand_array->FillUnique(1.0f); + auto operand_literal = Literal::CreateR2FromArray2D(*operand_array); + + HloInstruction* operand = b.AddInstruction( + HloInstruction::CreateConstant(std::move(operand_literal))); + + auto start_indices = b.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR1({2, 1}))); + + Shape shape = ShapeUtil::MakeShape(F32, {2, 3}); + b.AddInstruction(HloInstruction::CreateDynamicSlice(shape, operand, + start_indices, {2, 3})); + HloModule module(TestName()); + auto computation = module.AddEntryComputation(b.Build()); + + std::unique_ptr result = + evaluator_->Evaluate(*computation, {}).ConsumeValueOrDie(); + + auto expected = Literal::CreateR2({ + {2, 3, 4}, + {6, 7, 8}, + }); + + LiteralTestUtil::ExpectEqual(*expected, *result); +} + +TEST_F(HloEvaluatorTest, DynamicSliceUpdate) { + HloComputation::Builder b(TestName()); + + // arg: + // f32[2,3] { + // { 1, 2, 3 }, + // { 5, 6, 7 }, + // } + auto operand_array = MakeUnique>(2, 3); + operand_array->FillUnique(1.0); + auto operand_literal = Literal::CreateR2FromArray2D(*operand_array); + + HloInstruction* operand = b.AddInstruction( + HloInstruction::CreateConstant(std::move(operand_literal))); + + auto start_indices = b.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR1({0, 1}))); + + auto update = b.AddInstruction(HloInstruction::CreateConstant( + Literal::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)); + HloModule module(TestName()); + auto computation = module.AddEntryComputation(b.Build()); + + std::unique_ptr result = + evaluator_->Evaluate(*computation, {}).ConsumeValueOrDie(); + + auto expected = Literal::CreateR2({ + {1, -2, -3}, + {5, -6, -7}, + }); + + LiteralTestUtil::ExpectEqual(*expected, *result); +} + +TEST_F(HloEvaluatorTest, SetAndGetTuples) { + HloComputation::Builder b(TestName()); + + // arg: + // f32[2,3] { + // { 1, 2, 3 }, + // { 5, 6, 7 }, + // } + auto operand_array = MakeUnique>(2, 3); + operand_array->FillUnique(1.0); + auto operand_literal2 = Literal::CreateR2FromArray2D(*operand_array); + + HloInstruction* operand2 = b.AddInstruction( + HloInstruction::CreateConstant(std::move(operand_literal2))); + HloInstruction* operand1 = b.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR1({0, 1}))); + + auto tuple = + b.AddInstruction(HloInstruction::CreateTuple({operand1, operand2})); + + Shape shape = ShapeUtil::MakeShape(F64, {2, 3}); + b.AddInstruction(HloInstruction::CreateGetTupleElement(shape, tuple, 1)); + + HloModule module(TestName()); + auto computation = module.AddEntryComputation(b.Build()); + std::unique_ptr result = + evaluator_->Evaluate(*computation, {}).ConsumeValueOrDie(); + + auto expected = Literal::CreateR2({ + {1, 2, 3}, + {5, 6, 7}, + }); + + LiteralTestUtil::ExpectEqual(*expected, *result); +} + +TEST_F(HloEvaluatorTest, SetAndGetNestedTuples) { + HloComputation::Builder b(TestName()); + + // arg: + // f32[2,3] { + // { 1, 2, 3 }, + // { 5, 6, 7 }, + // } + auto operand_array = MakeUnique>(2, 3); + operand_array->FillUnique(1.0); + + HloInstruction* operand2 = b.AddInstruction(HloInstruction::CreateConstant( + Literal::CreateR2FromArray2D(*operand_array))); + HloInstruction* operand1 = b.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR1({0, 1}))); + + auto tuple1 = + b.AddInstruction(HloInstruction::CreateTuple({operand1, operand2})); + auto tuple2 = + b.AddInstruction(HloInstruction::CreateTuple({operand2, operand2})); + + auto outer_tuple = + b.AddInstruction(HloInstruction::CreateTuple({tuple1, tuple2})); + + b.AddInstruction( + HloInstruction::CreateGetTupleElement(tuple2->shape(), outer_tuple, 1)); + + HloModule module(TestName()); + auto computation = module.AddEntryComputation(b.Build()); + std::unique_ptr result = + evaluator_->Evaluate(*computation, {}).ConsumeValueOrDie(); + + auto result_inner_literal = + Literal::CreateR2FromArray2D(*operand_array); + auto expected = Literal::MakeTuple({ + result_inner_literal.get(), + result_inner_literal.get(), + }); + + LiteralTestUtil::ExpectEqual(*expected, *result); +} + +TEST_F(HloEvaluatorTest, Reverse) { + HloComputation::Builder b(TestName()); + + // Input shape is float[4x3x2x1]. + // clang-format off + Array4D input({ + {{{1.0f}, {2.0f}}, + {{3.0f}, {4.0f}}, + {{5.0f}, {6.0f}}}, + {{{7.0f}, {8.0f}}, + {{9.0f}, {10.0f}}, + {{11.0f}, {12.0f}}}, + {{{13.0f}, {14.0f}}, + {{15.0f}, {16.0f}}, + {{17.0f}, {18.0f}}}, + {{{19.0f}, {20.0f}}, + {{21.0f}, {22.0f}}, + {{23.0f}, {24.0f}}}, + }); + // clang-format on + auto operand_literal = Literal::CreateR4FromArray4D(input); + HloInstruction* operand = b.AddInstruction( + HloInstruction::CreateConstant(std::move(operand_literal))); + + const Shape shape = ShapeUtil::MakeShape(F32, {4, 3, 2, 1}); + b.AddInstruction(HloInstruction::CreateReverse(shape, operand, {0, 1})); + HloModule module(TestName()); + auto computation = module.AddEntryComputation(b.Build()); + + std::unique_ptr result = + evaluator_->Evaluate(*computation, {}).ConsumeValueOrDie(); + + // clang-format off + auto expected = Literal::CreateR4FromArray4D({ + {{{23.0f}, {24.0f}}, + {{21.0f}, {22.0f}}, + {{19.0f}, {20.0f}}}, + + {{{17.0f}, {18.0f}}, + {{15.0f}, {16.0f}}, + {{13.0f}, {14.0f}}}, + + {{{11.0f}, {12.0f}}, + {{9.0f}, {10.0f}}, + {{7.0f}, {8.0f}}}, + + {{{5.0f}, {6.0f}}, + {{3.0f}, {4.0f}}, + {{1.0f}, {2.0f}}}, + }); + // clang-format on + + LiteralTestUtil::ExpectEqual(*expected, *result); +} + +} // namespace +} // namespace xla diff --git a/tensorflow/compiler/xla/service/hlo_execution_profile.cc b/tensorflow/compiler/xla/service/hlo_execution_profile.cc index 447892c8dec9ea0549a35c9ea2b20303c52b9aa2..0fe7c9fe1b25a1846f3f8b509312dfcb20a162ca 100644 --- a/tensorflow/compiler/xla/service/hlo_execution_profile.cc +++ b/tensorflow/compiler/xla/service/hlo_execution_profile.cc @@ -19,14 +19,11 @@ limitations under the License. #include #include -#include "tensorflow/compiler/xla/metric_table_report.h" #include "tensorflow/compiler/xla/service/hlo_instruction.h" #include "tensorflow/compiler/xla/service/hlo_module.h" +#include "tensorflow/compiler/xla/service/human_readable_profile_builder.h" #include "tensorflow/compiler/xla/types.h" #include "tensorflow/compiler/xla/util.h" -#include "tensorflow/core/lib/strings/numbers.h" -#include "tensorflow/core/lib/strings/strcat.h" -#include "tensorflow/core/lib/strings/stringprintf.h" namespace xla { @@ -47,95 +44,28 @@ uint64 HloExecutionProfile::GetProfileResult(const HloInstruction& hlo) const { string HloExecutionProfile::ToString( const HloComputation& computation, const DeviceDescription& device_description, - const HloCostAnalysis::ShapeSizeFunction& shape_size) const { - HloCostAnalysis cost_analysis(shape_size); + HloCostAnalysis* cost_analysis) const { tensorflow::Status analysis_status = - computation.root_instruction()->Accept(&cost_analysis); + computation.root_instruction()->Accept(cost_analysis); if (!analysis_status.ok()) { return ""; } - using Item = std::pair; - std::vector items; - for (Item item : hlo_to_cycles_taken_) { - // Only include the HLOs which are part of the desired computation. - if (item.first->parent() == &computation) { - items.push_back(item); - } - } - auto custom_less = [](const Item& lhs, const Item& rhs) { - return lhs.second > rhs.second; - }; - std::sort(items.begin(), items.end(), custom_less); - string result; - const int64 total_cycles = total_cycles_executed(computation); - double clock_rate_ghz = device_description.clock_rate_ghz(); - - const auto cycles_to_microseconds = [&](double cycles) { - return cycles / clock_rate_ghz / 1000.0; - }; - - auto append_item = [&](int64 cycles, int64 flops, int64 bytes_accessed, - const string& name) { - double nsecs = cycles / clock_rate_ghz; - string bytes_per_sec; - string bytes_per_cycle; - if (bytes_accessed >= 0) { - bytes_per_sec = tensorflow::strings::HumanReadableNumBytes( - bytes_accessed / (nsecs / 1e9)); - bytes_per_cycle = - tensorflow::strings::HumanReadableNumBytes(bytes_accessed / cycles); - } else { - bytes_per_sec = ""; - bytes_per_cycle = ""; - } - - tensorflow::strings::StrAppend( - &result, - tensorflow::strings::Printf( - "%15lld cycles (%6.2f%%) :: %12.1f usec @ f_nom :: %18s :: %12s/s " - ":: " - "%12s/cycle :: " - "%s", - cycles, cycles / static_cast(total_cycles) * 100, - cycles_to_microseconds(cycles), - flops <= 0 ? "" : HumanReadableNumFlops(flops, nsecs).c_str(), - bytes_per_sec.c_str(), bytes_per_cycle.c_str(), name.c_str())); - }; - tensorflow::strings::StrAppend( - &result, tensorflow::strings::Printf( - "HLO execution profile for %s: (%s @ f_nom)\n\t", - computation.name().c_str(), - tensorflow::strings::HumanReadableElapsedTime( - total_cycles / clock_rate_ghz / 1e9) - .c_str())); - - append_item(total_cycles, -1, -1, "[total]"); - for (const auto& item : items) { + HumanReadableProfileBuilder builder(computation.name(), + total_cycles_executed(computation), + device_description.clock_rate_ghz()); + for (const auto& item : hlo_to_cycles_taken_) { const HloInstruction* hlo = item.first; - tensorflow::strings::StrAppend(&result, "\n\t"); - int64 flops = hlo == nullptr ? -1 : cost_analysis.flop_count(*hlo); - int64 bytes_accessed = - hlo == nullptr ? -1 : cost_analysis.bytes_accessed(*hlo); - string display = hlo == nullptr ? "" : hlo->ToString(); - append_item(item.second, flops, bytes_accessed, display); - } - - MetricTableReport table; - table.SetMetricName("microseconds"); - table.SetEntryName("ops"); - table.SetShowCategoryTable(); - for (const auto& item : items) { - MetricTableReport::Entry entry; - entry.text = item.first->ToString(); - entry.short_text = item.first->ToString(/*compact_operands=*/true); - entry.category_text = item.first->ToCategory(); - entry.metric = cycles_to_microseconds(item.second); - table.AddEntry(std::move(entry)); + int64 cycles = item.second; + + builder.AddOp(/*op_name=*/hlo->ToString(), + /*short_name=*/hlo->ToString(/*compact_operands=*/true), + hlo->ToCategory(), cycles, cost_analysis->flop_count(*hlo), + cost_analysis->transcendental_count(*hlo), + cost_analysis->bytes_accessed(*hlo), + cost_analysis->seconds(*hlo)); } - result += table.MakeReport(cycles_to_microseconds(total_cycles)); - - return result; + return builder.ToString(); } } // namespace xla diff --git a/tensorflow/compiler/xla/service/hlo_execution_profile.h b/tensorflow/compiler/xla/service/hlo_execution_profile.h index 70b94a3f95020bad9b9009d679a4a27a4ba1f34d..a980c1617f395fc6668b8f8739e04d18fd1b689e 100644 --- a/tensorflow/compiler/xla/service/hlo_execution_profile.h +++ b/tensorflow/compiler/xla/service/hlo_execution_profile.h @@ -60,12 +60,12 @@ class HloExecutionProfile { // 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 - // for the operations in a given computation. - // Returns an empty string if it wasn't possible to generate a printable - // version. + // for the operations in a given computation. Returns an empty string if it + // wasn't possible to generate a printable version. cost_analysis should be a + // clean analysis that can be used to visit the computation. string ToString(const HloComputation& computation, const DeviceDescription& device_description, - const HloCostAnalysis::ShapeSizeFunction& shape_size) const; + HloCostAnalysis* cost_analysis) const; // Returns the computations we have profiled. std::unordered_set profiled_computations() const { diff --git a/tensorflow/compiler/xla/service/hlo_graph_dumper.cc b/tensorflow/compiler/xla/service/hlo_graph_dumper.cc index 0af4c99d0a51ab6e4d3048abae1b9c3fb6dca5e6..07b3369d5c1276f0a62af4d3882fed70277f9a91 100644 --- a/tensorflow/compiler/xla/service/hlo_graph_dumper.cc +++ b/tensorflow/compiler/xla/service/hlo_graph_dumper.cc @@ -16,16 +16,25 @@ limitations under the License. #include "tensorflow/compiler/xla/service/hlo_graph_dumper.h" #include +#include +#include +#include +#include +#include #include +#include +#include +#include #include "tensorflow/compiler/xla/layout_util.h" -#include "tensorflow/compiler/xla/legacy_flags/hlo_graph_dumper_flags.h" #include "tensorflow/compiler/xla/literal_util.h" #include "tensorflow/compiler/xla/service/hlo_module.h" +#include "tensorflow/compiler/xla/service/hlo_tfgraph_builder.h" #include "tensorflow/compiler/xla/shape_util.h" #include "tensorflow/compiler/xla/types.h" #include "tensorflow/compiler/xla/window_util.h" #include "tensorflow/core/lib/core/status.h" +#include "tensorflow/core/lib/gtl/optional.h" #include "tensorflow/core/lib/io/path.h" #include "tensorflow/core/lib/strings/numbers.h" #include "tensorflow/core/lib/strings/str_util.h" @@ -33,355 +42,911 @@ limitations under the License. #include "tensorflow/core/lib/strings/stringprintf.h" #include "tensorflow/core/platform/env.h" #include "tensorflow/core/platform/protobuf.h" +#include "tensorflow/core/platform/regexp.h" using ::tensorflow::Env; -using ::tensorflow::WriteStringToFile; +using ::tensorflow::gtl::nullopt; +using ::tensorflow::gtl::optional; using ::tensorflow::io::JoinPath; -using ::tensorflow::strings::Appendf; -using ::tensorflow::strings::Printf; using ::tensorflow::strings::StrAppend; using ::tensorflow::strings::StrCat; using ::tensorflow::str_util::Join; +using ::tensorflow::str_util::StringReplace; +using ::tensorflow::WriteStringToFile; namespace xla { namespace hlo_graph_dumper { namespace { -// Returns the dot graph identifier for the given instruction. -string InstructionId(const HloInstruction* instruction) { - return Printf("%lld", reinterpret_cast(instruction)); +// Helpers for Printf and Appendf. +template +struct PrintfConvert { + const T& operator()(const T& t) const { return t; } +}; +template <> +struct PrintfConvert { + const char* operator()(const string& s) const { return s.c_str(); } +}; + +// Like tensorflow::strings::Printf/Appendf, but you don't need to call c_str() +// on strings. +template +string Printf(const char* fmt, const Ts&... ts) { + return tensorflow::strings::Printf(fmt, PrintfConvert()(ts)...); +} +template +void Appendf(string* s, const char* fmt, const Ts&... ts) { + tensorflow::strings::Appendf(s, fmt, PrintfConvert()(ts)...); } -// Returns the dot graph identifier for the given computation. -string ComputationId(const HloComputation* computation) { - return Printf("%lld", reinterpret_cast(computation)); +// Used to indicate how we should treat a given HLOInstruction in the graph. +// should we treat it like normal, hide it, and so on? +enum NodeFilterResult { + kNormalNode, + kHideNode, + // Make the node easy to find in the final graph. + kHighlightNode, + // "Gray out" the node to indicate that some of its operands have been + // omitted. + kSomeOperandsOmitted, + // Style the node the same as kSomeOperandsOmitted, but also don't connect it + // to its operands, even if they're present in the graph. + kOmitNodeOperands, + // Same style as kSomeOperandsOmitted, but used to indicate that some of the + // node's *users* have been omitted. + kSomeUsersOmitted, +}; + +// NodeFilter is essentially a map from HloInstruction*s to NodeFilterResult. +// It lets callers tell the graph-drawing routines which nodes they want to be +// shown, hidden, or highlighted. +class NodeFilter { + public: + NodeFilter() : filter_([](const HloInstruction*) { return kNormalNode; }) {} + + explicit NodeFilter( + std::function filter) + : filter_(std::move(filter)) {} + + bool Show(const HloInstruction* instr) const { + return filter_(instr) != kHideNode; + } + bool Highlight(const HloInstruction* instr) const { + return filter_(instr) == kHighlightNode; + } + bool OmitOperands(const HloInstruction* instr) const { + return filter_(instr) == kOmitNodeOperands; + } + bool SomeOrAllOperandsOmitted(const HloInstruction* instr) const { + auto result = filter_(instr); + return result == kOmitNodeOperands || result == kSomeOperandsOmitted; + } + bool Deemphasized(const HloInstruction* instr) const { + auto result = filter_(instr); + return result == kOmitNodeOperands || result == kSomeOperandsOmitted || + result == kSomeUsersOmitted; + } + + bool ShowFusionSubcomputation(const HloInstruction* instr) const { + CHECK_EQ(instr->opcode(), HloOpcode::kFusion); + return Show(instr) && !SomeOrAllOperandsOmitted(instr); + } + + private: + std::function filter_; +}; + +// Node color schemes, used by NodeColorAttributes. +enum ColorScheme { + kBlue, + kBrown, + kDarkBlue, + kDarkGreen, + kDarkRed, + kGray, + kGreen, + kOrange, + kPurple, + kRed, + kWhite, + kYellow, + + // Causes the node's border to be a dashed line, and its content to be gray + // text on a white background, suggesting that this is an "unimportant" node. + kDashedBorder, +}; + +// Given a ColorScheme, returns an attribute string for a node of that color. +// Sets the node's style and fill/stroke/text colors. +// +// Colors are from https://material.io/color. +string NodeColorAttributes(ColorScheme color) { + using std::make_tuple; + + const char *style, *fill_color, *stroke_color, *font_color; + std::tie(style, fill_color, stroke_color, font_color) = [color] { + switch (color) { + case kBlue: + return make_tuple("filled", "#bbdefb", "#8aacc8", "black"); + case kBrown: + return make_tuple("filled", "#bcaaa4", "#8c7b75", "black"); + case kDarkBlue: + return make_tuple("filled", "#1565c0", "#003c8f", "white"); + case kDarkGreen: + return make_tuple("filled", "#2e7d32", "#005005", "white"); + case kDarkRed: + return make_tuple("filled", "#b71c1c", "#7f0000", "white"); + case kGray: + return make_tuple("filled", "#cfd8dc", "#9ea7aa", "black"); + case kGreen: + return make_tuple("filled", "#c8e6c9", "#97b498", "black"); + case kOrange: + return make_tuple("filled", "#ffe0b2", "#cbae82", "black"); + case kPurple: + return make_tuple("filled", "#e1bee7", "#af8eb5", "black"); + case kRed: + return make_tuple("filled", "#ffcdd2", "#cb9ca1", "black"); + case kWhite: + return make_tuple("filled", "white", "black", "black"); + case kYellow: + return make_tuple("filled", "#fff9c4", "#cbc693", "black"); + case kDashedBorder: + // "filled,dashed" looks the same as "dashed", since we have a white + // background. But we use "filled,dashed" so that when you hover over + // any part of the node (not just the text inside the node), our css + // :hover rule is triggered. + return make_tuple("filled,dashed", "white", "#757575", "#757575"); + } + }(); + + return Printf( + R"(style="%s", fontcolor="%s", color="%s", fillcolor="%s")", style, + font_color, stroke_color, fill_color); } -// Returns the dot graph edges and nodes for the given instruction sequence. -// Edges which extend between computations are added to the vector -// intercomputation_edges. This is necessary because graphviz does not render -// the graph properly unless these inter-computation edges appear after all -// subgraph statements. -string InstructionSequenceGraph( - const std::list>& instructions, - bool show_addresses, bool show_layouts, - std::vector* intercomputation_edges, - const HloExecutionProfile* hlo_execution_profile) { - string graph_body; - - // Create a single "record" node for the parameters. This node is a - // partitioned rectangle with one partition per parameter node. The keeps - // all the parameter instructions together. - std::vector param_instructions; - for (auto& instruction : instructions) { - if (instruction->opcode() == HloOpcode::kParameter) { - size_t param_number = instruction->parameter_number(); - - if (param_instructions.size() < param_number + 1) { - param_instructions.resize(param_number + 1, nullptr); - } - param_instructions[param_number] = instruction.get(); - } - } - string param_node_name; - if (!param_instructions.empty()) { - std::vector param_ports; - param_node_name = - StrCat("parameters_", InstructionId(param_instructions[0])); - for (auto& param : param_instructions) { - string label = StrCat(param->parameter_name(), "\\n", - ShapeUtil::HumanString(param->shape())); - if (show_addresses) { - Appendf(&label, "\\n[%p]", param); - } - if (show_layouts) { - StrAppend(&label, "\\nlayout=\\{", - Join(param->shape().layout().minor_to_major(), ","), "\\}"); - } - param_ports.push_back( - Printf("<%s> %s", InstructionId(param).c_str(), label.c_str())); - } - StrAppend(&graph_body, param_node_name, - " [shape=record,style=filled,fillcolor=\"lightblue1\",", - "label=\"{parameters | {", Join(param_ports, "|"), "}}\"];\n"); - } - - for (auto& instruction : instructions) { - string color = "peachpuff"; - string shape = "ellipse"; - string name = instruction->ExtendedOpcodeStr(); - if (HloOpcode::kConvolution == instruction->opcode()) { - name += ":\\n" + instruction->ConvolutionDimensionNumbersToString() + - "\\n" + window_util::ToString(instruction->window()); - } - - name += "\\n" + instruction->name(); - if (!instruction->metadata().op_type().empty()) { - StrAppend(&name, "\\n", instruction->metadata().op_type()); - } - if (!instruction->metadata().op_name().empty()) { - StrAppend(&name, "\\n", instruction->metadata().op_name()); - } - if (!instruction->metadata().source_file().empty() && - instruction->metadata().source_line() != 0) { - StrAppend(&name, "\\n", instruction->metadata().source_file(), ":", - instruction->metadata().source_line()); - } - - // Pick different colors or shapes for instructions which are particularly - // expensive (eg, dot) and those which are unusual in some way or unique - // (eg, parameter). - switch (instruction->opcode()) { - // "Normal" instructions. Mostly cheap and elementwise. No call to - // embedded computations. In this case, use default color, shape and - // label. - case HloOpcode::kAbs: - case HloOpcode::kAdd: - case HloOpcode::kCeil: - case HloOpcode::kClamp: - case HloOpcode::kConcatenate: - case HloOpcode::kConvert: - case HloOpcode::kDivide: - case HloOpcode::kDynamicSlice: - case HloOpcode::kDynamicUpdateSlice: - case HloOpcode::kEq: - case HloOpcode::kExp: - case HloOpcode::kFloor: +// Replaces <> with <>, so that this string is safe(er) for use in a +// graphviz HTML-like string. +string HtmlLikeStringSanitize(tensorflow::StringPiece s) { + return StringReplace(StringReplace(s, "<", "<", /*replace_all=*/true), ">", + ">", /*replace_all=*/true); +} + +// Tries to generates a human-readable one-word description of the given +// computation. +// +// Currently we support: +// +// "return param0 + param1;" --> "add" +// "return param0 * param1;" --> "multiply" +// "return min(param0, param1);" --> "min" +// "return max(param0, param1);" --> "max" +// "return param0 <= param1;" --> "less-or-equal" +// "return param0 >= param1;" --> "greater-or-equal" +// "return param0 > param1;" --> "greater-than" +// "return param0 < param1;" --> "less-than" +// "return param0 == param1;" --> "equal-to" +// "return param0 != param1;" --> "not-equal-to" +// +// where param0 and param1 are effective scalars. For the ops that are +// commutative, we also support them with param0 and param1 swapped. +// +// This is useful primarily for reduce and map nodes. These take a +// subcomputation which is almost always one of the four above, and pattern +// matching 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) { + 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) { + 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) { + switch (root->opcode()) { + case HloOpcode::kLe: case HloOpcode::kGe: case HloOpcode::kGt: - case HloOpcode::kIndex: - case HloOpcode::kIsFinite: - case HloOpcode::kLe: - case HloOpcode::kLog: - case HloOpcode::kLogicalAnd: - case HloOpcode::kLogicalNot: - case HloOpcode::kLogicalOr: case HloOpcode::kLt: - case HloOpcode::kMaximum: - case HloOpcode::kMinimum: - case HloOpcode::kMultiply: - case HloOpcode::kNe: - case HloOpcode::kNegate: - case HloOpcode::kPad: - case HloOpcode::kPower: - case HloOpcode::kRemainder: - case HloOpcode::kReshape: - case HloOpcode::kReverse: - case HloOpcode::kSelect: - case HloOpcode::kSign: - case HloOpcode::kSlice: - case HloOpcode::kSort: - case HloOpcode::kSubtract: - case HloOpcode::kTanh: - case HloOpcode::kTuple: - case HloOpcode::kUpdate: + return nullopt; + default: break; + } + } + + // 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: + return "add"; + case HloOpcode::kMultiply: + return "multiply"; + case HloOpcode::kMinimum: + return "min"; + case HloOpcode::kMaximum: + return "max"; + case HloOpcode::kLe: + return "less-or-equal"; + case HloOpcode::kGe: + return "greater-or-equal"; + case HloOpcode::kGt: + return "greater-than"; + case HloOpcode::kLt: + return "less-than"; + case HloOpcode::kEq: + return "equal-to"; + case HloOpcode::kNe: + return "not-equal-to"; + default: + return nullopt; + } +} + +// Encapsulates logic for dumping an HLO module to DOT (i.e. graphviz syntax). +class HloDotDumper { + public: + HloDotDumper(const HloComputation* computation, tensorflow::StringPiece label, + bool show_addresses, const HloExecutionProfile* profile, + NodeFilter filter) + : computation_(computation), + label_(label.ToString()), + show_addresses_(show_addresses), + profile_(profile), + filter_(std::move(filter)) {} + + string Dump(); + + private: + // Returns the dot graph identifier for the given instruction. + string InstructionId(const HloInstruction* instruction) { + return StrCat(reinterpret_cast(instruction)); + } + + // Returns the dot graph identifier for the given computation. + string SubcomputationId(const HloComputation* computation) { + return StrCat("cluster_", reinterpret_cast(computation)); + } + + // Generates graph header/footer. These should be called *after* dumping all + // of the instructions and subcomputations for the graph, as they both use + // data generated while dumping the graph. + string Header(); + string Footer(); + + // Maps HloComputations we should dump to their parent instruction in the + // outer computation. + std::unordered_map + SubcomputationsToDump(); + + string DumpSubcomputation(const HloComputation* subcomp, + const HloInstruction* parent_instr); + string DumpComputation(const HloComputation* comp); + string DumpInstruction(const HloInstruction* instr); + ColorScheme GetInstructionColor(const HloInstruction* instr); + string GetInstructionNodeShape(const HloInstruction* instr); + string GetInstructionNodeLabel(const HloInstruction* instr); + string GetInstructionNodeExtraInfo(const HloInstruction* instr); + string GetInstructionNodeInlinedConstants(const HloInstruction* instr); + void AddInstructionIncomingEdges(const HloInstruction* instr); + + // If instr has just one computation and it's trivial (e.g. "return param0 + + // param1"), returns a string you can put into the node's body that names the + // subcomputation, e.g. "Subcomputation: add". + string GetInstructionTrivialComputationStr(const HloInstruction* instr); + + const HloComputation* computation_; // never null + const string label_; // overall name for the graph + const bool show_addresses_; + const HloExecutionProfile* profile_; // may be null + const NodeFilter filter_; + + // 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_; + // Each (from, to) edge gets a monotonically-increasing ID. This is a + // multimap because it's possible for the same edge to appear multiple times + // in the graph (e.g. x^2 may be represented as mul(x, x)). + int64 next_edge_id_ = 1; + std::unordered_multimap< + std::pair, int64, + tensorflow::hash>> + edge_ids_; + + // Each HloComputation that's emitted gets a monotonically-increasing ID. + int64 next_cluster_id_ = 1; + std::unordered_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 + // appears before both the inner computation and the destination node are + // defined. + std::vector edges_; +}; + +string HloDotDumper::Dump() { + string body; + for (const auto& kv : SubcomputationsToDump()) { + const HloComputation* subcomp = kv.first; + const HloInstruction* parent = kv.second; + StrAppend(&body, DumpSubcomputation(subcomp, parent)); + } + StrAppend(&body, DumpComputation(computation_)); + + // By contract, Header() and Footer() have to be called after we've dumped all + // our instructions, because they use state generated during that process. + string g = Header(); + StrAppend(&g, body); + StrAppend(&g, Footer()); + return g; +} + +string HloDotDumper::Header() { + const char* fmt = R"(digraph G { +rankdir = TB; +compound = true; +label = <%s>; +labelloc = t; +// Disable the tooltip. Interestingly, "" doesn't work! +tooltip = " "; +// DOT graphs accept a stylesheet as a URI. So naturally, an inline +// stylesheet is a data URI! +stylesheet=" + data:text/css, + @import url(https://fonts.googleapis.com/css?family=Roboto:400,700); + svg text { + font-family: 'Roboto'; + font-size: 12px; + } + + %s +" + +)"; + + string graph_label = StrCat(label_, "
", computation_->name()); + if (profile_ != nullptr) { + auto cycles = profile_->total_cycles_executed(*computation_); + Appendf(&graph_label, "
total cycles = %lld (%s)", cycles, + tensorflow::strings::HumanReadableNum(cycles)); + } + + // Create CSS rules that say, when you hover over the given node or cluster, + // turn the given edge the given color. + // + // We rely on a few properties of how graphviz generates SVGs: + // + // - Nodes are named "nodeN", where N corresponds to the 1-based index of + // the node in our DOT (i.e. the first node in the DOT is "node1", etc.). + // Edges are similarly named "edgeN", and clusters are named "clustN". + // - Nodes come before their in- and out-edges in the SVG. We need this + // because the "X ~ Y" CSS selector finds a sibling of X that *comes + // after X in the DOM* and matches Y. + std::vector edge_css_rules; + const char* kBlue = "#1976d2"; + const char* kRed = "#d32f2f"; + for (const auto& kv : edge_ids_) { + const HloInstruction* from_node = kv.first.first; + const HloInstruction* to_node = kv.first.second; + int64 edge_id = kv.second; + + auto add_hover_css_rule = [&](string elem_type, int64 elem_id, + const char* color) { + // One could imagine other ways of writing this CSS rule that involve less + // duplication, but this way seems to be relatively performant. + edge_css_rules.push_back(Printf( + " #%s%d:hover ~ #edge%lld text { fill: %s; }\n" + " #%s%d:hover ~ #edge%lld path { stroke: %s; stroke-width: .2em; }\n" + " #%s%d:hover ~ #edge%lld polygon { " + "fill: %s; stroke: %s; stroke-width: .2em; }\n", + elem_type, elem_id, edge_id, color, // + elem_type, elem_id, edge_id, color, // + elem_type, elem_id, edge_id, color, color)); + }; + + int64 from_node_id = node_ids_.at(from_node); + int64 to_node_id = node_ids_.at(to_node); + add_hover_css_rule("node", from_node_id, kBlue); + add_hover_css_rule("node", to_node_id, kRed); + + // If this edge crosses a fusion cluster boundary, highlight it when the + // cluster is hovered over. + if (from_node->IsFused() && + from_node->parent()->root_instruction() == from_node) { + int64 cluster_id = cluster_ids_.at(from_node->parent()); + add_hover_css_rule("clust", cluster_id, kBlue); + } + if (to_node->IsFused() && to_node->opcode() == HloOpcode::kParameter) { + int64 cluster_id = cluster_ids_.at(to_node->parent()); + add_hover_css_rule("clust", cluster_id, kRed); + } + } + + return Printf(fmt, graph_label, Join(edge_css_rules, "\n")); +} + +string HloDotDumper::Footer() { return StrCat(Join(edges_, "\n"), "\n}"); } + +std::unordered_map +HloDotDumper::SubcomputationsToDump() { + // Dump the subcomputations of each instruction that's shown and doesn't have + // its operands omitted. If an instruction has just one subcomputation and + // it's trivial, omit it: We'll display that subcomputation inlined into the + // instruction's node when we draw it. + std::unordered_map to_dump; + for (const auto& instr : computation_->instructions()) { + if (!filter_.Show(instr.get()) || + filter_.SomeOrAllOperandsOmitted(instr.get())) { + continue; + } + if (instr->opcode() == HloOpcode::kFusion) { + to_dump[instr->fused_instructions_computation()] = instr.get(); + } + + for (const HloComputation* comp : instr->called_computations()) { + if (!MatchTrivialComputation(comp)) { + to_dump[comp] = instr.get(); + } + } + } + return to_dump; +} + +string HloDotDumper::DumpSubcomputation(const HloComputation* subcomp, + const HloInstruction* parent_instr) { + const char* computation_fmt = R"(subgraph %s { +%s +label = <%s>; +labelloc = t; +tooltip = " "; +%s +} // %s + +)"; + + cluster_ids_[subcomp] = next_cluster_id_++; + + string id = SubcomputationId(subcomp); + + string subcomp_label, style; + if (parent_instr->opcode() == HloOpcode::kFusion) { + subcomp_label = Printf("Fused expression for %s
%s", + HtmlLikeStringSanitize(parent_instr->name()), + HtmlLikeStringSanitize(parent_instr->ToCategory())); + string extra_info = GetInstructionNodeExtraInfo(parent_instr); + if (!extra_info.empty()) { + StrAppend(&subcomp_label, "
", extra_info); + } + + // Subcomputation's fill/stroke color is light/dark red/gray, depending on + // whether or not the subcomputation's fusion node is highlighted. + bool highlight = filter_.Highlight(parent_instr); + const char* fillcolor = highlight ? "#ffcdd2" : "#f5f5f5"; + const char* strokecolor = highlight ? "#b71c1c" : "#c2c2c2"; + style = + Printf(R"(style="rounded,filled,bold"; fillcolor="%s"; color="%s;")", + fillcolor, strokecolor); + } else { + subcomp_label = Printf("Subcomputation for %s
%s", + HtmlLikeStringSanitize(parent_instr->name()), + HtmlLikeStringSanitize(subcomp->name())); + style = "style=rounded; color=black;"; + } + + string comp_body = DumpComputation(subcomp); + + if (parent_instr->opcode() == HloOpcode::kFusion) { + // Dump any nested fusion nodes. + for (const auto& subcomp_instr : subcomp->instructions()) { + if (subcomp_instr->opcode() == HloOpcode::kFusion) { + StrAppend( + &comp_body, + DumpSubcomputation(subcomp_instr->fused_instructions_computation(), + subcomp_instr.get())); + } + } + } else { + // Add an edge from the subcomputation to its parent node. If subcomp + // belongs to a fusion node, it's drawn in place of the fusion instruction, + // so there's no need to link those. + edge_ids_.insert( + {{subcomp->root_instruction(), parent_instr}, next_edge_id_++}); + const char* edge_fmt = + R"(%s -> %s [ltail="%s", style="dashed" tooltip="%s -> %s"];)"; + edges_.push_back( + Printf(edge_fmt, InstructionId(subcomp->root_instruction()), + InstructionId(parent_instr), SubcomputationId(subcomp), + subcomp->name(), parent_instr->name())); + } + + string computation = + Printf(computation_fmt, id, style, subcomp_label, comp_body, id); + + return computation; +} + +string HloDotDumper::DumpComputation(const HloComputation* comp) { + string g; + for (const auto& instr : comp->instructions()) { + if (!filter_.Show(instr.get())) { + continue; + } + StrAppend(&g, DumpInstruction(instr.get())); + } + return g; +} + +string HloDotDumper::DumpInstruction(const HloInstruction* instr) { + // We don't display constants as separate nodes; they're merged into their + // users. + if (instr->opcode() == HloOpcode::kConstant) { + return ""; + } + // Omit the fusion node if its subcomputation is drawn, since the + // subcomputation will be drawn inline. + if (instr->opcode() == HloOpcode::kFusion && + filter_.ShowFusionSubcomputation(instr)) { + return ""; + } + + node_ids_[instr] = next_node_id_++; + + ColorScheme color = GetInstructionColor(instr); + string node_shape = GetInstructionNodeShape(instr); + string node_label = GetInstructionNodeLabel(instr); + string extra_info = GetInstructionNodeExtraInfo(instr); + string inlined_constants = GetInstructionNodeInlinedConstants(instr); + string trivial_subcomputation = GetInstructionTrivialComputationStr(instr); + AddInstructionIncomingEdges(instr); + + // Override the node's styling if it should be (de-)emphasized. + if (filter_.Deemphasized(instr)) { + color = kDashedBorder; + } + if (filter_.Highlight(instr)) { + node_shape = "diamond"; + color = kDarkRed; + } + + // Build the text that will be displayed inside the node. + string node_body = node_label; + for (const string& s : + {trivial_subcomputation, extra_info, inlined_constants}) { + if (!s.empty()) { + StrAppend(&node_body, "
", s); + } + } + + return Printf(R"(%s [label=<%s>, shape=%s, tooltip=" ", %s];)" + "\n", + InstructionId(instr), node_body, node_shape, + NodeColorAttributes(color)); +} + +string HloDotDumper::GetInstructionNodeInlinedConstants( + const HloInstruction* instr) { + auto stringify_constant = [](const HloInstruction* constant) { + if (ShapeUtil::IsEffectiveScalar(constant->shape())) { + auto elem_idx = IndexUtil::LinearIndexToMultidimensionalIndex( + constant->shape(), /*linear_index=*/0); + return Printf("%s (%s)", constant->literal().GetAsString(elem_idx), + ShapeUtil::HumanString(constant->shape())); + } + if (tensorflow::StringPiece(constant->name()).starts_with("%constant")) { + return constant->name(); + } + return StrCat("constant ", constant->name()); + }; + + // Special case: If instr is a parameter to a fusion node, check whether the + // corresponding operand to the fusion node is a constant. + if (instr->opcode() == HloOpcode::kParameter && instr->IsFused()) { + const HloInstruction* fusion = instr->parent()->FusionInstruction(); + const HloInstruction* operand = fusion->operand(instr->parameter_number()); + if (operand->opcode() != HloOpcode::kConstant) { + return ""; + } + return StrCat("constant ", stringify_constant(operand)); + } + + std::vector lines; + for (int64 i = 0; i < instr->operand_count(); ++i) { + const HloInstruction* operand = instr->operand(i); + if (operand->opcode() != HloOpcode::kConstant) { + continue; + } + lines.push_back( + Printf("operand %lld = %s", i, stringify_constant(operand))); + } + return Join(lines, "
"); +} + +ColorScheme HloDotDumper::GetInstructionColor(const HloInstruction* instr) { + // Pick different colors or shapes for instructions which are particularly + // expensive (eg, dot) and those which are unusual in some way or unique + // (eg, parameter). + switch (instr->opcode()) { + case HloOpcode::kAbs: + case HloOpcode::kAdd: + case HloOpcode::kCeil: + case HloOpcode::kClamp: + case HloOpcode::kConvert: + case HloOpcode::kCos: + case HloOpcode::kDivide: + case HloOpcode::kEq: + case HloOpcode::kExp: + case HloOpcode::kFloor: + case HloOpcode::kGe: + case HloOpcode::kGt: + case HloOpcode::kIndex: + case HloOpcode::kIsFinite: + case HloOpcode::kLe: + case HloOpcode::kLog: + case HloOpcode::kLogicalAnd: + case HloOpcode::kLogicalNot: + case HloOpcode::kLogicalOr: + case HloOpcode::kLt: + case HloOpcode::kMaximum: + case HloOpcode::kMinimum: + case HloOpcode::kMultiply: + case HloOpcode::kNe: + case HloOpcode::kNegate: + case HloOpcode::kPower: + case HloOpcode::kRemainder: + case HloOpcode::kSelect: + case HloOpcode::kSign: + case HloOpcode::kSin: + case HloOpcode::kSlice: + case HloOpcode::kSort: + case HloOpcode::kSubtract: + case HloOpcode::kTanh: + case HloOpcode::kRng: + case HloOpcode::kBroadcast: + case HloOpcode::kTranspose: + return kYellow; + case HloOpcode::kBitcast: + case HloOpcode::kTuple: + case HloOpcode::kTrace: + case HloOpcode::kGetTupleElement: + return kWhite; + case HloOpcode::kConcatenate: + case HloOpcode::kCopy: + case HloOpcode::kDynamicSlice: + case HloOpcode::kDynamicUpdateSlice: + case HloOpcode::kPad: + case HloOpcode::kReshape: + case HloOpcode::kReverse: + case HloOpcode::kUpdate: + return kGreen; + case HloOpcode::kConvolution: + case HloOpcode::kDot: + return kDarkBlue; + case HloOpcode::kReducePrecision: + return kRed; + case HloOpcode::kParameter: + return kOrange; + case HloOpcode::kBatchNormTraining: + case HloOpcode::kBatchNormInference: + case HloOpcode::kBatchNormGrad: + case HloOpcode::kReduce: + case HloOpcode::kSelectAndScatter: + case HloOpcode::kReduceWindow: + return kPurple; + case HloOpcode::kMap: + case HloOpcode::kFusion: + return kGray; + case HloOpcode::kSend: + case HloOpcode::kRecv: + case HloOpcode::kInfeed: + case HloOpcode::kOutfeed: + case HloOpcode::kCrossReplicaSum: + return kBrown; + case HloOpcode::kCustomCall: + case HloOpcode::kWhile: + case HloOpcode::kCall: + return kDarkGreen; + case HloOpcode::kConstant: + LOG(FATAL) << "Constants don't get their own nodes in the graph."; + } +} + +string HloDotDumper::GetInstructionNodeShape(const HloInstruction* instr) { + // Give while loops a different shape so they're easier to pick out. + switch (instr->opcode()) { + case HloOpcode::kWhile: + return "ellipse"; + default: + return "rect"; + } +} + +string HloDotDumper::GetInstructionNodeLabel(const HloInstruction* instr) { + // If we have a parameter, put the param number in the name. + if (instr->opcode() == HloOpcode::kParameter) { + return Printf("Parameter %lld", instr->parameter_number()); + } + + // The HLO instruction name contains usually the opcode, e.g. "%add.42" is + // an add instruction. In this case we render just the name. + if (tensorflow::StringPiece(instr->name()) + .starts_with(StrCat("%", HloOpcodeString(instr->opcode())))) { + return Printf("%s", HtmlLikeStringSanitize(instr->name())); + } + + // If the name does not contain the opcode, render both. + return Printf("%s
%s", + HtmlLikeStringSanitize(instr->ExtendedOpcodeStr()), + HtmlLikeStringSanitize(instr->name())); +} + +string HloDotDumper::GetInstructionNodeExtraInfo(const HloInstruction* instr) { + string opcode_specific_info = [&]() -> string { + switch (instr->opcode()) { + case HloOpcode::kRng: + return RandomDistribution_Name(instr->random_distribution()); + case HloOpcode::kConvolution: + return StrCat( + HtmlLikeStringSanitize( + instr->ConvolutionDimensionNumbersToString()), + "
", + HtmlLikeStringSanitize(window_util::ToString(instr->window()))); case HloOpcode::kBroadcast: case HloOpcode::kTranspose: - StrAppend(&name, "\\n", "dims={", Join(instruction->dimensions(), ","), - "}"); - break; + case HloOpcode::kReduce: + return Printf("dims={%s}", Join(instr->dimensions(), ",")); case HloOpcode::kGetTupleElement: - StrAppend(&name, "\\nindex=", instruction->tuple_index()); - break; - case HloOpcode::kRng: - StrAppend(&name, "\\n", - RandomDistribution_Name(instruction->random_distribution())); - break; - case HloOpcode::kConstant: - shape = "boxed"; - color = "palegreen"; - if (ShapeUtil::IsScalar(instruction->shape())) { - StrAppend(&name, "\\n", "value=", LiteralUtil::GetAsString( - instruction->literal(), {})); - } - break; - case HloOpcode::kBitcast: - case HloOpcode::kCopy: - color = "white"; - break; - case HloOpcode::kCall: - color = "tomato"; - break; + return Printf("index=%lld", instr->tuple_index()); + case HloOpcode::kBatchNormTraining: + case HloOpcode::kBatchNormGrad: + return Printf("feature_index=%lld", instr->feature_index()); case HloOpcode::kCustomCall: - color = "tomato4"; - StrAppend(&name, "\\n", - "custom_call_target=", instruction->custom_call_target()); - break; - case HloOpcode::kDot: - color = "slateblue"; - break; - case HloOpcode::kSend: - color = "purple"; - break; - case HloOpcode::kRecv: - color = "orange"; - break; - case HloOpcode::kMap: - color = "palevioletred"; - break; - case HloOpcode::kParameter: - // A single record node is created for all the parameter nodes with a - // port for each parameter instruction. No need to emit anything in this - // case. - continue; - case HloOpcode::kReduce: - StrAppend(&name, " dims=", Join(instruction->dimensions(), ",")); - color = "lightsalmon"; - break; - case HloOpcode::kSelectAndScatter: - case HloOpcode::kReduceWindow: - color = "lightsalmon"; - break; - case HloOpcode::kTrace: - color = "white"; - break; - case HloOpcode::kWhile: - color = "forestgreen"; - break; - case HloOpcode::kFusion: - color = "gray"; - break; - case HloOpcode::kConvolution: - color = "red"; - break; - case HloOpcode::kCrossReplicaSum: - color = "turquoise"; - break; - case HloOpcode::kInfeed: - case HloOpcode::kOutfeed: - color = "blue"; - break; + return Printf("custom_call_target=%s", instr->custom_call_target()); + default: + return ""; } + }(); - // Create instruction node with appropriate label, shape, and color. - string label = - StrCat(name, "\\n", ShapeUtil::HumanString(instruction->shape())); - if (show_addresses) { - Appendf(&label, "\\n[%p]", instruction.get()); - } - if (show_layouts && LayoutUtil::HasLayout(instruction->shape())) { - string layout_string; - if (ShapeUtil::IsTuple(instruction->shape())) { - // For tuples, emit the full shape because the layout of a tuple is not - // represented in a single Layout field. - layout_string = ShapeUtil::HumanStringWithLayout(instruction->shape()); - } else { - layout_string = - Join(instruction->shape().layout().minor_to_major(), ","); - } - StrAppend(&label, "\\nlayout={", layout_string, "}"); - } - if (hlo_execution_profile != nullptr) { - auto hlo_cycles_executed = - hlo_execution_profile->GetProfileResult(*instruction); - auto total_cycles_executed = - hlo_execution_profile->total_cycles_executed(*instruction->parent()); - if (hlo_cycles_executed > 0 && total_cycles_executed > 0) { - Appendf(&label, "\\n%% of cycles executed=%.2f", - (static_cast(hlo_cycles_executed) / - static_cast(total_cycles_executed)) * - 100); - } + std::vector lines; + if (!opcode_specific_info.empty()) { + lines.push_back(opcode_specific_info); + } + + // Show the shape and layout of the instruction, unless it's an inlined fusion + // node -- there the shape and layout is present in the output node. + if (instr->opcode() != HloOpcode::kFusion || + !filter_.ShowFusionSubcomputation(instr)) { + string instr_shape = ShapeUtil::HumanString(instr->shape()); + + // Show layout of non-tuple shapes with more than one dimension. + if (LayoutUtil::HasLayout(instr->shape()) && + instr->shape().dimensions_size() > 1 && + !ShapeUtil::IsTuple(instr->shape())) { + StrAppend(&instr_shape, "{", + Join(instr->shape().layout().minor_to_major(), ","), "}"); } - Appendf(&graph_body, - "%s [label=\"%s\", shape=%s, style=filled, fillcolor=%s];\n", - InstructionId(instruction.get()).c_str(), label.c_str(), - shape.c_str(), color.c_str()); - - // Create edges from the instruction's operands to the instruction. - int64 operand_number = 0; - for (auto* operand : instruction->operands()) { - string src; - if (operand->opcode() == HloOpcode::kParameter) { - // If operand is a parameter, then select the proper partition (port) in - // the unified parameter node. - src = param_node_name + ":" + InstructionId(operand); - } else { - src = InstructionId(operand); - } - Appendf(&graph_body, "%s -> %s", src.c_str(), - InstructionId(instruction.get()).c_str()); - if (instruction->operand_count() > 1) { - Appendf(&graph_body, " [headlabel=\"%lld\",labeldistance=2]", - operand_number); - } - StrAppend(&graph_body, ";\n"); - ++operand_number; - } - - // Fusion nodes are handled specially because they contain nested - // expressions. - if (instruction->opcode() == HloOpcode::kFusion) { - string cluster_name = - StrCat("cluster_", InstructionId(instruction.get())); - StrAppend(&graph_body, "subgraph ", cluster_name, " {\n"); - StrAppend(&graph_body, - "label=\"fused expression\";\nstyle=filled;\n" - "color=lightgrey;\n"); - StrAppend(&graph_body, InstructionSequenceGraph( - instruction->fused_instructions(), - show_addresses, show_layouts, - intercomputation_edges, hlo_execution_profile), - "}\n"); - string fusion_edge = - StrCat(InstructionId(instruction->fused_expression_root()), " -> ", - InstructionId(instruction.get()), - " [ style = \"dotted\", arrowsize=0.0, ltail=", cluster_name, - " ];\n"); - intercomputation_edges->push_back(fusion_edge); - } else { - // Add a dotted edge between the instruction and any computations that the - // instruction calls. - for (const HloComputation* computation : - instruction->called_computations()) { - string cluster_name = StrCat("cluster_", ComputationId(computation)); - string call_edge = Printf( - "%s -> %s [ style=dashed; ltail=%s ];\n", - InstructionId(computation->root_instruction()).c_str(), - InstructionId(instruction.get()).c_str(), cluster_name.c_str()); - intercomputation_edges->push_back(call_edge); - } + + // Some instructions have giant tuples as their shapes, so truncate the + // HLO's shape to kMaxShapeLen characters. + constexpr int kMaxShapeLen = 64; + if (instr_shape.length() > kMaxShapeLen) { + instr_shape = StrCat( + tensorflow::StringPiece(instr_shape).substr(0, kMaxShapeLen - 3), + "..."); + } + lines.push_back(instr_shape); + } + + if (show_addresses_) { + lines.push_back(Printf("[%p]", instr)); + } + if (profile_ != nullptr) { + double hlo_cycles_executed = profile_->GetProfileResult(*instr); + double total_cycles_executed = + profile_->total_cycles_executed(*instr->parent()); + if (hlo_cycles_executed > 0 && total_cycles_executed > 0) { + lines.push_back( + Printf("%% of cycles executed=%.2f", + 100 * hlo_cycles_executed / total_cycles_executed)); } } - return graph_body; + return Join(lines, "
"); } -string ComputationToDotGraph(const HloComputation& computation, - const string& label, bool show_addresses, - bool show_layouts, - const HloExecutionProfile* hlo_execution_profile) { - string graph_label = StrCat(label, "\\n", computation.name()); - if (hlo_execution_profile != nullptr) { - auto cycles = hlo_execution_profile->total_cycles_executed(computation); - Appendf(&graph_label, "\\ntotal cycles = %lld (%s)", cycles, - tensorflow::strings::HumanReadableNum(cycles).c_str()); +void HloDotDumper::AddInstructionIncomingEdges(const HloInstruction* instr) { + auto add_edge = [&](const HloInstruction* from, const HloInstruction* to, + int64 operand_num, bool control_edge = false) { + // Fusion nodes' subcomputations are displayed inline, so if 'from' is a + // fusion node and the node's subcomputation is shown, we draw our edge + // starting at the fusion node's root instead of at the fusion node itself. + if (from->opcode() == HloOpcode::kFusion && + filter_.ShowFusionSubcomputation(from)) { + from = from->fused_expression_root(); + } + if (!filter_.Show(from) || from->opcode() == HloOpcode::kConstant) { + return; + } + edge_ids_.insert({{from, to}, next_edge_id_++}); + + string edge_label; + if (instr->operand_count() > 1 && !control_edge) { + edge_label = Printf(R"( headlabel="%lld", labeldistance=2)", operand_num); + } else if (control_edge) { + edge_label = "style=\"dotted\" color=\"gray\" label=\"ctrl\""; + } + const char* kEdgeFmt = R"(%s -> %s [tooltip="%s -> %s" %s];)"; + edges_.push_back(Printf(kEdgeFmt, InstructionId(from), InstructionId(to), + from->name(), to->name(), edge_label)); + }; + + // Add edges from instr's operands to instr. Parameters within fusion + // expressions are handled specially -- we draw an edge from the corresponding + // operand on the fusion node itself to the parameter. + if (instr->opcode() == HloOpcode::kParameter && instr->IsFused()) { + const HloInstruction* fusion = instr->parent()->FusionInstruction(); + add_edge(fusion->operand(instr->parameter_number()), instr, + /*operand_num=*/0); + } else { + for (int64 i = 0; i < instr->operand_count(); ++i) { + add_edge(instr->operand(i), instr, i); + } + for (const HloInstruction* pred : instr->control_predecessors()) { + add_edge(pred, instr, /*operand_num=*/0, /*control_edge=*/true); + } } - string graph = - Printf("digraph G {\nrankdir=TB;\ncompound=true;\nlabel=\"%s\"\n", - graph_label.c_str()); - - // Emit embedded computations as subgraph clusters. - std::vector intercomputation_edges; - for (auto embedded : computation.MakeEmbeddedComputationsList()) { - string graph_body = InstructionSequenceGraph( - embedded->instructions(), show_addresses, show_layouts, - &intercomputation_edges, hlo_execution_profile); - Appendf(&graph, "subgraph cluster_%s {\nlabel=\"%s\";\n%s}\n", - ComputationId(embedded).c_str(), embedded->name().c_str(), - graph_body.c_str()); - } - StrAppend(&graph, - InstructionSequenceGraph(computation.instructions(), show_addresses, - show_layouts, &intercomputation_edges, - hlo_execution_profile)); - - // Edges between computations (subgraph clusters) must be emitted last for the - // graph to be rendered properly for some reason. - StrAppend(&graph, Join(intercomputation_edges, "\n"), "}\n"); - - return graph; +} + +string HloDotDumper::GetInstructionTrivialComputationStr( + const HloInstruction* instr) { + // called_computations() on a fusion node "inherits" any called computations + // of the fused root, which isn't what we want. Just ignore fusion nodes + // here; they're handled separately. + if (instr->opcode() == HloOpcode::kFusion) { + return ""; + } + + std::vector lines; + for (int64 i = 0; i < instr->called_computations().size(); ++i) { + optional computation_type = + MatchTrivialComputation(instr->called_computations()[i]); + if (!computation_type) { + continue; + } + if (instr->called_computations().size() == 1) { + lines.push_back(Printf("Subcomputation: %s", + HtmlLikeStringSanitize(*computation_type))); + } else { + lines.push_back(Printf("Subcomputation %lld: %s", i, + HtmlLikeStringSanitize(*computation_type))); + } + } + return Join(lines, "
"); } tensorflow::mutex& RendererMutex() { @@ -414,14 +979,23 @@ namespace { class FileGraphRenderer : public GraphRendererInterface { public: - string RenderGraph(const string& graph) override { + string RenderGraph(const string& graph, GraphKind graph_kind, + const DebugOptions& debug_options) override { static std::atomic output_num(0); - legacy_flags::HloGraphDumperFlags* flags = - legacy_flags::GetHloGraphDumperFlags(); - string path = StrCat(flags->xla_hlo_dump_graph_path, "hlo_graph_", - output_num++, ".XXXXXX.dot"); + string file_extension; + switch (graph_kind) { + case DOT_GRAPH: + file_extension = ".dot"; + break; + case TF_GRAPHDEF: + file_extension = ".pbtxt"; + break; + } + string path = + JoinPath(debug_options.xla_hlo_graph_path(), + StrCat("hlo_graph_", output_num++, ".XXXXXX", file_extension)); auto status = Status::OK(); - int fd = mkstemps(&path[0], 4); + int fd = mkstemps(&path[0], file_extension.length()); if (fd < 0) { status = Status(tensorflow::error::Code::UNKNOWN, @@ -439,22 +1013,155 @@ class FileGraphRenderer : public GraphRendererInterface { } }; +// Gets a NodeFilter that includes roughly all instructions whose distance from +// root is <= radius. +NodeFilter MakeNodeFilter(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; + std::deque> worklist; + worklist.push_back({root, 0}); + while (!worklist.empty()) { + const HloInstruction* instr; + int64 depth; + std::tie(instr, depth) = worklist.front(); + worklist.pop_front(); + + nodes[instr] = kNormalNode; + if (depth == radius) { + continue; + } + + // Traverse into instr's operands. + // + // Don't traverse into tuples' operands unless the tuple is the root. + // Usually a tuple is the bottommost node in the graph, and so its operands + // 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)) { + worklist.push_back({operand, depth + 1}); + } + } + } + + // Traverse into instr's users, unless: + // + // - there are a ton of them, in which case they're probably not + // interesting (and anyway, rendering them all would make the graph + // unreadable), or + // - instr is a constant, in which case its users are probably not + // interesting. + if (instr->opcode() == HloOpcode::kConstant) { + continue; + } + constexpr int kMaxUsersToRender = 16; + if (instr->user_count() > kMaxUsersToRender) { + // If we're going to skip this node's users, style it as such. + nodes[instr] = kSomeUsersOmitted; + continue; + } + for (const HloInstruction* user : instr->users()) { + if (!nodes.count(user)) { + worklist.push_back({user, depth + 1}); + } + } + } + + auto is_displayed = [&](const HloInstruction* instr) { + // Constants are displayed inline with their users; they're never omitted. + return nodes.count(instr) > 0 || instr->opcode() == HloOpcode::kConstant; + }; + + // Make a second pass over 'nodes' to fix up the NodeFilterResults now that we + // know which nodes will be included in the graph. + for (auto& kv : nodes) { + const HloInstruction* instr = kv.first; + 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)) { + // Mark nodes with some operands omitted appropriately. + filter_result = kSomeOperandsOmitted; + } else if (!operands.empty() && + std::none_of(operands.begin(), operands.end(), is_displayed)) { + // Mark nodes with *all* operands omitted appropriately. + filter_result = kOmitNodeOperands; + } + + // 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)) { + filter_result = kNormalNode; + } + } + + // Highlight the root node. + nodes[root] = kHighlightNode; + + return NodeFilter([=](const HloInstruction* instr) { + auto it = nodes.find(instr); + if (it != nodes.end()) { + return it->second; + } + // Show all nodes in subcomputations. + if (instr->parent() != root->parent()) { + return kNormalNode; + } + return kHideNode; + }); +} + XLA_REGISTER_GRAPH_RENDERER(FileGraphRenderer, 0); } // namespace string DumpGraph(const HloComputation& computation, const string& label, - bool show_addresses, bool show_layouts, + const DebugOptions& debug_options, const HloExecutionProfile* hlo_execution_profile) { - string graph = ComputationToDotGraph(computation, label, show_addresses, - show_layouts, hlo_execution_profile); - - string graph_url = GetGraphRenderer()->RenderGraph(graph); + string graph; + string graph_url; + if (debug_options.xla_hlo_dump_as_graphdef()) { + HloTfGraphBuilder builder; + TF_CHECK_OK(builder.AddComputation(computation)); + CHECK(tensorflow::protobuf::TextFormat::PrintToString(builder.GetGraphDef(), + &graph)); + // TODO(b/37198616): Use the default registered renderers when all + // renderers support rendering GraphDefs. Always dump GraphDefs to files + // for now. + graph_url = FileGraphRenderer().RenderGraph( + graph, GraphRendererInterface::TF_GRAPHDEF, debug_options); + } else { + graph = + HloDotDumper(&computation, label, + /*show_addresses=*/debug_options.xla_hlo_graph_addresses(), + hlo_execution_profile, NodeFilter()) + .Dump(); + graph_url = GetGraphRenderer()->RenderGraph( + graph, GraphRendererInterface::DOT_GRAPH, debug_options); + } LOG(INFO) << "computation " << computation.name() << " [" << label << "]: " << graph_url; return graph_url; } +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); + string graph = + HloDotDumper(node.parent(), label, + /*show_addresses=*/debug_options.xla_hlo_graph_addresses(), + /*profile=*/nullptr, filter) + .Dump(); + return GetGraphRenderer()->RenderGraph( + 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(); @@ -464,8 +1171,31 @@ void DumpText(const HloModule& module, const string& label, do_prefix ? StrCat(prefix, "-", label, ".txt") : StrCat(label, ".txt"); string path = JoinPath(directory_path, filename); TF_CHECK_OK(WriteStringToFile(env, path, module.ToString())); + LOG(INFO) << "dumping module '" << module.name() << "' to " << path; } -} // namespace hlo_graph_dumper +string MaybeDumpHloModule(const HloModule& module, const string& label, + const HloExecutionProfile* profile) { + VLOG(2) << "MaybeDumpHloModule called on module " << module.name(); + string graph_url; + const DebugOptions& debug_options = module.config().debug_options(); + if (!debug_options.xla_generate_hlo_graph().empty() && + RE2::PartialMatch(module.name(), + debug_options.xla_generate_hlo_graph())) { + graph_url = + DumpGraph(*module.entry_computation(), label, debug_options, profile); + } + if (!debug_options.xla_log_hlo_text().empty() && + RE2::PartialMatch(module.name(), debug_options.xla_log_hlo_text())) { + LOG(INFO) << "HLO for module " << module.name(); + LOG(INFO) << "Label: " << label; + XLA_LOG_LINES(2, module.ToString()); + } + if (!debug_options.xla_generate_hlo_text_to().empty()) { + DumpText(module, label, debug_options.xla_generate_hlo_text_to()); + } + return graph_url; +} +} // 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 5f841da1f35c40042fde54dbc03eb7682a8d31cb..0100d50c050a30a2464b912fcf3688426618513e 100644 --- a/tensorflow/compiler/xla/service/hlo_graph_dumper.h +++ b/tensorflow/compiler/xla/service/hlo_graph_dumper.h @@ -21,16 +21,49 @@ limitations under the License. #include "tensorflow/compiler/xla/service/hlo_computation.h" #include "tensorflow/compiler/xla/service/hlo_execution_profile.h" #include "tensorflow/compiler/xla/types.h" +#include "tensorflow/compiler/xla/xla.pb.h" namespace xla { namespace hlo_graph_dumper { -// Dumps a graph of the computation to the GraphViz server and returns -// a description of the rendered graph (e.g., a URL). +// Abstract interface for classes that render HLO graphs (e.g. DOT graph, +// tensorflow GraphDef). +class GraphRendererInterface { + public: + enum GraphKind { + DOT_GRAPH, + TF_GRAPHDEF, + }; + + virtual ~GraphRendererInterface() = default; + + // Renders a DOT graph, returning a description of the rendered output + // (e.g., a URL) + virtual string RenderGraph(const string& graph, GraphKind graph_kind, + const DebugOptions& debug_options) = 0; +}; + +// Dump the given HLO module if a dump is requested in its debug options. Based +// on the debug options, either a graph dump, a text dump or both may be +// generated. If a graph dump is generated, the description (e.g. an URL) is +// returned; otherwise an empty string is returned. +string MaybeDumpHloModule(const HloModule& module, const string& label, + const HloExecutionProfile* profile = nullptr); + +// Dumps a graph of the computation and returns a description of the rendered +// graph (e.g., a URL) based on the renderer. The "best" renderer in the +// registry is used. string DumpGraph(const HloComputation& computation, const string& label, - bool show_addresses, bool show_layouts, + const DebugOptions& debug_options, const HloExecutionProfile* hlo_execution_profile = nullptr); +// Like DumpGraph, but renders only nodes "near" the given node in the graph. +// +// The number of nodes dumped is controlled by the radius parameter, which +// (roughly) corresponds to the max distance a node may be from the primary node +// before it's omitted from the graph. +string DumpNeighborhoodAround(const HloInstruction& node, int radius); + // Dumps the HloModule::ToString() as a file into the provided directory path // suffixed with the provided label. // @@ -40,16 +73,6 @@ string DumpGraph(const HloComputation& computation, const string& label, void DumpText(const HloModule& module, const string& label, const string& directory_path, bool do_prefix = true); -// Abstract interface for classes that render DOT graphs. -class GraphRendererInterface { - public: - virtual ~GraphRendererInterface() = default; - - // Renders a DOT graph, returning a description of the rendered output - // (e.g., a URL) - virtual string RenderGraph(const string& graph) = 0; -}; - // 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. diff --git a/tensorflow/compiler/xla/service/hlo_instruction.cc b/tensorflow/compiler/xla/service/hlo_instruction.cc index 905647c2ed9f30dca31dccd07b6bcf99479ae2aa..75b88aeb1280b22f4a356bf7067c727fc20d8c54 100644 --- a/tensorflow/compiler/xla/service/hlo_instruction.cc +++ b/tensorflow/compiler/xla/service/hlo_instruction.cc @@ -17,6 +17,7 @@ limitations under the License. #include #include +#include #include #include #include @@ -27,6 +28,8 @@ limitations under the License. #include "tensorflow/compiler/xla/ptr_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_module.h" +#include "tensorflow/compiler/xla/service/name_uniquer.h" #include "tensorflow/compiler/xla/shape_util.h" #include "tensorflow/compiler/xla/status_macros.h" #include "tensorflow/compiler/xla/types.h" @@ -36,14 +39,13 @@ limitations under the License. #include "tensorflow/core/lib/gtl/flatmap.h" #include "tensorflow/core/lib/strings/str_util.h" #include "tensorflow/core/lib/strings/strcat.h" -#include "tensorflow/core/lib/strings/stringprintf.h" #include "tensorflow/core/platform/logging.h" namespace xla { -using ::tensorflow::strings::StrAppend; using ::tensorflow::str_util::Join; -using ::tensorflow::strings::Printf; +using ::tensorflow::strings::StrAppend; +using ::tensorflow::strings::StrCat; /* static */ std::unique_ptr HloInstruction::CreateParameter( int64 parameter_number, const Shape& shape, const string& name) { @@ -61,7 +63,7 @@ using ::tensorflow::strings::Printf; WrapUnique(new HloInstruction(HloOpcode::kTrace, ShapeUtil::MakeNil())); instruction->operands_.push_back(operand); instruction->literal_.reset(new Literal); - *instruction->literal_->mutable_u8s() += tag; + instruction->literal_->append_u8s(tag); return instruction; } @@ -118,6 +120,7 @@ HloInstruction::CreateGetTupleElement(const Shape& shape, case HloOpcode::kBitcast: case HloOpcode::kCeil: case HloOpcode::kCopy: + case HloOpcode::kCos: case HloOpcode::kExp: case HloOpcode::kFloor: case HloOpcode::kIsFinite: @@ -125,6 +128,7 @@ HloInstruction::CreateGetTupleElement(const Shape& shape, case HloOpcode::kLogicalNot: case HloOpcode::kNegate: case HloOpcode::kSign: + case HloOpcode::kSin: case HloOpcode::kSort: case HloOpcode::kTanh: break; @@ -209,10 +213,10 @@ HloInstruction::CreateGetTupleElement(const Shape& shape, auto instruction = WrapUnique(new HloInstruction(HloOpcode::kConvolution, shape)); if (window_util::HasBaseDilation(window)) { - instruction->set_name(instruction->name() + "-base-dilated"); + instruction->name_ = instruction->name() + "-base-dilated"; } if (window_util::HasWindowDilation(window)) { - instruction->set_name(instruction->name() + "-window-dilated"); + instruction->name_ = instruction->name() + "-window-dilated"; } instruction->AppendOperand(lhs); instruction->AppendOperand(rhs); @@ -222,6 +226,19 @@ HloInstruction::CreateGetTupleElement(const Shape& shape, return instruction; } +/* static */ std::unique_ptr +HloInstruction::CreateReducePrecision(const Shape& shape, + HloInstruction* operand, + const int exponent_bits, + const int mantissa_bits) { + auto instruction = + WrapUnique(new HloInstruction(HloOpcode::kReducePrecision, shape)); + instruction->AppendOperand(operand); + instruction->exponent_bits_ = exponent_bits; + instruction->mantissa_bits_ = mantissa_bits; + return instruction; +} + /* static */ std::unique_ptr HloInstruction::CreateCrossReplicaSum(const Shape& shape, HloInstruction* operand) { @@ -288,11 +305,19 @@ HloInstruction::CreateCrossReplicaSum(const Shape& shape, /* static */ std::unique_ptr HloInstruction::CreateSlice( const Shape& shape, HloInstruction* operand, tensorflow::gtl::ArraySlice start_indices, - tensorflow::gtl::ArraySlice limit_indices) { + tensorflow::gtl::ArraySlice limit_indices, + tensorflow::gtl::ArraySlice strides) { auto instruction = WrapUnique(new HloInstruction(HloOpcode::kSlice, shape)); instruction->AppendOperand(operand); instruction->slice_starts_.assign(start_indices.begin(), start_indices.end()); instruction->slice_limits_.assign(limit_indices.begin(), limit_indices.end()); + instruction->slice_strides_.assign(strides.begin(), strides.end()); + // For backward compatibility with old serialized computations: if there are + // no strides, assume all strides are 1. + // TODO(b/63317920): remove this code. + if (instruction->slice_strides_.empty()) { + instruction->slice_strides_ = std::vector(start_indices.size(), 1LL); + } return instruction; } @@ -365,6 +390,57 @@ HloInstruction::CreateDynamicUpdateSlice(const Shape& shape, return instruction; } +/* static */ std::unique_ptr +HloInstruction::CreateBatchNormTraining(const Shape& shape, + HloInstruction* operand, + HloInstruction* scale, + HloInstruction* offset, float epsilon, + int64 feature_index) { + auto instruction = + WrapUnique(new HloInstruction(HloOpcode::kBatchNormTraining, shape)); + instruction->AppendOperand(operand); + instruction->AppendOperand(scale); + instruction->AppendOperand(offset); + instruction->epsilon_ = epsilon; + instruction->feature_index_ = feature_index; + return instruction; +} + +/* static */ std::unique_ptr +HloInstruction::CreateBatchNormInference( + const Shape& shape, HloInstruction* operand, HloInstruction* scale, + HloInstruction* offset, HloInstruction* mean, HloInstruction* variance, + float epsilon, int64 feature_index) { + auto instruction = + WrapUnique(new HloInstruction(HloOpcode::kBatchNormInference, shape)); + instruction->AppendOperand(operand); + instruction->AppendOperand(scale); + instruction->AppendOperand(offset); + instruction->AppendOperand(mean); + instruction->AppendOperand(variance); + instruction->epsilon_ = epsilon; + instruction->feature_index_ = feature_index; + return instruction; +} + +/* static */ std::unique_ptr +HloInstruction::CreateBatchNormGrad(const Shape& shape, HloInstruction* operand, + HloInstruction* scale, HloInstruction* mean, + HloInstruction* variance, + HloInstruction* grad_output, float epsilon, + int64 feature_index) { + auto instruction = + WrapUnique(new HloInstruction(HloOpcode::kBatchNormGrad, shape)); + instruction->AppendOperand(operand); + instruction->AppendOperand(scale); + instruction->AppendOperand(mean); + instruction->AppendOperand(variance); + instruction->AppendOperand(grad_output); + instruction->epsilon_ = epsilon; + instruction->feature_index_ = feature_index; + return instruction; +} + /* static */ std::unique_ptr HloInstruction::CreateSelectAndScatter( const Shape& shape, HloInstruction* operand, HloComputation* select, @@ -406,7 +482,9 @@ HloInstruction::CreateSelectAndScatter( /* static */ std::unique_ptr HloInstruction::CreateReshape( const Shape& shape, HloInstruction* operand) { CHECK_EQ(ShapeUtil::ElementsIn(shape), - ShapeUtil::ElementsIn(operand->shape())); + ShapeUtil::ElementsIn(operand->shape())) + << "shape: " << ShapeUtil::HumanString(shape) + << " operand: " << ShapeUtil::HumanString(operand->shape()); auto instruction = WrapUnique(new HloInstruction(HloOpcode::kReshape, shape)); instruction->AppendOperand(operand); return instruction; @@ -432,6 +510,7 @@ HloInstruction::CreateSelectAndScatter( auto instruction = WrapUnique(new HloInstruction(HloOpcode::kFusion, shape)); instruction->fusion_kind_ = fusion_kind; instruction->set_parent(fused_root->parent()); + instruction->set_metadata(fused_root->metadata()); instruction->CloneAndFuseInternal(fused_root); instruction->CheckFusionInstruction(); return instruction; @@ -453,19 +532,24 @@ void HloInstruction::MergeFusionInstruction( HloInstruction* instruction_to_merge) { CHECK_EQ(opcode_, HloOpcode::kFusion); CHECK_EQ(instruction_to_merge->opcode(), HloOpcode::kFusion); + CHECK(std::find(operands().begin(), operands().end(), instruction_to_merge) != + operands().end()); // Clone the instruction from which to merge fused instructions. std::unique_ptr clone = instruction_to_merge->Clone(); // Replace uses of fused parameters with the corresponding operand of the - // fusion. - // Add all non-parameter fused instructions to 'unfused_instructions' to be - // merged into 'this'. + // fusion. Add all non-parameter fused instructions to 'unfused_instructions' + // to be merged into 'this'. This is done in reverse post order. std::vector unfused_instructions; - for (auto& fused_instruction : clone->fused_instructions()) { + auto fused_instructions = + clone->fused_instructions_computation()->MakeInstructionPostOrder(); + for (auto fused_it = fused_instructions.rbegin(); + fused_it != fused_instructions.rend(); ++fused_it) { + auto fused_instruction = *fused_it; if (fused_instruction->opcode() == HloOpcode::kParameter) { TF_CHECK_OK(fused_instruction->ReplaceAllUsesWith( clone->mutable_operand(fused_instruction->parameter_number()))); } else { - unfused_instructions.push_back(fused_instruction.get()); + unfused_instructions.push_back(fused_instruction); } } CHECK(unfused_instructions.front() == clone->fused_expression_root()); @@ -479,39 +563,111 @@ void HloInstruction::MergeFusionInstruction( } CHECK_EQ(0, clone->user_count()); clone->DetachFromOperands(); + TF_CHECK_OK(parent()->parent()->RemoveEmbeddedComputation( + clone->fused_instructions_computation())); } -HloInstruction* HloInstruction::FuseInstruction( - HloInstruction* instruction_to_fuse) { +void HloInstruction::MergeFusionInstructionIntoMultiOutput( + HloInstruction* instruction_to_merge) { + CHECK_EQ(opcode_, HloOpcode::kFusion); + CHECK_EQ(instruction_to_merge->opcode(), HloOpcode::kFusion); + // Add all non-parameter fused instructions to 'unfused_instructions' to be + // merged into 'this'. `old_to_new' maps the instructions in the fused node + // to the disaseembled fusion instructions. + // Note that we add the unfused instructions to this->parent_ computation. + // This is necessary because the unique_id needs for an instruction and + // it's only added when inserting to the computation. + tensorflow::gtl::FlatMap old_to_new; + std::vector unfused_instructions; + auto computation_to_merge = + instruction_to_merge->fused_instructions_computation(); + auto post_order = computation_to_merge->MakeInstructionPostOrder(); + for (auto rit = post_order.rbegin(); rit != post_order.rend(); ++rit) { + auto fused_instruction = *rit; + if (fused_instruction->opcode() == HloOpcode::kParameter) { + InsertOrDie(&old_to_new, fused_instruction, + instruction_to_merge->mutable_operand( + fused_instruction->parameter_number())); + continue; + } + + // Here we clone the insertion and call FuseInstructionIntoMultiOutput() + // which clones again. This can be improved. + auto cloned_instruction = + parent_->AddInstruction(fused_instruction->Clone()); + unfused_instructions.push_back(cloned_instruction); + InsertOrDie(&old_to_new, fused_instruction, cloned_instruction); + } + for (auto unfused_instruction : unfused_instructions) { + for (int64 index = 0; index < unfused_instruction->operand_count(); + index++) { + auto new_operand = + FindOrDie(old_to_new, unfused_instruction->mutable_operand(index)); + TF_CHECK_OK(unfused_instruction->ReplaceOperandWith(index, new_operand)); + } + } + + HloInstruction* unfused_root = unfused_instructions.front(); + TF_CHECK_OK(instruction_to_merge->ReplaceAllUsesWith(unfused_root)); + + TF_CHECK_OK( + instruction_to_merge->parent()->RemoveInstruction(instruction_to_merge)); + if (GetModule()) { + TF_CHECK_OK(GetModule()->RemoveEmbeddedComputation(computation_to_merge)); + } + + // Fuse the root instruction and generate multiple outputs. + FuseInstructionIntoMultiOutput(unfused_root); + // The rest instructions are of normal fusing. + for (int64 i = 1; i < unfused_instructions.size(); i++) { + auto instruction = unfused_instructions[i]; + FuseInstruction(instruction); + } +} + +HloInstruction* HloInstruction::FuseInstructionInternal( + HloInstruction* instruction_to_fuse, bool add_output) { CHECK_EQ(opcode_, HloOpcode::kFusion); - // This fusion instruction must be a user of instruction_to_fuse. - CHECK(IsUserOf(instruction_to_fuse)); - HloInstruction* fused_instruction = CloneAndFuseInternal(instruction_to_fuse); + // When add_output is false, this fusion instruction must be a user of + // instruction_to_fuse. + if (!add_output) { + CHECK(IsUserOf(instruction_to_fuse)); + } + HloInstruction* fused_instruction = + CloneAndFuseInternal(instruction_to_fuse, add_output); CheckFusionInstruction(); return fused_instruction; } HloInstruction* HloInstruction::CloneAndFuseInternal( - HloInstruction* instruction_to_fuse) { + HloInstruction* instruction_to_fuse, bool add_output) { CHECK_EQ(opcode_, HloOpcode::kFusion); CHECK(instruction_to_fuse->IsFusable()); - - bool new_fusion_instruction = fused_instructions_.empty(); - fused_instructions_.emplace_back(instruction_to_fuse->Clone()); - HloInstruction* clone = fused_instructions_.back().get(); - clone->parent_fusion_instruction_ = this; - - if (new_fusion_instruction) { - fused_root_ = clone; + VLOG(3) << "CloneAndFuseInternal:\n" << instruction_to_fuse->ToString(); + HloInstruction* clone = nullptr; + if (called_computations_.empty()) { + // New fusion instruction. It should not be a multioutput instruction. + CHECK(!add_output); + auto builder = HloComputation::Builder("fused_computation", this); + builder.AddInstruction(instruction_to_fuse->Clone(/*suffix=*/"")); + called_computations_.push_back( + CHECK_NOTNULL(GetModule())->AddEmbeddedComputation(builder.Build())); + clone = fused_expression_root(); } else { - // instruction_to_fuse is necessarily an operand of the fusion instruction. - // After fusion this will no longer be the case. Remove the operand from the - // operand list and remove its corresponding fused parameter - // instruction. Renumber parameters as necessary to make parameter numbers - // consistent with their index in the fused_parameter_ vector. - CHECK(std::find(operands_.begin(), operands_.end(), instruction_to_fuse) != - operands_.end()); + clone = fused_instructions_computation()->AddInstruction( + instruction_to_fuse->Clone(/*suffix=*/"")); + // When add_output is false, instruction_to_fuse is necessarily an operand + // of the fusion instruction. After fusion this will no longer be the case. + // Remove the operand from the operand list and remove its 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(); + CHECK(add_output || in_operand_list); + const std::vector& fused_parameters_ = + fused_instructions_computation()->parameter_instructions(); for (int64 operand_num = 0; operand_num < operand_count(); ++operand_num) { if (instruction_to_fuse == operands_[operand_num]) { // replace the fused parameter instruction's uses with the clone. @@ -520,30 +676,28 @@ HloInstruction* HloInstruction::CloneAndFuseInternal( // Remove the corresponding fused parameter and operand from their // respective vectors. - fused_parameters_.erase(fused_parameters_.begin() + operand_num); + TF_CHECK_OK( + fused_instructions_computation()->RemoveParameter(operand_num)); operands_.erase(operands_.begin() + operand_num); - - // Renumber fused parameter numbers to match the vector index. - while (operand_num < fused_parameters_.size()) { - fused_parameters_[operand_num]->parameter_number_ = operand_num; - operand_num++; - } - // Throw removed fused parameter instruction away. - auto inst_it = - std::find_if(fused_instructions_.begin(), fused_instructions_.end(), - [=](const std::unique_ptr& inst) { - return inst.get() == fused_parameter; - }); - CHECK(inst_it != fused_instructions_.end()); - fused_instructions_.erase(inst_it); break; } } // We've cloned instruction_to_fuse into this fusion instruction, so this // fusion instruction is no longer a use of instruction_to_fuse. - instruction_to_fuse->RemoveUser(this); + if (in_operand_list) { + instruction_to_fuse->RemoveUser(this); + // When the instruction_to_fuse does not have other users, we don't need + // to generate a multioutput fusion instruction. + if (instruction_to_fuse->user_count() == 0) { + add_output = false; + } + } } + // Reread the parameters in the computation. + const std::vector& fused_parameters_ = + fused_instructions_computation()->parameter_instructions(); + // Add each operand of the clone as an operand of the fusion instruction. A // complication is that some clone operands may already be operands of the // fusion instruction. @@ -566,28 +720,80 @@ HloInstruction* HloInstruction::CloneAndFuseInternal( // instruction. Add it as an operand and add a corresponding fused // parameter instruction. int64 param_no = fused_parameters_.size(); - std::unique_ptr param_instruction = - CreateParameter(param_no, operand->shape(), "fusion_param"); - - param_instruction->set_parent(parent()); - param_instruction->parent_fusion_instruction_ = this; - fused_parameters_.push_back(param_instruction.get()); - fused_instructions_.push_back(std::move(param_instruction)); + // Name the parameter after the instruction it represents in the outer + // (non-fusion) computation. Strip the leading "%" from the operand name + // to avoid a double %%. + string param_name = + StrCat(operand->name().substr(1), ".param_", param_no); + fused_param = fused_instructions_computation()->AddParameter( + CreateParameter(param_no, operand->shape(), param_name)); AppendOperand(operand); - - fused_param = fused_instructions_.back().get(); } TF_CHECK_OK(clone->ReplaceOperandWith(operand_num, fused_param)); } - for (HloComputation* computation : - instruction_to_fuse->called_computations()) { - if (std::find(called_computations_.begin(), called_computations_.end(), - computation) == called_computations_.end()) { - called_computations_.push_back(computation); + if (add_output) { + CHECK_GT(instruction_to_fuse->user_count(), 0); + // If this is already a multioutput fusion instruction, expand the root + // tuple by 1. + HloInstruction* fused_root = fused_expression_root(); + HloInstruction::InstructionVector tuple_elements; + bool newly_created_tuple_instr = false; + if (fused_root->opcode() == HloOpcode::kTuple) { + tuple_elements = fused_root->operands(); + } else { + tuple_elements.push_back(fused_root); + newly_created_tuple_instr = true; + } + if (clone->opcode() == HloOpcode::kTuple) { + for (auto inst : clone->operands()) { + tuple_elements.push_back(inst); + } + } else { + tuple_elements.push_back(clone); + } + HloInstruction* new_root = fused_instructions_computation()->AddInstruction( + HloInstruction::CreateTuple(tuple_elements)); + fused_instructions_computation()->set_root_instruction(new_root); + shape_ = new_root->shape(); + if (fused_root->opcode() == HloOpcode::kTuple) { + TF_CHECK_OK( + fused_instructions_computation()->RemoveInstruction(fused_root)); + } + + // If this is a newly created multioutput instruction, we need to update + // the use of the original fusion instruction. + if (newly_created_tuple_instr) { + HloInstruction* new_instr = parent_->AddInstruction( + HloInstruction::CreateGetTupleElement(fused_root->shape(), this, 0)); + TF_CHECK_OK(ReplaceAllUsesWith(new_instr)); + } + int64 index = tuple_elements.size(); + if (instruction_to_fuse->opcode() == HloOpcode::kTuple) { + index -= instruction_to_fuse->operand_count(); + std::vector to_be_removed; + for (auto old_gte : instruction_to_fuse->users()) { + CHECK_EQ(old_gte->opcode(), HloOpcode::kGetTupleElement); + int64 old_tuple_index = old_gte->tuple_index(); + HloInstruction* new_gte = + parent_->AddInstruction(HloInstruction::CreateGetTupleElement( + old_gte->shape(), this, index + old_tuple_index)); + TF_CHECK_OK(old_gte->ReplaceAllUsesWith(new_gte)); + to_be_removed.push_back(old_gte); + } + for (auto old_gte : to_be_removed) { + TF_CHECK_OK(parent_->RemoveInstruction(old_gte)); + } + TF_CHECK_OK(fused_instructions_computation()->RemoveInstruction(clone)); + } else { + HloInstruction* new_gte = + parent_->AddInstruction(HloInstruction::CreateGetTupleElement( + clone->shape(), this, index - 1)); + TF_CHECK_OK(instruction_to_fuse->ReplaceAllUsesWith(new_gte)); } } + VLOG(2) << "New clone:\n" << clone->ToString(); return clone; } @@ -596,22 +802,40 @@ RandomDistribution HloInstruction::random_distribution() const { return distribution_; } +bool HloInstruction::HasSideEffect() const { + switch (opcode_) { + case HloOpcode::kSend: + case HloOpcode::kRecv: + case HloOpcode::kInfeed: + case HloOpcode::kOutfeed: + case HloOpcode::kTrace: + return true; + default: { + // Check if any of the called computations has a side effect. + for (const auto& computation : called_computations()) { + if (computation->HasSideEffect()) { + return true; + } + } + return false; + } + } +} + void HloInstruction::CheckFusionInstruction() const { CHECK_EQ(opcode_, HloOpcode::kFusion); - // All instructions owned by this fusion instruction must be fused, and the - // parent fusion instruction of the fused instructions must be 'this'. - for (auto& instruction : fused_instructions_) { - CHECK(instruction->IsFused()); - CHECK_EQ(this, instruction->fusion_instruction()); - CHECK_EQ(parent(), instruction->parent()) << instruction->ToString(); - } + // The parent fusion instruction of the fusion computation must be 'this'. + HloComputation* fused_computation = fused_instructions_computation(); + CHECK_EQ(this, fused_computation->FusionInstruction()); // Fused root instruction and fused parameters must all be owned by the fusion - // instruction. + // computation. bool root_owned = false; + const std::vector& fused_parameters_ = fused_parameters(); + const HloInstruction* fused_root_ = fused_expression_root(); std::vector parameter_owned(fused_parameters_.size(), false); - for (auto& instruction : fused_instructions_) { + for (auto& instruction : fused_computation->instructions()) { if (fused_root_ == instruction.get()) { CHECK(!root_owned); root_owned = true; @@ -632,14 +856,13 @@ void HloInstruction::CheckFusionInstruction() const { // Fused root must have no users. CHECK_EQ(0, fused_root_->user_count()); - // All uses of fused instructions must be in the fusion instruction, and every + // All uses of fused instructions must be in the fusion computation, and every // non-root instruction must have at least one use. - for (auto& instruction : fused_instructions_) { + for (auto& instruction : fused_instructions_computation()->instructions()) { if (instruction.get() != fused_root_) { CHECK_GT(instruction->user_count(), 0); for (auto& user : instruction->users()) { - CHECK(user->IsFused()); - CHECK_EQ(this, user->fusion_instruction()); + CHECK_EQ(fused_computation, user->parent()); } } } @@ -702,7 +925,14 @@ void HloInstruction::CheckFusionInstruction() const { } std::unique_ptr HloInstruction::CloneWithNewOperands( - const Shape& shape, tensorflow::gtl::ArraySlice operands) { + const Shape& shape, + tensorflow::gtl::ArraySlice new_operands) { + VLOG(3) << "CloneWithNewOperands:\n " << ToString(); + VLOG(3) << " new operands:"; + for (const HloInstruction* new_operand : new_operands) { + VLOG(3) << " " << new_operand->name(); + } + // Explicitly call the factory for the instruction type. This is more robust // in the face of code changes than copying fields explicitly. This also // properly sets the user fields of the operands. @@ -712,6 +942,7 @@ std::unique_ptr HloInstruction::CloneWithNewOperands( case HloOpcode::kBitcast: case HloOpcode::kCeil: case HloOpcode::kCopy: + case HloOpcode::kCos: case HloOpcode::kExp: case HloOpcode::kIsFinite: case HloOpcode::kFloor: @@ -719,10 +950,11 @@ std::unique_ptr HloInstruction::CloneWithNewOperands( case HloOpcode::kLogicalNot: case HloOpcode::kNegate: case HloOpcode::kSign: + case HloOpcode::kSin: case HloOpcode::kSort: case HloOpcode::kTanh: - CHECK_EQ(operands.size(), 1); - return CreateUnary(shape, opcode_, operands[0]); + CHECK_EQ(new_operands.size(), 1); + return CreateUnary(shape, opcode_, new_operands[0]); // Binary ops. case HloOpcode::kAdd: case HloOpcode::kDivide: @@ -741,107 +973,165 @@ std::unique_ptr HloInstruction::CloneWithNewOperands( case HloOpcode::kRemainder: case HloOpcode::kLogicalAnd: case HloOpcode::kLogicalOr: - CHECK_EQ(operands.size(), 2); - return CreateBinary(shape, opcode_, operands[0], operands[1]); + CHECK_EQ(new_operands.size(), 2); + return CreateBinary(shape, opcode_, new_operands[0], new_operands[1]); // Ternary ops. case HloOpcode::kClamp: case HloOpcode::kSelect: - CHECK_EQ(operands.size(), 3); - return CreateTernary(shape, opcode_, operands[0], operands[1], - operands[2]); + CHECK_EQ(new_operands.size(), 3); + return CreateTernary(shape, opcode_, new_operands[0], new_operands[1], + new_operands[2]); // Other supported ops. case HloOpcode::kBroadcast: - CHECK_EQ(operands.size(), 1); - return CreateBroadcast(shape, operands[0], dimensions_); + CHECK_EQ(new_operands.size(), 1); + return CreateBroadcast(shape, new_operands[0], dimensions_); case HloOpcode::kCall: - return CreateCall(shape, operands, to_apply()); + return CreateCall(shape, new_operands, to_apply()); case HloOpcode::kCustomCall: - return CreateCustomCall(shape, operands, custom_call_target_); + return CreateCustomCall(shape, new_operands, custom_call_target_); case HloOpcode::kConcatenate: - return CreateConcatenate(shape, operands, dimensions(0)); + return CreateConcatenate(shape, new_operands, dimensions(0)); case HloOpcode::kConvert: - CHECK_EQ(operands.size(), 1); - return CreateConvert(shape, operands[0]); + CHECK_EQ(new_operands.size(), 1); + return CreateConvert(shape, new_operands[0]); + case HloOpcode::kReducePrecision: + CHECK_EQ(new_operands.size(), 1); + return CreateReducePrecision(shape, new_operands[0], exponent_bits_, + mantissa_bits_); case HloOpcode::kConvolution: - CHECK_EQ(operands.size(), 2); - return CreateConvolve(shape, operands[0], operands[1], *window_, + CHECK_EQ(new_operands.size(), 2); + return CreateConvolve(shape, new_operands[0], new_operands[1], *window_, *convolution_dimension_numbers_); case HloOpcode::kCrossReplicaSum: - CHECK_EQ(operands.size(), 1); - return CreateCrossReplicaSum(shape, operands[0]); + CHECK_EQ(new_operands.size(), 1); + return CreateCrossReplicaSum(shape, new_operands[0]); case HloOpcode::kGetTupleElement: - CHECK_EQ(operands.size(), 1); - return CreateGetTupleElement(shape, operands[0], tuple_index()); + CHECK_EQ(new_operands.size(), 1); + return CreateGetTupleElement(shape, new_operands[0], tuple_index()); case HloOpcode::kMap: - return CreateMap(shape, operands, to_apply()); + return CreateMap(shape, new_operands, to_apply()); case HloOpcode::kPad: - CHECK_EQ(operands.size(), 2); - return CreatePad(shape, operands[0], operands[1], *padding_config_); + CHECK_EQ(new_operands.size(), 2); + return CreatePad(shape, new_operands[0], new_operands[1], + *padding_config_); case HloOpcode::kReduce: - CHECK_EQ(operands.size(), 2); - return CreateReduce(shape, operands[0], operands[1], dimensions_, + CHECK_EQ(new_operands.size(), 2); + return CreateReduce(shape, new_operands[0], new_operands[1], dimensions_, to_apply()); case HloOpcode::kReduceWindow: - CHECK_EQ(operands.size(), 2); - return CreateReduceWindow(shape, operands[0], operands[1], *window_, - to_apply()); + CHECK_EQ(new_operands.size(), 2); + return CreateReduceWindow(shape, new_operands[0], new_operands[1], + *window_, to_apply()); case HloOpcode::kSelectAndScatter: - CHECK_EQ(operands.size(), 3); - return CreateSelectAndScatter(shape, operands[0], select(), *window_, - operands[1], operands[2], scatter()); - case HloOpcode::kRecv: - CHECK_EQ(operands.size(), 0); - return CreateRecv(shape, channel_id_); + CHECK_EQ(new_operands.size(), 3); + return CreateSelectAndScatter(shape, new_operands[0], select(), *window_, + new_operands[1], new_operands[2], + scatter()); case HloOpcode::kReverse: - CHECK_EQ(operands.size(), 1); - return CreateReverse(shape, operands[0], dimensions_); + CHECK_EQ(new_operands.size(), 1); + return CreateReverse(shape, new_operands[0], dimensions_); case HloOpcode::kRng: - return CreateRng(shape, distribution_, operands); + return CreateRng(shape, distribution_, new_operands); case HloOpcode::kReshape: - CHECK_EQ(operands.size(), 1); - return CreateReshape(shape, operands[0]); - case HloOpcode::kSend: - CHECK_EQ(operands.size(), 1); - return CreateSend(operands[0], channel_id_); + CHECK_EQ(new_operands.size(), 1); + return CreateReshape(shape, new_operands[0]); case HloOpcode::kSlice: - CHECK_EQ(operands.size(), 1); - return CreateSlice(shape, operands[0], slice_starts_, slice_limits_); + CHECK_EQ(new_operands.size(), 1); + return CreateSlice(shape, new_operands[0], slice_starts_, slice_limits_, + slice_strides_); case HloOpcode::kDynamicSlice: - return CreateDynamicSlice(shape, operands[0], operands[1], + return CreateDynamicSlice(shape, new_operands[0], new_operands[1], dynamic_slice_sizes_); case HloOpcode::kDynamicUpdateSlice: - CHECK_EQ(operands.size(), 3); - return CreateDynamicUpdateSlice(shape, operands[0], operands[1], - operands[2]); + CHECK_EQ(new_operands.size(), 3); + return CreateDynamicUpdateSlice(shape, new_operands[0], new_operands[1], + new_operands[2]); case HloOpcode::kTranspose: - CHECK_EQ(operands.size(), 1); - return CreateTranspose(shape, operands[0], dimensions_); + CHECK_EQ(new_operands.size(), 1); + return CreateTranspose(shape, new_operands[0], dimensions_); case HloOpcode::kTuple: - return CreateTuple(operands_); + return CreateTuple(new_operands); case HloOpcode::kWhile: - CHECK_EQ(operands.size(), 1); - return CreateWhile(shape, while_condition(), while_body(), operands[0]); + CHECK_EQ(new_operands.size(), 1); + return CreateWhile(shape, while_condition(), while_body(), + new_operands[0]); case HloOpcode::kConstant: - return CreateConstant(LiteralUtil::CloneToUnique(*literal_)); + return CreateConstant(literal_->CloneToUnique()); case HloOpcode::kFusion: - return CloneFusionWithNewOperands(shape, operands); + return CloneFusionWithNewOperands(shape, new_operands); case HloOpcode::kParameter: return CreateParameter(parameter_number_, shape, parameter_name_); - // Unsupported ops for cloning. - case HloOpcode::kUpdate: - case HloOpcode::kIndex: + case HloOpcode::kBatchNormTraining: + CHECK_EQ(new_operands.size(), 3); + return CreateBatchNormTraining(shape, new_operands[0], new_operands[1], + new_operands[2], epsilon(), + feature_index()); + + case HloOpcode::kBatchNormInference: + CHECK_EQ(new_operands.size(), 5); + return CreateBatchNormInference( + shape, new_operands[0], new_operands[1], new_operands[2], + new_operands[3], new_operands[4], epsilon(), feature_index()); case HloOpcode::kInfeed: + CHECK_EQ(new_operands.size(), 0); + return CreateInfeed(shape, infeed_config()); case HloOpcode::kOutfeed: + CHECK_EQ(new_operands.size(), 1); + return CreateOutfeed(outfeed_shape_, new_operands[0], outfeed_config()); + case HloOpcode::kBatchNormGrad: + CHECK_EQ(new_operands.size(), 5); + return CreateBatchNormGrad(shape, new_operands[0], new_operands[1], + new_operands[2], new_operands[3], + new_operands[4], epsilon(), feature_index()); + case HloOpcode::kRecv: + case HloOpcode::kSend: + case HloOpcode::kUpdate: + case HloOpcode::kIndex: case HloOpcode::kTrace: LOG(FATAL) << "Not yet implemented, clone: " << HloOpcodeString(opcode_); } } +HloInstruction::~HloInstruction() {} + std::unique_ptr HloInstruction::Clone(const string& suffix) { std::unique_ptr clone = CloneWithNewOperands(shape_, operands_); - clone->name_ = name() + "." + suffix; + if (suffix.empty()) { + clone->name_ = name(); + } else { + // If an instruction is cloned multiple times avoid names like + // foo.suffix.suffix.suffix. Instead of repeating the suffix add a numeric + // suffix. Specifically, the clone of foo.suffix is named foo.suffix2, the + // clone of foo.suffix2 is named foo.suffix3 and so on. + const string dot_suffix = "." + suffix; + size_t index = name().rfind(dot_suffix); + if (index == string::npos) { + // Existing name does not include ".suffix". + clone->name_ = name() + dot_suffix; + } else { + // Existing name includes ".suffix". Determine if substring after + // ".suffix" is numeric and should be replaced with an incremented number. + string after_suffix = name().substr(index + dot_suffix.size()); + if (after_suffix.empty()) { + // Existing name ends in ".suffix". New name should end in ".suffix2". + clone->name_ = name() + "2"; + } else { + // If names ends with .suffix[0-9]+ then replace with a suffix with the + // numeric value incremented. + int64 numeric_suffix; + if (tensorflow::strings::safe_strto64(after_suffix, &numeric_suffix)) { + clone->name_ = + StrCat(name().substr(0, index), dot_suffix, numeric_suffix + 1); + } else { + // Substring after ".suffix" is non-numeric. + clone->name_ = name() + dot_suffix; + } + } + } + } clone->set_parent(parent()); + clone->set_metadata(metadata_); return clone; } @@ -861,18 +1151,14 @@ std::unique_ptr HloInstruction::CloneFusionWithNewOperands( std::list> new_fused_instructions; // Create the list of fused parameters by mapping through the cloned, // fused instructions. - std::vector new_fused_parameters; - for (HloInstruction* old_fused_parameter : fused_parameters_) { + for (HloInstruction* old_fused_parameter : + fused_instructions_computation()->parameter_instructions()) { new_fused_instructions.push_back(old_fused_parameter->Clone()); HloInstruction* new_fusion_parameter = new_fused_instructions.back().get(); - new_fusion_parameter->parent_fusion_instruction_ = new_instruction.get(); - new_fused_parameters.push_back(new_fusion_parameter); InsertOrDie(&old_to_new, old_fused_parameter, new_fusion_parameter); } - for (auto old_fused_instruction_iter = fused_instructions_.rbegin(); - old_fused_instruction_iter != fused_instructions_.rend(); - ++old_fused_instruction_iter) { - HloInstruction* old_fused_instruction = old_fused_instruction_iter->get(); + for (auto old_fused_instruction : + fused_instructions_computation()->MakeInstructionPostOrder()) { if (old_fused_instruction->opcode() == HloOpcode::kParameter) { FindOrDie(old_to_new, old_fused_instruction); continue; @@ -889,16 +1175,24 @@ std::unique_ptr HloInstruction::CloneFusionWithNewOperands( old_fused_instruction->shape(), new_operands)); HloInstruction* new_fused_instruction = new_fused_instructions.back().get(); new_fused_instruction->set_parent(parent()); - new_fused_instruction->parent_fusion_instruction_ = new_instruction.get(); InsertOrDie(&old_to_new, old_fused_instruction, new_fused_instruction); } + new_instruction->fusion_kind_ = fusion_kind_; + auto computation_builder = HloComputation::Builder( + fused_instructions_computation()->name() + ".clone", + new_instruction.get()); // We iterated the fusion instructions in reverse post order which means // that we must reverse our new list of fusion instructions. - std::reverse(new_fused_instructions.begin(), new_fused_instructions.end()); - new_instruction->fusion_kind_ = fusion_kind_; - new_instruction->fused_instructions_ = std::move(new_fused_instructions); - new_instruction->fused_parameters_ = std::move(new_fused_parameters); - new_instruction->fused_root_ = FindOrDie(old_to_new, fused_root_); + for (auto new_fused_instruction_iter = new_fused_instructions.rbegin(); + new_fused_instruction_iter != new_fused_instructions.rend(); + ++new_fused_instruction_iter) { + computation_builder.AddInstruction(std::move(*new_fused_instruction_iter)); + } + auto fused_root_ = fused_expression_root(); + new_instruction->called_computations_.push_back( + CHECK_NOTNULL(GetModule()) + ->AddEmbeddedComputation( + computation_builder.Build(FindOrDie(old_to_new, fused_root_)))); new_instruction->set_parent(parent()); new_instruction->CheckFusionInstruction(); return new_instruction; @@ -975,7 +1269,7 @@ Status HloInstruction::RemoveControlDependencyTo(HloInstruction* instruction) { auto pred_it = std::find(instruction->control_predecessors_.begin(), instruction->control_predecessors_.end(), this); TF_RET_CHECK(pred_it != instruction->control_predecessors_.end()); - instruction->control_predecessors_.erase(succ_it); + instruction->control_predecessors_.erase(pred_it); return Status::OK(); } @@ -1005,25 +1299,10 @@ bool HloInstruction::HasConstantOperand() const { return false; } -bool HloInstruction::Identical( +bool HloInstruction::IdenticalSlowPath( const HloInstruction& other, - std::function - eq_operands, std::function eq_computations) const { - // An instruction is always identical to itself. - if (this == &other) { - return true; - } - - // Identical instruction must have the same opcode and identical operands. In - // general, there is no need to check shape because shape is inferred from the - // shape of the operands. - if (opcode() != other.opcode() || - !ContainersEqual(operands(), other.operands(), eq_operands)) { - return false; - } - // Perform opcode specific checks. switch (opcode()) { // The result of these instructions only depend upon their opcode and @@ -1033,6 +1312,7 @@ bool HloInstruction::Identical( case HloOpcode::kCeil: case HloOpcode::kClamp: case HloOpcode::kCopy: + case HloOpcode::kCos: case HloOpcode::kCrossReplicaSum: case HloOpcode::kDivide: case HloOpcode::kDot: @@ -1057,6 +1337,7 @@ bool HloInstruction::Identical( case HloOpcode::kRemainder: case HloOpcode::kSelect: case HloOpcode::kSign: + case HloOpcode::kSin: case HloOpcode::kSubtract: case HloOpcode::kTanh: case HloOpcode::kTuple: @@ -1075,15 +1356,26 @@ bool HloInstruction::Identical( // different HloComputations. ShapeUtil::Compatible(shape(), other.shape()); + case HloOpcode::kBatchNormTraining: + case HloOpcode::kBatchNormInference: + case HloOpcode::kBatchNormGrad: + return feature_index() == other.feature_index() && + epsilon() == other.epsilon(); + // A constant is defined by the value in the literal. case HloOpcode::kConstant: - return LiteralUtil::Equal(literal(), other.literal()); + return literal().Equal(other.literal()); // A convert result is determined by the primitive type that the operand is // converted into. case HloOpcode::kConvert: return shape().element_type() == other.shape().element_type(); + // A reduce-precision operation is determined by the bit sizes. + case HloOpcode::kReducePrecision: + return exponent_bits() == other.exponent_bits() && + mantissa_bits() == other.mantissa_bits(); + // Convolution has a window and dimensions. case HloOpcode::kConvolution: return protobuf_util::ProtobufEquals(window(), other.window()) && @@ -1242,8 +1534,9 @@ Status HloInstruction::ReplaceAllUsesWith(HloInstruction* new_producer) { } void HloInstruction::DetachFromOperands() { + VLOG(3) << "DetachFromOperands:\n " << ToString(); CHECK_EQ(0, user_count()); - // An intruction may be repeated as an operand. To avoid calling RemoveUser + // An instruction may be repeated as an operand. To avoid calling RemoveUser // twice on the same operand, keep a set of already detached operands. std::set detached_operands; for (int64 operand_num = 0; operand_num < operand_count(); ++operand_num) { @@ -1355,8 +1648,7 @@ string HloInstruction::SignatureString() const { Join(operands_, ", ", [](string* out, HloInstruction* operand) { StrAppend(out, ShapeUtil::HumanString(operand->shape())); }); - return tensorflow::strings::StrCat("(", operands, ") -> ", - ShapeUtil::HumanString(shape())); + return StrCat("(", operands, ") -> ", ShapeUtil::HumanString(shape())); } string HloInstruction::ExtendedOpcodeStr() const { @@ -1368,28 +1660,33 @@ string HloInstruction::ExtendedOpcodeStr() const { return opc_name; } -string HloInstruction::ToString(bool compact_operands) const { +string HloInstruction::ToString(bool compact_operands, + bool include_metadata) const { string operands; if (opcode() == HloOpcode::kConstant) { // For constants, show the actual value in place of an empty operand list. - if (ShapeUtil::ElementsIn(shape()) <= 10) { - // LiteralUtil::ToString emits multidimensional arrays over multiple + if (!ShapeUtil::IsTuple(shape()) && ShapeUtil::ElementsIn(shape()) <= 10) { + // Literal::ToString emits multidimensional arrays over multiple // lines. Compact this into one line by stripping out white space. - string tmp = LiteralUtil::ToString(literal()); + string tmp = literal().ToString(); std::replace(tmp.begin(), tmp.end(), '\n', ' '); std::vector v = tensorflow::str_util::Split(tmp, ' '); bool first = true; // Concatenate elements in "v" with spaces separating them, but ignoring // empty entries. for (const auto& s : v) { - if (s.empty()) continue; + if (s.empty()) { + continue; + } StrAppend(&operands, (first ? "" : " "), s); first = false; } } else { - // Do not show large constants. + // Do not show large constants or tuples. operands = "{...}"; } + } else if (opcode() == HloOpcode::kParameter) { + StrAppend(&operands, parameter_number_); } else { tensorflow::gtl::ArraySlice slice(operands_); const int64 kMaxOperandsToShowIfCompact = 4; @@ -1419,9 +1716,10 @@ string HloInstruction::ToString(bool compact_operands) const { } if (!slice_starts_.empty() && !slice_limits_.empty()) { std::vector bounds; + bounds.reserve(slice_starts_.size()); for (int i = 0; i < slice_starts_.size(); ++i) { - bounds.push_back(tensorflow::strings::StrCat("[", slice_starts_[i], ":", - slice_limits_[i], "]")); + bounds.push_back( + StrCat("[", slice_starts_[i], ":", slice_limits_[i], "]")); } StrAppend(&extra, ", slice={", Join(bounds, ", "), "}"); } @@ -1447,23 +1745,75 @@ string HloInstruction::ToString(bool compact_operands) const { if (opcode() == HloOpcode::kGetTupleElement) { StrAppend(&extra, ", index=", tuple_index()); } - if (!metadata_.op_type().empty() || !metadata_.op_name().empty() || - !metadata_.source_file().empty()) { + if (!control_successors_.empty()) { + StrAppend( + &extra, ", control-successors=", + Join(control_successors_, ", ", [](string* out, HloInstruction* succ) { + StrAppend(out, succ->name()); + })); + } + if (include_metadata && + (!metadata_.op_type().empty() || !metadata_.op_name().empty() || + !metadata_.source_file().empty())) { StrAppend(&extra, " # metadata=", metadata_.ShortDebugString()); } - return Printf("%s = %s %s(%s)%s", name().c_str(), - ShapeUtil::HumanStringWithLayout(shape()).c_str(), - ExtendedOpcodeStr().c_str(), operands.c_str(), extra.c_str()); + + return StrCat(name(), " = ", ShapeUtil::HumanStringWithLayout(shape()), " ", + ExtendedOpcodeStr(), "(", operands, ")", extra); } string HloInstruction::ToShortString() const { - return Printf("%s = %s(%s)", name().c_str(), - HloOpcodeString(opcode()).c_str(), + return StrCat(name(), " = ", HloOpcodeString(opcode()), "(", Join(operands_, ", ", [](string* out, HloInstruction* operand) { StrAppend(out, operand->name()); - }) - .c_str()); + }), + ")"); +} + +HloInstructionProto HloInstruction::ToProto() const { + HloInstructionProto proto; + proto.set_name(name_); + proto.set_opcode(HloOpcodeString(opcode_)); + *proto.mutable_shape() = shape_; + for (const HloInstruction* operand : operands_) { + *proto.add_operand_names() = operand->name(); + } + for (const HloInstruction* control : control_predecessors_) { + *proto.add_control_predecessor_names() = control->name(); + } + for (const HloComputation* computation : called_computations_) { + *proto.add_called_computation_names() = computation->name(); + } + *proto.mutable_metadata() = metadata_; + switch (opcode_) { + case HloOpcode::kConstant: + *proto.mutable_literal() = literal_->ToProto(); + break; + case HloOpcode::kParameter: + proto.set_parameter_number(parameter_number_); + proto.set_parameter_name(parameter_name_); + break; + case HloOpcode::kFusion: { + HloComputationProto* proto_fused_computation = + proto.mutable_fused_instructions_computation(); + proto_fused_computation->set_name(name()); + + // Fill in fused instructions in post order. + auto fused_instructions = + fused_instructions_computation()->MakeInstructionPostOrder(); + for (auto fused_instruction : fused_instructions) { + HloInstructionProto fused_proto = fused_instruction->ToProto(); + proto_fused_computation->add_instructions()->Swap(&fused_proto); + } + break; + } + case HloOpcode::kGetTupleElement: + proto.set_tuple_index(tuple_index_); + break; + default: {} // Nothing to do + } + return proto; } string HloInstruction::ToCategory() const { @@ -1503,12 +1853,16 @@ string HloInstruction::ToCategory() const { return "non-elementwise fusion"; } case FusionKind::kInput: - return "reduce fusion"; + return "input fusion"; + case FusionKind::kOutput: + return "output fusion"; case FusionKind::kTransposeDot: return "dot fusion"; case FusionKind::kConvBackwardFilter: case FusionKind::kConvBackwardInput: return "convolution fusion"; + case FusionKind::kCustom: + return "custom fusion"; } } @@ -1519,30 +1873,19 @@ string HloInstruction::ToCategory() const { return HloOpcodeString(opcode()); } -string HloInstruction::FullyQualifiedName() const { - if (IsFused()) { - return tensorflow::strings::StrCat(fusion_instruction()->parent()->name(), - "::", fusion_instruction()->name(), - "::", name_); - } - return tensorflow::strings::StrCat(parent_->name(), "::", name_); -} - HloInstruction* HloInstruction::tracing() const { return trace_instruction_; } void HloInstruction::set_tracing(HloInstruction* trace_instruction) { trace_instruction_ = trace_instruction; } -const string& HloInstruction::tracing_tag() const { +string HloInstruction::TracingTag() const { CHECK_EQ(HloOpcode::kTrace, opcode()); CHECK(literal_ != nullptr); - return literal_->u8s(); + return literal_->u8s_string(); } -bool HloInstruction::IsFused() const { - return parent_fusion_instruction_ != nullptr; -} +bool HloInstruction::IsFused() const { return parent_->IsFusionComputation(); } bool HloInstruction::IsFusable() const { // Instructions which are traced should not be fused. @@ -1552,7 +1895,6 @@ bool HloInstruction::IsFusable() const { // Some kinds of instructions don't make sense to fuse. switch (opcode_) { - case HloOpcode::kFusion: case HloOpcode::kInfeed: case HloOpcode::kOutfeed: case HloOpcode::kParameter: @@ -1561,44 +1903,50 @@ bool HloInstruction::IsFusable() const { case HloOpcode::kRecv: return false; // Only fuse Rng if it is used once, otherwise the random numbers generated - // will be different in each fusion. + // will be different in each fusion. If it is the root (user count = 0) + // then it is the equivalent of having one user. case HloOpcode::kRng: - return users_.size() == 1; + return users_.size() <= 1; default: return true; } } -HloInstruction* HloInstruction::fusion_instruction() const { - CHECK(IsFused()); - return parent_fusion_instruction_; +HloComputation* HloInstruction::fused_instructions_computation() const { + CHECK_EQ(opcode_, HloOpcode::kFusion); + CHECK(!called_computations_.empty()); + auto* fused_instructions_computation = called_computations_.front(); + CHECK(fused_instructions_computation->IsFusionComputation()); + return fused_instructions_computation; } HloInstruction* HloInstruction::fused_expression_root() const { CHECK_EQ(opcode_, HloOpcode::kFusion); - return fused_root_; + return fused_instructions_computation()->root_instruction(); } HloInstruction* HloInstruction::fused_parameter(int64 parameter_number) const { CHECK_EQ(opcode_, HloOpcode::kFusion); - CHECK_GE(parameter_number, 0); - CHECK_LT(parameter_number, fused_parameters_.size()); - return fused_parameters_[parameter_number]; + return fused_instructions_computation()->parameter_instruction( + parameter_number); } const std::vector& HloInstruction::fused_parameters() const { CHECK_EQ(opcode_, HloOpcode::kFusion); - return fused_parameters_; + return fused_instructions_computation()->parameter_instructions(); } const std::list>& HloInstruction::fused_instructions() const { CHECK_EQ(opcode_, HloOpcode::kFusion); - return fused_instructions_; + return fused_instructions_computation()->instructions(); } HloInstruction::HloInstruction(HloOpcode opcode, const Shape& shape) - : shape_(shape), opcode_(opcode), name_("%" + HloOpcodeString(opcode)) { + : unique_id_(-1), + opcode_(opcode), + shape_(shape), + name_("%" + HloOpcodeString(opcode)) { TF_DCHECK_OK(ShapeUtil::ValidateShapeWithOptionalLayout(shape_)); } @@ -1606,6 +1954,12 @@ Status HloInstruction::Visit(DfsHloVisitor* visitor) { switch (opcode_) { case HloOpcode::kAbs: return visitor->HandleAbs(this, operands_[0]); + case HloOpcode::kBatchNormTraining: + return visitor->HandleBatchNormTraining(this); + case HloOpcode::kBatchNormInference: + return visitor->HandleBatchNormInference(this); + case HloOpcode::kBatchNormGrad: + return visitor->HandleBatchNormGrad(this); case HloOpcode::kSign: return visitor->HandleSign(this, operands_[0]); case HloOpcode::kConstant: @@ -1628,9 +1982,9 @@ Status HloInstruction::Visit(DfsHloVisitor* visitor) { case HloOpcode::kSubtract: return visitor->HandleSubtract(this, operands_[0], operands_[1]); case HloOpcode::kMaximum: - return visitor->HandleMaximum(this, operands_[0], operands_[1]); + return visitor->HandleMaximum(this); case HloOpcode::kMinimum: - return visitor->HandleMinimum(this, operands_[0], operands_[1]); + return visitor->HandleMinimum(this); case HloOpcode::kLogicalAnd: return visitor->HandleLogicalAnd(this, operands_[0], operands_[1]); case HloOpcode::kLogicalOr: @@ -1638,9 +1992,9 @@ Status HloInstruction::Visit(DfsHloVisitor* visitor) { case HloOpcode::kConcatenate: return visitor->HandleConcatenate(this, operands_); case HloOpcode::kConvert: - return visitor->HandleConvert(this, operands_[0]); + return visitor->HandleConvert(this); case HloOpcode::kCopy: - return visitor->HandleCopy(this, operands_[0]); + return visitor->HandleCopy(this); case HloOpcode::kMultiply: return visitor->HandleMultiply(this, operands_[0], operands_[1]); case HloOpcode::kDot: @@ -1684,6 +2038,10 @@ Status HloInstruction::Visit(DfsHloVisitor* visitor) { return visitor->HandleLog(this, operands_[0]); case HloOpcode::kTanh: return visitor->HandleTanh(this, operands_[0]); + case HloOpcode::kCos: + return visitor->HandleCos(this, operands_[0]); + case HloOpcode::kSin: + return visitor->HandleSin(this, operands_[0]); case HloOpcode::kIsFinite: return visitor->HandleIsFinite(this, operands_[0]); case HloOpcode::kLogicalNot: @@ -1700,10 +2058,12 @@ Status HloInstruction::Visit(DfsHloVisitor* visitor) { return visitor->HandleTranspose(this); case HloOpcode::kReverse: return visitor->HandleReverse(this, operands_[0]); + case HloOpcode::kReducePrecision: + return visitor->HandleReducePrecision(this); case HloOpcode::kSlice: return visitor->HandleSlice(this, operands_[0]); case HloOpcode::kDynamicSlice: - return visitor->HandleDynamicSlice(this, operands_); + return visitor->HandleDynamicSlice(this, operands_[0], operands_[1]); case HloOpcode::kDynamicUpdateSlice: return visitor->HandleDynamicUpdateSlice(this, operands_[0], operands_[1], operands_[2]); @@ -1716,12 +2076,11 @@ Status HloInstruction::Visit(DfsHloVisitor* visitor) { case HloOpcode::kRng: return visitor->HandleRng(this, distribution_); case HloOpcode::kWhile: - return visitor->HandleWhile(this, operands_[0], while_condition(), - while_body()); + return visitor->HandleWhile(this); case HloOpcode::kFusion: return visitor->HandleFusion(this); case HloOpcode::kCall: - return visitor->HandleCall(this, operands_, to_apply()); + return visitor->HandleCall(this); case HloOpcode::kCustomCall: return visitor->HandleCustomCall(this, operands_, custom_call_target_); case HloOpcode::kSend: @@ -1739,62 +2098,112 @@ Status HloInstruction::Visit(DfsHloVisitor* visitor) { HloOpcodeString(opcode_).c_str()); } -Status HloInstruction::AcceptInternal(DfsHloVisitor* visitor, - const CompareFunction* operand_order) { - // Do not visit this HLO node again if it is already visited. - if (visitor->DidVisit(*this)) { - VLOG(3) << "Not visiting HLO " << name() << " as it was already visited."; - return Status::OK(); - } +using DFSStack = + tensorflow::gtl::InlinedVector, 16>; - // If the instruction is in the visiting state, it means a cycle. - if (visitor->IsVisiting(*this)) { - return FailedPrecondition( - "A cycle is detected while visiting instruction %s", - ToString().c_str()); - } - visitor->SetVisiting(*this); +// Push "child" onto the dfs_stack if not already visited. Returns false if a +// cycle was detected, and true otherwise. +inline bool PushDFSChild(DfsHloVisitor* visitor, DFSStack* dfs_stack, + HloInstruction* child) { + CHECK(child != nullptr); + const int id = child->unique_id(); + CHECK_GE(id, 0) << "instruction may not have a parent computation"; + switch (visitor->GetVisitState(id)) { + case DfsHloVisitor::kVisiting: + return false; - // Sort operands and control predecessors, if an ordering was provided. Note - // that 'temp_sorted_operands' must live at this scope, since 'operands' will - // point to it if the operands are sorted. The point of the 'operands' - // pointer is to avoid copying the operands in the common case where the - // operands are not sorted. - std::vector* operands = &operands_; - std::vector temp_sorted_operands; - std::vector predecessors(control_predecessors_.begin(), - control_predecessors_.end()); - if (operand_order != nullptr) { - temp_sorted_operands = operands_; - std::sort(temp_sorted_operands.begin(), temp_sorted_operands.end(), - *operand_order); - std::sort(predecessors.begin(), predecessors.end(), *operand_order); - operands = &temp_sorted_operands; - } + case DfsHloVisitor::kVisited: + // Nothing to do + return true; - for (auto operand : *operands) { - VLOG(3) << "Going to visit HLO " << operand->name() << " as operand of HLO " - << name(); - TF_RETURN_IF_ERROR(operand->AcceptInternal(visitor, operand_order)); + case DfsHloVisitor::kNotVisited: + dfs_stack->push_back(std::make_pair(id, child)); + return true; } +} - for (auto control_predecessor : predecessors) { - VLOG(3) << "Going to visit HLO " << control_predecessor->name() - << " as a control predecessor of HLO " << name(); - TF_RETURN_IF_ERROR( - control_predecessor->AcceptInternal(visitor, operand_order)); - } +using InternalCompareFunction = + std::function, + std::pair)>; +static Status PostOrderDFS(HloInstruction* root, DfsHloVisitor* visitor, + const InternalCompareFunction* operand_order, + bool ignore_control_predecessors) { + visitor->ReserveVisitStates(root->GetModule()->NumUniqueInstructionIds()); + + // dfs_stack holds pairs of unique_id(), HloInstruction*>. + // + // We need to keep track of both the id and the instruction because + // instructions can get deleted while they are on the stack, so we + // can't always use the (potentiall dead) instruction object to grab + // its id. + DFSStack dfs_stack; + dfs_stack.emplace_back(root->unique_id(), root); + + do { + DCHECK(!dfs_stack.empty()); + + int current_id = dfs_stack.back().first; + HloInstruction* current_node = dfs_stack.back().second; + CHECK_GE(current_id, 0) << current_id << ": " << current_node + << ": instruction may not have parent computation"; + DfsHloVisitor::VisitState visit_state = visitor->GetVisitState(current_id); + if (visit_state == DfsHloVisitor::kVisited) { + dfs_stack.pop_back(); + VLOG(3) << "Not visiting HLO " << current_node->name() + << " as it was already visited."; + continue; + } + + if (visit_state == DfsHloVisitor::kVisiting) { + dfs_stack.pop_back(); + + TF_RETURN_IF_ERROR(visitor->Preprocess(current_node)); + VLOG(2) << "Visiting HLO " << current_node->name(); + TF_RETURN_IF_ERROR(current_node->Visit(visitor)); + visitor->SetVisitState(current_id, DfsHloVisitor::kVisited); + TF_RETURN_IF_ERROR(visitor->Postprocess(current_node)); + continue; + } + + visitor->SetVisitState(current_id, DfsHloVisitor::kVisiting); + + const size_t old_dfs_stack_size = dfs_stack.size(); + for (HloInstruction* child : current_node->operands()) { + if (!TF_PREDICT_TRUE(PushDFSChild(visitor, &dfs_stack, child))) { + return FailedPrecondition( + "A cycle is detected while visiting instruction %s", + current_node->ToString().c_str()); + } + } + + if (!ignore_control_predecessors) { + for (HloInstruction* child : current_node->control_predecessors()) { + if (!TF_PREDICT_TRUE(PushDFSChild(visitor, &dfs_stack, child))) { + return FailedPrecondition( + "A cycle is detected while visiting instruction %s", + current_node->ToString().c_str()); + } + } + } + + if (operand_order != nullptr) { + std::sort(dfs_stack.begin() + old_dfs_stack_size, dfs_stack.end(), + *operand_order); + } - TF_RETURN_IF_ERROR(visitor->Preprocess(this)); - VLOG(2) << "Visiting HLO " << name(); - TF_RETURN_IF_ERROR(Visit(visitor)); - visitor->SetVisited(*this); - return visitor->Postprocess(this); + // This makes the traversal order the same as what you'd expect + // out of a recursive algorithm. + std::reverse(dfs_stack.begin() + old_dfs_stack_size, dfs_stack.end()); + } while (!dfs_stack.empty()); + + return Status::OK(); } -Status HloInstruction::Accept(DfsHloVisitor* visitor, bool call_finish_visit) { - VLOG(2) << "HloInstruction::Accept(" << name() << ")"; - TF_RETURN_IF_ERROR(AcceptInternal(visitor, nullptr)); +Status HloInstruction::Accept(DfsHloVisitor* visitor, bool call_finish_visit, + bool ignore_control_predecessors) { + VLOG(3) << "HloInstruction::Accept(" << name() << ")"; + TF_RETURN_IF_ERROR( + PostOrderDFS(this, visitor, nullptr, ignore_control_predecessors)); if (call_finish_visit) { TF_RETURN_IF_ERROR(visitor->FinishVisit(this)); } @@ -1805,17 +2214,28 @@ Status HloInstruction::AcceptWithOperandOrder( DfsHloVisitor* visitor, const CompareFunction& operand_order, bool call_finish_visit) { VLOG(2) << "HloInstruction::AcceptWithOperandOrder(" << name() << ")"; - TF_RETURN_IF_ERROR(AcceptInternal(visitor, &operand_order)); + InternalCompareFunction func = [&operand_order]( + std::pair a, + std::pair b) { + // Call the client's comparison function on the actual HloInstruction* + // objects (ignoring the internal ids we also have in our stack entries) + return operand_order(a.second, b.second); + }; + TF_RETURN_IF_ERROR(PostOrderDFS(this, visitor, &func, + /*ignore_control_predecessors=*/false)); if (call_finish_visit) { + VLOG(3) << "HloInstruction::AcceptWithOperandOrder BEFORE FINISH VISIT"; TF_RETURN_IF_ERROR(visitor->FinishVisit(this)); + VLOG(3) << "HloInstruction::AcceptWithOperandOrder AFTER FINISH VISIT"; } + VLOG(2) << "HloInstruction::AcceptWithOperandOrder EXIT"; return Status::OK(); } namespace { -// Returns true if the given order is a topological sort of the instructions it -// contains. +// 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; @@ -1826,8 +2246,8 @@ bool OrderIsTopologicalSort(const std::vector& order) { } } // 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). + // 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) || @@ -1909,6 +2329,32 @@ std::vector HloInstruction::OperandIndices( return result; } +bool HloInstruction::IsElementwiseBinary() const { + switch (opcode_) { + // Binary elementwise operations. If you update this, please update + // IsElementwise() accordingly. + case HloOpcode::kAdd: + case HloOpcode::kDivide: + case HloOpcode::kEq: + case HloOpcode::kGe: + case HloOpcode::kGt: + case HloOpcode::kLe: + case HloOpcode::kLt: + case HloOpcode::kMaximum: + case HloOpcode::kMinimum: + case HloOpcode::kMultiply: + case HloOpcode::kNe: + case HloOpcode::kPower: + case HloOpcode::kRemainder: + case HloOpcode::kSubtract: + case HloOpcode::kLogicalAnd: + case HloOpcode::kLogicalOr: + return true; + default: + return false; + } +} + bool HloInstruction::IsElementwise() const { switch (opcode_) { // Nullary elementwise operations. @@ -1920,17 +2366,21 @@ bool HloInstruction::IsElementwise() const { case HloOpcode::kCeil: case HloOpcode::kConvert: case HloOpcode::kCopy: + case HloOpcode::kCos: case HloOpcode::kExp: case HloOpcode::kFloor: case HloOpcode::kIsFinite: case HloOpcode::kLog: case HloOpcode::kLogicalNot: case HloOpcode::kNegate: + case HloOpcode::kReducePrecision: case HloOpcode::kSign: + case HloOpcode::kSin: case HloOpcode::kTanh: return true; - // Binary elementwise operations. + // Binary elementwise operations, the same as in IsElementwiseBinary(). + // If you update this, please update IsElementwiseBinary() accordingly. case HloOpcode::kAdd: case HloOpcode::kDivide: case HloOpcode::kEq: @@ -2010,6 +2460,7 @@ bool HloInstruction::IsElementwiseOnOperand(int64 operand_idx) const { HloInstruction* operand = worklist.front(); worklist.pop_front(); for (HloInstruction* user : operand->users()) { + CHECK_GE(user->unique_id(), 0); if (ContainsKey(visited, user)) { continue; } @@ -2025,6 +2476,70 @@ bool HloInstruction::IsElementwiseOnOperand(int64 operand_idx) const { return true; } +// A helper class for memoized, recursive computation of HloOpcode::kFusion +// in HloInstruction::OperandElementUse below. +class HloInstruction::FusionReusesParamElements { + public: + using UseKind = HloInstruction::UseKind; + + // We could rather iterate backwards thru fused_instructions_ here, as it is + // in reverse postorder, and compute whether each fused instruction reuses + // the value of this parameter, which would save stack space but not allow + // us to finish early if we find a reuse. + static UseKind Compute(int64 i, const HloInstruction& hlo) { + tensorflow::gtl::FlatMap memoization_cache; + return ComputeInternal(i, hlo, &memoization_cache); + } + + private: + static UseKind ComputeInternal( + int64 i, const HloInstruction& hlo, + tensorflow::gtl::FlatMap* cache) { + if (hlo.opcode_ == HloOpcode::kParameter && hlo.parameter_number_ == i) { + return UseKind::kUse; + } + + auto p = cache->emplace(&hlo, UseKind{}); + auto value_it = p.first; + const bool key_is_new = p.second; + + if (key_is_new) { + for (int64 j = 0; j < hlo.operands_.size(); ++j) { + UseKind old_val = value_it->second; + + // The next operation invalidates iterators. + UseKind new_val = + Plus(old_val, std::min(hlo.OperandElementUse(j), + ComputeInternal(i, *hlo.operand(j), cache))); + + // Re-acquire the iterator. We could work harder to do this only if + // absolutely necessary, but this code is not hot enough to warrant + // that. + value_it = cache->find(&hlo); + value_it->second = new_val; + } + } + return value_it->second; + } + + // Fold operation for UseKinds. + static UseKind Plus(UseKind a, UseKind b) { + if (a == UseKind::kNoUse) { + return b; + } else if (b == UseKind::kNoUse) { + return a; + } else if (a == UseKind::kReuse || b == UseKind::kReuse) { + return UseKind::kReuse; + } else if (a == UseKind::kUsePermutingElements || + b == UseKind::kUsePermutingElements) { + return UseKind::kReuse; + } else { + CHECK(a == UseKind::kUse && b == UseKind::kUse); + return UseKind::kUse; + } + } +}; + HloInstruction::UseKind HloInstruction::OperandElementUse(int64 i) const { switch (opcode_) { case HloOpcode::kBitcast: @@ -2039,45 +2554,21 @@ HloInstruction::UseKind HloInstruction::OperandElementUse(int64 i) const { // Pad reuses the padding value but not the padded array elements. // Reduce reuses the init value but not the operand array elements. return i > 0 ? UseKind::kReuse : UseKind::kUsePermutingElements; - case HloOpcode::kFusion: { - tensorflow::gtl::FlatMap cache; - // We could rather iterate backwards thru fused_instructions_ here, as it - // is in reverse postorder, and compute whether each fused instruction - // reuses the value of this parameter, which would save stack space but - // not allow us to finish early if we find a reuse. - std::function reuses_parameter_elements = - [i, &cache, &reuses_parameter_elements](const HloInstruction& hlo) { - auto plus = [](const UseKind& a, const UseKind& b) { - if (a == UseKind::kNoUse) return b; - if (b == UseKind::kNoUse) return a; - if (a == UseKind::kReuse || b == UseKind::kReuse) { - return UseKind::kReuse; - } - if (a == UseKind::kUsePermutingElements || - b == UseKind::kUsePermutingElements) { - return UseKind::kReuse; - } - CHECK(UseKind::kUse == a && UseKind::kUse == b); - return UseKind::kUse; - }; - - if (hlo.opcode_ == HloOpcode::kParameter && - hlo.parameter_number_ == i) { - return UseKind::kUse; - } - if (!ContainsKey(cache, &hlo)) { - for (int64 j = 0; j < hlo.operands_.size(); ++j) { - UseKind old = cache[&hlo]; - UseKind updated = plus( - old, std::min(hlo.OperandElementUse(j), - reuses_parameter_elements(*hlo.operand(j)))); - cache[&hlo] = updated; - } - } - return cache[&hlo]; - }; - return reuses_parameter_elements(*fused_root_); - } + case HloOpcode::kFusion: + // Uses the memoizing, recursive computation defined above. + return FusionReusesParamElements::Compute(i, *fused_expression_root()); + case HloOpcode::kDot: + // Dot operations with inputs [A,B] * [B,1] do not re-use + // elements on their left operand. + // Dot operations with inputs [1,A] * [A,B] do not re-use + // elements on their right operand. + if (shape().dimensions_size() == 2) { + if ((i == 0 && shape().dimensions(1) == 1) || + (i == 1 && shape().dimensions(0) == 1)) { + return UseKind::kUse; + } + } + return UseKind::kReuse; default: return IsElementwise() ? UseKind::kUse : UseKind::kReuse; } @@ -2098,15 +2589,23 @@ string ToString(HloInstruction::FusionKind kind) { return "kLoop"; case HloInstruction::FusionKind::kInput: return "kInput"; + case HloInstruction::FusionKind::kOutput: + return "kOutput"; case HloInstruction::FusionKind::kTransposeDot: return "kTransposeDot"; case HloInstruction::FusionKind::kConvBackwardFilter: return "kConvBackwardFilter"; case HloInstruction::FusionKind::kConvBackwardInput: return "kConvBackwardInput"; + case HloInstruction::FusionKind::kCustom: + return "kCustom"; } } +std::ostream& operator<<(std::ostream& os, HloInstruction::FusionKind kind) { + return os << ToString(kind); +} + string HloInstruction::ConvolutionDimensionNumbersToString() const { string result; if (convolution_dimension_numbers_ == nullptr) { @@ -2133,15 +2632,14 @@ string HloInstruction::ConvolutionDimensionNumbersToString() const { lhs_dims[dnums.batch_dimension()] = 'b'; lhs_dims[dnums.feature_dimension()] = 'f'; for (int64 i = 0; i < dnums.spatial_dimensions().size(); ++i) { - lhs_dims[dnums.spatial_dimensions(i)] = tensorflow::strings::StrCat(i); + lhs_dims[dnums.spatial_dimensions(i)] = StrCat(i); } std::vector rhs_dims(2 + dnums.kernel_spatial_dimensions().size()); rhs_dims[dnums.kernel_input_feature_dimension()] = "i"; rhs_dims[dnums.kernel_output_feature_dimension()] = "o"; for (int64 i = 0; i < dnums.spatial_dimensions().size(); ++i) { - rhs_dims[dnums.kernel_spatial_dimensions(i)] = - tensorflow::strings::StrCat(i); + rhs_dims[dnums.kernel_spatial_dimensions(i)] = StrCat(i); } result += "dim_labels="; @@ -2164,4 +2662,21 @@ bool HloInstruction::CouldBeBitcast() const { } } +HloModule* HloInstruction::GetModule() const { + if (parent_) { + return parent_->parent(); + } + return nullptr; +} + +void HloInstruction::UniquifyName(NameUniquer* name_uniquer) { + string parent_str = parent() == nullptr ? "noparent" : parent()->name(); + name_ = name_uniquer->GetUniqueName(name_); +} + +void HloInstruction::set_outer_dimension_partitions( + const std::vector& outer_dimension_partitions) { + outer_dimension_partitions_ = outer_dimension_partitions; +} + } // namespace xla diff --git a/tensorflow/compiler/xla/service/hlo_instruction.h b/tensorflow/compiler/xla/service/hlo_instruction.h index 6557ca9116312c4bc31b9f0ba734edd11106d1e7..e393e05c344491ee3ba88572df00a93cdb142d64 100644 --- a/tensorflow/compiler/xla/service/hlo_instruction.h +++ b/tensorflow/compiler/xla/service/hlo_instruction.h @@ -22,6 +22,7 @@ limitations under the License. #define TENSORFLOW_COMPILER_XLA_SERVICE_HLO_INSTRUCTION_H_ #include +#include #include #include #include @@ -29,15 +30,19 @@ limitations under the License. #include #include +#include "tensorflow/compiler/xla/literal_util.h" #include "tensorflow/compiler/xla/map_util.h" #include "tensorflow/compiler/xla/service/dfs_hlo_visitor.h" #include "tensorflow/compiler/xla/service/dfs_hlo_visitor_with_default.h" +#include "tensorflow/compiler/xla/service/hlo.pb.h" #include "tensorflow/compiler/xla/service/hlo_opcode.h" +#include "tensorflow/compiler/xla/service/name_uniquer.h" #include "tensorflow/compiler/xla/types.h" #include "tensorflow/compiler/xla/xla_data.pb.h" #include "tensorflow/core/lib/core/status.h" #include "tensorflow/core/lib/core/stringpiece.h" #include "tensorflow/core/lib/gtl/array_slice.h" +#include "tensorflow/core/lib/gtl/inlined_vector.h" #include "tensorflow/core/platform/logging.h" #include "tensorflow/core/platform/macros.h" #include "tensorflow/core/platform/types.h" @@ -45,18 +50,26 @@ limitations under the License. namespace xla { class HloComputation; +class HloModule; // HLO instructions are the IR used by the high-level compiler. class HloInstruction { public: enum class FusionKind { kLoop, // Fused into a loop. - kInput, // Fused into a reduction kernel. + kInput, // Op's input is fused into the op itself. + kOutput, // Op's output is fused into the op itself. + // REQUIRES: At least one operand buffer must be able + // to alias the output buffer. kTransposeDot, // Fused into a dot with transposed operands. kConvBackwardFilter, // Fused into a backward filter convolution. kConvBackwardInput, // Fused into a backward input convolution. + + kCustom, // Custom category for backend-specific fusions that + // do not match any of the more specific ones. }; + ~HloInstruction(); // Creates a parameter-retrieving instruction. static std::unique_ptr CreateParameter(int64 parameter_number, const Shape& shape, @@ -122,6 +135,13 @@ class HloInstruction { const Window& window, const ConvolutionDimensionNumbers& dimension_numbers); + // Creates a reduce-precision op, where operand is the data to reduce in + // precision, and exponent_bits and mantissa_bits describe the precision to + // reduce it to. + static std::unique_ptr CreateReducePrecision( + const Shape& shape, HloInstruction* operand, const int exponent_bits, + const int mantissa_bits); + // Creates a cross replica sum op. static std::unique_ptr CreateCrossReplicaSum( const Shape& shape, HloInstruction* operand); @@ -158,7 +178,8 @@ class HloInstruction { static std::unique_ptr CreateSlice( const Shape& shape, HloInstruction* operand, tensorflow::gtl::ArraySlice start_indices, - tensorflow::gtl::ArraySlice limit_indices); + tensorflow::gtl::ArraySlice limit_indices, + tensorflow::gtl::ArraySlice strides); // Creates a slice instruction, where the first operand is sliced by // start indices specified in the second operand, and by size specfied in @@ -199,6 +220,23 @@ class HloInstruction { const Shape& shape, HloInstruction* operand, HloInstruction* init_value, const Window& window, HloComputation* reduce_computation); + // Creates a batch-norm-training instruction. + static std::unique_ptr CreateBatchNormTraining( + const Shape& shape, HloInstruction* operand, HloInstruction* scale, + HloInstruction* offset, float epsilon, int64 feature_index); + + // Creates a batch-norm-inference instruction. + static std::unique_ptr CreateBatchNormInference( + const Shape& shape, HloInstruction* operand, HloInstruction* scale, + HloInstruction* offset, HloInstruction* mean, HloInstruction* variance, + float epsilon, int64 feature_index); + + // Creates a batch-norm-grad instruction. + static std::unique_ptr CreateBatchNormGrad( + const Shape& shape, HloInstruction* operand, HloInstruction* scale, + HloInstruction* mean, HloInstruction* variance, + HloInstruction* grad_output, float epsilon, int64 feature_index); + // Creates a scatter computation that scatters the `source` array to the // selected indices of each window. static std::unique_ptr CreateSelectAndScatter( @@ -278,6 +316,11 @@ class HloInstruction { // Returns the opcode for this instruction. HloOpcode opcode() const { return opcode_; } + // Returns true if this instruction has a side effect. An instruction has a + // side effect if it uses certain opcodes or calls a computation with a side + // effect. + bool HasSideEffect() const; + // Returns the result shape of this instruction. const Shape& shape() const; @@ -294,7 +337,8 @@ class HloInstruction { int64 operand_count() const { return operands_.size(); } // Returns the vector of operands of this instruction. - const std::vector& operands() const { return operands_; } + using InstructionVector = tensorflow::gtl::InlinedVector; + const InstructionVector& operands() const { return operands_; } // Returns the index of 'target' in the operands sequence. // Precondition: target must be an operand (or a fatal error will occur). @@ -343,7 +387,23 @@ class HloInstruction { std::function eq_operands = std::equal_to(), std::function - eq_computations = std::equal_to()) const; + eq_computations = std::equal_to()) const { + // An instruction is always identical to itself. + if (this == &other) { + return true; + } + + // Identical instruction must have the same opcode and identical operands. + // In general, there is no need to check shape because shape is inferred + // from the shape of the operands. + if (opcode() != other.opcode() || + !ContainersEqual(operands(), other.operands(), + std::move(eq_operands))) { + return false; + } + + return IdenticalSlowPath(other, eq_computations); + } // Returns whether the instruction has a constant operand. bool HasConstantOperand() const; @@ -371,8 +431,12 @@ class HloInstruction { // Performs a postorder DFS visit using this node as the root. If // call_finish_visit is true, then DfsHloVisitor::FinishVisit is called when - // complete. - Status Accept(DfsHloVisitor* visitor, bool call_finish_visit = true); + // complete. If ignore_control_predecessors is true, instructions only + // reachable via control dependencies will not be visited, and the postorder + // will not take control dependencies into account. It is as if the control + // dependencies didn't exist in the graph at all. + Status Accept(DfsHloVisitor* visitor, bool call_finish_visit = true, + bool ignore_control_predecessors = false); // Same as Accept() above, but the order of operand and control predecessor // visitation is determined by the given operand order; if compare(A, B) == @@ -418,6 +482,11 @@ class HloInstruction { return parameter_name_; } + void set_parameter_name(const string& str) { + CHECK_EQ(HloOpcode::kParameter, opcode_); + parameter_name_ = str; + } + // Returns the dimension sizes or numbers associated with this instruction. // // Precondition: opcode() is one of: concatenate, reduce, broadcast, reshape, @@ -476,18 +545,21 @@ class HloInstruction { string SignatureString() const; // Returns a debugging string that represents this instruction. - string ToString(bool compact_operands = false) const; + string ToString(bool compact_operands = false, + bool include_metadata = true) const; + + string ToStringNoMetadata() const { return ToString(false, false); } // As ToString, but returns a shorter string. string ToShortString() const; + // Returns a serialized representation of this instruction. + HloInstructionProto ToProto() const; + // Returns a category for the HLO. This could be something like "convolution" // or "elementwise". string ToCategory() const; - // Returns the string concatenation of parent name and this instructions name. - string FullyQualifiedName() const; - // Returns a logging instruction, if the output of this instruction is logged. // // Postcondition: retval == nullptr || retval->opcode() == HloOpcode::kTrace @@ -501,6 +573,18 @@ class HloInstruction { // Precondition: opcode() == HloOpcode::kSend or HloOpcode::kRecv int64 channel_id() const { return channel_id_; } + // Returns feature_index field associated with the instruction. The index + // represents the index of the feature dimension. + // + // Precondition: opcode() == HloOpcode::kBatchNormTraining + int64 feature_index() const { return feature_index_; } + + // Returns a epsilon value associated with the instruction. The is a small + // number added to the variance to avoid divide-by-zero error. + // + // Precondition: opcode() == HloOpcode::kBatchNormTraining + float epsilon() const { return epsilon_; } + // Returns the infeed configuration string. The infeed configuration includes // any metadata needed for the backend compiler (e.g., infeed buffer address) // and is target-dependent. @@ -510,7 +594,7 @@ class HloInstruction { // Returns a tag to be used in tracing. // // Precondition: opcode() == HloOpcode::kTrace - const string& tracing_tag() const; + string TracingTag() const; // Returns whether the instruction is a constant. bool IsConstant() const; @@ -519,24 +603,22 @@ class HloInstruction { // instruction. bool IsFused() const; + // Returns the computation for this fused instruction. + // + // Precondition: opcode() == HloOpcode::kFusion + HloComputation* fused_instructions_computation() const; + // Returns true if this instruction can be legally fused into a fusion // instruction. bool IsFusable() const; - // Returns the fusion instruction that contains this instruction. - // - // Note: only valid if this instruction is fused into a fusion instruction. - HloInstruction* fusion_instruction() const; - // Returns the root instruction of the fused expression contained within this // fusion instruction. // // Precondition: opcode() == HloOpcode::kFusion HloInstruction* fused_expression_root() const; - // Returns the vector of fused instructions inside this fusion - // instruction. The order is a reverse postorder of the fused expression (root - // is first in the order). + // Returns the list of fused instructions inside this fusioninstruction. // // Note: although the list itself is const, the instructions contained in the // list returned here are mutable. @@ -555,11 +637,23 @@ class HloInstruction { // Precondition: opcode() == HloOpcode::kFusion const std::vector& fused_parameters() const; + // Returns true if this instruction is a fusion instruction that generates + // multiple outputs. + const bool IsMultiOutputFusion() const { + return (opcode() == HloOpcode::kFusion && + fused_expression_root()->opcode() == HloOpcode::kTuple); + } + FusionKind fusion_kind() const { CHECK_EQ(HloOpcode::kFusion, opcode_); return fusion_kind_; } + void set_fusion_kind(FusionKind kind) { + CHECK_EQ(HloOpcode::kFusion, opcode_); + fusion_kind_ = kind; + } + // Merges the fused instructions from 'instruction_to_merge' into the // fused instruction set of 'this', updating operands as necessary. // @@ -567,6 +661,16 @@ class HloInstruction { // Predondition: 'instruction_to_merge' must be an operand of 'this'. void MergeFusionInstruction(HloInstruction* instruction_to_merge); + // Merges the fused instructions from 'instruction_to_merge' into the + // fused instruction set of 'this' and generate multioutput fusion + // instructions. All the user of instruction_to_merge will be redirected + // to 'this' instruction. `instruction_to_merge' will be removed from its + // parent computation. + // + // Precondition: opcode() == HloOpcode::kFusion + void MergeFusionInstructionIntoMultiOutput( + HloInstruction* instruction_to_merge); + // Fuses the given instruction in this fusion instruction. instruction_to_fuse // is cloned and the clone is placed in the fusion // instruction. instruction_to_fuse is unchanged. Instruction is cloned rather @@ -576,7 +680,21 @@ class HloInstruction { // and significantly complicate code generation. // // Precondition: this->opcode() == HloOpcode::kFusion - HloInstruction* FuseInstruction(HloInstruction* instruction_to_fuse); + HloInstruction* FuseInstruction(HloInstruction* instruction_to_fuse) { + return FuseInstructionInternal(instruction_to_fuse); + } + + // Fuses the given instruction in this fusion instruction and generate + // multioutput fusion instruction. A clone of the instruction_to_fuse will + // be part of the output of fusion instructions. The users of + // instruction_to_fuse will be redirected to this fusion instructions. + // instruction_to_fuse will be removed from its parent computation. + // + // Precondition: this->opcode() == HloOpcode::kFusion + HloInstruction* FuseInstructionIntoMultiOutput( + HloInstruction* instruction_to_fuse) { + return FuseInstructionInternal(instruction_to_fuse, /* add_output */ true); + } // Returns the start index in the given dimension for a slice node. // @@ -600,6 +718,15 @@ class HloInstruction { return slice_limits_; } + // Returns the stride in the given dimension for a slice node. + // + // Precondition: opcode() == HloOpcode::kSlice + int64 slice_strides(int64 dimension) const { + CHECK_EQ(HloOpcode::kSlice, opcode_); + return slice_strides_[dimension]; + } + const std::vector& slice_strides() const { return slice_strides_; } + // Returns the size of the slice in the given dimension for a dynamic // slice node. // @@ -613,6 +740,22 @@ class HloInstruction { return dynamic_slice_sizes_; } + // Returns the number of exponent bits for a reduce-precision node. + // + // Precondition: opcode() == HloOpcode::kReducePrecision + int32 exponent_bits() const { + CHECK_EQ(HloOpcode::kReducePrecision, opcode_); + return exponent_bits_; + } + + // Returns the number of mantissa bits for a reduce-precision node. + // + // Precondition: opcode() == HloOpcode::kReducePrecision + int32 mantissa_bits() const { + CHECK_EQ(HloOpcode::kReducePrecision, opcode_); + return mantissa_bits_; + } + // Returns data on the window in a windowed operation such as // convolution. const Window& window() const { @@ -654,12 +797,21 @@ class HloInstruction { const Shape& shape, tensorflow::gtl::ArraySlice operands); - // Returns the computations this instruction calls (if any). This includes - // computations called by fused instructions inside of a fusion instruction. + // Returns the computations this instruction directly calls (if any). const std::vector& called_computations() const { return called_computations_; } + // Replaces all called computations based on a map function. This is needed + // when we clone hlo_computations and want to let the instructions to point + // to the newly cloned nodes. + void ReplaceCalledComputations( + std::function map_function) { + for (int64 i = 0; i < called_computations_.size(); ++i) { + called_computations_[i] = map_function(called_computations_[i]); + } + } + // Returns true if this instruction performs an elementwise operation on // `operand_idx`-th operand. An instruction is elementwise on an operand iff, // after performing necessary implicit broadcast @@ -675,6 +827,9 @@ class HloInstruction { // Returns true if this instruction is elementwise on all its operands. bool IsElementwise() const; + // Returns true if this instruction is binary and elementwise. + bool IsElementwiseBinary() const; + // Returns whether this instruction may reuse elements of its `i`th operand. bool ReusesOperandElements(int64 i) const { return OperandElementUse(i) == UseKind::kReuse; @@ -694,9 +849,9 @@ class HloInstruction { std::tuple, std::vector> ReshapeMerelyInsertsOrDeletes1SizedDimensions() const; - // Returns the opcode string for this instruction. Compared with - // HloOpcodeString method, this wrapper dumps additional information - // such as fusion kind. + // Returns the opcode string for this instruction. This is the result from + // HloOpcodeString plus, for fusion nodes, the fusion kind, separated by a + // ':'. string ExtendedOpcodeStr() const; // Returns a string identifier for this instruction. If no string identifier @@ -704,8 +859,20 @@ class HloInstruction { // this instruction. const string& name() const { return name_; } - // Sets the string identifier for this instruction. - void set_name(const string& name) { name_ = name; } + // Use the given NameUniquer to select a unique name for the instruction based + // on the instruction's existing name. + void UniquifyName(NameUniquer* name_uniquer); + + // Set the unique id for this instruction to "id" + void SetUniqueId(int id) { + CHECK_EQ(unique_id_, -1); // Should not be assigned already + CHECK_GE(id, 0); + unique_id_ = id; + } + + // Return the unique ID assigned to this node via SetUniqueId (or -1 + // if no id has been assigned yet). + int unique_id() const { return unique_id_; } // Sets the debug metadata for this instruction. void set_metadata(const OpMetadata& metadata) { metadata_ = metadata; } @@ -718,13 +885,39 @@ class HloInstruction { const HloComputation* parent() const { return parent_; } HloComputation* parent() { return parent_; } + // Returns the module for this instruction. + HloModule* GetModule() const; + // Returns whether we could assign input and output layouts to this // instruction to make it a bitcast. bool CouldBeBitcast() const; + // CHECKs various invariants of a fusion instruction. + void CheckFusionInstruction() const; + + // Get/Set the number of partitions per outer dimension (in order, starting + // with outer-most dimension first). Currently used by the parallel cpu + // backend to partition HLOs into parallel tasks. + // TODO(b/62783254) Replace these methods with a more general way to + // annotate HLOs with backend-specific information. + const std::vector& outer_dimension_partitions() const { + return outer_dimension_partitions_; + } + void set_outer_dimension_partitions( + const std::vector& outer_dimension_partitions); + private: enum class UseKind { kNoUse, kReuse, kUsePermutingElements, kUse }; + // Helper class for computing OperandElementUse for kFusion. + class FusionReusesParamElements; + + // See comments on Identical(). + bool IdenticalSlowPath( + const HloInstruction& other, + std::function + eq_computations) const; + // Creates an n-ary elementwise operation. static std::unique_ptr CreateNary( const Shape& shape, HloOpcode opcode, @@ -744,25 +937,34 @@ class HloInstruction { // by factory methods. HloInstruction(HloOpcode opcode, const Shape& shape); + // Fuses the given instruction into this fusion instruction. When add_output + // is false (which is the default), instruction_to_fuse is cloned and the + // clone is placed in the fusion instruction. instruction_to_fuse is + // unchanged. + // + // When add_output is true, a clone of the instruction_to_fuse will be part + // of the output of fusion instructions. The users of instruction_to_fuse + // will be redirected to this fusion instructions. instruction_to_fuse will + // be removed from its parent computation. + // + // Precondition: this->opcode() == HloOpcode::kFusion + HloInstruction* FuseInstructionInternal(HloInstruction* instruction_to_fuse, + bool add_output = false); + // Clones the given instruction_to_fuse and insert the clone into this fusion - // instruction. + // instruction. If add_output is true, a clone of instruction_to_fuse will + // be in the output of the this fusion instruction (part of the tuple of the + // fusion root). // // Precondition: opcode() == HloOpcode::kFusion - HloInstruction* CloneAndFuseInternal(HloInstruction* instruction_to_fuse); + HloInstruction* CloneAndFuseInternal(HloInstruction* instruction_to_fuse, + bool add_output = false); // Clones a fusion instruction with a new shape and operands. std::unique_ptr CloneFusionWithNewOperands( const Shape& shape, tensorflow::gtl::ArraySlice operands); - // Inner DFS traversal function -- this function being called (rather than - // Accept above) allows us to distinguish the root of the traversal. - Status AcceptInternal(DfsHloVisitor* visitor, - const CompareFunction* operand_order); - - // CHECKs various invariants of a fusion instruction. - void CheckFusionInstruction() const; - // Returns true if this instruction can legally have the dimensions field // set. Used for checking precondition of dimensions field accessors. bool CanHaveDimensionsField() const; @@ -770,15 +972,36 @@ class HloInstruction { // Returns how this instruction uses elements of its `i`th operand. UseKind OperandElementUse(int64 i) const; + int unique_id_; // Unique to this HloInstruction within a HloModule + + // Opcode for this instruction. + HloOpcode opcode_; + + // Instruction operands. + InstructionVector operands_; + + // The set of control predecessors of this instruction. + std::vector control_predecessors_; + + // The users of this instruction. Users are HLOs where this instruction is an + // operand. The vector users_ and the set user_set_ contain identical + // members. The set enables fast membership testing and the vector enables + // fast, stable iteration. + std::vector users_; + std::unordered_set user_set_; + + // The set of control successors of this instruction. + std::vector control_successors_; + + // The computation in which this instruction is contained. + HloComputation* parent_ = nullptr; + // Shape of outfeed request. Shape outfeed_shape_; // Result shape of this instruction. Shape shape_; - // Opcode for this instruction. - HloOpcode opcode_; - // Literal, only present for kConstant. std::unique_ptr literal_; @@ -798,6 +1021,11 @@ class HloInstruction { // Describes the [begin, end) index range for a slice. std::vector slice_starts_; std::vector slice_limits_; + std::vector slice_strides_; + + // The bit sizes for a reduce-precision operation. + int32 exponent_bits_; + int32 mantissa_bits_; // Describes the [start, start + size) range size for a dynamic slice // ('start' is specified dynamically in the second operand of the operation). @@ -807,22 +1035,6 @@ class HloInstruction { // padding of this pad instruction. Only set for pad instructions. std::unique_ptr padding_config_; - // The set of instruction fused into this fusion instruction. Only set for - // fusion instructions. - std::list> fused_instructions_; - - // If this instruction is fused into a fusion instruction, this field points - // to the fusion instruction. - HloInstruction* parent_fusion_instruction_ = nullptr; - - // The vector of parameter instructions inside this fusion instruction. The - // index of the vector is the parameter_number of the parameter instruction. - // This vector is non-empty only for fusion instructions. - std::vector fused_parameters_; - - // The root of the expression fused into this fusion instruction. - HloInstruction* fused_root_ = nullptr; - // The type of the fusion. Used by kFusion only. FusionKind fusion_kind_; @@ -851,22 +1063,6 @@ class HloInstruction { // Outfeed configuration information, only present for kOutfeed. string outfeed_config_; - // Instruction operands. - std::vector operands_; - - // The users of this instruction. Users are HLOs where this instruction is an - // operand. The vector users_ and the set user_set_ contain identical - // members. The set enables fast membership testing and the vector enables - // fast, stable iteration. - std::vector users_; - std::unordered_set user_set_; - - // The set of control predecessors of this instruction. - std::vector control_predecessors_; - - // The set of control successors of this instruction. - std::vector control_successors_; - // A trace instruction that consumes this instruction. // // Invariant: if trace_instruction_ != nullptr, trace_instruction has this as @@ -877,6 +1073,14 @@ class HloInstruction { // Only present for kRng. RandomDistribution distribution_; + // A small float number added to the variance to avoid divide-by-zero error. + // Only present for kBatchNormTraining. + float epsilon_; + + // An integer value representing the index of the feature dimension. + // Only present for kBatchNormTraining. + int64 feature_index_; + // Represents a unique identifier for each Send/Recv instruction pair. // Only present for kSend or kRecv. int64 channel_id_ = -1; @@ -887,17 +1091,20 @@ class HloInstruction { // String identifier for instruction. string name_; - // The computation in which this instruction is contained. - HloComputation* parent_ = nullptr; - // Metadata for debugging. OpMetadata metadata_; + // The number of partitions per outer dimension (listed in order from + // outer-most dimension first). + std::vector outer_dimension_partitions_; + TF_DISALLOW_COPY_AND_ASSIGN(HloInstruction); }; string ToString(HloInstruction::FusionKind kind); +std::ostream& operator<<(std::ostream& os, HloInstruction::FusionKind kind); + } // namespace xla #endif // TENSORFLOW_COMPILER_XLA_SERVICE_HLO_INSTRUCTION_H_ diff --git a/tensorflow/compiler/xla/service/hlo_instruction_test.cc b/tensorflow/compiler/xla/service/hlo_instruction_test.cc index 8eabaa1c474aa068c423099919d3382f04c7591c..2e1eeee36b58826045f2aeabf74497b019aa1764 100644 --- a/tensorflow/compiler/xla/service/hlo_instruction_test.cc +++ b/tensorflow/compiler/xla/service/hlo_instruction_test.cc @@ -21,9 +21,11 @@ limitations under the License. #include #include "tensorflow/compiler/xla/literal_util.h" +#include "tensorflow/compiler/xla/protobuf_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/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/util.h" @@ -31,6 +33,9 @@ limitations under the License. namespace xla { namespace { +using ::testing::ElementsAre; +using ::testing::UnorderedElementsAre; + class HloInstructionTest : public HloTestBase { protected: HloInstructionTest() {} @@ -143,22 +148,27 @@ TEST_F(HloInstructionTest, UserWithTwoOperands) { // [Param foo]-----> |-----| // | Add | // [Param bar]-----> |-----| - auto foo = HloInstruction::CreateParameter(0, r0f32_, "foo"); - auto bar = HloInstruction::CreateParameter(1, r0f32_, "bar"); - auto add = HloInstruction::CreateBinary(r0f32_, HloOpcode::kAdd, foo.get(), - bar.get()); - - ExpectEqOrdered(add->operands(), {foo.get(), bar.get()}); - ExpectEqUnordered(foo->users(), {add.get()}); - ExpectEqUnordered(bar->users(), {add.get()}); + HloComputation::Builder builder(TestName()); + auto foo = + builder.AddInstruction(HloInstruction::CreateParameter(0, r0f32_, "foo")); + auto bar = + builder.AddInstruction(HloInstruction::CreateParameter(1, r0f32_, "bar")); + auto add = builder.AddInstruction( + HloInstruction::CreateBinary(r0f32_, HloOpcode::kAdd, foo, bar)); + HloModule module(TestName()); + module.AddEntryComputation(builder.Build()); + + EXPECT_THAT(add->operands(), UnorderedElementsAre(foo, bar)); + EXPECT_THAT(foo->users(), UnorderedElementsAre(add)); + EXPECT_THAT(bar->users(), UnorderedElementsAre(add)); OpAndUserCollectingVisitor visitor; ASSERT_IS_OK(add->Accept(&visitor)); - EXPECT_EQ(2, visitor.NumOperands(add.get())); - EXPECT_EQ(0, visitor.NumUsers(add.get())); - EXPECT_EQ(1, visitor.NumUsers(foo.get())); - EXPECT_EQ(1, visitor.NumUsers(bar.get())); + EXPECT_EQ(2, visitor.NumOperands(add)); + EXPECT_EQ(0, visitor.NumUsers(add)); + EXPECT_EQ(1, visitor.NumUsers(foo)); + EXPECT_EQ(1, visitor.NumUsers(bar)); } TEST_F(HloInstructionTest, MultipleUsers) { @@ -171,12 +181,19 @@ TEST_F(HloInstructionTest, MultipleUsers) { // ------- ------- ----------- // | exp | | exp | | add | // ------- ------- ----------- - auto foo = HloInstruction::CreateParameter(0, r0f32_, "foo"); - auto bar = HloInstruction::CreateParameter(1, r0f32_, "bar"); - auto exp1 = HloInstruction::CreateUnary(r0f32_, HloOpcode::kExp, foo.get()); - auto exp2 = HloInstruction::CreateUnary(r0f32_, HloOpcode::kExp, foo.get()); - auto add = HloInstruction::CreateBinary(r0f32_, HloOpcode::kAdd, foo.get(), - bar.get()); + HloComputation::Builder builder(TestName()); + auto foo = + builder.AddInstruction(HloInstruction::CreateParameter(0, r0f32_, "foo")); + auto bar = + builder.AddInstruction(HloInstruction::CreateParameter(1, r0f32_, "bar")); + auto exp1 = builder.AddInstruction( + HloInstruction::CreateUnary(r0f32_, HloOpcode::kExp, foo)); + auto exp2 = builder.AddInstruction( + HloInstruction::CreateUnary(r0f32_, HloOpcode::kExp, foo)); + auto add = builder.AddInstruction( + HloInstruction::CreateBinary(r0f32_, HloOpcode::kAdd, foo, bar)); + HloModule module(TestName()); + module.AddEntryComputation(builder.Build()); EXPECT_EQ(3, foo->user_count()); EXPECT_EQ(1, bar->user_count()); @@ -187,8 +204,8 @@ TEST_F(HloInstructionTest, MultipleUsers) { OpAndUserCollectingVisitor visitor; ASSERT_IS_OK(add->Accept(&visitor)); - EXPECT_EQ(2, visitor.NumOperands(add.get())); - EXPECT_EQ(3, visitor.NumUsers(foo.get())); + EXPECT_EQ(2, visitor.NumOperands(add)); + EXPECT_EQ(3, visitor.NumUsers(foo)); } TEST_F(HloInstructionTest, RepeatedUser) { @@ -203,9 +220,14 @@ TEST_F(HloInstructionTest, RepeatedUser) { // ------- // | add | // ------- - auto foo = HloInstruction::CreateParameter(0, r0f32_, "foo"); - auto add = HloInstruction::CreateBinary(r0f32_, HloOpcode::kAdd, foo.get(), - foo.get()); + HloComputation::Builder builder(TestName()); + auto foo = + builder.AddInstruction(HloInstruction::CreateParameter(0, r0f32_, "foo")); + auto add = builder.AddInstruction( + HloInstruction::CreateBinary(r0f32_, HloOpcode::kAdd, foo, foo)); + HloModule module(TestName()); + module.AddEntryComputation(builder.Build()); + EXPECT_EQ(1, foo->user_count()); // But 'add' still has two operands, even if both are the same HLO. @@ -225,23 +247,29 @@ TEST_F(HloInstructionTest, MultipleUsersAndOperands) { // \ ------- / // ---->| add |<---- // ------- - auto param0 = HloInstruction::CreateParameter(0, r0f32_, "param0"); - auto param1 = HloInstruction::CreateParameter(1, r0f32_, "param1"); - auto c0 = HloInstruction::CreateConstant(LiteralUtil::CreateR0(1.1f)); - auto addleft = HloInstruction::CreateBinary(r0f32_, HloOpcode::kAdd, - param0.get(), c0.get()); - auto addright = HloInstruction::CreateBinary(r0f32_, HloOpcode::kAdd, - c0.get(), param1.get()); - auto addtotal = HloInstruction::CreateBinary(r0f32_, HloOpcode::kAdd, - addleft.get(), addright.get()); + HloComputation::Builder builder(TestName()); + auto param0 = builder.AddInstruction( + HloInstruction::CreateParameter(0, r0f32_, "param0")); + auto param1 = builder.AddInstruction( + HloInstruction::CreateParameter(1, r0f32_, "param1")); + auto c0 = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(1.1f))); + auto addleft = builder.AddInstruction( + HloInstruction::CreateBinary(r0f32_, HloOpcode::kAdd, param0, c0)); + auto addright = builder.AddInstruction( + HloInstruction::CreateBinary(r0f32_, HloOpcode::kAdd, c0, param1)); + auto addtotal = builder.AddInstruction( + HloInstruction::CreateBinary(r0f32_, HloOpcode::kAdd, addleft, addright)); + HloModule module(TestName()); + module.AddEntryComputation(builder.Build()); OpAndUserCollectingVisitor visitor; ASSERT_IS_OK(addtotal->Accept(&visitor)); - EXPECT_EQ(2, visitor.NumUsers(c0.get())); - EXPECT_EQ(2, visitor.NumOperands(addleft.get())); - EXPECT_EQ(2, visitor.NumOperands(addright.get())); - EXPECT_EQ(2, visitor.NumOperands(addtotal.get())); + EXPECT_EQ(2, visitor.NumUsers(c0)); + EXPECT_EQ(2, visitor.NumOperands(addleft)); + EXPECT_EQ(2, visitor.NumOperands(addright)); + EXPECT_EQ(2, visitor.NumOperands(addtotal)); } TEST_F(HloInstructionTest, MultipleUsersAndOperandsWithUnaryOps) { @@ -264,29 +292,36 @@ TEST_F(HloInstructionTest, MultipleUsersAndOperandsWithUnaryOps) { // ------- // | neg | // ------- - auto param0 = HloInstruction::CreateParameter(0, r0f32_, "param0"); - auto param1 = HloInstruction::CreateParameter(1, r0f32_, "param1"); - auto c0 = HloInstruction::CreateConstant(LiteralUtil::CreateR0(1.1f)); - auto neg1 = HloInstruction::CreateUnary(r0f32_, HloOpcode::kNegate, c0.get()); - auto addleft = HloInstruction::CreateBinary(r0f32_, HloOpcode::kAdd, - param0.get(), neg1.get()); - auto addright = HloInstruction::CreateBinary(r0f32_, HloOpcode::kAdd, - neg1.get(), param1.get()); - auto addtotal = HloInstruction::CreateBinary(r0f32_, HloOpcode::kAdd, - addleft.get(), addright.get()); - auto neg2 = - HloInstruction::CreateUnary(r0f32_, HloOpcode::kNegate, addtotal.get()); + HloComputation::Builder builder(TestName()); + auto param0 = builder.AddInstruction( + HloInstruction::CreateParameter(0, r0f32_, "param0")); + auto param1 = builder.AddInstruction( + HloInstruction::CreateParameter(1, r0f32_, "param1")); + auto c0 = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(1.1f))); + auto neg1 = builder.AddInstruction( + HloInstruction::CreateUnary(r0f32_, HloOpcode::kNegate, c0)); + auto addleft = builder.AddInstruction( + HloInstruction::CreateBinary(r0f32_, HloOpcode::kAdd, param0, neg1)); + auto addright = builder.AddInstruction( + HloInstruction::CreateBinary(r0f32_, HloOpcode::kAdd, neg1, param1)); + auto addtotal = builder.AddInstruction( + HloInstruction::CreateBinary(r0f32_, HloOpcode::kAdd, addleft, addright)); + auto neg2 = builder.AddInstruction( + HloInstruction::CreateUnary(r0f32_, HloOpcode::kNegate, addtotal)); + HloModule module(TestName()); + module.AddEntryComputation(builder.Build()); OpAndUserCollectingVisitor visitor; ASSERT_IS_OK(neg2->Accept(&visitor)); - EXPECT_EQ(1, visitor.NumUsers(c0.get())); - EXPECT_EQ(2, visitor.NumUsers(neg1.get())); - EXPECT_EQ(2, visitor.NumOperands(addleft.get())); - EXPECT_EQ(2, visitor.NumOperands(addright.get())); - EXPECT_EQ(2, visitor.NumOperands(addtotal.get())); - EXPECT_EQ(1, visitor.NumOperands(neg2.get())); - EXPECT_EQ(0, visitor.NumUsers(neg2.get())); + EXPECT_EQ(1, visitor.NumUsers(c0)); + EXPECT_EQ(2, visitor.NumUsers(neg1)); + EXPECT_EQ(2, visitor.NumOperands(addleft)); + EXPECT_EQ(2, visitor.NumOperands(addright)); + EXPECT_EQ(2, visitor.NumOperands(addtotal)); + EXPECT_EQ(1, visitor.NumOperands(neg2)); + EXPECT_EQ(0, visitor.NumUsers(neg2)); } TEST_F(HloInstructionTest, TrivialMap) { @@ -296,29 +331,33 @@ TEST_F(HloInstructionTest, TrivialMap) { // Shape r0f32 = ShapeUtil::MakeShape(F32, {}); Shape f32a100x10 = ShapeUtil::MakeShape(F32, {100, 10}); + HloModule module(TestName()); // Builds an x+1.0 computation to use in a Map. - auto builder = HloComputation::Builder("f32+1"); - auto param = - builder.AddInstruction(HloInstruction::CreateParameter(0, r0f32, "x")); - auto value = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(1.0))); - builder.AddInstruction( + auto embedded_builder = HloComputation::Builder("f32+1"); + auto param = embedded_builder.AddInstruction( + HloInstruction::CreateParameter(0, r0f32, "x")); + auto value = embedded_builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(1.0))); + embedded_builder.AddInstruction( HloInstruction::CreateBinary(r0f32, HloOpcode::kAdd, param, value)); - auto add_f32 = builder.Build(); + auto add_f32 = module.AddEmbeddedComputation(embedded_builder.Build()); // Builds a parameter and feeds it to the map. - auto param0 = HloInstruction::CreateParameter(1, f32a100x10, ""); - auto map = - HloInstruction::CreateMap(f32a100x10, {param0.get()}, add_f32.get()); + HloComputation::Builder builder(TestName()); + auto param0 = builder.AddInstruction( + HloInstruction::CreateParameter(0, f32a100x10, "")); + auto map = builder.AddInstruction( + HloInstruction::CreateMap(f32a100x10, {param0}, add_f32)); + module.AddEntryComputation(builder.Build()); OpAndUserCollectingVisitor visitor; ASSERT_IS_OK(map->Accept(&visitor)); // Check counts. We aren't walking the mapper computation yet. - EXPECT_EQ(1, visitor.NumUsers(param0.get())); - EXPECT_EQ(0, visitor.NumUsers(map.get())); - EXPECT_EQ(1, visitor.NumOperands(map.get())); + EXPECT_EQ(1, visitor.NumUsers(param0)); + EXPECT_EQ(0, visitor.NumUsers(map)); + EXPECT_EQ(1, visitor.NumOperands(map)); // TODO(dehnert): Add walking and counters for the wrapped computation. } @@ -333,164 +372,197 @@ TEST_F(HloInstructionTest, TrivialReduce) { Shape f32a100x10 = ShapeUtil::MakeShape(F32, {100, 10}); // Builds an x+y computation to use in a Reduce. - auto builder = HloComputation::Builder("f32+f32"); - auto paramx = - builder.AddInstruction(HloInstruction::CreateParameter(0, r0f32, "x")); - auto paramy = - builder.AddInstruction(HloInstruction::CreateParameter(1, r0f32, "y")); - builder.AddInstruction( + auto embedded_builder = HloComputation::Builder("f32+f32"); + auto paramx = embedded_builder.AddInstruction( + HloInstruction::CreateParameter(0, r0f32, "x")); + auto paramy = embedded_builder.AddInstruction( + HloInstruction::CreateParameter(1, r0f32, "y")); + embedded_builder.AddInstruction( HloInstruction::CreateBinary(r0f32, HloOpcode::kAdd, paramx, paramy)); - auto add_f32 = builder.Build(); + HloModule module(TestName()); + auto add_f32 = module.AddEmbeddedComputation(embedded_builder.Build()); // Builds a parameter and an initial value and feeds them to the reduce. - auto param0 = HloInstruction::CreateParameter(0, f32a100x10, ""); - auto const0 = - HloInstruction::CreateConstant(LiteralUtil::CreateR0(0.0f)); - auto c0 = HloInstruction::CreateConstant(LiteralUtil::CreateR0(1.1f)); - auto reduce = - HloInstruction::CreateReduce(f32v100, param0.get(), const0.get(), - /*dimensions_to_reduce=*/{1}, add_f32.get()); + HloComputation::Builder builder(TestName()); + auto param0 = builder.AddInstruction( + HloInstruction::CreateParameter(0, f32a100x10, "")); + auto const0 = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(0.0f))); + builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(1.1f))); + auto reduce = builder.AddInstruction( + HloInstruction::CreateReduce(f32v100, param0, const0, + /*dimensions_to_reduce=*/{1}, add_f32)); + module.AddEntryComputation(builder.Build()); OpAndUserCollectingVisitor visitor; ASSERT_IS_OK(reduce->Accept(&visitor)); // Check counts. We aren't walking the reducer computation. - EXPECT_EQ(1, visitor.NumUsers(param0.get())); - EXPECT_EQ(1, visitor.NumUsers(const0.get())); - EXPECT_EQ(0, visitor.NumUsers(reduce.get())); - EXPECT_EQ(2, visitor.NumOperands(reduce.get())); + EXPECT_EQ(1, visitor.NumUsers(param0)); + EXPECT_EQ(1, visitor.NumUsers(const0)); + EXPECT_EQ(0, visitor.NumUsers(reduce)); + EXPECT_EQ(2, visitor.NumOperands(reduce)); } TEST_F(HloInstructionTest, ReplaceUseInBinaryOps) { // Construct a graph of a few binary ops using two different // parameters. Replace one of the parameters with the other parameter in one // of the instructions. - auto foo = HloInstruction::CreateParameter(0, r0f32_, "foo"); - auto bar = HloInstruction::CreateParameter(1, r0f32_, "bar"); - auto add_foobar = HloInstruction::CreateBinary(r0f32_, HloOpcode::kAdd, - foo.get(), bar.get()); - auto add_foofoo = HloInstruction::CreateBinary(r0f32_, HloOpcode::kAdd, - foo.get(), foo.get()); - auto sum = HloInstruction::CreateBinary(r0f32_, HloOpcode::kAdd, - add_foobar.get(), add_foofoo.get()); + HloComputation::Builder builder(TestName()); + auto foo = + builder.AddInstruction(HloInstruction::CreateParameter(0, r0f32_, "foo")); + auto bar = + builder.AddInstruction(HloInstruction::CreateParameter(1, r0f32_, "bar")); + auto add_foobar = builder.AddInstruction( + HloInstruction::CreateBinary(r0f32_, HloOpcode::kAdd, foo, bar)); + auto add_foofoo = builder.AddInstruction( + HloInstruction::CreateBinary(r0f32_, HloOpcode::kAdd, foo, foo)); + builder.AddInstruction(HloInstruction::CreateBinary(r0f32_, HloOpcode::kAdd, + add_foobar, add_foofoo)); + HloModule module(TestName()); + module.AddEntryComputation(builder.Build()); EXPECT_EQ(2, foo->user_count()); EXPECT_EQ(1, bar->user_count()); // Replace the use of foo in add_foofoo with bar. - ASSERT_IS_OK(foo->ReplaceUseWith(add_foofoo.get(), bar.get())); + ASSERT_IS_OK(foo->ReplaceUseWith(add_foofoo, bar)); EXPECT_EQ(1, foo->user_count()); EXPECT_EQ(2, bar->user_count()); - ExpectEqUnordered(foo->users(), {add_foobar.get()}); - ExpectEqOrdered(add_foobar->operands(), {foo.get(), bar.get()}); + EXPECT_THAT(foo->users(), UnorderedElementsAre(add_foobar)); + EXPECT_THAT(add_foobar->operands(), ElementsAre(foo, bar)); - ExpectEqUnordered(bar->users(), {add_foobar.get(), add_foofoo.get()}); - ExpectEqOrdered(add_foobar->operands(), {foo.get(), bar.get()}); - ExpectEqOrdered(add_foofoo->operands(), {bar.get(), bar.get()}); + EXPECT_THAT(bar->users(), UnorderedElementsAre(add_foobar, add_foofoo)); + EXPECT_THAT(add_foobar->operands(), ElementsAre(foo, bar)); + EXPECT_THAT(add_foofoo->operands(), ElementsAre(bar, bar)); } TEST_F(HloInstructionTest, ReplaceUseInVariadicOp) { // Construct a tuple containing several parameters. Replace one parameter with // another in the tuple. - auto foo = HloInstruction::CreateParameter(0, r0f32_, "foo"); - auto bar = HloInstruction::CreateParameter(1, r0f32_, "bar"); - auto baz = HloInstruction::CreateParameter(2, r0f32_, "baz"); + HloComputation::Builder builder(TestName()); + auto foo = + builder.AddInstruction(HloInstruction::CreateParameter(0, r0f32_, "foo")); + auto bar = + builder.AddInstruction(HloInstruction::CreateParameter(1, r0f32_, "bar")); + auto baz = + builder.AddInstruction(HloInstruction::CreateParameter(2, r0f32_, "baz")); auto tuple = - HloInstruction::CreateTuple({foo.get(), bar.get(), baz.get(), foo.get()}); - auto add_foobar = HloInstruction::CreateBinary(r0f32_, HloOpcode::kAdd, - foo.get(), bar.get()); + builder.AddInstruction(HloInstruction::CreateTuple({foo, bar, baz, foo})); + auto add_foobar = builder.AddInstruction( + HloInstruction::CreateBinary(r0f32_, HloOpcode::kAdd, foo, bar)); + HloModule module(TestName()); + module.AddEntryComputation(builder.Build()); EXPECT_EQ(2, foo->user_count()); - ExpectEqUnordered(foo->users(), {tuple.get(), add_foobar.get()}); + EXPECT_THAT(foo->users(), UnorderedElementsAre(tuple, add_foobar)); // Replace the use of foo in tuple with bar. - ASSERT_IS_OK(foo->ReplaceUseWith(tuple.get(), bar.get())); + ASSERT_IS_OK(foo->ReplaceUseWith(tuple, bar)); - ExpectEqUnordered(foo->users(), {add_foobar.get()}); + EXPECT_THAT(foo->users(), UnorderedElementsAre(add_foobar)); // Both uses of foo in tuple should have been replaced with bar. - ExpectEqOrdered(tuple->operands(), - {bar.get(), bar.get(), baz.get(), bar.get()}); + EXPECT_THAT(tuple->operands(), ElementsAre(bar, bar, baz, bar)); } TEST_F(HloInstructionTest, ReplaceUseInUnaryOp) { // Construct a couple unary instructions which use a parameter. Replace the // use of a parameter in one of the unary ops with the other parameter. - auto foo = HloInstruction::CreateParameter(0, r0f32_, "foo"); - auto bar = HloInstruction::CreateParameter(1, r0f32_, "bar"); - - auto exp = HloInstruction::CreateUnary(r0f32_, HloOpcode::kExp, foo.get()); - auto log = HloInstruction::CreateUnary(r0f32_, HloOpcode::kLog, foo.get()); + HloComputation::Builder builder(TestName()); + auto foo = + builder.AddInstruction(HloInstruction::CreateParameter(0, r0f32_, "foo")); + auto bar = + builder.AddInstruction(HloInstruction::CreateParameter(1, r0f32_, "bar")); + + auto exp = builder.AddInstruction( + HloInstruction::CreateUnary(r0f32_, HloOpcode::kExp, foo)); + auto log = builder.AddInstruction( + HloInstruction::CreateUnary(r0f32_, HloOpcode::kLog, foo)); + HloModule module(TestName()); + module.AddEntryComputation(builder.Build()); EXPECT_EQ(2, foo->user_count()); - ExpectEqUnordered(foo->users(), {exp.get(), log.get()}); + EXPECT_THAT(foo->users(), UnorderedElementsAre(exp, log)); EXPECT_EQ(0, bar->user_count()); // Replace the use of foo in exp with bar. - ASSERT_IS_OK(foo->ReplaceUseWith(exp.get(), bar.get())); + ASSERT_IS_OK(foo->ReplaceUseWith(exp, bar)); // The use of foo in log should not have been affected. EXPECT_EQ(1, foo->user_count()); - ExpectEqUnordered(foo->users(), {log.get()}); - ExpectEqOrdered(log->operands(), {foo.get()}); + EXPECT_THAT(foo->users(), UnorderedElementsAre(log)); + EXPECT_THAT(log->operands(), ElementsAre(foo)); // Bar should now be used in exp. EXPECT_EQ(1, bar->user_count()); - EXPECT_EQ(*bar->users().begin(), exp.get()); + EXPECT_EQ(*bar->users().begin(), exp); EXPECT_EQ(1, exp->operands().size()); - EXPECT_EQ(*exp->operands().begin(), bar.get()); + EXPECT_EQ(*exp->operands().begin(), bar); } TEST_F(HloInstructionTest, ReplaceAllUsesWithInBinaryOps) { // Construct a simple graph of a few binary ops using two different // parameters. Replace all uses of one of the parameters with the other // parameter. - auto foo = HloInstruction::CreateParameter(0, r0f32_, "foo"); - auto bar = HloInstruction::CreateParameter(1, r0f32_, "bar"); - auto add_foobar = HloInstruction::CreateBinary(r0f32_, HloOpcode::kAdd, - foo.get(), bar.get()); - auto add_foofoo = HloInstruction::CreateBinary(r0f32_, HloOpcode::kAdd, - foo.get(), foo.get()); - auto sum = HloInstruction::CreateBinary(r0f32_, HloOpcode::kAdd, - add_foobar.get(), add_foofoo.get()); + HloComputation::Builder builder(TestName()); + auto foo = + builder.AddInstruction(HloInstruction::CreateParameter(0, r0f32_, "foo")); + auto bar = + builder.AddInstruction(HloInstruction::CreateParameter(1, r0f32_, "bar")); + auto add_foobar = builder.AddInstruction( + HloInstruction::CreateBinary(r0f32_, HloOpcode::kAdd, foo, bar)); + auto add_foofoo = builder.AddInstruction( + HloInstruction::CreateBinary(r0f32_, HloOpcode::kAdd, foo, foo)); + builder.AddInstruction(HloInstruction::CreateBinary(r0f32_, HloOpcode::kAdd, + add_foobar, add_foofoo)); + HloModule module(TestName()); + module.AddEntryComputation(builder.Build()); EXPECT_EQ(2, foo->user_count()); EXPECT_EQ(1, bar->user_count()); // Replace all uses of foo with bar. - ASSERT_IS_OK(foo->ReplaceAllUsesWith(bar.get())); + ASSERT_IS_OK(foo->ReplaceAllUsesWith(bar)); EXPECT_EQ(0, foo->user_count()); EXPECT_EQ(2, bar->user_count()); - ExpectEqUnordered(bar->users(), {add_foobar.get(), add_foofoo.get()}); + EXPECT_THAT(bar->users(), UnorderedElementsAre(add_foobar, add_foofoo)); } TEST_F(HloInstructionTest, ReplaceAllUsesInMultipleOps) { // Construct a graph containing several ops (a unary, binary, and variadic) // which use two parameters. Replace all uses of one of the parameters with // the other parameter. - auto foo = HloInstruction::CreateParameter(0, r0f32_, "foo"); - auto bar = HloInstruction::CreateParameter(1, r0f32_, "bar"); - - auto add_foobar = HloInstruction::CreateBinary(r0f32_, HloOpcode::kAdd, - foo.get(), bar.get()); - auto exp = HloInstruction::CreateUnary(r0f32_, HloOpcode::kExp, foo.get()); - auto tuple = HloInstruction::CreateTuple({foo.get(), bar.get()}); + HloComputation::Builder builder(TestName()); + auto foo = + builder.AddInstruction(HloInstruction::CreateParameter(0, r0f32_, "foo")); + auto bar = + builder.AddInstruction(HloInstruction::CreateParameter(1, r0f32_, "bar")); + + auto add_foobar = builder.AddInstruction( + HloInstruction::CreateBinary(r0f32_, HloOpcode::kAdd, foo, bar)); + auto exp = builder.AddInstruction( + HloInstruction::CreateUnary(r0f32_, HloOpcode::kExp, foo)); + auto tuple = builder.AddInstruction(HloInstruction::CreateTuple({foo, bar})); + HloModule module(TestName()); + module.AddEntryComputation(builder.Build()); EXPECT_EQ(3, foo->user_count()); EXPECT_EQ(2, bar->user_count()); // Replace all uses of foo with bar. - ASSERT_IS_OK(foo->ReplaceAllUsesWith(bar.get())); + ASSERT_IS_OK(foo->ReplaceAllUsesWith(bar)); EXPECT_EQ(0, foo->user_count()); EXPECT_EQ(3, bar->user_count()); - ExpectEqUnordered(bar->users(), {add_foobar.get(), exp.get(), tuple.get()}); + EXPECT_THAT(bar->users(), UnorderedElementsAre(add_foobar, exp, tuple)); } // Simple visitor that collects and post-processes each node in the graph. @@ -534,11 +606,17 @@ TEST_F(HloInstructionTest, PostProcessAllVisitedNodes) { // /--> exp --\ // foo add // \--> log --/ - auto foo = HloInstruction::CreateParameter(0, r0f32_, "foo"); - auto exp = HloInstruction::CreateUnary(r0f32_, HloOpcode::kExp, foo.get()); - auto log = HloInstruction::CreateUnary(r0f32_, HloOpcode::kLog, foo.get()); - auto add = HloInstruction::CreateBinary(r0f32_, HloOpcode::kAdd, exp.get(), - log.get()); + HloComputation::Builder builder(TestName()); + auto foo = + builder.AddInstruction(HloInstruction::CreateParameter(0, r0f32_, "foo")); + auto exp = builder.AddInstruction( + HloInstruction::CreateUnary(r0f32_, HloOpcode::kExp, foo)); + auto log = builder.AddInstruction( + HloInstruction::CreateUnary(r0f32_, HloOpcode::kLog, foo)); + auto add = builder.AddInstruction( + HloInstruction::CreateBinary(r0f32_, HloOpcode::kAdd, exp, log)); + HloModule module(TestName()); + module.AddEntryComputation(builder.Build()); NodeCollectorAndPostProcessor visitor; ASSERT_IS_OK(add->Accept(&visitor)); @@ -549,95 +627,148 @@ TEST_F(HloInstructionTest, PostProcessAllVisitedNodes) { } TEST_F(HloInstructionTest, SingletonFusionOp) { + HloComputation::Builder builder(TestName()); // Create a fusion instruction containing a single unary operation. - auto constant = - HloInstruction::CreateConstant(LiteralUtil::CreateR0(1.1f)); - auto exp = - HloInstruction::CreateUnary(r0f32_, HloOpcode::kExp, constant.get()); - - auto fusion = HloInstruction::CreateFusion( - r0f32_, HloInstruction::FusionKind::kLoop, exp.get()); - - ExpectEqOrdered(fusion->operands(), {constant.get()}); - ExpectEqUnordered(constant->users(), {fusion.get(), exp.get()}); + auto constant = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(1.1f))); + auto exp = builder.AddInstruction( + HloInstruction::CreateUnary(r0f32_, HloOpcode::kExp, constant)); + HloModule module(TestName()); + auto* computation = module.AddEntryComputation(builder.Build()); + auto* fusion = computation->CreateFusionInstruction( + {exp}, HloInstruction::FusionKind::kLoop); + + EXPECT_THAT(fusion->operands(), ElementsAre(constant)); + EXPECT_THAT(constant->users(), ElementsAre(fusion)); } TEST_F(HloInstructionTest, BinaryFusionOp) { + HloComputation::Builder builder(TestName()); // Create a fusion instruction containing a single binary operation. - auto constant1 = - HloInstruction::CreateConstant(LiteralUtil::CreateR0(1.1f)); - auto constant2 = - HloInstruction::CreateConstant(LiteralUtil::CreateR0(42.1f)); - auto add = HloInstruction::CreateBinary(r0f32_, HloOpcode::kAdd, - constant1.get(), constant2.get()); - - auto fusion = HloInstruction::CreateFusion( - r0f32_, HloInstruction::FusionKind::kLoop, add.get()); - - ExpectEqOrdered(fusion->operands(), {constant1.get(), constant2.get()}); - ExpectEqUnordered(constant1->users(), {fusion.get(), add.get()}); - ExpectEqUnordered(constant2->users(), {fusion.get(), add.get()}); + auto constant1 = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(1.1f))); + auto constant2 = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(42.1f))); + auto add = builder.AddInstruction(HloInstruction::CreateBinary( + r0f32_, HloOpcode::kAdd, constant1, constant2)); + HloModule module(TestName()); + auto* computation = module.AddEntryComputation(builder.Build()); + auto* fusion = computation->CreateFusionInstruction( + {add}, HloInstruction::FusionKind::kLoop); + + EXPECT_THAT(fusion->operands(), ElementsAre(constant1, constant2)); + EXPECT_THAT(constant1->users(), ElementsAre(fusion)); + EXPECT_THAT(constant2->users(), ElementsAre(fusion)); } TEST_F(HloInstructionTest, ChainFusionOp) { + HloComputation::Builder builder(TestName()); // Create a chain of fused unary ops. - auto constant = - HloInstruction::CreateConstant(LiteralUtil::CreateR0(1.1f)); - auto exp1 = - HloInstruction::CreateUnary(r0f32_, HloOpcode::kExp, constant.get()); - auto exp2 = HloInstruction::CreateUnary(r0f32_, HloOpcode::kExp, exp1.get()); - auto exp3 = HloInstruction::CreateUnary(r0f32_, HloOpcode::kExp, exp2.get()); + auto constant = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(1.1f))); + auto exp1 = builder.AddInstruction( + HloInstruction::CreateUnary(r0f32_, HloOpcode::kExp, constant)); + auto exp2 = builder.AddInstruction( + HloInstruction::CreateUnary(r0f32_, HloOpcode::kExp, exp1)); + auto exp3 = builder.AddInstruction( + HloInstruction::CreateUnary(r0f32_, HloOpcode::kExp, exp2)); + + HloModule module(TestName()); + auto* computation = module.AddEntryComputation(builder.Build()); + auto* fusion = computation->CreateFusionInstruction( + {exp3, exp2, exp1}, HloInstruction::FusionKind::kLoop); + + EXPECT_THAT(fusion->operands(), ElementsAre(constant)); + EXPECT_THAT(constant->users(), ElementsAre(fusion)); +} - auto fusion = HloInstruction::CreateFusion( - r0f32_, HloInstruction::FusionKind::kLoop, exp3.get()); - fusion->FuseInstruction(exp2.get()); - fusion->FuseInstruction(exp1.get()); +TEST_F(HloInstructionTest, PreserveMetadataInFusionAndClone) { + HloComputation::Builder builder(TestName()); + // Create a chain of fused unary ops. + auto constant = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(1.1f))); + auto exp1 = builder.AddInstruction( + HloInstruction::CreateUnary(r0f32_, HloOpcode::kExp, constant)); + auto exp2 = builder.AddInstruction( + HloInstruction::CreateUnary(r0f32_, HloOpcode::kExp, exp1)); + OpMetadata metadata; + metadata.set_op_name("tf_op"); + exp1->set_metadata(metadata); + exp2->set_metadata(metadata); + + HloModule module(TestName()); + auto* computation = module.AddEntryComputation(builder.Build()); + auto* fusion = computation->CreateFusionInstruction( + {exp2, exp1}, HloInstruction::FusionKind::kLoop); + + EXPECT_TRUE(protobuf_util::ProtobufEquals(metadata, fusion->metadata())); + EXPECT_TRUE(protobuf_util::ProtobufEquals( + metadata, fusion->fused_expression_root()->metadata())); + EXPECT_TRUE(protobuf_util::ProtobufEquals( + metadata, fusion->fused_expression_root()->operand(0)->metadata())); +} - ExpectEqOrdered(fusion->operands(), {constant.get()}); - ExpectEqUnordered(constant->users(), {fusion.get(), exp1.get()}); +TEST_F(HloInstructionTest, PreserveOutfeedShapeThroughClone) { + HloComputation::Builder builder(TestName()); + auto constant = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR2({ + {1, 2}, + {3, 4}, + }))); + auto shape10 = ShapeUtil::MakeShapeWithLayout(F32, {2, 3}, {1, 0}); + auto shape01 = ShapeUtil::MakeShapeWithLayout(F32, {2, 3}, {0, 1}); + auto outfeed10 = builder.AddInstruction( + HloInstruction::CreateOutfeed(shape10, constant, "")); + auto outfeed01 = builder.AddInstruction( + HloInstruction::CreateOutfeed(shape01, constant, "")); + + auto clone01 = builder.AddInstruction(outfeed01->Clone()); + auto clone10 = builder.AddInstruction(outfeed10->Clone()); + + EXPECT_TRUE(ShapeUtil::Equal(clone01->outfeed_shape(), shape01)); + EXPECT_TRUE(ShapeUtil::Equal(clone10->outfeed_shape(), shape10)); } TEST_F(HloInstructionTest, FusionOpWithCalledComputations) { // Create a fusion instruction containing a single unary operation. const Shape scalar_shape = ShapeUtil::MakeShape(F32, {}); + HloModule module(TestName()); auto make_map_computation = [&]() { auto builder = HloComputation::Builder("FusionMap"); builder.AddInstruction( HloInstruction::CreateParameter(0, scalar_shape, "param")); - return builder.Build(); + return module.AddEmbeddedComputation(builder.Build()); }; - std::unique_ptr computation_x = make_map_computation(); - std::unique_ptr computation_y = make_map_computation(); - - auto constant = - HloInstruction::CreateConstant(LiteralUtil::CreateR0(1.1f)); - auto map_1_x = - HloInstruction::CreateMap(scalar_shape, {constant.get()}, - computation_x.get(), /*static_operands=*/{}); - auto map_2_x = - HloInstruction::CreateMap(scalar_shape, {map_1_x.get()}, - computation_x.get(), /*static_operands=*/{}); - auto map_3_y = - HloInstruction::CreateMap(scalar_shape, {map_2_x.get()}, - computation_y.get(), /*static_operands=*/{}); - - auto fusion = HloInstruction::CreateFusion( - scalar_shape, HloInstruction::FusionKind::kLoop, map_3_y.get()); - - ASSERT_EQ(fusion->called_computations().size(), 1); - EXPECT_EQ(fusion->called_computations()[0], computation_y.get()); - - fusion->FuseInstruction(map_2_x.get()); - ASSERT_EQ(fusion->called_computations().size(), 2); - EXPECT_EQ(fusion->called_computations()[1], computation_x.get()); - - fusion->FuseInstruction(map_1_x.get()); - ASSERT_EQ(fusion->called_computations().size(), 2); + HloComputation* computation_x = make_map_computation(); + HloComputation* computation_y = make_map_computation(); + + HloComputation::Builder builder(TestName()); + auto constant = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(1.1f))); + auto map_1_x = builder.AddInstruction(HloInstruction::CreateMap( + scalar_shape, {constant}, computation_x, /*static_operands=*/{})); + auto map_2_x = builder.AddInstruction(HloInstruction::CreateMap( + scalar_shape, {map_1_x}, computation_x, /*static_operands=*/{})); + auto map_3_y = builder.AddInstruction(HloInstruction::CreateMap( + scalar_shape, {map_2_x}, computation_y, /*static_operands=*/{})); + auto* computation = module.AddEntryComputation(builder.Build()); + + auto* fusion = computation->CreateFusionInstruction( + {map_3_y}, HloInstruction::FusionKind::kLoop); + auto* fused_computation = fusion->fused_instructions_computation(); + EXPECT_THAT(fusion->called_computations(), ElementsAre(fused_computation)); + + fusion->FuseInstruction(map_2_x); + EXPECT_THAT(fusion->called_computations(), ElementsAre(fused_computation)); + + fusion->FuseInstruction(map_1_x); + EXPECT_THAT(fusion->called_computations(), ElementsAre(fused_computation)); } TEST_F(HloInstructionTest, ComplexFusionOp) { + HloComputation::Builder builder(TestName()); // Fuse all instructions in complicated expression: // // add = Add(C1, C2) @@ -649,34 +780,35 @@ TEST_F(HloInstructionTest, ComplexFusionOp) { // // Notable complexities are repeated operands in a same instruction, different // shapes, use of value in different expressions. - auto c1 = HloInstruction::CreateConstant(LiteralUtil::CreateR0(1.1f)); - auto c2 = HloInstruction::CreateConstant(LiteralUtil::CreateR0(2.1f)); - auto c3 = HloInstruction::CreateConstant(LiteralUtil::CreateR0(9.0f)); - - auto add = - HloInstruction::CreateBinary(r0f32_, HloOpcode::kAdd, c1.get(), c2.get()); - auto clamp = HloInstruction::CreateTernary(r0f32_, HloOpcode::kClamp, - c2.get(), add.get(), add.get()); - auto exp = HloInstruction::CreateUnary(r0f32_, HloOpcode::kExp, add.get()); - auto mul = HloInstruction::CreateBinary(r0f32_, HloOpcode::kMultiply, - exp.get(), c3.get()); - auto sub = HloInstruction::CreateBinary(r0f32_, HloOpcode::kSubtract, - mul.get(), clamp.get()); + auto c1 = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(1.1f))); + auto c2 = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(2.1f))); + auto c3 = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(9.0f))); + + auto add = builder.AddInstruction( + HloInstruction::CreateBinary(r0f32_, HloOpcode::kAdd, c1, c2)); + auto clamp = builder.AddInstruction( + HloInstruction::CreateTernary(r0f32_, HloOpcode::kClamp, c2, add, add)); + auto exp = builder.AddInstruction( + HloInstruction::CreateUnary(r0f32_, HloOpcode::kExp, add)); + auto mul = builder.AddInstruction( + HloInstruction::CreateBinary(r0f32_, HloOpcode::kMultiply, exp, c3)); + auto sub = builder.AddInstruction( + HloInstruction::CreateBinary(r0f32_, HloOpcode::kSubtract, mul, clamp)); auto tuple = - HloInstruction::CreateTuple({sub.get(), sub.get(), mul.get(), c1.get()}); + builder.AddInstruction(HloInstruction::CreateTuple({sub, sub, mul, c1})); - auto fusion = HloInstruction::CreateFusion( - r0f32_, HloInstruction::FusionKind::kLoop, tuple.get()); - fusion->FuseInstruction(sub.get()); - fusion->FuseInstruction(mul.get()); - fusion->FuseInstruction(exp.get()); - fusion->FuseInstruction(clamp.get()); - fusion->FuseInstruction(add.get()); + HloModule module(TestName()); + auto* computation = module.AddEntryComputation(builder.Build()); + auto* fusion = computation->CreateFusionInstruction( + {tuple, sub, mul, exp, clamp, add}, HloInstruction::FusionKind::kLoop); // Operands in the fusion instruction's operands() vector should be in the // order in which their users were added fused. - ExpectEqOrdered(fusion->operands(), {c1.get(), c3.get(), c2.get()}); - ExpectEqUnordered(c1->users(), {add.get(), tuple.get(), fusion.get()}); + EXPECT_THAT(fusion->operands(), ElementsAre(c1, c3, c2)); + EXPECT_THAT(c1->users(), ElementsAre(fusion)); } // Convenience function for comparing two HloInstructions inside of @@ -699,11 +831,11 @@ TEST_F(HloInstructionTest, IdenticalInstructions) { // Create a set of random constant operands to use below. Make them matrices // so dimensions are interesting. auto operand1 = HloInstruction::CreateConstant( - LiteralUtil::CreateR2({{1.0, 2.0}, {3.0, 4.0}})); + Literal::CreateR2({{1.0, 2.0}, {3.0, 4.0}})); auto operand2 = HloInstruction::CreateConstant( - LiteralUtil::CreateR2({{10.0, 20.0}, {30.0, 40.0}})); - auto vector_operand = HloInstruction::CreateConstant( - LiteralUtil::CreateR1({42.0, 123.0})); + Literal::CreateR2({{10.0, 20.0}, {30.0, 40.0}})); + auto vector_operand = + HloInstruction::CreateConstant(Literal::CreateR1({42.0, 123.0})); Shape shape = operand1->shape(); // Convenient short names for the operands. @@ -760,12 +892,17 @@ TEST_F(HloInstructionTest, FunctionVisitor) { // \ / // add const Shape f32 = ShapeUtil::MakeShape(F32, {}); - auto param = HloInstruction::CreateParameter(0, f32, "0"); - auto negate = - HloInstruction::CreateUnary(f32, HloOpcode::kNegate, param.get()); - auto exp = HloInstruction::CreateUnary(f32, HloOpcode::kExp, param.get()); - auto add = HloInstruction::CreateBinary(f32, HloOpcode::kAdd, negate.get(), - exp.get()); + HloComputation::Builder builder(TestName()); + auto param = + builder.AddInstruction(HloInstruction::CreateParameter(0, f32, "0")); + auto negate = builder.AddInstruction( + HloInstruction::CreateUnary(f32, HloOpcode::kNegate, param)); + auto exp = builder.AddInstruction( + HloInstruction::CreateUnary(f32, HloOpcode::kExp, param)); + auto add = builder.AddInstruction( + HloInstruction::CreateBinary(f32, HloOpcode::kAdd, negate, exp)); + HloModule module(TestName()); + module.AddEntryComputation(builder.Build()); int visit_num = 0; std::unordered_map visit_order; @@ -776,21 +913,26 @@ TEST_F(HloInstructionTest, FunctionVisitor) { return Status::OK(); })); - EXPECT_EQ(0, visit_order.at(param.get())); + EXPECT_EQ(0, visit_order.at(param)); // negate and exp can be visited in an arbitrary order. - EXPECT_TRUE(visit_order.at(exp.get()) == 1 || visit_order.at(exp.get()) == 2); - EXPECT_TRUE(visit_order.at(negate.get()) == 1 || - visit_order.at(negate.get()) == 2); - EXPECT_NE(visit_order.at(exp.get()), visit_order.at(negate.get())); - EXPECT_EQ(3, visit_order.at(add.get())); + EXPECT_TRUE(visit_order.at(exp) == 1 || visit_order.at(exp) == 2); + EXPECT_TRUE(visit_order.at(negate) == 1 || visit_order.at(negate) == 2); + EXPECT_NE(visit_order.at(exp), visit_order.at(negate)); + EXPECT_EQ(3, visit_order.at(add)); } TEST_F(HloInstructionTest, FullyElementwise) { const Shape r1f32 = ShapeUtil::MakeShape(F32, {5}); - auto x = HloInstruction::CreateParameter(0, r1f32, "x"); - auto y = HloInstruction::CreateParameter(1, r1f32, "y"); - auto add = - HloInstruction::CreateBinary(r1f32, HloOpcode::kAdd, x.get(), y.get()); + HloComputation::Builder builder(TestName()); + auto x = + builder.AddInstruction(HloInstruction::CreateParameter(0, r1f32, "x")); + auto y = + builder.AddInstruction(HloInstruction::CreateParameter(1, r1f32, "y")); + auto add = builder.AddInstruction( + HloInstruction::CreateBinary(r1f32, HloOpcode::kAdd, x, y)); + HloModule module(TestName()); + module.AddEntryComputation(builder.Build()); + EXPECT_TRUE(add->IsElementwise()); for (int i = 0; i < add->operand_count(); ++i) { EXPECT_TRUE(add->IsElementwiseOnOperand(i)); @@ -832,7 +974,8 @@ TEST_F(HloInstructionTest, PartiallyElementwise) { HloInstruction* max = builder.AddInstruction( HloInstruction::CreateBinary(r2f32, HloOpcode::kMaximum, div, broadcast)); - auto computation = builder.Build(); + HloModule module(TestName()); + auto* computation = module.AddEntryComputation(builder.Build()); HloInstruction* fusion = computation->CreateFusionInstruction( {max, broadcast, div, mul}, HloInstruction::FusionKind::kLoop); EXPECT_FALSE(fusion->IsElementwise()); @@ -874,7 +1017,8 @@ TEST_F(HloInstructionTest, PartiallyElementwiseWithReuse) { HloInstruction* sub = builder.AddInstruction(HloInstruction::CreateBinary( r1f32, HloOpcode::kSubtract, min, broadcast)); - auto computation = builder.Build(); + HloModule module(TestName()); + auto* computation = module.AddEntryComputation(builder.Build()); HloInstruction* fusion = computation->CreateFusionInstruction( {sub, broadcast, min}, HloInstruction::FusionKind::kLoop); EXPECT_FALSE(fusion->IsElementwise()); @@ -913,7 +1057,8 @@ TEST_F(HloInstructionTest, CloneOfFusionPreservesShape) { HloInstruction* dot = builder.AddInstruction( HloInstruction::CreateBinary(sout, HloOpcode::kDot, x, reshape)); - auto computation = builder.Build(); + HloModule module(TestName()); + auto* computation = module.AddEntryComputation(builder.Build()); HloInstruction* fusion = computation->CreateFusionInstruction( {dot, reshape}, HloInstruction::FusionKind::kTransposeDot); @@ -929,5 +1074,127 @@ TEST_F(HloInstructionTest, CloneOfFusionPreservesShape) { root2->operand(1)->operand(0)->shape())); } +TEST_F(HloInstructionTest, IsRandomFusable) { + auto shape = ShapeUtil::MakeShape(F32, {2, 2}); + { + auto builder = HloComputation::Builder(TestName()); + auto hlo_module = CreateNewModule(); + auto const0 = builder.AddInstruction(HloInstruction::CreateConstant( + Literal::CreateR0(0.0))); + auto const1 = builder.AddInstruction(HloInstruction::CreateConstant( + Literal::CreateR0(1.0))); + auto rng = builder.AddInstruction(HloInstruction::CreateRng( + shape, RandomDistribution::RNG_NORMAL, {const0, const1})); + + auto* computation = hlo_module->AddEntryComputation(builder.Build()); + computation->CreateFusionInstruction({rng, const0, const1}, + HloInstruction::FusionKind::kLoop); + + auto* root = computation->root_instruction(); + + EXPECT_EQ(HloOpcode::kFusion, root->opcode()); + } + { + auto builder = HloComputation::Builder(TestName()); + auto hlo_module = CreateNewModule(); + auto const0 = builder.AddInstruction(HloInstruction::CreateConstant( + Literal::CreateR0(0.0))); + auto const1 = builder.AddInstruction(HloInstruction::CreateConstant( + Literal::CreateR0(1.0))); + auto rng = builder.AddInstruction(HloInstruction::CreateRng( + shape, RandomDistribution::RNG_NORMAL, {const0, const1})); + builder.AddInstruction(HloInstruction::CreateUnary( + shape, HloOpcode::kNegate, rng)); + auto* computation = hlo_module->AddEntryComputation(builder.Build()); + computation->CreateFusionInstruction({rng, const0, const1}, + HloInstruction::FusionKind::kLoop); + + auto* root = computation->root_instruction(); + + EXPECT_EQ(HloOpcode::kFusion, root->operand(0)->opcode()); + } +} + + +TEST_F(HloInstructionTest, CloneSuffixNames) { + // Test that the suffix string added to cloned instructions is not + // duplicated. Rather a numeric incrementing value should be appended. That + // is, we want "foo.clone2", not "foo.clone.clone". + + // Test cloning the same instruction multiple times. + auto foo = + HloInstruction::CreateParameter(0, ShapeUtil::MakeShape(F32, {}), "foo"); + EXPECT_EQ(foo->Clone()->name(), "%foo.clone"); + EXPECT_EQ(foo->Clone()->Clone()->name(), "%foo.clone2"); + EXPECT_EQ(foo->Clone()->Clone()->Clone()->name(), "%foo.clone3"); + + // Test custom suffixes. + EXPECT_EQ(foo->Clone("bar")->name(), "%foo.bar"); + EXPECT_EQ(foo->Clone("bar")->Clone("bar")->name(), "%foo.bar2"); + EXPECT_EQ(foo->Clone("bar")->Clone("bar")->Clone()->name(), + "%foo.bar2.clone"); + + // Test instruction name with a dot. + auto foo_baz = HloInstruction::CreateParameter( + 0, ShapeUtil::MakeShape(F32, {}), "foo.baz"); + EXPECT_EQ(foo_baz->Clone()->name(), "%foo.baz.clone"); + + // Test incrementing a large number after the suffix. + auto foo_clone234 = HloInstruction::CreateParameter( + 0, ShapeUtil::MakeShape(F32, {}), "foo.clone234"); + EXPECT_EQ(foo_clone234->Clone()->name(), "%foo.clone235"); + + // Test a non-numeric string after the cloning suffix. + auto foo_clonexyz = HloInstruction::CreateParameter( + 0, ShapeUtil::MakeShape(F32, {}), "foo.clonexyz"); + EXPECT_EQ(foo_clonexyz->Clone()->name(), "%foo.clonexyz.clone"); + + // Test a name with multiple appearances of the suffix. + auto foo_clone_clone3 = HloInstruction::CreateParameter( + 0, ShapeUtil::MakeShape(F32, {}), "foo.clone.clone3"); + EXPECT_EQ(foo_clone_clone3->Clone()->name(), "%foo.clone.clone4"); +} + +TEST_F(HloInstructionTest, Stringification) { + // Tests stringification of a simple op, fusion, and while. + const Shape s1 = ShapeUtil::MakeShape(F32, {5, 10}); + const Shape s2 = ShapeUtil::MakeShape(F32, {20, 10}); + const Shape s2t = ShapeUtil::MakeShape(F32, {10, 20}); + const Shape sout = ShapeUtil::MakeShape(F32, {5, 20}); + + HloComputation::Builder builder("TransposeDot"); + HloInstruction* x = + builder.AddInstruction(HloInstruction::CreateParameter(0, s1, "x")); + HloInstruction* y = + builder.AddInstruction(HloInstruction::CreateParameter(1, s2, "y")); + HloInstruction* reshape = + builder.AddInstruction(HloInstruction::CreateTranspose(s2t, y, {1, 0})); + HloInstruction* dot = builder.AddInstruction( + HloInstruction::CreateBinary(sout, HloOpcode::kDot, x, reshape)); + + EXPECT_EQ(dot->ToString(false, false), + "%dot = f32[5,20]{1,0} dot(f32[5,10]{1,0} %x, f32[10,20]{1,0} " + "%transpose)"); + + HloModule module(TestName()); + auto* computation = module.AddEntryComputation(builder.Build()); + HloInstruction* fusion = computation->CreateFusionInstruction( + {dot, reshape}, HloInstruction::FusionKind::kTransposeDot); + + EXPECT_EQ(fusion->ToString(false, false), + "%fusion = f32[5,20]{1,0} fusion:kTransposeDot(f32[5,10]{1,0} %x, " + "f32[20,10]{1,0} %y), calls=fused_computation"); + + HloInstruction* loop = builder.AddInstruction( + HloInstruction::CreateWhile(sout, computation, computation, x)); + EXPECT_EQ(loop->ToString(false, false), + "%while = f32[5,20]{1,0} while(f32[5,10]{1,0} %x), " + "condition=TransposeDot, body=TransposeDot"); +} + } // namespace } // namespace xla + +int main(int argc, char** argv) { + return xla::ParseDebugOptionsFlagsAndRunTests(argc, argv); +} diff --git a/tensorflow/compiler/xla/service/hlo_matchers.cc b/tensorflow/compiler/xla/service/hlo_matchers.cc new file mode 100644 index 0000000000000000000000000000000000000000..e022c4836d87866925ab7e56c2250d87d0f5dfec --- /dev/null +++ b/tensorflow/compiler/xla/service/hlo_matchers.cc @@ -0,0 +1,77 @@ +/* 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/hlo_matchers.h" + +#include "tensorflow/compiler/xla/service/hlo_instruction.h" +#include "tensorflow/compiler/xla/test.h" + +namespace xla { +namespace testing { + +bool HloMatcher::MatchAndExplain( + const HloInstruction* instruction, + ::testing::MatchResultListener* listener) const { + // These cases are self-explanatory from the printed value. + if (!instruction || instruction->opcode() != opcode_) { + return false; + } + // Special case: no operand matchers means don't verify. + if (operands_.empty()) { + return true; + } + const auto& operands = instruction->operands(); + if (operands.size() != operands_.size()) { + *listener << "has too " + << (operands.size() > operands_.size() ? "many" : "few") + << " operands (got " << operands.size() << ", want " + << operands_.size() << ")"; + return false; + } + for (int index = 0; index < operands.size(); index++) { + ::testing::StringMatchResultListener inner_listener; + if (!operands_[index].MatchAndExplain(operands[index], &inner_listener)) { + if (listener->IsInterested()) { + *listener << "\noperand " << index << ":\n\t" + << operands[index]->ToString() + << "\ndoesn't match expected:\n\t"; + operands_[index].DescribeTo(listener->stream()); + string explanation = inner_listener.str(); + if (!explanation.empty()) { + *listener << ", " << explanation; + } + } + return false; + } + } + return true; +} + +void HloMatcher::DescribeTo(::std::ostream* os) const { + *os << opcode_; + if (!operands_.empty()) { + *os << "("; + for (int i = 0; i < operands_.size(); i++) { + if (i > 0) { + *os << ", "; + } + operands_[i].DescribeTo(os); + } + *os << ")"; + } +} + +} // namespace testing +} // namespace xla diff --git a/tensorflow/compiler/xla/service/hlo_matchers.h b/tensorflow/compiler/xla/service/hlo_matchers.h new file mode 100644 index 0000000000000000000000000000000000000000..79f17bbb6bd9bfc0c6ed48c68599ef51fbd27af8 --- /dev/null +++ b/tensorflow/compiler/xla/service/hlo_matchers.h @@ -0,0 +1,143 @@ +/* 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_MATCHERS_H_ +#define TENSORFLOW_COMPILER_XLA_SERVICE_HLO_MATCHERS_H_ + +#include "tensorflow/compiler/xla/service/hlo_instruction.h" +#include "tensorflow/compiler/xla/test.h" + +namespace xla { +namespace testing { + +class HloMatcher : public ::testing::MatcherInterface { + public: + HloMatcher(HloOpcode opcode, + std::vector<::testing::Matcher> operands) + : opcode_(opcode), operands_(operands) {} + + bool MatchAndExplain(const HloInstruction* instruction, + ::testing::MatchResultListener* listener) const override; + + void DescribeTo(::std::ostream* os) const override; + + private: + HloOpcode opcode_; + std::vector<::testing::Matcher> operands_; +}; + +// HloInstruction* matchers for opcode and operands. Example: +// namespace op = xla::opcode_matchers; +// EXPECT_THAT(instruction, +// op::Add(op::Reshape(), op::Add(op::Reshape(), _))); +namespace opcode_matchers { +#define HLO_MATCHER(opcode) \ + template \ + ::testing::Matcher opcode(M... operands) { \ + return ::testing::MakeMatcher(new ::xla::testing::HloMatcher( \ + ::xla::HloOpcode::k##opcode, {operands...})); \ + } +HLO_MATCHER(Abs); +HLO_MATCHER(Add); +HLO_MATCHER(Bitcast); +HLO_MATCHER(Broadcast); +HLO_MATCHER(Call); +HLO_MATCHER(Ceil); +HLO_MATCHER(Clamp); +HLO_MATCHER(Concatenate); +HLO_MATCHER(Constant); +HLO_MATCHER(Convert); +HLO_MATCHER(Convolution); +HLO_MATCHER(Copy); +HLO_MATCHER(CrossReplicaSum); +HLO_MATCHER(CustomCall); +HLO_MATCHER(Divide); +HLO_MATCHER(Dot); +HLO_MATCHER(DynamicSlice); +HLO_MATCHER(DynamicUpdateSlice); +HLO_MATCHER(Eq); +HLO_MATCHER(Exp); +HLO_MATCHER(Floor); +HLO_MATCHER(Fusion); +HLO_MATCHER(Ge); +HLO_MATCHER(GetTupleElement); +HLO_MATCHER(Gt); +HLO_MATCHER(Index); +HLO_MATCHER(Infeed); +HLO_MATCHER(IsFinite); +HLO_MATCHER(Le); +HLO_MATCHER(Log); +HLO_MATCHER(LogicalAnd); +HLO_MATCHER(LogicalNot); +HLO_MATCHER(LogicalOr); +HLO_MATCHER(Lt); +HLO_MATCHER(Map); +HLO_MATCHER(Maximum); +HLO_MATCHER(Minimum); +HLO_MATCHER(Multiply); +HLO_MATCHER(Ne); +HLO_MATCHER(Negate); +HLO_MATCHER(Outfeed); +HLO_MATCHER(Pad); +HLO_MATCHER(Parameter); +HLO_MATCHER(Power); +HLO_MATCHER(Recv); +HLO_MATCHER(Reduce); +HLO_MATCHER(ReducePrecision); +HLO_MATCHER(ReduceWindow); +HLO_MATCHER(Remainder); +HLO_MATCHER(Reshape); +HLO_MATCHER(Reverse); +HLO_MATCHER(Rng); +HLO_MATCHER(Select); +HLO_MATCHER(SelectAndScatter); +HLO_MATCHER(Send); +HLO_MATCHER(Sign); +HLO_MATCHER(Slice); +HLO_MATCHER(Sort); +HLO_MATCHER(Subtract); +HLO_MATCHER(Tanh); +HLO_MATCHER(Trace); +HLO_MATCHER(Transpose); +HLO_MATCHER(Tuple); +HLO_MATCHER(Update); +HLO_MATCHER(While); +#undef HLO_MATCHER +} // namespace opcode_matchers + +// Helper to convert smart to raw pointers for matching. +template +std::vector Pointers(const Container& container) { + std::vector result; + result.reserve(container.size()); + for (const auto& entry : container) result.push_back(entry.get()); + return result; +} + +} // namespace testing + +// 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) { + *os << (inst ? inst->ToString() : "nullptr"); +} + +void PrintTo(HloInstruction* inst, ::std::ostream* os) { + PrintTo(const_cast(inst), os); +} + +} // namespace xla + +#endif // TENSORFLOW_COMPILER_XLA_SERVICE_HLO_MATCHERS_H_ diff --git a/tensorflow/compiler/xla/service/hlo_matchers_test.cc b/tensorflow/compiler/xla/service/hlo_matchers_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..1465d1cacdc971a04c620bc48bed33239a67a955 --- /dev/null +++ b/tensorflow/compiler/xla/service/hlo_matchers_test.cc @@ -0,0 +1,71 @@ +/* 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/hlo_matchers.h" +#include "tensorflow/compiler/xla/shape_util.h" + +namespace op = xla::testing::opcode_matchers; +using ::testing::_; +using ::testing::Eq; + +namespace xla { +namespace { + +template +string Explain(const T& t, const M& m) { + ::testing::StringMatchResultListener listener; + EXPECT_THAT(t, ::testing::Not(m)); // For the error message. + EXPECT_FALSE(m.MatchAndExplain(t, &listener)); + return listener.str(); +} + +TEST(HloMatchersTest, Test) { + auto shape = ShapeUtil::MakeShape(F32, {1}); + auto param = HloInstruction::CreateParameter(0, shape, "param"); + auto mul = HloInstruction::CreateBinary(shape, HloOpcode::kMultiply, + param.get(), param.get()); + auto add = HloInstruction::CreateBinary(shape, HloOpcode::kAdd, param.get(), + mul.get()); + + EXPECT_THAT(add.get(), op::Add()); + EXPECT_THAT(add.get(), op::Add(op::Parameter(), op::Multiply())); + EXPECT_THAT(add.get(), + op::Add(op::Parameter(), op::Multiply(_, op::Parameter()))); + + // Negative matches: check the explanation string. + EXPECT_THAT(Explain(add.get(), op::Parameter()), Eq("")); + EXPECT_THAT(Explain(add.get(), op::Add(op::Parameter())), + Eq("has too many operands (got 2, want 1)")); + EXPECT_THAT( + Explain(add.get(), op::Add(op::Parameter(), op::Parameter())), + Eq("\noperand 1:\n\t" + "%multiply = f32[1]{0} multiply(f32[1]{0} %param, f32[1]{0} %param)\n" + "doesn't match expected:\n\t" + "parameter")); + EXPECT_THAT( + Explain(add.get(), + op::Add(op::Parameter(), op::Multiply(op::Add(), op::Add()))), + Eq("\noperand 1:\n\t" + "%multiply = f32[1]{0} multiply(f32[1]{0} %param, f32[1]{0} %param)\n" + "doesn't match expected:\n\t" + "multiply(add, add), \n" + "operand 0:\n\t" + "%param = f32[1]{0} parameter(0)\n" + "doesn't match expected:\n\t" + "add")); +} + +} // namespace +} // namespace xla diff --git a/tensorflow/compiler/xla/service/hlo_module.cc b/tensorflow/compiler/xla/service/hlo_module.cc index 8ed672aa9b8fb73cc120f55d93530b3124519fcb..9e172d940a6b6c384c64be28b5bc6cee0a6fbb11 100644 --- a/tensorflow/compiler/xla/service/hlo_module.cc +++ b/tensorflow/compiler/xla/service/hlo_module.cc @@ -23,6 +23,7 @@ limitations under the License. #include #include "tensorflow/compiler/xla/map_util.h" +#include "tensorflow/compiler/xla/ptr_util.h" #include "tensorflow/compiler/xla/shape_util.h" #include "tensorflow/compiler/xla/types.h" #include "tensorflow/core/lib/gtl/map_util.h" @@ -32,22 +33,22 @@ limitations under the License. namespace xla { HloModule::HloModule(const string& name, - const VersionedComputationHandle& entry_computation_handle) + const VersionedComputationHandle& entry_computation_handle, + const HloModuleConfig& config) : name_(name), - entry_computation_(nullptr), + config_(config), has_entry_computation_handle_(true), - entry_computation_handle_(entry_computation_handle), - computation_name_uniquer_(/*separator=*/".") {} + entry_computation_handle_(entry_computation_handle) {} -HloModule::HloModule(const string& name) - : name_(name), - entry_computation_(nullptr), - computation_name_uniquer_(/*separator=*/".") {} +HloModule::HloModule(const string& name) : name_(name) {} HloComputation* HloModule::AddComputationInternal( std::unique_ptr computation) { - computation->set_name( - computation_name_uniquer_.GetUniqueName(computation->name())); + computation->UniquifyName(&computation_name_uniquer_); + for (auto& instruction : computation->instructions()) { + instruction->UniquifyName(&instruction_name_uniquer_); + instruction->SetUniqueId(NewUniqueInstructionId()); + } computation->set_parent(this); computations_.push_back(std::move(computation)); return computations_.back().get(); @@ -57,9 +58,27 @@ HloComputation* HloModule::AddEntryComputation( std::unique_ptr computation) { CHECK_EQ(nullptr, entry_computation_); entry_computation_ = computation.get(); + + // If the module configuration has no entry layout computation set, create a + // default one based on the program shape. + if (!config_.has_entry_computation_layout()) { + config_.SetDefaultComputationLayout( + entry_computation_->ComputeProgramShape()); + } return AddComputationInternal(std::move(computation)); } +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; + }); + TF_RET_CHECK(it->get() == to_remove); + computations_.erase(it); + return Status::OK(); +} + HloComputation* HloModule::AddEmbeddedComputation( std::unique_ptr computation) { return AddComputationInternal(std::move(computation)); @@ -141,6 +160,17 @@ string HloModule::ToString() const { return s.str(); } +HloModuleProto HloModule::ToProto() const { + HloModuleProto proto; + proto.set_name(name_); + proto.set_entry_computation_name(entry_computation_->name()); + for (const HloComputation* computation : MakeComputationPostOrder()) { + HloComputationProto computation_proto = computation->ToProto(); + proto.add_computations()->Swap(&computation_proto); + } + return proto; +} + namespace { // Returns whether `hlo` is used outside the given subcomputation. // `instructions_in_subcomputation` is the instruction set of the given @@ -281,6 +311,36 @@ std::list HloModule::MakeComputationPostOrder() const { return post_order; } +std::unique_ptr HloModule::Clone(const string& suffix) { + VLOG(1) << "Cloning module :" << name_ << " --> " << suffix << "\n"; + auto module = MakeUnique(name_ + "-" + suffix); + module->config_ = config_; + module->entry_computation_handle_ = entry_computation_handle_; + module->has_entry_computation_handle_ = has_entry_computation_handle_; + + std::unordered_map clone_map; + for (auto& computation : computations_) { + auto cloned_computation = computation->Clone(suffix); + InsertOrDie(&clone_map, computation.get(), cloned_computation.get()); + + if (entry_computation_ == computation.get()) { + module->AddEntryComputation(std::move(cloned_computation)); + } else { + module->AddEmbeddedComputation(std::move(cloned_computation)); + } + } + + for (auto& cloned_computation : module->computations_) { + for (auto& instruction : cloned_computation->instructions()) { + // Rewrite instruction's called_computation to point to the cloned + // computations. + instruction->ReplaceCalledComputations( + [&](HloComputation* hlo) { return FindOrDie(clone_map, hlo); }); + } + } + return module; +} + uint64 HloModule::RandomNew64() const { tensorflow::mutex_lock l(rng_mutex_); return rng_(); diff --git a/tensorflow/compiler/xla/service/hlo_module.h b/tensorflow/compiler/xla/service/hlo_module.h index 1ff5c5dacb810593e9f1e251867fdb58f20e57e0..d149500bcad1e4c7725abf84b8aefff28b5082eb 100644 --- a/tensorflow/compiler/xla/service/hlo_module.h +++ b/tensorflow/compiler/xla/service/hlo_module.h @@ -23,8 +23,10 @@ limitations under the License. #include #include +#include "tensorflow/compiler/xla/service/hlo.pb.h" #include "tensorflow/compiler/xla/service/hlo_computation.h" #include "tensorflow/compiler/xla/service/hlo_instruction.h" +#include "tensorflow/compiler/xla/service/hlo_module_config.h" #include "tensorflow/compiler/xla/service/name_uniquer.h" #include "tensorflow/compiler/xla/service/versioned_computation_handle.h" #include "tensorflow/compiler/xla/types.h" @@ -43,12 +45,13 @@ namespace xla { class HloModule { public: HloModule(const string& name, - const VersionedComputationHandle& entry_computation_handle); + const VersionedComputationHandle& entry_computation_handle, + const HloModuleConfig& config); // Constructor without a versioned computation handle. This constructor should // only be used for HloModules used outside of the XLA service (eg // tests). The versioned handle is used by the service in the compilation - // cache. + // cache. A default configuration is created for this module. explicit HloModule(const string& name); // Adds an entry computation to the module. A module can only have one entry @@ -60,6 +63,9 @@ class HloModule { HloComputation* AddEmbeddedComputation( std::unique_ptr computation); + // Removes an embedded computation. + Status RemoveEmbeddedComputation(HloComputation* to_remove); + // Replaces all uses of computations that are keys of 'replacements' with // the corresponding values in 'replacements'. Replaces the entry computation, // if applicable. @@ -72,12 +78,19 @@ class HloModule { const string& name() const { return name_; } + // Returns a deep copy of this module including all computations. + std::unique_ptr Clone(const string& suffix = "clone"); + // Return a pointer to the entry computation of the module.. HloComputation* entry_computation() const { CHECK_NE(nullptr, entry_computation_); return entry_computation_; } + ComputationLayout* mutable_entry_computation_layout() { + return config_.mutable_entry_computation_layout(); + } + const VersionedComputationHandle& entry_computation_handle() const { return entry_computation_handle_; } @@ -91,7 +104,10 @@ class HloModule { // computation B, then A will appear after B in the sort. std::list MakeComputationPostOrder() const; + const HloModuleConfig& config() const { return config_; } + string ToString() const; + HloModuleProto ToProto() const; // Outlines the given expression from the given computation. // instructions_to_outline contains the instructions that form the expression. @@ -106,12 +122,32 @@ class HloModule { // Returns a randomly generated uint64. uint64 RandomNew64() const; + // Returns the unique name for a computation in this module. + string GetUniqueCompuationName(const string& prefix) { + return computation_name_uniquer_.GetUniqueName(prefix); + } + + // Returns the NameUniquer for uniquing instruction names in this module. + NameUniquer& instruction_name_uniquer() { return instruction_name_uniquer_; } + + // Assign a new unique dense id for an instruction + int NewUniqueInstructionId() { + int result = next_unique_id_; + next_unique_id_++; + return result; + } + + // Returns the number of unique intruction ids given out. All ids up to + // this point are guaranteed to be in the range [0..NumUniqueInstructionIds()) + int NumUniqueInstructionIds() const { return next_unique_id_; } + private: HloComputation* AddComputationInternal( std::unique_ptr computation); const string name_; - HloComputation* entry_computation_; + HloModuleConfig config_; + HloComputation* entry_computation_ = nullptr; std::vector> computations_; // Random number generator engine to use when generating random numbers per @@ -125,8 +161,11 @@ class HloModule { bool has_entry_computation_handle_ = false; VersionedComputationHandle entry_computation_handle_; - // Unique name generator for computation names, which are unique per module. - NameUniquer computation_name_uniquer_; + // Unique name generator for computation and instruction names, which are + // unique per module. + NameUniquer computation_name_uniquer_{/*separator=*/"."}; + NameUniquer instruction_name_uniquer_{/*separator=*/"."}; + int next_unique_id_ = 0; }; } // namespace xla diff --git a/tensorflow/compiler/xla/service/hlo_module_config.cc b/tensorflow/compiler/xla/service/hlo_module_config.cc index c129ad1b3924c94110bc356197105f01cdcbf677..8974deb530c2e4561b5ab57f43c65fd525db3617 100644 --- a/tensorflow/compiler/xla/service/hlo_module_config.cc +++ b/tensorflow/compiler/xla/service/hlo_module_config.cc @@ -18,6 +18,7 @@ limitations under the License. #include #include +#include "tensorflow/compiler/xla/ptr_util.h" #include "tensorflow/compiler/xla/shape_layout.h" #include "tensorflow/compiler/xla/types.h" #include "tensorflow/core/lib/strings/str_util.h" @@ -27,20 +28,27 @@ namespace xla { using tensorflow::strings::StrAppend; +HloModuleConfig::HloModuleConfig() {} + HloModuleConfig::HloModuleConfig(const ProgramShape& program_shape) : entry_computation_layout_(program_shape) {} +void HloModuleConfig::SetDefaultComputationLayout( + const ProgramShape& program_shape) { + entry_computation_layout_ = ComputationLayout(program_shape); +} + string HloModuleConfig::compilation_cache_key() const { string key = tensorflow::strings::StrCat("profiling=", hlo_profiling_enabled_, "::hybrid=", has_hybrid_result_); StrAppend(&key, "::("); std::vector params; for (const ShapeLayout& param_layout : - entry_computation_layout_.parameter_layouts()) { + entry_computation_layout_->parameter_layouts()) { params.push_back(param_layout.shape().DebugString()); } StrAppend(&key, tensorflow::str_util::Join(params, ", "), ") => ", - entry_computation_layout_.result_shape().SerializeAsString()); + entry_computation_layout_->result_shape().SerializeAsString()); if (seed() != 0) { // TODO(b/32083678): force recompilation to reset global state. static std::atomic counter{0}; @@ -49,7 +57,11 @@ string HloModuleConfig::compilation_cache_key() const { if (replica_count() != 1) { StrAppend(&key, "::replica_count=", replica_count()); } - StrAppend(&key, "::fast_math_disabled=", fast_math_disabled_); + StrAppend(&key, debug_options_.DebugString()); + if (intra_op_parallelism_threads() > 0) { + StrAppend(&key, "::intra_op_parallelism_threads=", + intra_op_parallelism_threads()); + } return key; } diff --git a/tensorflow/compiler/xla/service/hlo_module_config.h b/tensorflow/compiler/xla/service/hlo_module_config.h index f9a61c1cd1c3cca90e09cbbb88554013772becac..2299200b5be969c065fded840709a3d6034efe47 100644 --- a/tensorflow/compiler/xla/service/hlo_module_config.h +++ b/tensorflow/compiler/xla/service/hlo_module_config.h @@ -22,6 +22,7 @@ limitations under the License. #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/gtl/optional.h" namespace xla { @@ -32,14 +33,34 @@ namespace xla { // executable. class HloModuleConfig { public: + // A configuration can be created either with, or without an entry + // ComputationLayout. The default ctor creates it without -- in this case + // accessing entry_computation_layout will CHECK-fail. The ctor accepting a + // ProgramShape creates a computation layout using this shape. + HloModuleConfig(); explicit HloModuleConfig(const ProgramShape& program_shape); - // Return a reference to the layout of the entry computation. + // Checks if this config has an entry computation layout already. + bool has_entry_computation_layout() const { + return entry_computation_layout_.has_value(); + } + + // Sets the entry computation layout for this config. If the entry computation + // layout already exists, it is silently replaced. + void SetDefaultComputationLayout(const ProgramShape& program_shape); + + // Returns a constant reference to the layout of the entry computation. + // Assumes the layout was set. const ComputationLayout& entry_computation_layout() const { - return entry_computation_layout_; + CHECK(entry_computation_layout_.has_value()); + return *entry_computation_layout_; } + + // Returns a mutable pointer to the layout of the entry computation. Assumes + // the layout was set. ComputationLayout* mutable_entry_computation_layout() { - return &entry_computation_layout_; + CHECK(entry_computation_layout_.has_value()); + return &(*entry_computation_layout_); } // Sets/returns whether to enable HLO-level profiling. @@ -60,23 +81,30 @@ class HloModuleConfig { } int64 replica_count() const { return replica_count_; } - // Sets/returns whether unsafe math optimizations are disabled for this - // module. Default is fast-math enabled. - // - // This is named fast_math_disabled rather than the more natural - // fast_math_enabled for consistency with the ExecutionOptions proto. - bool fast_math_disabled() const { return fast_math_disabled_; } - void set_fast_math_disabled(bool disabled) { fast_math_disabled_ = disabled; } - // Return a string which unambiguously represents all the fields of this data // structure. Used for generating a cache key for storing the compiled // executable. string compilation_cache_key() const; + const DebugOptions& debug_options() const { return debug_options_; } + + void set_debug_options(const DebugOptions& debug_options) { + debug_options_ = debug_options; + } + + // Sets/returns the number of intra op threads for this module. + void set_intra_op_parallelism_threads( + const int intra_op_parallelism_threads) { + intra_op_parallelism_threads_ = intra_op_parallelism_threads; + } + int64 intra_op_parallelism_threads() const { + return intra_op_parallelism_threads_; + } + private: // If you add new members, be sure to update compilation_cache_key. - ComputationLayout entry_computation_layout_; + tensorflow::gtl::optional entry_computation_layout_; // Whether to enable HLO-level profiling. bool hlo_profiling_enabled_ = false; @@ -97,7 +125,11 @@ class HloModuleConfig { // The number of replicas to compile this binary for. int64 replica_count_ = 1; - bool fast_math_disabled_ = false; + // The target maximum parallelism at which to partition HLOs for parallel + // execution on the CPU backend. + int64 intra_op_parallelism_threads_ = -1; + + DebugOptions debug_options_; }; } // namespace xla diff --git a/tensorflow/compiler/xla/service/hlo_module_test.cc b/tensorflow/compiler/xla/service/hlo_module_test.cc index dba9731e2a0dd7f227cf16d77c598b17dbaf7307..56dc5632035c625445018becfd25d69557e6232a 100644 --- a/tensorflow/compiler/xla/service/hlo_module_test.cc +++ b/tensorflow/compiler/xla/service/hlo_module_test.cc @@ -23,7 +23,7 @@ limitations under the License. #include "tensorflow/compiler/xla/tests/hlo_test_base.h" #include "tensorflow/compiler/xla/xla_data.pb.h" -#include "tensorflow/compiler/xla/test_helpers.h" +#include "tensorflow/compiler/xla/test.h" #include "tensorflow/core/lib/gtl/array_slice.h" namespace xla { @@ -38,7 +38,7 @@ class HloModuleTest : public HloTestBase { std::unique_ptr CreateConstantComputation() { auto builder = HloComputation::Builder("Constant"); builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(42.0f))); + HloInstruction::CreateConstant(Literal::CreateR0(42.0f))); return builder.Build(); } @@ -58,22 +58,22 @@ class HloModuleTest : public HloTestBase { TEST_F(HloModuleTest, OneComputationPostOrder) { // Create a module with a single computation. - auto module = MakeUnique(TestName()); + auto module = CreateNewModule(); auto computation = module->AddEntryComputation(CreateConstantComputation()); - EXPECT_EQ(module->MakeComputationPostOrder().front(), computation); + EXPECT_THAT(module->MakeComputationPostOrder(), + ::testing::ElementsAre(computation)); } TEST_F(HloModuleTest, TwoComputationsPostOrder) { // Create a module with two unconnected computations. - auto module = MakeUnique(TestName()); + auto module = CreateNewModule(); auto computation1 = module->AddEntryComputation(CreateConstantComputation()); auto computation2 = module->AddEmbeddedComputation(CreateConstantComputation()); - EXPECT_MATCH( - testing::ListToVec(module->MakeComputationPostOrder()), - testing::UnorderedMatcher(computation1, computation2)); + EXPECT_THAT(module->MakeComputationPostOrder(), + ::testing::UnorderedElementsAre(computation1, computation2)); // We specified the same name for both computations, but the HloModule should // have made the names unique. @@ -81,9 +81,33 @@ TEST_F(HloModuleTest, TwoComputationsPostOrder) { EXPECT_EQ(computation2->name(), "Constant.1"); } +TEST_F(HloModuleTest, CloneTest) { + // Create and copy a module with a diamond call graph of computations. + auto module = CreateNewModule(); + auto computation1 = + module->AddEmbeddedComputation(CreateConstantComputation()); + auto computation2 = + module->AddEmbeddedComputation(CreateCallComputation({computation1})); + auto computation3 = + module->AddEmbeddedComputation(CreateCallComputation({computation1})); + module->AddEntryComputation( + CreateCallComputation({computation2, computation3})); + + auto post_order = module->MakeComputationPostOrder(); + auto cloned_module = module->Clone("copy"); + auto post_order_copied = cloned_module->MakeComputationPostOrder(); + + EXPECT_EQ(post_order.size(), post_order_copied.size()); + for (auto origin = post_order.begin(), copied = post_order_copied.begin(); + origin != post_order.end() && copied != post_order_copied.end(); + ++origin, ++copied) { + EXPECT_EQ((*origin)->name() + "copy", (*copied)->name()); + } +} + TEST_F(HloModuleTest, DiamondComputationsPostOrder) { // Create a module with a diamond call graph of computations. - auto module = MakeUnique(TestName()); + auto module = CreateNewModule(); auto computation1 = module->AddEmbeddedComputation(CreateConstantComputation()); auto computation2 = @@ -94,9 +118,9 @@ TEST_F(HloModuleTest, DiamondComputationsPostOrder) { CreateCallComputation({computation2, computation3})); auto post_order = module->MakeComputationPostOrder(); - EXPECT_MATCH(testing::ListToVec(post_order), - testing::UnorderedMatcher( - computation1, computation2, computation3, computation4)); + EXPECT_THAT(post_order, + ::testing::UnorderedElementsAre(computation1, computation2, + computation3, computation4)); EXPECT_EQ(post_order.back(), computation4); EXPECT_EQ(post_order.front(), computation1); } @@ -104,3 +128,7 @@ TEST_F(HloModuleTest, DiamondComputationsPostOrder) { } // namespace } // namespace xla + +int main(int argc, char** argv) { + return xla::ParseDebugOptionsFlagsAndRunTests(argc, argv); +} diff --git a/tensorflow/compiler/xla/service/hlo_opcode.cc b/tensorflow/compiler/xla/service/hlo_opcode.cc index 616b239a9310bc13e14c861184b7efebe7da6b2f..314512d0a8d32b59ea000bdaa6f7399b0194dd03 100644 --- a/tensorflow/compiler/xla/service/hlo_opcode.cc +++ b/tensorflow/compiler/xla/service/hlo_opcode.cc @@ -19,11 +19,24 @@ limitations under the License. namespace xla { string HloOpcodeString(HloOpcode opcode) { + // Note: Do not use ':' in opcode strings. It is used as a special character + // in these places: + // - In extended opcode strings (HloInstruction::ExtendedOpcodeString()), to + // separate the opcode from the fusion kind + // - In fully qualified names (HloInstruction::FullyQualifiedName()), to + // separate the qualifiers (name of the computation and potentially the + // fusion instruction) from the name switch (opcode) { case HloOpcode::kAbs: return "abs"; case HloOpcode::kAdd: return "add"; + case HloOpcode::kBatchNormTraining: + return "batch-norm-training"; + case HloOpcode::kBatchNormInference: + return "batch-norm-inference"; + case HloOpcode::kBatchNormGrad: + return "batch-norm-grad"; case HloOpcode::kBitcast: return "bitcast"; case HloOpcode::kBroadcast: @@ -40,6 +53,8 @@ string HloOpcodeString(HloOpcode opcode) { return "convert"; case HloOpcode::kConvolution: return "convolution"; + case HloOpcode::kCos: + return "cosine"; case HloOpcode::kCrossReplicaSum: return "cross-replica-sum"; case HloOpcode::kCustomCall: @@ -112,6 +127,8 @@ string HloOpcodeString(HloOpcode opcode) { return "recv"; case HloOpcode::kReduce: return "reduce"; + case HloOpcode::kReducePrecision: + return "reduce-precision"; case HloOpcode::kReduceWindow: return "reduce-window"; case HloOpcode::kRemainder: @@ -130,6 +147,8 @@ string HloOpcodeString(HloOpcode opcode) { return "send"; case HloOpcode::kSign: return "sign"; + case HloOpcode::kSin: + return "sine"; case HloOpcode::kSlice: return "slice"; case HloOpcode::kSort: @@ -165,4 +184,17 @@ bool HloOpcodeIsComparison(HloOpcode opcode) { } } +bool HloOpcodeIsVariadic(HloOpcode opcode) { + switch (opcode) { + case HloOpcode::kCall: + case HloOpcode::kConcatenate: + case HloOpcode::kFusion: + case HloOpcode::kMap: + case HloOpcode::kTuple: + return true; + default: + return false; + } +} + } // namespace xla diff --git a/tensorflow/compiler/xla/service/hlo_opcode.h b/tensorflow/compiler/xla/service/hlo_opcode.h index 978ed5e79b90c3c12f31b4d4e3d3314849fed75c..c4d5efad9035c32d8ecdf63ab385bf80faabb965 100644 --- a/tensorflow/compiler/xla/service/hlo_opcode.h +++ b/tensorflow/compiler/xla/service/hlo_opcode.h @@ -30,6 +30,9 @@ namespace xla { enum class HloOpcode { kAbs, kAdd, + kBatchNormTraining, + kBatchNormInference, + kBatchNormGrad, kBitcast, kBroadcast, kCall, @@ -40,6 +43,7 @@ enum class HloOpcode { kConvert, kConvolution, kCopy, + kCos, kCrossReplicaSum, kCustomCall, kDivide, @@ -74,6 +78,7 @@ enum class HloOpcode { kPower, kRecv, kReduce, + kReducePrecision, kReduceWindow, kRemainder, kReshape, @@ -83,6 +88,7 @@ enum class HloOpcode { kSelectAndScatter, kSend, kSign, + kSin, kSlice, kSort, kSubtract, @@ -104,6 +110,14 @@ inline std::ostream& operator<<(std::ostream& os, HloOpcode opcode) { // Returns true iff the given opcode is a comparison operation. bool HloOpcodeIsComparison(HloOpcode opcode); +// Returns true iff the given opcode has variadic operands. +bool HloOpcodeIsVariadic(HloOpcode opcode); + +// Returns the number of HloOpcode values. +inline const uint32_t HloOpcodeCount() { + return static_cast(HloOpcode::kWhile) + 1; +} + } // namespace xla #endif // TENSORFLOW_COMPILER_XLA_SERVICE_HLO_OPCODE_H_ diff --git a/tensorflow/compiler/xla/service/hlo_opcode_test.cc b/tensorflow/compiler/xla/service/hlo_opcode_test.cc index 0b64c16fdc6639a0288b4a69698a600b09ba32f7..892c89f9df209f2e39005a4901feae6699ce4d0b 100644 --- a/tensorflow/compiler/xla/service/hlo_opcode_test.cc +++ b/tensorflow/compiler/xla/service/hlo_opcode_test.cc @@ -15,8 +15,8 @@ limitations under the License. #include "tensorflow/compiler/xla/service/hlo_opcode.h" +#include "tensorflow/compiler/xla/test.h" #include "tensorflow/compiler/xla/types.h" -#include "tensorflow/core/platform/test.h" namespace xla { namespace { diff --git a/tensorflow/compiler/xla/service/hlo_ordering.cc b/tensorflow/compiler/xla/service/hlo_ordering.cc index b3168ed40ece3ea65c6b26b96250f2ea77969953..08f572bb2aba6c972ca0e8ee826c2ffac2e739c2 100644 --- a/tensorflow/compiler/xla/service/hlo_ordering.cc +++ b/tensorflow/compiler/xla/service/hlo_ordering.cc @@ -15,13 +15,11 @@ limitations under the License. #include "tensorflow/compiler/xla/service/hlo_ordering.h" -#include #include #include -#include "tensorflow/compiler/xla/service/heap_simulator.h" #include "tensorflow/compiler/xla/service/hlo_computation.h" -#include "tensorflow/compiler/xla/service/logical_buffer.h" +#include "tensorflow/compiler/xla/service/liveness_util.h" #include "tensorflow/compiler/xla/shape_util.h" #include "tensorflow/compiler/xla/status_macros.h" #include "tensorflow/compiler/xla/statusor.h" @@ -34,17 +32,218 @@ limitations under the License. namespace xla { -PredecessorHloOrdering::PredecessorHloOrdering(const HloModule* module) - : module_(module) {} +bool HloOrdering::ExecutesBefore(const HloInstruction* a, + const HloInstruction* b) const { + // 'a' and 'b' may be in different computations. In this case, find the + // callgraph ancestor instructions which call (potentially transitively) the + // computations containing 'a' and 'b' and use these ancestor instructions to + // compare order. + const HloInstruction* a_ancestor; + const HloInstruction* b_ancestor; + std::tie(a_ancestor, b_ancestor) = + call_graph_->NearestAncestorsInSameComputation( + const_cast(a), const_cast(b)); + + if (a_ancestor == nullptr) { + // Ancestors in a common computation could not be found so consider the + // instructions 'a' and 'b' to be unordered. + return false; + } + // a_ancestor and b_ancestor must be either both null or both non-null. + CHECK_NE(b_ancestor, nullptr); + CHECK_EQ(a_ancestor->parent(), b_ancestor->parent()); + + // If the common ancestor is a while instruction there is an additional + // ordering criteria which may apply. The condition computation is considered + // to execute before the body computation so if 'a' is in the condition and + // 'b' is in the body, then 'a' executes before 'b'. + if (a_ancestor == b_ancestor && a_ancestor->opcode() == HloOpcode::kWhile) { + const HloComputation* body = a_ancestor->while_body(); + const HloComputation* condition = a_ancestor->while_condition(); + if (call_graph_->InstructionIsNestedIn(a, condition) && + call_graph_->InstructionIsNestedIn(b, body)) { + return true; + } + } + + return ExecutesBeforeInSameComputation(a_ancestor, b_ancestor); +} + +bool HloOrdering::IsDefinedBefore(const HloValue& a, const HloValue& b) const { + // If 'b' is an entry param then 'a' cannot be defined before 'b' because 'b' + // is live into the module. + const HloModule* module = b.defining_instruction()->parent()->parent(); + if (b.defining_instruction()->parent() == module->entry_computation() && + b.defining_instruction()->opcode() == HloOpcode::kParameter) { + return false; + } + + // Phi values require special handling. Because XLA does not have a phi + // instruction, the definition instruction of the phis values are + // placeholders: either the subcomputation parameter (body or condition) or + // the while instruction. However, the program point where these values are + // logically defined does not necessarily coincide exactly with program point + // of these place-holder instructions. So we explicitly define the following + // order for phi values: + // + // body/condition parameter phi: + // Defined before all values defined in its computation excepting other + // phis. + // + // while phi: + // defined after all values defined in the condition or body. + // + auto is_body_or_condition_phi = [](const HloValue& v) { + return v.is_phi() && + v.defining_instruction()->opcode() == HloOpcode::kParameter; + }; + if (is_body_or_condition_phi(a) && !is_body_or_condition_phi(b) && + call_graph_->InstructionIsNestedIn(b.defining_instruction(), + a.defining_instruction()->parent())) { + return true; + } + if (is_body_or_condition_phi(b) && + call_graph_->InstructionIsNestedIn(a.defining_instruction(), + b.defining_instruction()->parent())) { + return false; + } + + // If 'b' is a while phi and 'a' is in the body or condition, then 'a' + // executes before 'b'. + if (b.is_phi() && b.defining_instruction()->opcode() == HloOpcode::kWhile && + (call_graph_->InstructionIsNestedIn( + a.defining_instruction(), b.defining_instruction()->while_body()) || + call_graph_->InstructionIsNestedIn( + a.defining_instruction(), + b.defining_instruction()->while_condition()))) { + return true; + } + + return ExecutesBefore(a.defining_instruction(), b.defining_instruction()); +} + +/* static */ +bool HloOrdering::UseIsBeforeValueDefinition(const HloUse& use, + const HloValue& value) const { + VLOG(4) << "UseIsBeforeValueDefinition(use=" << use + << ", value=" << value.ToShortString() << ")"; + if (ExecutesBefore(use.instruction, value.defining_instruction())) { + VLOG(4) << " use instruction executes before value-defining instruction"; + return true; + } + + // If the use is at the instruction where the value is defined, then the use + // is before the def if the instruction allows buffer sharing (in place + // computation). + if (use.instruction == value.defining_instruction() && + CanShareOperandBufferWithUser( + use.instruction->mutable_operand(use.operand_number), + use.operand_index, value.defining_instruction(), + value.defining_index())) { + VLOG(4) << " use is value def, and instruction can share use buffer"; + return true; + } + + // The use at a while is an input to a phi, and logically occurs before values + // are defined in the body or condition computations. + if (use.instruction->opcode() == HloOpcode::kWhile) { + const HloInstruction* xla_while = use.instruction; + if (call_graph_->InstructionIsNestedIn(value.defining_instruction(), + xla_while->while_body()) || + call_graph_->InstructionIsNestedIn(value.defining_instruction(), + xla_while->while_condition())) { + VLOG(4) << " use is while " << use.instruction->name() + << " and def is in condition or body"; + return true; + } + } + + // Similarly if the value is defined at a while, it logically occurs after any + // uses in the body or condition computations. + if (value.defining_instruction()->opcode() == HloOpcode::kWhile) { + CHECK(value.is_phi()); + const HloInstruction* xla_while = value.defining_instruction(); + if (call_graph_->InstructionIsNestedIn(use.instruction, + xla_while->while_body()) || + call_graph_->InstructionIsNestedIn(use.instruction, + xla_while->while_condition())) { + VLOG(4) << " value is while " << value.defining_instruction()->name() + << " and use is in condition or body"; + return true; + } + } + VLOG(4) << " use is not before while"; + return false; +} + +bool HloOrdering::LiveRangeStrictlyBefore(const HloValue& a, + const HloValue& b) const { + VLOG(4) << "LiveRangeStrictlyBefore(a = " << a.ToShortString() + << ", b = " << b.ToShortString() << ")"; + if (!IsDefinedBefore(a, b)) { + VLOG(4) << "a not defined before b"; + return false; + } -bool PredecessorHloOrdering::ExecutesBefore(const HloInstruction* a, - const HloInstruction* b) const { - // Instructions in different computations are unordered. - if (a->parent() != b->parent()) { + // Live-out values from the module can never have ranges strictly before any + // other value. + if (a.live_out_of_module()) { + VLOG(4) << "a is live out of module"; return false; } + + // Live-out values of computations can never have ranges strictly before any + // other value in the computation (including values nested in + // subcomputations). + if (a.live_out_of_computation() && + call_graph_->InstructionIsNestedIn(b.defining_instruction(), + a.defining_instruction()->parent())) { + VLOG(4) << "a is live out of computation containing b"; + return false; + } + + // All uses of 'a' must be before 'b' is defined. + for (const HloUse& use : a.uses()) { + if (!UseIsBeforeValueDefinition(use, b)) { + VLOG(4) << "use of a (" << use << ") not before b is defined"; + return false; + } + } + + return true; +} + +bool HloOrdering::MayInterfere(const HloValue& a, const HloValue& b) const { + // Buffers without disjoint liveness may interfere. + return !LiveRangeStrictlyBefore(a, b) && !LiveRangeStrictlyBefore(b, a); +} + +HloOrderingProto HloOrdering::ToProto() const { + HloOrderingProto proto; + for (const auto& computation : module_->computations()) { + const std::vector* sequence = + SequentialOrder(*computation); + if (sequence != nullptr) { + HloOrderingProto::SequentialComputation* proto_computation = + proto.add_sequential_computations(); + proto_computation->set_computation_name(computation->name()); + for (const HloInstruction* instruction : *sequence) { + *proto_computation->add_instruction_names() = instruction->name(); + } + } + } + return proto; +} + +PredecessorHloOrdering::PredecessorHloOrdering(const HloModule* module) + : HloOrdering(module) {} + +bool PredecessorHloOrdering::ExecutesBeforeInSameComputation( + const HloInstruction* a, const HloInstruction* b) const { + CHECK_EQ(a->parent(), b->parent()); + // 'a' executes before 'b' if 'a' is in the strict predecessor set of 'b'. - return strict_predecessors_.at(b->parent())->IsReachable(b, a); + return a != b && predecessors_.at(a->parent())->IsReachable(a, b); } string PredecessorHloOrdering::ToStringHelper(const string& name) const { @@ -56,10 +255,10 @@ string PredecessorHloOrdering::ToStringHelper(const string& name) const { const auto all = computation->MakeInstructionPostOrder(); for (auto instruction : all) { pieces.push_back(tensorflow::strings::Printf( - " %s strict predecessors:", instruction->name().c_str())); + " %s predecessors:", instruction->name().c_str())); for (auto predecessor : all) { - if (strict_predecessors_.at(computation.get()) - ->IsReachable(instruction, predecessor)) { + if (predecessors_.at(computation.get()) + ->IsReachable(predecessor, instruction)) { pieces.push_back( tensorflow::strings::Printf(" %s", predecessor->name().c_str())); } @@ -75,8 +274,11 @@ 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->computations()) { - strict_predecessors_.emplace(computation.get(), - computation->ComputeTransitiveOperands()); + if (computation->IsFusionComputation()) { + continue; + } + predecessors_.emplace(computation.get(), + computation->ComputeReachability()); } } @@ -86,7 +288,7 @@ string DependencyHloOrdering::ToString() const { SequentialHloOrdering::SequentialHloOrdering( const HloModule* module, const HloModuleSequence& module_sequence) - : module_(module), module_sequence_(module_sequence) { + : HloOrdering(module), module_sequence_(module_sequence) { // Create a map from instruction to its order position. for (auto computation_order : module_sequence_) { const std::vector& order = computation_order.second; @@ -97,12 +299,9 @@ SequentialHloOrdering::SequentialHloOrdering( } } -bool SequentialHloOrdering::ExecutesBefore(const HloInstruction* a, - const HloInstruction* b) const { - // Instructions in different computations are unordered. - if (a->parent() != b->parent()) { - return false; - } +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) { return false; @@ -144,360 +343,6 @@ string SequentialHloOrdering::ToString() const { return tensorflow::str_util::Join(pieces, "\n"); } -namespace { -StatusOr MinimumMemoryForSequence( - const HloComputation& computation, - const std::vector& sequence, - const TuplePointsToAnalysis& points_to_analysis, - const LogicalBuffer::SizeFunction& size_function) { - // The absolute minimum memory required for a given sequence of instructions - // is determined by the sequence of Alloc and Free calls on a simulated heap, - // ignoring fragmentation. - TF_ASSIGN_OR_RETURN( - HeapSimulator::Result result, - HeapSimulator::Run(MakeUnique(), sequence, - computation, points_to_analysis, size_function)); - return result.heap_size; -} -} // namespace - -StatusOr MinimumMemoryForSequence( - const SequentialHloOrdering::HloModuleSequence& module_sequence, - const LogicalBuffer::SizeFunction& size_function) { - if (module_sequence.empty()) { - return 0; - } - - const HloModule* module = module_sequence.begin()->first->parent(); - TF_ASSIGN_OR_RETURN(std::unique_ptr points_to_analysis, - TuplePointsToAnalysis::Run(module)); - - int64 total_memory = 0; - for (const auto& pair : module_sequence) { - const HloComputation* computation = pair.first; - const std::vector& sequence = pair.second; - TF_ASSIGN_OR_RETURN( - const int64 memory, - MinimumMemoryForSequence(*computation, sequence, *points_to_analysis, - size_function)); - total_memory += memory; - } - return total_memory; -} - -namespace { - -// Class implementing a list scheduler of HLO instructions which produces a -// sequence which minimizes memory usage. -class ListScheduler { - public: - // Construct and return a memory-minimizing sequence of HLO instructions - // containing the given HLO computation. - static StatusOr> Run( - const HloComputation& computation, - const TuplePointsToAnalysis& points_to_analysis, - const LogicalBuffer::SizeFunction& size_function) { - ListScheduler scheduler(computation, points_to_analysis, size_function); - return scheduler.CreateSchedule(); - } - - private: - // The scheduling priority of an instruction is first the number of bytes - // freed by scheduling the instruction, and second (tie-breaker) by the number - // of users. This is represented as a std::pair containing these two values - // (first element is the bytes freed). std::pair provides the necessary - // comparison operators. - using Priority = std::pair; - - ListScheduler(const HloComputation& computation, - const TuplePointsToAnalysis& points_to_analysis, - const LogicalBuffer::SizeFunction& size_function) - : computation_(computation), - points_to_analysis_(points_to_analysis), - size_function_(size_function) { - // Create a map containing the LogicalBuffer uses for each HLO - // 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()) { - buffer_uses_.insert( - {instruction.get(), std::unordered_set()}); - for (auto* operand : instruction->operands()) { - for (const LogicalBuffer* buffer : - points_to_analysis.GetBuffersDefinedByInstruction(operand)) { - buffer_uses_[instruction.get()].insert(buffer); - } - } - } - - // Create map containing the number of unscheduled uses (hlo instructions) - // of each logical buffer. - for (auto& instruction : computation.instructions()) { - for (auto* buffer : points_to_analysis.GetBuffersDefinedByInstruction( - instruction.get())) { - unscheduled_use_count_[buffer] = 0; - } - } - for (auto& instruction : computation.instructions()) { - for (const LogicalBuffer* buffer : buffer_uses_.at(instruction.get())) { - ++unscheduled_use_count_[buffer]; - } - } - - // 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()) - .CreateFlattenedSet()) { - ++unscheduled_use_count_[live_out_buffer]; - } - } - - // Returns whether the memory used by the given buffer should be ignored by - // the scheduling heuristic. - bool IgnoreBuffer(const LogicalBuffer& buffer) { - return buffer.instruction()->opcode() == HloOpcode::kParameter || - buffer.instruction()->opcode() == HloOpcode::kConstant; - } - - // Return the number of bytes freed if the HLO instruction is scheduled. - int64 BytesFreedIfScheduled(const HloInstruction* instruction) { - int64 freed_bytes = 0; - // Sum the total size of the values last used by this instruction. - for (auto* buffer : buffer_uses_.at(instruction)) { - if (IgnoreBuffer(*buffer)) { - continue; - } - CHECK_GE(unscheduled_use_count_.at(buffer), 1); - if (unscheduled_use_count_.at(buffer) == 1) { - // This is the last use of the logical buffer. - freed_bytes += size_function_(*buffer); - } - } - // Then subtract the size of the value(s) defined by this instruction. - for (auto* buffer : - points_to_analysis_.GetBuffersDefinedByInstruction(instruction)) { - if (!IgnoreBuffer(*buffer)) { - freed_bytes -= size_function_(*buffer); - } - } - return freed_bytes; - } - - // Construct the scheduling priority of the given instruciton. - Priority GetPriority(const HloInstruction* instruction) { - return {BytesFreedIfScheduled(instruction), instruction->user_count()}; - } - - std::vector CreateSchedule() { - std::vector schedule; - - // Populate the ready list with instructions which have no operands or - // control predecessors. - std::unordered_map unscheduled_pred_count; - std::list ready_list; - 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()) { - unscheduled_pred_count[user]++; - } - for (const HloInstruction* succ : instruction->control_successors()) { - unscheduled_pred_count[succ]++; - } - } - for (auto& instruction : computation_.instructions()) { - // Instruction with no operands or control predecessors will - // not be in the map. - if (unscheduled_pred_count.count(instruction.get()) == 0) { - ready_list.push_back(instruction.get()); - } - } - - while (!ready_list.empty()) { - // Select the highest priority HLO instruction from the ready list. - auto best_it = ready_list.begin(); - Priority best_priority = GetPriority(*best_it); - for (auto ready_it = std::next(ready_list.begin()); - ready_it != ready_list.end(); ++ready_it) { - Priority priority = GetPriority(*ready_it); - if (priority > best_priority) { - best_it = ready_it; - best_priority = priority; - } - } - - // Remove the selected instruction from the ready list and add it to the - // schedule. - const HloInstruction* best = *best_it; - ready_list.erase(best_it); - schedule.push_back(best); - scheduled_instructions_.insert(best); - - // Update the unscheduled uses of the logical buffers. - for (const LogicalBuffer* buffer : buffer_uses_.at(best)) { - CHECK_GT(unscheduled_use_count_.at(buffer), 0); - --unscheduled_use_count_[buffer]; - } - - // Add new instructions to ready list. - auto update_pred_count = [&unscheduled_pred_count, - &ready_list](HloInstruction* inst) { - int64 pred_count = --unscheduled_pred_count.at(inst); - CHECK_GE(pred_count, 0); - if (pred_count == 0) { - ready_list.push_back(inst); - } - }; - // TODO(b/34466113): Replace this and above with successors() or - // predecessors() when these methods are added to HloInstruction. - for (HloInstruction* user : best->users()) { - update_pred_count(user); - } - for (HloInstruction* succ : best->control_successors()) { - update_pred_count(succ); - } - } - CHECK_EQ(schedule.size(), computation_.instructions().size()); - CHECK_EQ(scheduled_instructions_.size(), - computation_.instructions().size()); - - return schedule; - } - - const HloComputation& computation_; - const TuplePointsToAnalysis& points_to_analysis_; - const LogicalBuffer::SizeFunction& size_function_; - - // A map containing the LogicalBuffers that each instruction uses. - std::unordered_map> - buffer_uses_; - - // A map containing the count of unscheduled HLOs which using a particular - // LogicalBuffer. - std::unordered_map unscheduled_use_count_; - - // Set of instructions which have been scheduled. - std::unordered_set scheduled_instructions_; -}; - -int64 SumLogicalBufferSizes(const std::vector& buffers, - const LogicalBuffer::SizeFunction& size_function) { - int64 size = 0; - for (const LogicalBuffer* buffer : buffers) { - size += size_function(*buffer); - } - return size; -} - -StatusOr> RunDFSMemoryScheduler( - const HloComputation& computation, - const TuplePointsToAnalysis& points_to_analysis, - const LogicalBuffer::SizeFunction& size_function) { - // This ordering is based on DFS post-order, with a heuristic to decide which - // operand to visit first. The heuristic is based on 'extra_users', which is - // simply users-1 for each instruction. By subtracting 1, we're saying that - // instructions with no users or a single user don't count; instructions with - // lots of fan-out will be visited earlier. - tensorflow::gtl::FlatMap extra_users; - tensorflow::gtl::FlatMap total_sizes; - for (const HloInstruction* hlo : computation.MakeInstructionPostOrder()) { - extra_users[hlo] = hlo->users().empty() ? 0 : hlo->users().size() - 1; - total_sizes[hlo] = SumLogicalBufferSizes( - points_to_analysis.GetBuffersDefinedByInstruction(hlo), size_function); - tensorflow::gtl::FlatSet unique_operands( - hlo->operands().begin(), hlo->operands().end()); - for (const HloInstruction* operand : unique_operands) { - extra_users[hlo] += extra_users[operand]; - total_sizes[hlo] += total_sizes[operand]; - } - } - CHECK_EQ(extra_users.size(), computation.instructions().size()); - CHECK_EQ(total_sizes.size(), computation.instructions().size()); - - // Construct a total order based on DFS post-order, visiting operands in - // decreasing cumulative extra user order, and next by cumulative size, with a - // tiebreaker by name for determinism. - std::vector sequence; - FunctionVisitor visitor([&sequence](HloInstruction* hlo) { - sequence.push_back(hlo); - return Status::OK(); - }); - TF_RETURN_IF_ERROR(computation.AcceptWithOperandOrder( - &visitor, [&extra_users, &total_sizes](const HloInstruction* a, - const HloInstruction* b) { - if (extra_users[a] != extra_users[b]) { - return extra_users[a] > extra_users[b]; - } - if (total_sizes[a] != total_sizes[b]) { - return total_sizes[a] > total_sizes[b]; - } - return a->name() < b->name(); - })); - CHECK_EQ(sequence.size(), computation.instructions().size()); - return sequence; -} - -StatusOr> CreateMemoryMinimizingSequence( - const HloComputation& computation, - const TuplePointsToAnalysis& points_to_analysis, - const LogicalBuffer::SizeFunction& size_function) { - // We try both a list-scheduler based ordering and a DFS based ordering, and - // choose whichever returns a lower min-memory, not accounting for - // fragmentation. - TF_ASSIGN_OR_RETURN( - std::vector list_sequence, - ListScheduler::Run(computation, points_to_analysis, size_function)); - TF_ASSIGN_OR_RETURN( - const int64 list_memory, - MinimumMemoryForSequence(computation, list_sequence, points_to_analysis, - size_function)); - VLOG(2) << "Min-memory list sequence: " << list_memory << " bytes"; - - TF_ASSIGN_OR_RETURN( - std::vector dfs_sequence, - RunDFSMemoryScheduler(computation, points_to_analysis, size_function)); - TF_ASSIGN_OR_RETURN( - const int64 dfs_memory, - MinimumMemoryForSequence(computation, dfs_sequence, points_to_analysis, - size_function)); - VLOG(2) << "Min-memory dfs sequence: " << dfs_memory << " bytes"; - - if (list_memory <= dfs_memory) { - VLOG(2) << "Chose min-memory list sequence: " << list_memory << " bytes"; - return list_sequence; - } else { - VLOG(2) << "Chose min-memory dfs sequence: " << dfs_memory << " bytes"; - return dfs_sequence; - } -} - -} // namespace - -StatusOr -CreateMemoryMinimizingSequence( - const HloModule& module, const LogicalBuffer::SizeFunction& size_function) { - SequentialHloOrdering::HloModuleSequence sequence; - TF_ASSIGN_OR_RETURN(std::unique_ptr points_to_analysis, - TuplePointsToAnalysis::Run(&module)); - for (const auto& computation : module.computations()) { - TF_ASSIGN_OR_RETURN(sequence[computation.get()], - CreateMemoryMinimizingSequence( - *computation, *points_to_analysis, size_function)); - } - return sequence; -} - -StatusOr> CreateMemoryMinimizingSequence( - const HloComputation& computation, - const LogicalBuffer::SizeFunction& size_function) { - TF_ASSIGN_OR_RETURN(std::unique_ptr points_to_analysis, - TuplePointsToAnalysis::Run(computation.parent())); - return CreateMemoryMinimizingSequence(computation, *points_to_analysis, - size_function); -} - std::ostream& operator<<( std::ostream& out, const SequentialHloOrdering::HloModuleSequence& module_sequence) { diff --git a/tensorflow/compiler/xla/service/hlo_ordering.h b/tensorflow/compiler/xla/service/hlo_ordering.h index e964c4c51ae14f89d1f1b0450990cfc50c8a74be..efb5fca188a756b1fadda25f90defd94d8e3cb1c 100644 --- a/tensorflow/compiler/xla/service/hlo_ordering.h +++ b/tensorflow/compiler/xla/service/hlo_ordering.h @@ -20,49 +20,80 @@ limitations under the License. #include #include +#include "tensorflow/compiler/xla/service/call_graph.h" +#include "tensorflow/compiler/xla/service/hlo.pb.h" #include "tensorflow/compiler/xla/service/hlo_instruction.h" #include "tensorflow/compiler/xla/service/hlo_module.h" -#include "tensorflow/compiler/xla/service/tuple_points_to_analysis.h" -#include "tensorflow/compiler/xla/statusor.h" +#include "tensorflow/compiler/xla/service/hlo_value.h" #include "tensorflow/compiler/xla/types.h" -#include "tensorflow/core/lib/core/status.h" #include "tensorflow/core/lib/gtl/flatmap.h" -#include "tensorflow/core/lib/gtl/flatset.h" namespace xla { -// Abstract base class for describing a partial ordering of HLO -// instructions. Used to determine live range overlap of HLO instruction output -// buffers. +// Base class for describing a partial ordering of HLO instructions. Used to +// determine live range overlap of HLO instruction output buffers. class HloOrdering { public: - HloOrdering() = default; + HloOrdering(const HloModule* module) + : module_(module), call_graph_(CallGraph::Build(module)) {} virtual ~HloOrdering() = default; // Returns true if instruction 'a' executes before instruction 'b'. This is // not reflexive, that is, an instruction does not execute before itself. - virtual bool ExecutesBefore(const HloInstruction* a, - const HloInstruction* b) const = 0; + bool ExecutesBefore(const HloInstruction* a, const HloInstruction* b) const; + + // Returns whether the value 'a' is defined before the value 'b' under the + // given ordering. + bool IsDefinedBefore(const HloValue& a, const HloValue& b) const; + + // Returns whether the given use is before the given value definition under + // the given ordering. + bool UseIsBeforeValueDefinition(const HloUse& use, + const HloValue& value) const; + // Returns whether the given values interfere. Two values interfere if they + // may both be simultaneously live. + bool MayInterfere(const HloValue& a, const HloValue& b) const; + + // Returns true if the live range of the given value 'a' is strictly before + // the live range of value 'b' using the given HLO ordering. + bool LiveRangeStrictlyBefore(const HloValue& a, const HloValue& b) const; // 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( const HloComputation& computation) const = 0; + // Return the call graph of the module used to compute ordering. + const CallGraph& call_graph() const { return *call_graph_; } + virtual string ToString() const = 0; + + // Returns the serialized representation of this ordering. + // Only sequential computation orders are represented. + HloOrderingProto ToProto() const; + + protected: + // Returns true if instruction 'a' executes before instruction 'b'. + // Precondition: 'a' and 'b' are in the same computation. + // + // Derived classes should implement this method for determining order of + // instructions in the same comptuation. ExecutesBefore() analyzes the + // callgraph and uses this method to determine ordering of instructions in + // different computations. + virtual bool ExecutesBeforeInSameComputation( + const HloInstruction* a, const HloInstruction* b) const = 0; + + const HloModule* module_; + + std::unique_ptr call_graph_; }; -// Base class for partial orderings implemented by a map of strict predecessors -// for each instruction. Subclasses should fill in strict_predecessors_. +// Base class for partial orderings implemented by a map of predecessors for +// each instruction. Subclasses should fill in predecessors_. class PredecessorHloOrdering : public HloOrdering { public: ~PredecessorHloOrdering() override = default; - // Returns true if instruction 'a' executes before instruction 'b'. - // Instructions in different computations are not ordered. - bool ExecutesBefore(const HloInstruction* a, - const HloInstruction* b) const override; - // Returns nullptr indicating the computation does not have a sequential // ordering. const std::vector* SequentialOrder( @@ -70,20 +101,28 @@ class PredecessorHloOrdering : public HloOrdering { return nullptr; } + HloReachabilityMap& reachability_map(const HloComputation* computation) { + return *predecessors_.at(computation); + } + const HloReachabilityMap& reachability_map( + const HloComputation* computation) const { + return *predecessors_.at(computation); + } + protected: explicit PredecessorHloOrdering(const HloModule* module); string ToStringHelper(const string& name) const; - const HloModule* module_; + bool ExecutesBeforeInSameComputation(const HloInstruction* a, + const HloInstruction* b) const override; - // For each each computation in the module, this is the set of the - // instruction's strict predecessors. An instruction is not an element of its - // own strict predecessor set. + // For each computation in the module, this is the set of the instruction's + // predecessors. An instruction is an element of its own predecessor set. // // Subclasses should fill this in to define the desired ordering. tensorflow::gtl::FlatMap> - strict_predecessors_; + std::unique_ptr> + predecessors_; }; // An HLO ordering based on data dependencies in the HLO graph. In this partial @@ -150,12 +189,6 @@ class SequentialHloOrdering : public HloOrdering { const HloModuleSequence& module_sequence); ~SequentialHloOrdering() override = default; - // Instruction 'a' executes before 'b' if 'a' appears before 'b' in the - // instruction sequence for the computation. Instructions in different - // computations are unordered. - bool ExecutesBefore(const HloInstruction* a, - const HloInstruction* b) const override; - // Returns the sequential instruction order for the given computation. const std::vector* SequentialOrder( const HloComputation& computation) const override; @@ -163,7 +196,9 @@ class SequentialHloOrdering : public HloOrdering { string ToString() const override; protected: - const HloModule* module_; + bool ExecutesBeforeInSameComputation(const HloInstruction* a, + const HloInstruction* b) const override; + const HloModuleSequence module_sequence_; // The position of every instruction in the HLO module in its respective @@ -179,24 +214,6 @@ std::ostream& operator<<( std::ostream& out, const SequentialHloOrdering::HloModuleSequence& module_sequence); -// Returns the minimum memory required to compute the given module sequence, -// assuming no fragmentation. -StatusOr MinimumMemoryForSequence( - const SequentialHloOrdering::HloModuleSequence& module_sequence, - const LogicalBuffer::SizeFunction& size_function); - -// Returns an HloModuleSequence which seeks to minimize the memory required for -// the computation. size_function is the function returning the number of bytes -// required for a LogicalBuffer. -StatusOr -CreateMemoryMinimizingSequence( - const HloModule& module, const LogicalBuffer::SizeFunction& size_function); - -// Overload of above that computes the sequence for a single computation. -StatusOr> CreateMemoryMinimizingSequence( - const HloComputation& computation, - const LogicalBuffer::SizeFunction& size_function); - } // namespace xla #endif // TENSORFLOW_COMPILER_XLA_SERVICE_HLO_ORDERING_H_ diff --git a/tensorflow/compiler/xla/service/hlo_ordering_test.cc b/tensorflow/compiler/xla/service/hlo_ordering_test.cc index 425bee601a8d6357e21d3d00f8ccf5d69af03862..c95e44bd5d9d2ed87992d630bed4c1fe5c161383 100644 --- a/tensorflow/compiler/xla/service/hlo_ordering_test.cc +++ b/tensorflow/compiler/xla/service/hlo_ordering_test.cc @@ -19,8 +19,10 @@ limitations under the License. #include #include "tensorflow/compiler/xla/service/hlo_computation.h" +#include "tensorflow/compiler/xla/service/hlo_dataflow_analysis.h" #include "tensorflow/compiler/xla/service/hlo_instruction.h" #include "tensorflow/compiler/xla/service/hlo_opcode.h" +#include "tensorflow/compiler/xla/service/hlo_scheduling.h" #include "tensorflow/compiler/xla/shape_util.h" #include "tensorflow/compiler/xla/tests/hlo_test_base.h" #include "tensorflow/compiler/xla/types.h" @@ -58,26 +60,256 @@ TEST_F(HloOrderingTest, LastUseScheduledFirst) { auto sub = builder.AddInstruction( HloInstruction::CreateBinary(vec, HloOpcode::kSubtract, add, negate)); - HloModule module(TestName()); - module.AddEntryComputation(builder.Build()); + auto module = CreateNewModule(); + module->AddEntryComputation(builder.Build()); - TF_ASSIGN_OR_ASSERT_OK( + TF_ASSERT_OK_AND_ASSIGN( SequentialHloOrdering::HloModuleSequence sequence, - CreateMemoryMinimizingSequence(module, [](const LogicalBuffer& buffer) { + CreateMemoryMinimizingSequence(*module, [](const LogicalBuffer& buffer) { return ShapeUtil::ByteSizeOf(buffer.shape()); })); // Verify that all instructions are in the sequence. - EXPECT_EQ(module.entry_computation()->instruction_count(), - sequence.at(module.entry_computation()).size()); + EXPECT_EQ(module->entry_computation()->instruction_count(), + sequence.at(module->entry_computation()).size()); // The first instruction should be the parameter and the last the root "sub". - EXPECT_EQ(param, sequence.at(module.entry_computation()).front()); - EXPECT_EQ(sub, sequence.at(module.entry_computation()).back()); + EXPECT_EQ(param, sequence.at(module->entry_computation()).front()); + EXPECT_EQ(sub, sequence.at(module->entry_computation()).back()); - SequentialHloOrdering ordering(&module, sequence); + SequentialHloOrdering ordering(module.get(), sequence); EXPECT_TRUE(ordering.ExecutesBefore(add, negate)); } -} // namespace +TEST_F(HloOrderingTest, InstructionsInDifferentComputations) { + // Tests the ordering of instructions in different computations using the + // following HLO code: + // + // Entry computation: + // %x = Call(A, {}) + // %y = Call(B, {%x}) + // + // Computation A: + // %a = Call(C, {}) + // + // Computation B: + // %b = Call(C, {}) + // + // Computation C: + // %c = Constant(42.0f) + // + // This results in a diamond-shaped callgraph. + auto module = CreateNewModule(); + const Shape scalar_shape = ShapeUtil::MakeShape(xla::F32, {}); + + auto builder_c = HloComputation::Builder("C"); + HloInstruction* c = builder_c.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(42.0f))); + HloComputation* computation_c = + module->AddEmbeddedComputation(builder_c.Build()); + + auto builder_b = HloComputation::Builder("B"); + builder_b.AddInstruction( + HloInstruction::CreateParameter(0, scalar_shape, "param")); + HloInstruction* b = builder_b.AddInstruction( + HloInstruction::CreateCall(scalar_shape, {}, computation_c)); + HloComputation* computation_b = + module->AddEmbeddedComputation(builder_b.Build()); + + auto builder_a = HloComputation::Builder("A"); + HloInstruction* a = builder_a.AddInstruction( + HloInstruction::CreateCall(scalar_shape, {}, computation_c)); + HloComputation* computation_a = + module->AddEmbeddedComputation(builder_a.Build()); + + auto builder = HloComputation::Builder(TestName()); + HloInstruction* x = builder.AddInstruction( + HloInstruction::CreateCall(scalar_shape, {}, computation_a)); + HloInstruction* y = builder.AddInstruction( + HloInstruction::CreateCall(scalar_shape, {x}, computation_b)); + module->AddEntryComputation(builder.Build()); + + DependencyHloOrdering ordering(module.get()); + EXPECT_TRUE(ordering.ExecutesBefore(x, y)); + EXPECT_FALSE(ordering.ExecutesBefore(y, x)); + + EXPECT_TRUE(ordering.ExecutesBefore(a, b)); + EXPECT_FALSE(ordering.ExecutesBefore(b, a)); + + EXPECT_FALSE(ordering.ExecutesBefore(a, x)); + EXPECT_TRUE(ordering.ExecutesBefore(a, y)); + EXPECT_FALSE(ordering.ExecutesBefore(x, a)); + EXPECT_FALSE(ordering.ExecutesBefore(y, a)); + + EXPECT_FALSE(ordering.ExecutesBefore(b, x)); + EXPECT_FALSE(ordering.ExecutesBefore(b, y)); + EXPECT_TRUE(ordering.ExecutesBefore(x, b)); + EXPECT_FALSE(ordering.ExecutesBefore(y, b)); + + // Instruction 'c' is called from multiple callsites and should be unordered + // relative to all other instructions in the module. + EXPECT_FALSE(ordering.ExecutesBefore(c, a)); + EXPECT_FALSE(ordering.ExecutesBefore(c, b)); + EXPECT_FALSE(ordering.ExecutesBefore(c, x)); + EXPECT_FALSE(ordering.ExecutesBefore(c, y)); + EXPECT_FALSE(ordering.ExecutesBefore(a, c)); + EXPECT_FALSE(ordering.ExecutesBefore(b, c)); + EXPECT_FALSE(ordering.ExecutesBefore(x, c)); + EXPECT_FALSE(ordering.ExecutesBefore(y, c)); +} + +TEST_F(HloOrderingTest, InstructionsInWhileComputations) { + // Tests the ordering of instructions in the body and condition of a while + // instruction. HLO code: + // + // body(F32[]) %param): + // %negate = Negate(%param) + // + // condition(F32[] %param): + // %convert = Convert(%param) + // + // entry: + // %constant = Constant(1.0) + // return While(%constant, body, condition) + // + auto module = CreateNewModule(); + const Shape scalar_shape = ShapeUtil::MakeShape(xla::F32, {}); + + auto body_builder = HloComputation::Builder("body"); + auto body_param = body_builder.AddInstruction( + HloInstruction::CreateParameter(0, scalar_shape, "body_param")); + auto negate = body_builder.AddInstruction(HloInstruction::CreateUnary( + scalar_shape, HloOpcode::kNegate, body_param)); + HloComputation* body = module->AddEmbeddedComputation(body_builder.Build()); + + auto cond_builder = HloComputation::Builder("condition"); + auto cond_param = cond_builder.AddInstruction( + HloInstruction::CreateParameter(0, scalar_shape, "cond_param")); + auto convert = cond_builder.AddInstruction(HloInstruction::CreateConvert( + ShapeUtil::MakeShape(xla::PRED, {}), cond_param)); + HloComputation* condition = + module->AddEmbeddedComputation(cond_builder.Build()); + + auto builder = HloComputation::Builder(TestName()); + auto constant = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(1.0))); + auto xla_while = builder.AddInstruction( + HloInstruction::CreateWhile(scalar_shape, condition, body, constant)); + module->AddEntryComputation(builder.Build()); + + DependencyHloOrdering ordering(module.get()); + EXPECT_TRUE(ordering.ExecutesBefore(constant, xla_while)); + EXPECT_TRUE(ordering.ExecutesBefore(constant, cond_param)); + EXPECT_TRUE(ordering.ExecutesBefore(constant, convert)); + EXPECT_TRUE(ordering.ExecutesBefore(constant, body_param)); + EXPECT_TRUE(ordering.ExecutesBefore(constant, negate)); + + // The while should be unordered relative to the body and condition + // instructions. + EXPECT_FALSE(ordering.ExecutesBefore(xla_while, body_param)); + EXPECT_FALSE(ordering.ExecutesBefore(xla_while, cond_param)); + EXPECT_FALSE(ordering.ExecutesBefore(body_param, xla_while)); + EXPECT_FALSE(ordering.ExecutesBefore(cond_param, xla_while)); + + // Condition instructions should be ordered before body instructions. + EXPECT_TRUE(ordering.ExecutesBefore(cond_param, body_param)); + EXPECT_TRUE(ordering.ExecutesBefore(convert, body_param)); + EXPECT_TRUE(ordering.ExecutesBefore(cond_param, negate)); + EXPECT_TRUE(ordering.ExecutesBefore(convert, negate)); + + EXPECT_FALSE(ordering.ExecutesBefore(body_param, cond_param)); +} + +TEST_F(HloOrderingTest, ValuesInWhileComputations) { + // Tests the ordering of values (defined by dataflow analysis) in the body and + // condition of a while instruction. HLO code: + // + // body(F32[]) %param): + // %negate = Negate(%param) + // + // condition(F32[] %param): + // %convert = Convert(%param) + // + // entry: + // %constant = Constant(1.0) + // %while = While(%constant, body, condition) + // %add = Add(%constant, %while) + // + auto module = CreateNewModule(); + const Shape scalar_shape = ShapeUtil::MakeShape(xla::F32, {}); + + auto body_builder = HloComputation::Builder("body"); + auto body_param = body_builder.AddInstruction( + HloInstruction::CreateParameter(0, scalar_shape, "body_param")); + auto negate = body_builder.AddInstruction(HloInstruction::CreateUnary( + scalar_shape, HloOpcode::kNegate, body_param)); + HloComputation* body = module->AddEmbeddedComputation(body_builder.Build()); + + auto cond_builder = HloComputation::Builder("condition"); + auto cond_param = cond_builder.AddInstruction( + HloInstruction::CreateParameter(0, scalar_shape, "cond_param")); + auto convert = cond_builder.AddInstruction(HloInstruction::CreateConvert( + ShapeUtil::MakeShape(xla::PRED, {}), cond_param)); + HloComputation* condition = + module->AddEmbeddedComputation(cond_builder.Build()); + auto builder = HloComputation::Builder(TestName()); + auto constant = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(1.0))); + auto xla_while = builder.AddInstruction( + HloInstruction::CreateWhile(scalar_shape, condition, body, constant)); + auto add = builder.AddInstruction(HloInstruction::CreateBinary( + scalar_shape, HloOpcode::kAdd, constant, xla_while)); + module->AddEntryComputation(builder.Build()); + + TF_ASSERT_OK_AND_ASSIGN( + auto dataflow, HloDataflowAnalysis::Run(module.get(), /*ssa_form=*/true)); + DependencyHloOrdering ordering(module.get()); + + // Init value is defined before the while, but live range is not before the + // while because of the use of the init value in the add. + EXPECT_TRUE(ordering.IsDefinedBefore(dataflow->GetValueDefinedAt(constant), + dataflow->GetValueDefinedAt(xla_while))); + EXPECT_FALSE( + ordering.LiveRangeStrictlyBefore(dataflow->GetValueDefinedAt(constant), + dataflow->GetValueDefinedAt(xla_while))); + EXPECT_TRUE(ordering.MayInterfere(dataflow->GetValueDefinedAt(constant), + dataflow->GetValueDefinedAt(xla_while))); + + // Any value defined in the body or condition is defined before the while, and + // has a live range strictly before the while. + EXPECT_TRUE(ordering.IsDefinedBefore(dataflow->GetValueDefinedAt(negate), + dataflow->GetValueDefinedAt(xla_while))); + EXPECT_TRUE( + ordering.LiveRangeStrictlyBefore(dataflow->GetValueDefinedAt(negate), + dataflow->GetValueDefinedAt(xla_while))); + EXPECT_FALSE(ordering.MayInterfere(dataflow->GetValueDefinedAt(negate), + dataflow->GetValueDefinedAt(xla_while))); + + EXPECT_TRUE(ordering.IsDefinedBefore(dataflow->GetValueDefinedAt(convert), + dataflow->GetValueDefinedAt(xla_while))); + EXPECT_TRUE( + ordering.LiveRangeStrictlyBefore(dataflow->GetValueDefinedAt(convert), + dataflow->GetValueDefinedAt(xla_while))); + EXPECT_FALSE(ordering.MayInterfere(dataflow->GetValueDefinedAt(convert), + dataflow->GetValueDefinedAt(xla_while))); + + // The live range of the while should be before the add. + EXPECT_TRUE(ordering.IsDefinedBefore(dataflow->GetValueDefinedAt(xla_while), + dataflow->GetValueDefinedAt(add))); + ASSERT_EQ(dataflow->GetValueDefinedAt(xla_while).uses().size(), 1); + + const HloUse& while_use = dataflow->GetValueDefinedAt(xla_while).uses()[0]; + EXPECT_EQ(while_use.instruction, add); + EXPECT_TRUE(ordering.UseIsBeforeValueDefinition( + while_use, dataflow->GetValueDefinedAt(add))); + EXPECT_TRUE( + ordering.LiveRangeStrictlyBefore(dataflow->GetValueDefinedAt(xla_while), + dataflow->GetValueDefinedAt(add))); +} + +} // namespace } // namespace xla + +int main(int argc, char** argv) { + return xla::ParseDebugOptionsFlagsAndRunTests(argc, argv); +} diff --git a/tensorflow/compiler/xla/service/hlo_pass_pipeline.cc b/tensorflow/compiler/xla/service/hlo_pass_pipeline.cc index 6e3c983071245c548914bd9eecd0d1e86bc64d99..eb3da111a24df759370936daf902f3263030ff04 100644 --- a/tensorflow/compiler/xla/service/hlo_pass_pipeline.cc +++ b/tensorflow/compiler/xla/service/hlo_pass_pipeline.cc @@ -17,7 +17,7 @@ limitations under the License. #include -#include "tensorflow/compiler/xla/legacy_flags/hlo_pass_pipeline_flags.h" +#include "tensorflow/compiler/xla/service/hlo_graph_dumper.h" #include "tensorflow/compiler/xla/status_macros.h" #include "tensorflow/compiler/xla/types.h" #include "tensorflow/compiler/xla/util.h" @@ -31,23 +31,32 @@ using ::tensorflow::strings::StrAppend; namespace xla { namespace { -void DumpModule(const Compiler::HloDumper& dumper_, const HloModule& module, +void DumpModule(const HloModule& module, + const string& message) { - dumper_(module, message); + hlo_graph_dumper::MaybeDumpHloModule(module, message); VLOG(2) << "HLO " << message << ":"; XLA_VLOG_LINES(2, module.ToString()); } } // namespace StatusOr HloPassPipeline::Run(HloModule* module) { - legacy_flags::HloPassPipelineFlags* flags = - legacy_flags::GetHloPassPipelineFlags(); - std::vector tmp = - tensorflow::str_util::Split(flags->xla_disable_hlo_passes, ','); - tensorflow::gtl::FlatSet disabled_passes(tmp.begin(), tmp.end()); + run_called_ = true; + + VLOG(1) << "Running HLO pass pipeline " << name(); + + auto repeated_field = + module->config().debug_options().xla_disable_hlo_passes(); + tensorflow::gtl::FlatSet disabled_passes(repeated_field.begin(), + repeated_field.end()); + if (!disabled_passes.empty()) { + VLOG(1) << "Passes disabled by --xla_disable_hlo_passes: " + << tensorflow::str_util::Join(disabled_passes, ", "); + } auto run_invariant_checkers = [this, module]() -> Status { for (auto& invariant_checker : invariant_checkers_) { + VLOG(1) << " Invariant checker " << invariant_checker->name(); TF_ASSIGN_OR_RETURN(bool changed, invariant_checker->Run(module)); TF_RET_CHECK(!changed) << "invariant checkers must not change the graph"; } @@ -58,15 +67,18 @@ StatusOr HloPassPipeline::Run(HloModule* module) { bool changed = false; string message; for (auto& pass : passes_) { - if (!disabled_passes.empty() && - disabled_passes.count(pass->name().ToString()) > 0) { + if (disabled_passes.count(pass->name().ToString()) > 0) { + VLOG(1) << " Skipping HLO pass " << pass->name() + << ", disabled by --xla_disable_hlo_passes"; continue; } + VLOG(1) << " HLO pass " << pass->name(); + // Emit label containing: "after foo-pass, before bar-pass". message.clear(); StrAppend(&message, prefix, ", before ", pass->name()); - DumpModule(dumper_, *module, message); + DumpModule(*module, message); TF_RETURN_IF_ERROR(run_invariant_checkers()); TF_ASSIGN_OR_RETURN(bool changed_this_pass, pass->Run(module)); @@ -76,7 +88,7 @@ StatusOr HloPassPipeline::Run(HloModule* module) { StrAppend(&prefix, name(), ": after ", pass->name()); } TF_RETURN_IF_ERROR(run_invariant_checkers()); - DumpModule(dumper_, *module, prefix + ", pipeline end"); + DumpModule(*module, prefix + ", pipeline end"); return changed; } diff --git a/tensorflow/compiler/xla/service/hlo_pass_pipeline.h b/tensorflow/compiler/xla/service/hlo_pass_pipeline.h index a8c2d518730b9fab8febaae35797ea4a315ab9b1..a42d7e59fed2d838dfe3cb7f99e6b946edfdb0b4 100644 --- a/tensorflow/compiler/xla/service/hlo_pass_pipeline.h +++ b/tensorflow/compiler/xla/service/hlo_pass_pipeline.h @@ -22,7 +22,6 @@ limitations under the License. #include #include "tensorflow/compiler/xla/ptr_util.h" -#include "tensorflow/compiler/xla/service/compiler.h" #include "tensorflow/compiler/xla/service/hlo_module.h" #include "tensorflow/compiler/xla/service/hlo_pass_interface.h" #include "tensorflow/compiler/xla/statusor.h" @@ -34,9 +33,7 @@ namespace xla { // Pipeline of HLO passes. class HloPassPipeline : public HloPassInterface { public: - explicit HloPassPipeline(const string& name, - const Compiler::HloDumper& dumper) - : name_(name), dumper_(dumper) {} + explicit HloPassPipeline(const string& name) : name_(name) {} tensorflow::StringPiece name() const override { return name_; } // Add a pass to the pipeline. It should be called with the arguments for the @@ -47,6 +44,7 @@ class HloPassPipeline : public HloPassInterface { // Returns a reference to the added pass. template T& AddPass(Args&&... args) { + CHECK(!run_called_) << "AddPass cannot be called after Run"; auto pass = new T(std::forward(args)...); passes_.push_back(std::unique_ptr(pass)); return *pass; @@ -57,6 +55,7 @@ class HloPassPipeline : public HloPassInterface { // (it is required to always return "false" from its Run() method). template T& AddInvariantChecker(Args&&... args) { + CHECK(!run_called_) << "AddInvariantChecker cannot be called after Run"; auto pass = new T(std::forward(args)...); invariant_checkers_.push_back(std::unique_ptr(pass)); return *pass; @@ -67,9 +66,9 @@ class HloPassPipeline : public HloPassInterface { private: const string name_; - Compiler::HloDumper dumper_; std::vector> passes_; std::vector> invariant_checkers_; + bool run_called_ = false; TF_DISALLOW_COPY_AND_ASSIGN(HloPassPipeline); }; diff --git a/tensorflow/compiler/xla/service/hlo_proto_util.cc b/tensorflow/compiler/xla/service/hlo_proto_util.cc new file mode 100644 index 0000000000000000000000000000000000000000..727ad0178c6227cd2e64c31a4618e781671b9393 --- /dev/null +++ b/tensorflow/compiler/xla/service/hlo_proto_util.cc @@ -0,0 +1,33 @@ +/* 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/hlo_proto_util.h" + +namespace xla { + +HloProto MakeHloProto(const HloModule& module, + const BufferAssignment& assignment) { + HloModuleProto proto_module = module.ToProto(); + HloOrderingProto proto_ordering = + assignment.liveness().hlo_ordering().ToProto(); + BufferAssignmentProto proto_assignment = assignment.ToProto(); + HloProto proto; + proto.mutable_hlo_module()->Swap(&proto_module); + proto.mutable_hlo_ordering()->Swap(&proto_ordering); + proto.mutable_buffer_assignment()->Swap(&proto_assignment); + return proto; +} + +} // namespace xla diff --git a/tensorflow/compiler/xla/service/hlo_proto_util.h b/tensorflow/compiler/xla/service/hlo_proto_util.h new file mode 100644 index 0000000000000000000000000000000000000000..603259a11fcdca59f58653d9a7a164c983711a57 --- /dev/null +++ b/tensorflow/compiler/xla/service/hlo_proto_util.h @@ -0,0 +1,36 @@ +/* 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. +==============================================================================*/ + +// Utilities to manipulate data in hlo.proto. + +#ifndef TENSORFLOW_COMPILER_XLA_SERVICE_HLO_PROTO_UTIL_H_ +#define TENSORFLOW_COMPILER_XLA_SERVICE_HLO_PROTO_UTIL_H_ + +#include + +#include "tensorflow/compiler/xla/service/buffer_assignment.h" +#include "tensorflow/compiler/xla/service/hlo.pb.h" +#include "tensorflow/compiler/xla/service/hlo_module.h" +#include "tensorflow/compiler/xla/status.h" + +namespace xla { + +// Returns a serialized representation of the HLO state. +HloProto MakeHloProto(const HloModule& module, + const BufferAssignment& assignment); + +} // namespace xla + +#endif // TENSORFLOW_COMPILER_XLA_SERVICE_HLO_PROTO_UTIL_H_ diff --git a/tensorflow/compiler/xla/service/hlo_query.cc b/tensorflow/compiler/xla/service/hlo_query.cc index 1556d1772f934ea02506aff27396034814d61698..d45038f1f4a2e4aa19234eec93fdc9a068a902e1 100644 --- a/tensorflow/compiler/xla/service/hlo_query.cc +++ b/tensorflow/compiler/xla/service/hlo_query.cc @@ -25,13 +25,23 @@ namespace hlo_query { bool IsConstantR0F32(HloInstruction* instruction, float* out) { if (instruction->opcode() == HloOpcode::kConstant && ShapeUtil::IsScalarF32(instruction->shape())) { - *out = LiteralUtil::Get(instruction->literal(), {}); + *out = instruction->literal().Get({}); return true; } return false; } +bool AllOperandsAreParametersOrConstants(const HloInstruction& instruction) { + for (const auto& operand : instruction.operands()) { + if (operand->opcode() != HloOpcode::kParameter && + operand->opcode() != HloOpcode::kConstant) { + return false; + } + } + return true; +} + bool AllOperandsAreParameters(const HloInstruction& instruction) { for (const auto& operand : instruction.operands()) { if (operand->opcode() != HloOpcode::kParameter) { @@ -41,6 +51,15 @@ bool AllOperandsAreParameters(const HloInstruction& instruction) { return true; } +bool AllOperandsAreConstants(const HloInstruction& instruction) { + for (const auto& operand : instruction.operands()) { + if (operand->opcode() != HloOpcode::kConstant) { + return false; + } + } + return true; +} + HloInstruction* GetMatchingOperand( std::function matcher, HloInstruction* instruction) { diff --git a/tensorflow/compiler/xla/service/hlo_query.h b/tensorflow/compiler/xla/service/hlo_query.h index 864f892e92047e6f39b2949854190522b2f4a906..c79347bbf9d6146943b7b787f713369cb37fadee 100644 --- a/tensorflow/compiler/xla/service/hlo_query.h +++ b/tensorflow/compiler/xla/service/hlo_query.h @@ -28,9 +28,16 @@ namespace hlo_query { // Precondition: out != nullptr bool IsConstantR0F32(HloInstruction* instruction, float* out); +// Returns whether all of an instruction's operands are of the types constants +// and parameters. +bool AllOperandsAreParametersOrConstants(const HloInstruction& instruction); + // Returns whether all of an instruction's operands are parameters. bool AllOperandsAreParameters(const HloInstruction& instruction); +// Returns whether all of an instruction's operands are constants. +bool AllOperandsAreConstants(const HloInstruction& instruction); + // Returns whether the instruction is a scalar constant. bool IsScalarConstant(const HloInstruction* instruction); diff --git a/tensorflow/compiler/xla/service/hlo_reachability.cc b/tensorflow/compiler/xla/service/hlo_reachability.cc new file mode 100644 index 0000000000000000000000000000000000000000..8e167633bb13476301fa0c4afa0b123c9b47e40d --- /dev/null +++ b/tensorflow/compiler/xla/service/hlo_reachability.cc @@ -0,0 +1,64 @@ +/* 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/hlo_reachability.h" + +namespace xla { + +HloReachabilityMap::HloReachabilityMap( + const std::list& instructions) + : size_(instructions.size()) { + bit_vectors_.reserve(size_); + for (const HloInstruction* hlo : instructions) { + indices_[hlo] = bit_vectors_.size(); + bit_vectors_.emplace_back(size_); + } + CHECK_EQ(size_, indices_.size()); // instructions should be unique +} + +bool HloReachabilityMap::SetReachabilityToUnion( + tensorflow::gtl::ArraySlice inputs, + const HloInstruction* instruction) { + BitVector& bit_vector = GetBitVector(instruction); + tmp_bit_vector_ = bit_vector; + + // If instruction is part of inputs, don't reset the bit_vector. + if (std::find(inputs.begin(), inputs.end(), instruction) == inputs.end()) { + bit_vector.SetToZero(); + } + bit_vector.Set(GetIndex(instruction)); + for (const HloInstruction* input : inputs) { + bit_vector.OrWith(GetBitVector(input)); + } + + return bit_vector != tmp_bit_vector_; +} + +void HloReachabilityMap::SetReachable(const HloInstruction* a, + const HloInstruction* b) { + GetBitVector(b).Set(GetIndex(a)); +} + +bool HloReachabilityMap::IsReachable(const HloInstruction* a, + const HloInstruction* b) const { + return GetBitVector(b).Get(GetIndex(a)); +} + +bool HloReachabilityMap::IsConnected(const HloInstruction* a, + const HloInstruction* b) const { + return IsReachable(a, b) || IsReachable(b, a); +} + +} // namespace xla diff --git a/tensorflow/compiler/xla/service/hlo_reachability.h b/tensorflow/compiler/xla/service/hlo_reachability.h new file mode 100644 index 0000000000000000000000000000000000000000..d7bdac9c86579f19afbba133772c2c50894853d1 --- /dev/null +++ b/tensorflow/compiler/xla/service/hlo_reachability.h @@ -0,0 +1,138 @@ +/* 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_REACHABILITY_H_ +#define TENSORFLOW_COMPILER_XLA_SERVICE_HLO_REACHABILITY_H_ + +#include +#include + +#include "tensorflow/compiler/xla/map_util.h" +#include "tensorflow/compiler/xla/types.h" +#include "tensorflow/core/lib/core/status.h" +#include "tensorflow/core/lib/gtl/array_slice.h" +#include "tensorflow/core/lib/gtl/flatmap.h" +#include "tensorflow/core/platform/types.h" + +namespace xla { + +class HloInstruction; + +// A class for computing and representing reachability between HloInstructions. +class HloReachabilityMap { + public: + // Sets up an empty reachable matrix for the full set of instructions + // specified in 'instructions'. + explicit HloReachabilityMap(const std::list& instructions); + + // 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 + // for some 'input' in 'inputs'. Also sets 'instruction' to be reachable from + // itself. Returns whether the reachability set of 'instruction' changed. + bool SetReachabilityToUnion( + tensorflow::gtl::ArraySlice inputs, + const HloInstruction* instruction); + + // Sets entry so that IsReachable(a, b) will return true + void SetReachable(const HloInstruction* a, const HloInstruction* b); + + // Returns true if "b" is reachable from "a" + bool IsReachable(const HloInstruction* a, const HloInstruction* b) const; + + // Returns true if "b" is reachable from "a" or "a" is reachable from "b" + bool IsConnected(const HloInstruction* a, const HloInstruction* b) const; + + private: + // A bit-vector implementation specialized for this use case which provides a + // fast bitwise OR operation not available in tensorflow::gtl::BitMap. + class BitVector { + public: + BitVector() = default; + BitVector(size_t size) + : size_(size), vector_((size + kBits - 1) / kBits, 0) {} + + // Return the bit at the given index. + bool Get(size_t index) const { + DCHECK(index >= 0 && index < size_); + return vector_[index / kBits] & (1ull << (index % kBits)); + } + + // Set the bit at the given index. + void Set(size_t index) { + DCHECK(index >= 0 && index < size_); + vector_[index / kBits] |= 1ull << (index % kBits); + } + + // Set this bitvector to the Logical OR of this bitvector and 'other'. + void OrWith(const BitVector& other) { + for (size_t i = 0; i < vector_.size(); ++i) { + vector_[i] |= other.vector_[i]; + } + } + + // Set the bitvector to all zeros. + void SetToZero() { std::fill(vector_.begin(), vector_.end(), 0); } + + bool operator==(const BitVector& other) const { + return vector_ == other.vector_; + } + bool operator!=(const BitVector& other) const { + return vector_ != other.vector_; + } + + private: + using Word = uint64; + static const size_t kBits = 64; + + // Number of bits in the bitvector. + size_t size_; + + std::vector vector_; + }; + + // Return the bitvector storing the reachability-to of the given instruction. + const BitVector& GetBitVector(const HloInstruction* instruction) const { + return bit_vectors_[GetIndex(instruction)]; + } + BitVector& GetBitVector(const HloInstruction* instruction) { + return bit_vectors_[GetIndex(instruction)]; + } + + // 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); + } + + // 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. + tensorflow::gtl::FlatMap indices_; + + // Bitvectors holding the reachability to each instruction. The bit vector for + // instruction X includes ones for each instruction which X is reachable from. + std::vector bit_vectors_; + + // A temporary used by SetReachabilityToUnion to avoid an allocation with each + // call to the method. + BitVector tmp_bit_vector_; +}; + +} // namespace xla + +#endif // TENSORFLOW_COMPILER_XLA_SERVICE_HLO_REACHABILITY_H_ diff --git a/tensorflow/compiler/xla/service/hlo_reachability_test.cc b/tensorflow/compiler/xla/service/hlo_reachability_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..657a9ee83d29e72b95660325f9139f44159d6508 --- /dev/null +++ b/tensorflow/compiler/xla/service/hlo_reachability_test.cc @@ -0,0 +1,86 @@ +/* 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/hlo_reachability.h" + +#include + +#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_test_base.h" + +namespace xla { + +namespace { + +class HloReachabilityTest : public HloTestBase {}; + +TEST_F(HloReachabilityTest, Reachability) { + // Construct and test a reachability graph of the following form: + /* + a + / \ + b c + \ / \ + d e + */ + auto builder = HloComputation::Builder(TestName()); + auto a = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(0.0f))); + auto b = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(0.0f))); + auto c = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(0.0f))); + auto d = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(0.0f))); + auto e = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(0.0f))); + builder.Build(); + + HloReachabilityMap reachability({a, b, c, d, e}); + reachability.SetReachable(a, a); + EXPECT_TRUE(reachability.SetReachabilityToUnion({a}, b)); + EXPECT_TRUE(reachability.SetReachabilityToUnion({a}, c)); + EXPECT_TRUE(reachability.SetReachabilityToUnion({b, c}, d)); + EXPECT_TRUE(reachability.SetReachabilityToUnion({c}, e)); + + EXPECT_TRUE(reachability.IsReachable(a, a)); + EXPECT_TRUE(reachability.IsReachable(a, b)); + EXPECT_TRUE(reachability.IsReachable(a, c)); + EXPECT_TRUE(reachability.IsReachable(a, d)); + EXPECT_TRUE(reachability.IsReachable(a, e)); + + EXPECT_FALSE(reachability.IsReachable(b, a)); + EXPECT_TRUE(reachability.IsReachable(b, b)); + EXPECT_FALSE(reachability.IsReachable(b, c)); + EXPECT_TRUE(reachability.IsReachable(b, d)); + EXPECT_FALSE(reachability.IsReachable(b, e)); + + EXPECT_FALSE(reachability.IsReachable(e, a)); + EXPECT_FALSE(reachability.IsReachable(e, b)); + EXPECT_FALSE(reachability.IsReachable(e, c)); + EXPECT_FALSE(reachability.IsReachable(e, d)); + EXPECT_TRUE(reachability.IsReachable(e, e)); + + // Recomputing the same reachability for a previously computed instruction + // should return false (no change). + EXPECT_FALSE(reachability.SetReachabilityToUnion({a}, b)); + EXPECT_FALSE(reachability.SetReachabilityToUnion({b, c}, d)); +} + +} // namespace + +} // namespace xla diff --git a/tensorflow/compiler/xla/service/hlo_rematerialization.cc b/tensorflow/compiler/xla/service/hlo_rematerialization.cc index 52a0181029ddb7eb373bb9e9f91e2899c3140c71..20152cf0cefa4abfbdacdd26744890b12ddce6ac 100644 --- a/tensorflow/compiler/xla/service/hlo_rematerialization.cc +++ b/tensorflow/compiler/xla/service/hlo_rematerialization.cc @@ -22,14 +22,16 @@ limitations under the License. #include "tensorflow/compiler/xla/map_util.h" #include "tensorflow/compiler/xla/primitive_util.h" +#include "tensorflow/compiler/xla/service/flatten_call_graph.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_module.h" #include "tensorflow/compiler/xla/service/hlo_opcode.h" #include "tensorflow/compiler/xla/service/hlo_ordering.h" +#include "tensorflow/compiler/xla/service/hlo_scheduling.h" +#include "tensorflow/compiler/xla/service/liveness_util.h" #include "tensorflow/compiler/xla/service/logical_buffer.h" -#include "tensorflow/compiler/xla/service/tuple_points_to_analysis.h" #include "tensorflow/compiler/xla/status_macros.h" #include "tensorflow/compiler/xla/statusor.h" #include "tensorflow/compiler/xla/types.h" @@ -45,416 +47,815 @@ namespace xla { namespace { -// Returns a vector of the operands of 'instruction' with repeated elements -// removed. -std::vector UniqueOperands(const HloInstruction* instruction) { - std::vector unique_operands; - for (HloInstruction* operand : instruction->operands()) { - if (std::find(unique_operands.begin(), unique_operands.end(), operand) == - unique_operands.end()) { - unique_operands.push_back(operand); - } - } - return unique_operands; -} +// Potential optimizations: +// . TODO(b/35244891): Avoid N^2 behavior by keeping a priority queue +// of candidates. +// . Cache IsRematerializable in Item? Only correct if control +// predecessors and successors don't change. // Returns true if the given instruction is rematerializable. bool IsRematerializable(const HloInstruction* instruction) { - // Don't rematerialize instructions with side effects, those with a cost that - // might not be captured by HloCostAnalysis, or instructions which cannot be - // cloned safely. + // 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::kCrossReplicaSum: case HloOpcode::kCustomCall: case HloOpcode::kOutfeed: case HloOpcode::kInfeed: + case HloOpcode::kParameter: case HloOpcode::kRecv: case HloOpcode::kSend: case HloOpcode::kTrace: case HloOpcode::kWhile: return false; default: - break; + return true; } +} - // Skip tuple shapes because we do not currently account for buffer aliasing - // properly which results in improperly accounting of rematerialization cost - // for these shapes. - if (ShapeUtil::IsTuple(instruction->shape())) { - return false; - } - for (auto* operand : instruction->operands()) { - if (ShapeUtil::IsTuple(operand->shape())) { - return false; - } - } +// Type holding a unique identifier for each Buffer object. +using BufferId = int64; +using BufferIdList = tensorflow::gtl::InlinedVector; - return true; -} +// We wrap HloInstruction* with an Item that holds auxiliary +// per-instruction state. +struct Item { + HloInstruction* instruction; + + // True once the instruction is marked as placed (when BeginInstruction + // has been called for this instruction). + bool placed = false; + + // To avoid an infinite loop rematerializing the same set of + // instructions ad infinitum, keep a blacklist of instructions + // which should not be rematerialized. + bool blacklisted = false; -// Class which maintains an ordered list of instructions with fast insertion and -// removal of arbitrary elements. + // The buffers defined by this instruction. + BufferIdList buffers_defined; + + // The buffers used by this instruction. + BufferIdList buffers_used; + + private: + friend class InstructionList; + + // Items are arranged in a doubly linked list. + Item* next; + Item* prev; + + // List is ordered by position, which can however be duplicated as + // new instructions are inserted. See InsertBeforeInstructions + // comment for details. + int64 position; +}; + +using ItemList = tensorflow::gtl::InlinedVector; + +// Class which maintains an ordered list of instructions with fast insertion +// before arbitrary elements. class InstructionList { public: - explicit InstructionList(const std::vector order) { + explicit InstructionList(const std::vector& order) { + int64 position = 0; + Item* last = nullptr; for (const HloInstruction* inst : order) { - instructions_.push_back(const_cast(inst)); - instruction_iterators_.insert({const_cast(inst), - std::next(instructions_.end(), -1)}); + // Add a new item to the linked list. + Item* item = new Item; + item->next = nullptr; + item->prev = last; + if (last == nullptr) { + first_ = item; + } else { + last->next = item; + } + last = item; + + // Initially position numbers are uniquely assigned in order. Later as + // instructions are added with InsertBefore* methods, some instructions + // may have duplicate position numbers, but the values will be guaranteed + // 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->position = position; + position++; + + item_map_[inst] = item; + } + } + + ~InstructionList() { + for (Item* item = first_; item != nullptr;) { + Item* next = item->next; + delete item; + item = next; } } - // Returns the list of instructions. - const std::list& instructions() const { - return instructions_; + size_t size() const { return item_map_.size(); } + + // For ordered iteration over items. + // for (auto item = q.first(); item != nullptr; item = q.next(item)) {...} + Item* first() const { return first_; } + Item* next(Item* item) const { return item->next; } + + // Creates an Item for the given instruction, but doesn't add it to the list. + // (Use InsertBeforeInstructions to add the Item to the list.) + Item* CreateItem(HloInstruction* inst) { + Item* item = new Item; + item->instruction = inst; + CHECK(item_map_.insert({inst, item}).second) << "inserting inst twice"; + return item; } - // Insert instruction 'to_insert' before instruction 'before' in the list. - Status InsertBefore(HloInstruction* to_insert, HloInstruction* before) { - auto it = instruction_iterators_.find(before); - TF_RET_CHECK(it != instruction_iterators_.end()); - instruction_iterators_.insert( - {to_insert, instructions_.insert(it->second, to_insert)}); - return Status::OK(); + // Return the Item corresponding to inst. + Item* GetItem(const HloInstruction* inst) const { + auto iter = item_map_.find(inst); + CHECK(iter != item_map_.end()) << "Did not find " << inst->name(); + return iter->second; } - // Removes instruction from the list. - Status Remove(HloInstruction* instruction) { - auto it = instruction_iterators_.find(instruction); - TF_RET_CHECK(it != instruction_iterators_.end()); - instructions_.erase(it->second); - instruction_iterators_.erase(it); - return Status::OK(); + // Insert instruction 'to_insert' immediately before the earliest instruction + // in 'before_instructions'. + // + // Each instruction gets a non-decreasing ordinal number. We use this to let + // InsertBeforeInstructions quickly insert an instruction before the earliest + // instruction in a set of instructions. If position_number_[a] < + // position_number_[b] then 'a' comes before 'b' in the list. If the position + // numbers are the same then nothing can be said about their order without + // examining the list. + // + // On object construction this ordinal is precisely the instruction's index + // in the list. Later, instructions inserted via InsertBefore receive + // duplicate values. However, monotonicity is preserved. + void InsertBeforeInstructions( + Item* to_insert, tensorflow::gtl::ArraySlice before_instructions) { + VLOG(3) << "InsertBeforeInstructions: " << to_insert->instruction->name() + << " before {" + << tensorflow::str_util::Join(before_instructions, ", ", + [](string* out, Item* item) { + tensorflow::strings::StrAppend( + out, item->instruction->name()); + }) + << "}"; + + // Find the minimal position number of any instruction in + // 'before_instructions'. + CHECK(!before_instructions.empty()); + Item* min_position_item = nullptr; + for (Item* item : before_instructions) { + if (min_position_item == nullptr || + item->position < min_position_item->position) { + min_position_item = item; + } + } + + // Because more than one instruction in 'before_instructions' may have a + // position number of 'min_position_number', find the first such instruction + // with position number 'min_position_number'. + + // First find first instruction with the min position. + while (min_position_item->prev != nullptr && + min_position_item->position == min_position_item->prev->position) { + min_position_item = min_position_item->prev; + } + + // 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()) { + min_position_item = min_position_item->next; + } + return InsertBefore(to_insert, min_position_item); + } + + void Blacklist(const HloInstruction* inst) { + GetItem(inst)->blacklisted = true; } private: - // List of instructions. - std::list instructions_; + // Insert instruction 'item' immediately before 'before' in the list. + void InsertBefore(Item* item, Item* before) { + VLOG(3) << "InsertBefore: " << item->instruction->name() << " before " + << before->instruction->name(); + // Insert new item into linked list. + item->prev = before->prev; + item->next = before; + before->prev = item; + if (item->prev != nullptr) { + item->prev->next = item; + } else { + first_ = item; + } - // Iterators for each instruction in the list. - tensorflow::gtl::FlatMap::iterator> - instruction_iterators_; + // Assign the same position number to the newly added instruction as + // 'before'. This guarantees monotonicity of the position numbers, but not + // uniqueness. + item->position = before->position; + } + + Item* first_; + + // Item for each instruction. + tensorflow::gtl::FlatMap item_map_; }; +// Return the items which use the given LogicalBuffer. Sets +// has_indirect_users to whether any of the uses is indirect. A use is indirect +// if the instruction defining logical_buffer is not an operand of the use. This +// can happen via buffer aliasing (eg, tuples). +ItemList GetUsers(const InstructionList& instruction_list, + const LogicalBuffer* logical_buffer, + const TuplePointsToAnalysis& points_to_analysis, + bool* has_indirect_users) { + ItemList users; + // To identify uses iterate through all HloInstruction users of the + // BufferAliases of the logical buffer. + *has_indirect_users = false; + for (const BufferAlias& buffer_alias : + points_to_analysis.GetBufferAliases(*logical_buffer)) { + for (const HloInstruction* user : buffer_alias.instruction()->users()) { + if (DoesNotUseOperandBuffer(buffer_alias.instruction(), + buffer_alias.index(), user, + points_to_analysis)) { + // The alias may be an operand of 'user', but the LogicalBuffer cannot + // possibly be used by the instruction so ignore 'user'. This is the + // case, for example, for the tuple element buffers in a GetTupleElement + // instruction (the GTE instruction only uses the pointer vector). + continue; + } + if (buffer_alias.instruction() != logical_buffer->instruction()) { + *has_indirect_users = true; + } + // 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()) { + users.push_back(user_item); + } + } + } + return users; +} + // Class for tracking memory usage of a computation as the instructions are -// placed sequentially. Memory usage is the sum of live values at the current -// point in the instruction sequence. +// placed sequentially. Memory usage is the sum of the sizes of live values +// (LogicalBuffers) at the current point in the instruction sequence. class MemoryUsageTracker { public: MemoryUsageTracker( const HloComputation* computation, - const HloRematerialization::ShapeSizeFunction& size_function) - : computation_(computation), size_function_(size_function) { - for (const std::unique_ptr& instruction : - computation->instructions()) { - // Initially only live-in values occupy memory. - if (IsLiveIn(instruction.get())) { - memory_usage_ += TotalSizeBytes(instruction->shape()); - } - } + const HloRematerialization::ShapeSizeFunction& size_function, + const TuplePointsToAnalysis& points_to_analysis, + const InstructionList& instruction_list); + + // Starts the placement of the given instruction. This adds the sizes of the + // LogicalBuffers defined by the instruction to the current memory + // usage. Placement is broken into two steps (BeginInstruction and + // EndInstruction) to accurately model memory usage. At BeginInstruction the + // memory for the output value(s) of the current instruction is allocated. At + // EndInstruction memory for dead operand(s) is freed. + Status BeginInstruction(Item* item); + + // Finishes the placement of the current instruction. This frees any dead + // operands or dead result of the instruction. This must be called after + // each call to BeginInstruction. + Status EndInstruction(); + + // Returns the number of bytes that the current memory usage will be reduced + // if the given instruction is rematerialized. + int64 MemoryReducedIfRematerialized(Item* item) const; + + // Adjusts memory usage to account for the rematerialization of + // original_item for all remaining unplaced uses. The rematerialization + // is remat_item. This method should be called after the HLO graph has + // been transformed (rematerialization instruction created and connected to + // uses). + Status AddRematerializedInstruction(Item* original_item, Item* remat_item); + + // Returns whether the given instruction has been placed (BeginInstruction + // has been called with 'instruction' as the argument). + bool IsPlaced(const HloInstruction* instruction) const { + return instruction_list_.GetItem(instruction)->placed; } - // Starts the placement of the given instruction. This adds the output size of - // the instruction to the current memory usage. Placement is broken into two - // steps (BeginInstruction and EndInstruction) to accurately model memory - // usage. At BeginInstruction the memory for the output value of the current - // instruction is allocated. At EndInstruction memory for dead operands is - // freed. - Status BeginInstruction(const HloInstruction* instruction) { - VLOG(3) << "BeginInstruction " << instruction->name(); - TF_RET_CHECK(in_progress_instruction_ == nullptr); - in_progress_instruction_ = instruction; + // Returns the current memory usage. This is the sum of sizes of all live + // values. + int64 memory_usage() const { return memory_usage_; } + + // Check invariants of the data structure. This is expensive to call. + bool Check() const; - // Add instruction to remaining_uses_. - TF_RET_CHECK(!ContainsKey(remaining_uses_, instruction)); - std::vector& instruction_uses = - remaining_uses_[instruction]; - instruction_uses.insert(instruction_uses.begin(), - instruction->users().begin(), - instruction->users().end()); + string ToString() const; - if (!IsLiveIn(instruction)) { - // Instruction was not previously live so add output size to memory usage. - memory_usage_ += TotalSizeBytes(instruction->shape()); + private: + // A Buffer represents a single LogicalBuffer in the computation including + // various metadata useful for tracking liveness of the value. A LogicalBuffer + // is not used directly because the HLO graph is transformed and + // TuplePointsToAnalysis which owns all LogicalBuffers cannot be updated after + // HLO graph transformations. + struct Buffer { + // The unique id of this Buffer. This value is equal to the buffer's index + // in the vector buffers_. + const BufferId id; + + // The instruction which defines this buffer. + Item* defining_instruction; + + // The materialized size of the buffer in bytes. + const int64 size; + + // Whether this buffer is live-out of the computation. + bool live_out; + + // Whether this buffer has indirect uses. Ie, an instruction which is not a + // user of defining_instruction uses this buffer. This can occur due to + // buffer aliasing (eg, tuples). + bool has_indirect_uses; + + // The instructions which use this buffer. + ItemList users; + + // The number of users (HloInstructions) of this buffer which have not yet + // been placed in the sequence. + int64 unfinished_user_count; + + string ToString() const { + return tensorflow::strings::StrCat( + "Buffer ", id, " (defined by ", + defining_instruction->instruction->name(), ", size ", size, + " bytes)"); } + }; + + // Creates a Buffer representing the given logical buffer. The buffer is added + // to buffers_ and a reference is returned. + Buffer& CreateBufferFromLogicalBuffer( + const LogicalBuffer* logical_buffer, + const TuplePointsToAnalysis& points_to_analysis, + const HloRematerialization::ShapeSizeFunction& size_function, + bool live_out) { + bool has_indirect_uses = false; + ItemList users = GetUsers(instruction_list_, logical_buffer, + points_to_analysis, &has_indirect_uses); + return NewBuffer(instruction_list_.GetItem(logical_buffer->instruction()), + size_function(logical_buffer->shape()), std::move(users), + live_out, has_indirect_uses); + } - VLOG(3) << " memory usage = " << memory_usage_; - VLOG(10) << ToString(); - return Status::OK(); + // Create a new buffer representing a rematerialization of given buffer for + // 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); + for (Item* use : rematerialized_uses) { + CHECK(!use->placed); + } + return NewBuffer(remat_item, original_buffer.size, + std::move(rematerialized_uses), /*live_out=*/false, + /*has_indirect_uses=*/false); } - // Finishes the placement of the current instruction. This frees any dead - // operands or dead result of the instruction. This must be called after each - // call to BeginInstruction. - Status EndInstruction() { - TF_RET_CHECK(in_progress_instruction_ != nullptr); - VLOG(3) << "EndInstruction " << in_progress_instruction_->name(); - - for (HloInstruction* operand : UniqueOperands(in_progress_instruction_)) { - TF_RET_CHECK(ContainsKey(remaining_uses_, operand)); - std::vector& uses = remaining_uses_.at(operand); - auto it = std::find(uses.begin(), uses.end(), in_progress_instruction_); - TF_RET_CHECK(it != uses.end()); - uses.erase(it); - - if (uses.empty()) { - // Operand is dead. - int64 operand_size = TotalSizeBytes(operand->shape()); - if (!IsLiveOut(operand)) { - VLOG(4) << operand->name() << " (" - << HumanReadableNumBytes(operand_size) << ") is dead"; - memory_usage_ -= operand_size; - TF_RET_CHECK(memory_usage_ >= 0); - } - } + // Return number of bytes allocated for the buffer with the given id. Buffers + // allocated by the calling computation (eg, parameter and output buffers) are + // considered to have zero bytes because the memory is accounted for in a + // different computation. + int64 AllocatedSize(BufferId buffer_id) const { + const Buffer& buffer = buffers_.at(buffer_id); + HloOpcode def_opcode = buffer.defining_instruction->instruction->opcode(); + if (buffer.live_out || def_opcode == HloOpcode::kParameter) { + return 0; + } else { + return buffer.size; } + } + + // Returns true if BeginInstruction and EndInstruction has been called for the + // given instruction. + bool IsFinished(Item* item) const { + return item->placed && item != in_progress_item_; + } - // Value is dead if the instruction has no uses and is not live out. - if (in_progress_instruction_->users().empty() && - !IsLiveOut(in_progress_instruction_)) { - memory_usage_ -= TotalSizeBytes(in_progress_instruction_->shape()); - TF_RET_CHECK(memory_usage_ >= 0); + // Returns whether the given buffer is being used by the in-progress + // instruction. + bool IsInUse(BufferId buffer_id) const { + if (in_progress_item_ == nullptr) { + 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(); + } - in_progress_instruction_ = nullptr; + // Returns whether the given instruction is live at the current program + // point. + bool IsCurrentlyLive(BufferId buffer_id) const { + const Buffer& buffer = buffers_[buffer_id]; + return (buffer.defining_instruction->placed && + buffer.unfinished_user_count > 0); + } - VLOG(3) << " memory usage = " << memory_usage_; - VLOG(10) << ToString(); - return Status::OK(); + // Create a new buffer, add it to buffers_, and return a reference. + Buffer& NewBuffer(Item* defining_instruction, int64 size, ItemList&& users, + bool live_out, bool has_indirect_uses) { + int buffer_id = buffers_.size(); + buffers_.push_back(Buffer{buffer_id, defining_instruction, size, live_out, + has_indirect_uses, users, + static_cast(users.size())}); + return buffers_.back(); } - // Adjusts memory usage to account for the rematerialization of - // original_instruction for the given use. The rematerialization is - // remat_instruction. This method should be called after the HLO graph has - // been transformed (rematerialization instruction created and connected to - // its use). - Status RematerializeInstructionForUse(HloInstruction* original_instruction, - HloInstruction* remat_instruction, - HloInstruction* use) { - VLOG(3) << "RematerializeInstructionForUse: original_instruction = " - << original_instruction->name() - << ", remat_instruction = " << remat_instruction->name() - << ", use = " << use->name(); - - TF_RET_CHECK(in_progress_instruction_ != nullptr); - TF_RET_CHECK(IsPlaced(original_instruction)); - TF_RET_CHECK(!IsPlaced(remat_instruction)); - TF_RET_CHECK(!IsPlaced(use)); - TF_RET_CHECK(IsCurrentlyLive(original_instruction)); - - // Remove 'use' from remaining uses of original_instruction. - auto it = std::find(remaining_uses_[original_instruction].begin(), - remaining_uses_[original_instruction].end(), use); - TF_RET_CHECK(it != remaining_uses_[original_instruction].end()); - remaining_uses_[original_instruction].erase(it); - - // If original_instruction is no longer live ('use' was its last use) then - // deduct original_instruction's memory usage. - if (!IsCurrentlyLive(original_instruction)) { - memory_usage_ -= TotalSizeBytes(original_instruction->shape()); - TF_RET_CHECK(memory_usage_ >= 0); - } - - // Add the new remat_instruction to the remaining uses of its operands. - for (auto* operand : UniqueOperands(remat_instruction)) { - // Rematerialization may extend the lifetime of the operand so account for - // this in memory_usage_. - TF_RET_CHECK(IsPlaced(operand)); - if (!IsCurrentlyLive(operand)) { - memory_usage_ += TotalSizeBytes(operand->shape()); + const HloComputation* computation_; + + // Instruction list containing the ordering of instructions in + // computation_. This is the order in which instructions are placed + // (BeginInstruction/EndInstruction calls). + const InstructionList& instruction_list_; + + // Memory usage at the currently placed instruction. + int64 memory_usage_ = 0; + + // The instruction currently being placed. This value is non-null only + // between the calling of BeginInstruction and EndInstruction. + Item* in_progress_item_ = nullptr; + + // All buffers in the computation. + std::vector buffers_; +}; + +MemoryUsageTracker::MemoryUsageTracker( + const HloComputation* computation, + const HloRematerialization::ShapeSizeFunction& size_function, + const TuplePointsToAnalysis& points_to_analysis, + const InstructionList& instruction_list) + : computation_(computation), instruction_list_(instruction_list) { + PointsToSet::BufferSet live_out_set = + points_to_analysis.GetPointsToSet(computation_->root_instruction()) + .CreateFlattenedSet(); + tensorflow::gtl::FlatMap + logical_buffer_to_buffer_id; + + for (auto* item = instruction_list_.first(); item != nullptr; + item = instruction_list_.next(item)) { + const HloInstruction* const instruction = item->instruction; + for (const LogicalBuffer* logical_buffer : + points_to_analysis.GetBuffersDefinedByInstruction(instruction)) { + Buffer* buffer; + if (instruction->opcode() == HloOpcode::kWhile) { + // The while instruction defines no new buffers. Instead it reuses the + // buffers of its operand. Find the Buffer of its operand at the + // proper ShapeIndex. + const PointsToSet& operand_points_to = + points_to_analysis.GetPointsToSet(instruction->operand(0)); + CHECK_EQ(operand_points_to.element(logical_buffer->index()).size(), 1); + const LogicalBuffer* source_logical_buffer = + operand_points_to.element(logical_buffer->index())[0]; + buffer = + &buffers_.at(logical_buffer_to_buffer_id.at(source_logical_buffer)); + + // Mark buffer as has indirect use and live out. + buffer->has_indirect_uses = true; + buffer->live_out = + buffer->live_out || ContainsKey(live_out_set, logical_buffer); + + // Add users of while to Buffer users. + 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()) { + buffer->users.push_back(user_item); + buffer->unfinished_user_count++; + user_item->buffers_used.push_back(buffer->id); + } + } + } else { + buffer = &CreateBufferFromLogicalBuffer( + logical_buffer, points_to_analysis, size_function, + ContainsKey(live_out_set, logical_buffer)); + item->buffers_defined.push_back(buffer->id); + for (Item* user : buffer->users) { + user->buffers_used.push_back(buffer->id); + } } - remaining_uses_.at(operand).push_back(remat_instruction); + + logical_buffer_to_buffer_id[logical_buffer] = buffer->id; } + } + XLA_VLOG_LINES(10, ToString()); + DCHECK(Check()); +} + +Status MemoryUsageTracker::BeginInstruction(Item* item) { + const HloInstruction* instruction = item->instruction; + VLOG(3) << "BeginInstruction " << instruction->name(); + TF_RET_CHECK(in_progress_item_ == nullptr); + in_progress_item_ = item; - VLOG(3) << " memory usage = " << memory_usage_; - VLOG(10) << ToString(); - return Status::OK(); + item->placed = true; + + // All buffers defined by this instruction need memory. + for (BufferId buffer_id : item->buffers_defined) { + VLOG(3) << " Buffer " << buffers_.at(buffer_id).ToString() + << " is now live."; + memory_usage_ += AllocatedSize(buffer_id); } - // Returns the number of bytes that the current memory usage will be reduced - // if the given instruction is rematerialized. - int64 MemoryReducedIfRematerialized(const HloInstruction* instruction) const { - // To reduce memory consumption 'instruction' must be currently live and - // rematerialization must make 'instruction' not live. - if (IsLiveIn(instruction) || IsLiveOut(instruction) || - !IsCurrentlyLive(instruction)) { - return 0; - } + // TODO(b/37686934): Elementwise instructions can share the buffer of a (dead) + // operand. Account for this potential reuse here. - // If the in-progress instruction is a user of 'instruction' (or - // 'instruction' itself) then rematerializing 'instruction' cannot reduce - // memory usage because the value is required to be live at this program - // point. - if (in_progress_instruction_ == instruction || - in_progress_instruction_->IsUserOf(instruction)) { - return 0; - } + VLOG(3) << " memory usage = " << memory_usage_; + VLOG(10) << ToString(); - // Compute the amount of memory reduced (if any) by rematerializing - // 'instruction'. 'instruction' will no longer be live at this program - // point, so initially set memory_reduced to the size of its output value. - int64 memory_reduced = TotalSizeBytes(instruction->shape()); + DCHECK(Check()); + return Status::OK(); +} - // Account for any operands whose live range must be extended across this - // program point. - for (const HloInstruction* operand : UniqueOperands(instruction)) { - if (!IsCurrentlyLive(operand)) { - // This operand of candidate is not live at this program - // point. Rematerializing 'instruction' will extend the operand's live - // range across this program point. - memory_reduced -= TotalSizeBytes(operand->shape()); - } +Status MemoryUsageTracker::EndInstruction() { + TF_RET_CHECK(in_progress_item_ != nullptr); + VLOG(3) << "EndInstruction " << in_progress_item_->instruction->name(); + + for (BufferId buffer_id : in_progress_item_->buffers_used) { + Buffer& buffer = buffers_.at(buffer_id); + buffer.unfinished_user_count--; + CHECK_GE(buffer.unfinished_user_count, 0) + << buffer.ToString() << " has negative unfinished use count."; + if (buffer.unfinished_user_count == 0) { + // Buffer is now dead. + VLOG(3) << " " << buffer.ToString() << " is now dead."; + memory_usage_ -= AllocatedSize(buffer_id); + CHECK_GE(memory_usage_, 0); } - return memory_reduced; } - // Returns the remaining unplaced uses of the given instruction. - const std::vector& RemainingUses( - const HloInstruction* instruction) const { - return remaining_uses_.at(instruction); + // If any buffer defined by this instruction has no uses, then memory can be + // reclaimed immediately. + for (BufferId buffer_id : in_progress_item_->buffers_defined) { + const Buffer& buffer = buffers_.at(buffer_id); + if (buffer.unfinished_user_count == 0) { + VLOG(3) << " " << buffer.ToString() << " is immediately dead."; + memory_usage_ -= AllocatedSize(buffer_id); + CHECK_GE(memory_usage_, 0); + } } - // Returns whether the given instruction has been placed (BeginInstruction has - // been called with 'instruction' as the argument). - bool IsPlaced(const HloInstruction* instruction) const { - return ContainsKey(remaining_uses_, instruction); + in_progress_item_ = nullptr; + + VLOG(3) << " memory usage = " << memory_usage_; + VLOG(10) << ToString(); + + DCHECK(Check()); + + return Status::OK(); +} + +int64 MemoryUsageTracker::MemoryReducedIfRematerialized(Item* item) const { + CHECK_NE(in_progress_item_, nullptr); + if (!item->placed || item == in_progress_item_) { + return 0; } - // Returns whether the given instruction is live at the current program point. - bool IsCurrentlyLive(const HloInstruction* instruction) const { - return (!IsPlaced(instruction) && IsLiveIn(instruction)) || - (IsPlaced(instruction) && - (!RemainingUses(instruction).empty() || IsLiveOut(instruction))); + // TODO(b/37687140): Rematerialization can increase peak memory consumption at + // an earlier point in the program if rematerialization extends the live range + // of the operand of the instruction being rematerialized across the live + // range of the value of instruction being rematerialized. Don't rematerialize + // in this case (ie, return 0 here). + + // Compute the amount of memory reduced (if any) by rematerializing + // 'instruction'. The LogicalBuffers defined by 'instruction' will no longer + // be live at this program point, so initially set memory_reduced to the + // size of its defined values. + int64 memory_reduced = 0; + for (BufferId buffer_id : item->buffers_defined) { + // Avoid rematerializing instructions with indirect uses as it is difficult + // to reason about liveness after rematerializing the instruction. + // TODO(b/37714814): Consider rematerialzing instructions with indirect + // uses. + if (buffers_.at(buffer_id).has_indirect_uses) { + return 0; + } + + if (IsCurrentlyLive(buffer_id) && !IsInUse(buffer_id)) { + memory_reduced += AllocatedSize(buffer_id); + } } - string ToString() const { - string output = tensorflow::strings::StrCat("MemoryUsageTracker for ", - computation_->name(), "\n"); - tensorflow::strings::StrAppend(&output, "memory usage = ", memory_usage(), - "\n"); - tensorflow::strings::StrAppend(&output, "Live values:\n"); - for (const auto& pair : remaining_uses_) { - const HloInstruction* instruction = pair.first; - const std::vector& uses = pair.second; - tensorflow::strings::StrAppend( - &output, " ", instruction->name(), "; remaining uses: ", - tensorflow::str_util::Join(uses, ", ", - [](string* out, HloInstruction* use) { - tensorflow::strings::StrAppend( - out, use->name()); - }), - "\n"); + // Account for any logical buffers whose live range must be extended across + // this program point. + for (BufferId buffer_id : item->buffers_used) { + if (!IsCurrentlyLive(buffer_id)) { + // This logical buffer is used by 'instruction' but is not live at this + // program point. Rematerializing 'instruction' will extend the buffer's + // live range across this program point. + memory_reduced -= AllocatedSize(buffer_id); } - return output; } - // Returns the current memory usage. This is the sum of sizes of all live - // values. - int64 memory_usage() const { return memory_usage_; } + return memory_reduced; +} - // Returns the current instruction being placed. - const HloInstruction* in_progress_instruction() const { - return in_progress_instruction_; +Status MemoryUsageTracker::AddRematerializedInstruction(Item* original_item, + Item* remat_item) { + VLOG(3) << "AddRematerializedInstruction: original_instruction = " + << original_item->instruction->name() + << ", 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); + + // Construct the list of buffers used and defined by the rematerialization. + remat_item->buffers_used = original_item->buffers_used; + + // Account for the additional buffer uses created by the new rematerialization + // instruction. Update memory usage if the rematerialization makes a dead + // buffer live again. + for (BufferId buffer_id : original_item->buffers_used) { + Buffer& buffer = buffers_.at(buffer_id); + if (buffer.unfinished_user_count == 0) { + // Buffer used by this instruction was dead, now is alive. + memory_usage_ += AllocatedSize(buffer.id); + } + + buffer.unfinished_user_count++; + buffer.users.push_back(remat_item); } - private: - // Returns the total size of the shape (including nested elements) in bytes. - int64 TotalSizeBytes(const Shape& shape) const { - int64 total_size = 0; - ShapeUtil::ForEachSubshape( - shape, - [this, &total_size](const Shape& subshape, - const ShapeIndex& /*index*/) { - total_size += size_function_(subshape); - return Status::OK(); - }) - .IgnoreError(); - return total_size; - } - - // Returns true if the value of given instruction is live into the - // computation. - bool IsLiveIn(const HloInstruction* instruction) const { - return instruction->opcode() == HloOpcode::kConstant || - instruction->opcode() == HloOpcode::kParameter; - } - - // Returns true if the value of given instruction is live out of the - // computation. - bool IsLiveOut(const HloInstruction* instruction) const { - return instruction->opcode() == HloOpcode::kConstant || - instruction->opcode() == HloOpcode::kParameter || - instruction == instruction->parent()->root_instruction(); + // Create a new set of Buffers defined by the new rematerialization + // instruction. Update the internal data structures and memory use to account + // for them. + for (BufferId old_buffer_id : original_item->buffers_defined) { + Buffer& old_buffer = buffers_.at(old_buffer_id); + + ItemList placed_users; + ItemList unplaced_users; + for (Item* user : old_buffer.users) { + if (user->placed) { + CHECK(IsFinished(user)); + placed_users.push_back(user); + } else { + unplaced_users.push_back(user); + } + } + old_buffer.users = std::move(placed_users); + old_buffer.unfinished_user_count = 0; + + // Buffer is now dead. + memory_usage_ -= AllocatedSize(old_buffer.id); + + Buffer& new_buffer = + RematerializeBuffer(old_buffer, remat_item, std::move(unplaced_users)); + + remat_item->buffers_defined.push_back(new_buffer.id); + for (Item* user : new_buffer.users) { + BufferIdList& buffers_used = user->buffers_used; + std::replace(buffers_used.begin(), buffers_used.end(), old_buffer_id, + new_buffer.id); + } } - const HloComputation* computation_; + VLOG(3) << " memory usage = " << memory_usage_; + XLA_VLOG_LINES(10, ToString()); - // Function which computes the size of the top-level buffer of a shape. - const HloRematerialization::ShapeSizeFunction size_function_; + DCHECK(Check()); - // Memory usage at the currently placed instruction. - int64 memory_usage_ = 0; + return Status::OK(); +} - // The instruction currently being placed. This value is non-null only between - // the calling of BeginInstruction and EndInstruction. - const HloInstruction* in_progress_instruction_ = nullptr; - - // remaining_uses is a vector of uses of the HLO instruction's value which - // have not yet been visited by in the rematerialization loop. Use to track - // liveness of HLO instructions. - // TODO(b/35212854): Track values using logical buffers rather than HLO - // instructions. Using HLO instructions over-estimates memory usage because - // buffer aliasing is ignored. - tensorflow::gtl::FlatMap> - remaining_uses_; -}; +string MemoryUsageTracker::ToString() const { + string output = tensorflow::strings::StrCat("MemoryUsageTracker for ", + computation_->name(), "\n"); + tensorflow::strings::StrAppend( + &output, "Memory usage: ", HumanReadableNumBytes(memory_usage()), " (", + memory_usage(), " bytes)"); + for (auto* item = instruction_list_.first(); item != nullptr; + item = instruction_list_.next(item)) { + const HloInstruction* instruction = item->instruction; + string inprogress = item == in_progress_item_ ? " in-progress" : ""; + string placed = item->placed ? " placed" : ""; + tensorflow::strings::StrAppend(&output, " ", instruction->name(), + inprogress, placed, "\n Defines:\n"); + for (BufferId buffer_id : item->buffers_defined) { + const Buffer& buffer = buffers_[buffer_id]; + string live = IsCurrentlyLive(buffer_id) ? " live" : ""; + tensorflow::strings::StrAppend(&output, " ", buffer.ToString(), live, + ", ", buffer.unfinished_user_count, + " unfinished uses\n"); + } + tensorflow::strings::StrAppend(&output, " Uses:\n"); + for (BufferId buffer_id : item->buffers_used) { + tensorflow::strings::StrAppend(&output, " ", + buffers_[buffer_id].ToString(), "\n"); + } + } + return output; +} -// Computes and returns the cost of rematerializing the given instruction. Cost -// per rematerialized instruction is defined as: -// -// (flop_count + transcendental_count + element_count) / memory_reduced +bool MemoryUsageTracker::Check() const { + auto elements_are_unique = [](const BufferIdList& vec) { + return vec.size() == std::set(vec.begin(), vec.end()).size(); + }; + + // Verify buffers_defined per instruction. + for (auto& instruction : computation_->instructions()) { + const BufferIdList& defined_buffers = + instruction_list_.GetItem(instruction.get())->buffers_defined; + CHECK(elements_are_unique(defined_buffers)) + << "Instruction " << instruction->name() + << " does not have unique defined buffers: " + << tensorflow::str_util::Join( + defined_buffers, ", ", [this](string* out, BufferId buffer_id) { + tensorflow::strings::StrAppend( + out, buffers_.at(buffer_id).ToString()); + }); + + for (const Buffer& buffer : buffers_) { + if (buffer.defining_instruction->instruction == instruction.get()) { + CHECK(std::find(defined_buffers.begin(), defined_buffers.end(), + buffer.id) != defined_buffers.end()) + << "Instruction " << instruction->name() + << " defined buffers is missing: " << buffer.ToString(); + } + } + } + + // Verify buffers_used per instruction. + for (auto& instruction : computation_->instructions()) { + const BufferIdList& used_buffers = + instruction_list_.GetItem(instruction.get())->buffers_used; + CHECK(elements_are_unique(used_buffers)) + << "Instruction " << instruction->name() + << " does not have unique used buffers: " + << tensorflow::str_util::Join( + used_buffers, ", ", [this](string* out, BufferId buffer_id) { + tensorflow::strings::StrAppend( + out, buffers_.at(buffer_id).ToString()); + }); + } + for (const Buffer& buffer : buffers_) { + 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()) + << "Instruction " << user->instruction->name() + << " used buffers is missing " << buffer.ToString(); + if (!IsFinished(user)) { + unfinished_uses++; + } + } + CHECK_EQ(buffer.unfinished_user_count, unfinished_uses) + << "Incorrect unplaced use count for " << buffer.ToString(); + } + + // Verify live set size against memory_usage_. + int64 live_size = 0; + for (const Buffer& buffer : buffers_) { + // The while instruction reuses its input buffers as output buffers so + // don't double count its buffers if it is currently executing. + if (IsCurrentlyLive(buffer.id) && + !(buffer.defining_instruction == in_progress_item_ && + in_progress_item_->instruction->opcode() == HloOpcode::kWhile)) { + live_size += AllocatedSize(buffer.id); + } + } + CHECK(live_size == memory_usage_) + << "Live set size " << live_size << " is not same as memory usage " + << memory_usage_ + << ". This could happen if some nodes defined in the " + "computation are not being used/executed."; + + return true; +} + +// Computes and returns the cost of rematerializing the given instruction. +// Cost per rematerialized instruction is defined as: // -// flop_count: from HloCostAnalysis -// transcendental_count: from HloCostAnalysis -// element_count: number of elements accessed in operands and output of -// instruction -// memory_reduced: The memory usage reduced by rematerializing the -// instruction. +// memory_limit_bytes / memory_reduced // -// This is a rough estimate of the extra execution time per byte saved by -// rematerializing this instruction for its remaining uses. In general, we want -// the most memory saving for the least latency penalty which is captured by -// this heuristic. +// The idea is to choose the operation that will save the most memory for +// rematerialization and do not worry about how much the compute costs since +// running out of memory is more harmful than taking longer to get the answer. int64 RematerializationCost(const HloInstruction* instruction, const MemoryUsageTracker& memory_tracker, - const HloCostAnalysis& cost_analysis, - int64 memory_reduced) { - const int64 bytes_accessed = cost_analysis.bytes_accessed(*instruction); - const int64 elements_accessed = - bytes_accessed / - ShapeUtil::ByteSizeOfPrimitiveType(instruction->shape().element_type()); - - // A duplicate of the rematerialized instruction will be created at each - // remaining use. - int64 duplication = memory_tracker.RemainingUses(instruction).size(); - if (duplication == instruction->users().size()) { - // All remaining uses of instruction are after this point so we can remove - // the original instruciton after rematerialization. - duplication -= 1; + int64 memory_reduced, int64 memory_limit_bytes) { + // 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); + })) { + return 0; } - CHECK_GT(memory_reduced, 0); - // Multiply by 256 to improve precision of cost. Without this factor, - // many instructions such as many elementwise instructions would have - // zero cost because the bytes reduced can be several times greater than - // the element count. - return 256 * duplication * - (cost_analysis.flop_count(*instruction) + - cost_analysis.transcendental_count(*instruction) + - elements_accessed) / - memory_reduced; + CHECK_GT(memory_reduced, 0); + // Return the inverse of the benefit of rematerialization. + return memory_limit_bytes / memory_reduced; } // Selects and returns the best candidate instruction for rematerialization. @@ -462,30 +863,30 @@ int64 RematerializationCost(const HloInstruction* instruction, // candidate which reduce memory use at the program point of the current // instruction as indicated by memory_tracker. nullptr is returned if no // candidate can be found. -HloInstruction* PickRematerializationCandidate( - const MemoryUsageTracker& memory_tracker, - const InstructionList& instruction_list, - const HloCostAnalysis& cost_analysis, - const tensorflow::gtl::FlatSet& remat_instructions) { - HloInstruction* best = nullptr; +Item* PickRematerializationCandidate(const MemoryUsageTracker& memory_tracker, + const InstructionList& instruction_list, + int64 memory_limit_bytes) { + Item* best_item = nullptr; int64 best_cost = 0; // TODO(b/35244891): This is currently quadratic in the number of HLO // instructions. - for (HloInstruction* candidate : instruction_list.instructions()) { - if (!memory_tracker.IsPlaced(candidate)) { - // Only iterate up to the currently placed instruction as indicated by - // memory_tracker. We are trying to reduce memory usage at the placed + for (auto* item = instruction_list.first(); item != nullptr; + item = instruction_list.next(item)) { + if (!item->placed) { + // Only iterate up to the currently placed instruction. + // We are trying to reduce memory usage at the placed // instruction so rematerializing later values is of no benefit. break; } + HloInstruction* candidate = item->instruction; VLOG(5) << "considering rematerialization candidate " << candidate->name(); - if (ContainsKey(remat_instructions, candidate)) { - // Skip instructions which are rematerialization clones to avoid infinite - // loops of rematerializing the same instruction(s) repeatedly. + if (item->blacklisted) { + // Skip instructions on the blacklist to avoid infinite loops of + // rematerializing the same instruction(s) repeatedly. VLOG(5) << "candidate " << candidate->name() - << " not viable: is a rematerialized instruction"; + << " is excluded from rematerialization"; continue; } @@ -495,28 +896,41 @@ HloInstruction* PickRematerializationCandidate( continue; } + // If any of the candidate's control successor has been placed, we need to + // skip this candidate. Otherwise we will violate control dependency. + bool control_successor_placed = + std::any_of(candidate->control_successors().begin(), + candidate->control_successors().end(), + [&memory_tracker](const HloInstruction* inst) { + return memory_tracker.IsPlaced(inst); + }); + + if (control_successor_placed) { + continue; + } + const int64 memory_reduced = - memory_tracker.MemoryReducedIfRematerialized(candidate); + memory_tracker.MemoryReducedIfRematerialized(item); if (memory_reduced <= 0) { VLOG(5) << "candidate " << candidate->name() - << " memory reduced = " << memory_reduced << " <= 0"; + << " memory reduced = " << memory_reduced << " <= 0"; continue; } const int cost = RematerializationCost(candidate, memory_tracker, - cost_analysis, memory_reduced); + memory_reduced, memory_limit_bytes); VLOG(5) << "candidate " << candidate->name() << ", memory reduced " << memory_reduced << ", cost per byte " << cost; - if (best == nullptr || cost < best_cost) { + if (best_item == nullptr || cost < best_cost) { VLOG(5) << "candidate " << candidate->name() << " now best"; - best = candidate; + best_item = item; best_cost = cost; } } - return best; + return best_item; } } // namespace @@ -524,10 +938,14 @@ HloInstruction* PickRematerializationCandidate( StatusOr HloRematerialization::ComputePeakMemory( const HloComputation* computation, const std::vector& order) const { - MemoryUsageTracker tracker(computation, size_function_); + InstructionList instruction_list(order); + MemoryUsageTracker tracker(computation, size_function_, *points_to_analysis_, + instruction_list); int64 peak_memory = tracker.memory_usage(); - for (const HloInstruction* instruction : order) { - TF_RETURN_IF_ERROR(tracker.BeginInstruction(instruction)); + for (auto* item = instruction_list.first(); item != nullptr; + item = instruction_list.next(item)) { + const HloInstruction* instruction = item->instruction; + TF_RETURN_IF_ERROR(tracker.BeginInstruction(item)); TF_ASSIGN_OR_RETURN(int64 callee_usage, CalledComputationsMemoryUsage(instruction)); peak_memory = @@ -541,9 +959,8 @@ StatusOr HloRematerialization::ComputePeakMemory( StatusOr HloRematerialization::CalledComputationsMemoryUsage( const HloInstruction* instruction) const { - TF_ASSIGN_OR_RETURN(const CallGraphNode* node, - call_graph_->GetNode(instruction->parent())); - const CallSite* callsite = node->GetCallSite(instruction); + const CallSite* callsite = + call_graph_->GetNode(instruction->parent()).GetCallSite(instruction); if (callsite == nullptr || callsite->context() == CallContext::kParallel) { return 0; } @@ -563,15 +980,19 @@ StatusOr HloRematerialization::RematerializeComputation( << " with limit " << HumanReadableNumBytes(memory_limit_bytes); VLOG(1) << "peak memory usage is " << HumanReadableNumBytes(computation_peak_memory_.at(computation)); + CHECK(!ContainsKey(rematerialized_computations_, computation)); InstructionList instruction_list(sequence->at(computation)); - MemoryUsageTracker memory_tracker(computation, size_function_); + MemoryUsageTracker memory_tracker(computation, size_function_, + *points_to_analysis_, instruction_list); bool changed = false; - // Set of instruction clones (not the originals) created during - // rematerialization. A record is kept to avoid rematerializing an instruction - // more than once to avoid looping infinitely during rematerialization. - tensorflow::gtl::FlatSet remat_instructions; + // If the rematerialization makes the source instruction dead, then the + // rematerialization is added to 'remat_move_instructions' (the + // rematerialization is essentially a move). If the next rematerialization of + // the instruction is also a move then the rematerialization is added to the + // blacklist. + tensorflow::gtl::FlatSet remat_move_instructions; // The peak memory of the computation at any point in the instruction // sequence. @@ -583,22 +1004,24 @@ StatusOr HloRematerialization::RematerializeComputation( // instructions which are dead. int64 net_instructions_added = 0; - TF_ASSIGN_OR_RETURN(const CallGraphNode* call_graph_node, - call_graph_->GetNode(computation)); + const CallGraphNode& call_graph_node = call_graph_->GetNode(computation); // Iterate through all instructions in the sequence. At each instruction // (program point) if memory_usage exceeds the specified limit then // rematerialize HLO instructions until memory_usage is reduced. - for (auto list_it = instruction_list.instructions().begin(); - list_it != instruction_list.instructions().end(); ++list_it) { - HloInstruction* instruction = *list_it; + int64 instruction_index = 0; + for (auto* item = instruction_list.first(); item != nullptr; + item = instruction_list.next(item)) { + const HloInstruction* instruction = item->instruction; TF_ASSIGN_OR_RETURN(int64 callee_usage, CalledComputationsMemoryUsage(instruction)); - TF_RETURN_IF_ERROR(memory_tracker.BeginInstruction(instruction)); + TF_RETURN_IF_ERROR(memory_tracker.BeginInstruction(item)); VLOG(2) << "Program point at " << instruction->name() << ", memory usage = " << memory_tracker.memory_usage() - << ", callee usage = " << callee_usage; + << ", callee usage = " << callee_usage << ", [" << instruction_index + << "/" << instruction_list.size() << "]"; + instruction_index++; while (memory_tracker.memory_usage() + callee_usage > memory_limit_bytes) { VLOG(2) << "Over memory limit at instruction " << instruction->name() @@ -607,10 +1030,10 @@ StatusOr HloRematerialization::RematerializeComputation( callee_usage) << ", limit is " << HumanReadableNumBytes(memory_limit_bytes); - HloInstruction* best = PickRematerializationCandidate( - memory_tracker, instruction_list, cost_analysis_, remat_instructions); + Item* best_item = PickRematerializationCandidate( + memory_tracker, instruction_list, memory_limit_bytes); - if (best == nullptr) { + if (best_item == nullptr) { VLOG(3) << "Unable to find rematerialization candidate at program " "point " << instruction->name() << ". Memory usage = " @@ -619,44 +1042,68 @@ StatusOr HloRematerialization::RematerializeComputation( break; } - VLOG(1) << "Rematerializing instruction " << best->name(); + HloInstruction* best = best_item->instruction; + VLOG(1) << "Rematerializing instruction " << best->name() << " (saving " + << memory_tracker.MemoryReducedIfRematerialized(best_item) << ")"; changed = true; remat_count++; - // Create a rematerialized copy of the candidate at each remaining use. - // Make a copy of remaining uses because RematerializeInstructionForUse - // modifies the remaining uses vector in memory_tracker. - // TODO(b/35213652): It may be profitable to share one rematerialized copy - // amongst more than one use. - std::vector remaining_uses_copy = - memory_tracker.RemainingUses(best); - for (HloInstruction* use : remaining_uses_copy) { - // Create a new rematerialized instruction in the HLO graph. - HloInstruction* remat = - computation->AddInstruction(best->Clone(/*suffix=*/"remat")); - - VLOG(3) << "Replacing use of " << best->name() << " in " << use->name() - << " with rematerialization " << remat->name(); + HloInstruction* remat = + computation->AddInstruction(best->Clone(/*suffix=*/"remat")); - TF_RETURN_IF_ERROR(best->ReplaceUseWith(use, remat)); - - // Account for the rematerialization in the memory tracker. - TF_RETURN_IF_ERROR( - memory_tracker.RematerializeInstructionForUse(best, remat, use)); - - // Insert rematerialized instruction right before its use. - TF_RETURN_IF_ERROR(instruction_list.InsertBefore(remat, use)); + // Add control dependencies to the new operation. + for (auto successor : best->control_successors()) { + TF_RETURN_IF_ERROR(remat->AddControlDependencyTo(successor)); + } + for (auto predecessor : best->control_predecessors()) { + TF_RETURN_IF_ERROR(predecessor->AddControlDependencyTo(remat)); + } - // Add rematerialized instruction to remat_instructions so the - // rematerialized instruction is not rematerialized again. - remat_instructions.insert(remat); + Item* remat_item = instruction_list.CreateItem(remat); - net_instructions_added++; + // Replace each remaining use of 'best' with the rematerialization. + std::vector best_users_copy = best->users(); + for (HloInstruction* user : best_users_copy) { + if (!memory_tracker.IsPlaced(user)) { + VLOG(2) << " Replacing use of " << best->name() << " in " + << user->name() << " with " << remat->name(); + TF_RETURN_IF_ERROR(best->ReplaceUseWith(user, remat)); + } } - // Original instruction should no longer be live at this point. All - // of its remaining uses are fed by rematerialized instructions. - TF_RET_CHECK(!memory_tracker.IsCurrentlyLive(best)); + // Account for the rematerialization in the memory tracker. + TF_RETURN_IF_ERROR( + memory_tracker.AddRematerializedInstruction(best_item, remat_item)); + + // Insert rematerialized instruction right before the earliest unplaced + // use of the instruction *and* the earliest unplaced last use of any + // operands of remat. Unplaced uses of the remat's operands are included + // because we don't want to extend the live range of remat's operands as + // this could increase memory usage. + ItemList place_before; + for (auto user : remat->users()) { + place_before.push_back(instruction_list.GetItem(user)); + } + for (auto* operand : remat->operands()) { + for (auto* operand_user : operand->users()) { + if (operand_user != remat) { + Item* operand_user_item = instruction_list.GetItem(operand_user); + if (!operand_user_item->placed) { + place_before.push_back(operand_user_item); + } + } + } + } + // Insert rematerialized instruction before any of its successors to + // preserve ordering regarding control dependency. + for (auto successor : remat->control_successors()) { + 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); + place_before.push_back(successor_item); + } + instruction_list.InsertBeforeInstructions(remat_item, place_before); // If the rematerialized instruction is dead then rematerialization is // essentially a move. Don't delete the instruction now because we don't @@ -664,15 +1111,24 @@ StatusOr HloRematerialization::RematerializeComputation( // transformation because we keep maps with HloInstruction* values as // keys. if (best->users().empty()) { - VLOG(3) << best->name() << " is now dead"; - net_instructions_added--; + VLOG(2) << best->name() << " is now dead"; + if (ContainsKey(remat_move_instructions, best)) { + // Previously, 'best' was a rematerialization which killed the + // instruction it was a copying of. Now 'remat' is a rematerialization + // of 'best' and kills 'best'. Stop rematerializing this instruction + // to avoid an infinite loop. + instruction_list.Blacklist(remat); + } + remat_move_instructions.insert(remat); + } else { + net_instructions_added++; } VLOG(3) << "memory_usage after rematerialization = " << memory_tracker.memory_usage(); } - const CallSite* callsite = call_graph_node->GetCallSite(instruction); + const CallSite* callsite = call_graph_node.GetCallSite(instruction); if (callsite != nullptr && callsite->context() == CallContext::kSequential && memory_tracker.memory_usage() + callee_usage > memory_limit_bytes) { @@ -686,21 +1142,22 @@ StatusOr HloRematerialization::RematerializeComputation( // Recompute callee usage to account for any rematerialization performed // in the callee computations. - callee_usage = 0; for (HloComputation* called_computation : callsite->called_computations()) { - // Memory limit for the subcomputation is the memory limit less the - // amount of memory used at this point in the computation. - int64 subcomputation_memory_limit_bytes = std::max( - 0, memory_limit_bytes - memory_tracker.memory_usage()); - TF_ASSIGN_OR_RETURN( - bool subcomputation_changed, - RematerializeComputation(called_computation, sequence, - subcomputation_memory_limit_bytes)); - changed |= subcomputation_changed; - - callee_usage += computation_peak_memory_.at(called_computation); + if (!ContainsKey(rematerialized_computations_, called_computation)) { + // Memory limit for the subcomputation is the memory limit less the + // amount of memory used at this point in the computation. + int64 subcomputation_memory_limit_bytes = std::max( + 0, memory_limit_bytes - memory_tracker.memory_usage()); + TF_ASSIGN_OR_RETURN( + bool subcomputation_changed, + RematerializeComputation(called_computation, sequence, + subcomputation_memory_limit_bytes)); + changed |= subcomputation_changed; + } } + TF_ASSIGN_OR_RETURN(callee_usage, + CalledComputationsMemoryUsage(instruction)); } peak_memory = std::max(peak_memory, @@ -710,36 +1167,35 @@ StatusOr HloRematerialization::RematerializeComputation( TF_RETURN_IF_ERROR(memory_tracker.EndInstruction()); } - if (peak_memory > memory_limit_bytes) { - LOG(WARNING) << "Can't reduce memory use of computation " - << computation->name() << " below " - << HumanReadableNumBytes(memory_limit_bytes) - << " by rematerialization (only reduced to " - << HumanReadableNumBytes(peak_memory) << ")"; - } - - // Verify that there are no more remaining uses. + // Verify some invariants on the memory tracker. + CHECK_EQ(memory_tracker.memory_usage(), 0); for (auto& instruction : computation->instructions()) { - auto& remaining_uses = memory_tracker.RemainingUses(instruction.get()); - CHECK(remaining_uses.empty()) - << instruction->name() << " has remaining uses: " - << tensorflow::str_util::Join( - remaining_uses, ", ", [](string* out, HloInstruction* inst) { - tensorflow::strings::StrAppend(out, inst->name()); - }); + CHECK(memory_tracker.IsPlaced(instruction.get())); } - VLOG(1) << "Rematerialized " << remat_count << " instructions; " - << net_instructions_added << " net instructions added"; - VLOG(1) << "peak memory usage now " << HumanReadableNumBytes(peak_memory); + VLOG(1) << "In computation " << computation->name() << " rematerialized " + << remat_count << " instructions; " << net_instructions_added + << " net instructions added"; + VLOG(1) << " peak memory usage now " << HumanReadableNumBytes(peak_memory) + << " (was " + << HumanReadableNumBytes(computation_peak_memory_.at(computation)) + << ")"; // Update peak memory used by computation. - computation_peak_memory_[computation] = peak_memory; + computation_peak_memory_.at(computation) = peak_memory; // Update order to include rematerialized instructions. - sequence->at(computation) - .assign(instruction_list.instructions().begin(), - instruction_list.instructions().end()); + auto& dst = sequence->at(computation); + dst.clear(); + for (auto* item = instruction_list.first(); item != nullptr; + item = instruction_list.next(item)) { + const HloInstruction* instruction = item->instruction; + dst.push_back(instruction); + } + rematerialized_computations_.insert(computation); + + instructions_rematerialized_ += remat_count; + net_instructions_added_ += net_instructions_added; return changed; } @@ -753,18 +1209,36 @@ StatusOr HloRematerialization::Run( VLOG(1) << "HloRematerialization() with memory limit of " << HumanReadableNumBytes(memory_limit_bytes); - XLA_VLOG_LINES(3, "Before HloRematerialization:\n" + module->ToString()); + TF_ASSIGN_OR_RETURN(points_to_analysis_, TuplePointsToAnalysis::Run(module)); + + // Adjust memory limit to account for the output of the entry + // computation. This is necessary because the per-computation accounting in + // MemoryUsageTracker do not include output as these are typically allocated + // by the caller. + int64 module_output_size = 0; + ShapeUtil::ForEachSubshape( + module->entry_computation()->root_instruction()->shape(), + [&module_output_size, this](const Shape& subshape, + const ShapeIndex& /*index*/) { + module_output_size += size_function_(subshape); + }); + + const int64 adjusted_memory_limit_bytes = + memory_limit_bytes - module_output_size; + VLOG(1) << "Adjusted memory limit accounting for output (" + << HumanReadableNumBytes(module_output_size) + << "): " << HumanReadableNumBytes(adjusted_memory_limit_bytes); + XLA_VLOG_LINES(3, "Before HloRematerialization:\n" + module->ToString()); // Create initial sequence of HLO instructions. TF_ASSIGN_OR_RETURN(*sequence, CreateMemoryMinimizingSequence( *module, [this](const LogicalBuffer& buffer) { return size_function_(buffer.shape()); })); - // Compute peak memory usage of all computations in the module called in a // sequential context. - TF_ASSIGN_OR_RETURN(call_graph_, CallGraph::Build(module)); + call_graph_ = CallGraph::Build(module); TF_RETURN_IF_ERROR(call_graph_->VisitNodes( [this, sequence](const CallGraphNode& node) -> Status { if (node.context() == CallContext::kSequential) { @@ -774,22 +1248,24 @@ StatusOr HloRematerialization::Run( sequence->at(node.computation()))); } return Status::OK(); - })); - + }, + /*visit_unreachable_nodes=*/false)); + + // The peak memory usage of the module equals the peak memory use of the entry + // computation plus the output size of the computation. This is because the + // peak memory for a computation does not include the output as this is + // typically accounted for in the caller. + const int64 before_peak_memory = + computation_peak_memory_.at(module->entry_computation()) + + module_output_size; VLOG(1) << "Peak memory usage of module (before): " - << HumanReadableNumBytes( - computation_peak_memory_[module->entry_computation()]); - - // Run cost analysis. Operation cost is used in the heuristic for selecting - // instructions for rematerialization. - TF_RETURN_IF_ERROR( - module->entry_computation()->root_instruction()->Accept(&cost_analysis_)); + << HumanReadableNumBytes(before_peak_memory); // Subcomputations called by the entry computation will also be // rematerialized. - TF_ASSIGN_OR_RETURN(bool changed, - RematerializeComputation(module->entry_computation(), - sequence, memory_limit_bytes)); + TF_ASSIGN_OR_RETURN(bool changed, RematerializeComputation( + module->entry_computation(), sequence, + adjusted_memory_limit_bytes)); // Rematerialization can introduce dead code. This occurs if all uses of an // instruction are replaced with rematerializations of the instruction. @@ -799,6 +1275,9 @@ StatusOr HloRematerialization::Run( // After DCE, the module sequence may include instructions which no longer // exist. for (const auto& computation : module->computations()) { + if (computation->IsFusionComputation()) { + continue; + } if (sequence->at(computation.get()).size() != computation->instruction_count()) { // A size mismatch between the computation instruction count and the size @@ -824,19 +1303,38 @@ StatusOr HloRematerialization::Run( computation->instruction_count()); } } - - VLOG(1) << "Peak memory usage of module (after): " - << HumanReadableNumBytes( - computation_peak_memory_[module->entry_computation()]); + VLOG(1) << "Rematerialized " << instructions_rematerialized_ + << " instructions in module " << module->name() << "; " + << net_instructions_added_ << " net instructions added"; + const int64 current_peak_memory = + computation_peak_memory_.at(module->entry_computation()) + + module_output_size; + VLOG(1) << "Peak memory usage of module now " + << HumanReadableNumBytes(current_peak_memory) << " (" + << current_peak_memory << " bytes), was " + << HumanReadableNumBytes(before_peak_memory) << " (" + << before_peak_memory << " bytes)"; + const int64 reduced_peak_memory = before_peak_memory - current_peak_memory; + VLOG(1) << "Reduced peak memory by " + << HumanReadableNumBytes(reduced_peak_memory) << " (" + << reduced_peak_memory << " bytes)"; XLA_VLOG_LINES(3, "After HloRematerialization:\n" + module->ToString()); + if (current_peak_memory > memory_limit_bytes) { + LOG(WARNING) << "Can't reduce memory use below " + << HumanReadableNumBytes(memory_limit_bytes) + << " by rematerialization (only reduced to " + << HumanReadableNumBytes(current_peak_memory) << ")"; + } + return changed; } /* static */ StatusOr HloRematerialization::RematerializeAndSchedule( - const ShapeSizeFunction& size_function, int64 memory_limit_bytes, - HloModule* hlo_module, SequentialHloOrdering::HloModuleSequence* sequence) { + const HloRematerialization::ShapeSizeFunction& size_function, + int64 memory_limit_bytes, HloModule* hlo_module, + SequentialHloOrdering::HloModuleSequence* sequence) { HloRematerialization remat(size_function); return remat.Run(hlo_module, sequence, memory_limit_bytes); } diff --git a/tensorflow/compiler/xla/service/hlo_rematerialization.h b/tensorflow/compiler/xla/service/hlo_rematerialization.h index 86e1998b89454f75b1c10d0de2118fd1034c134d..42c279d440b78d90b9f19b92155c52787156e4b7 100644 --- a/tensorflow/compiler/xla/service/hlo_rematerialization.h +++ b/tensorflow/compiler/xla/service/hlo_rematerialization.h @@ -18,9 +18,9 @@ #include "tensorflow/compiler/xla/service/buffer_liveness.h" #include "tensorflow/compiler/xla/service/call_graph.h" #include "tensorflow/compiler/xla/service/hlo_computation.h" -#include "tensorflow/compiler/xla/service/hlo_cost_analysis.h" #include "tensorflow/compiler/xla/service/hlo_instruction.h" #include "tensorflow/compiler/xla/service/hlo_module.h" +#include "tensorflow/compiler/xla/service/tuple_points_to_analysis.h" namespace xla { @@ -60,7 +60,7 @@ class HloRematerialization { protected: HloRematerialization(const ShapeSizeFunction& size_function) - : size_function_(size_function), cost_analysis_(size_function_) {} + : size_function_(size_function) {} ~HloRematerialization() {} // Runs rematerialization on the given module. Returns whether the module was @@ -99,15 +99,29 @@ class HloRematerialization { // Call graph of the hlo_module. std::unique_ptr call_graph_; - // Analysis used for computing the rematerialization cost of instructions. - HloCostAnalysis cost_analysis_; - // The peak memory usage of each computation. The map contains only those // computations called from sequential context // (CallContext::kSequential). These values are updated as rematerialization // occurs. tensorflow::gtl::FlatMap computation_peak_memory_; + + std::unique_ptr points_to_analysis_; + + // Set of computations which have had rematerialization + // applied. Rematerialization is only applied once per computation. + tensorflow::gtl::FlatSet rematerialized_computations_; + + // Count of the total instructions rematerialized. + int64 instructions_rematerialized_ = 0; + + // Count of the net instructions added to the HLO module by + // rematerialization. This can be different than instructions_rematerialized_ + // because some rematerializations are effectively moves in the HLO + // schedule. In these cases, the rematerialization instruction replaces all + // uses of the original instruction and the original instruction is + // dead. Hence, no net instructions were added. + int64 net_instructions_added_ = 0; }; } // namespace xla diff --git a/tensorflow/compiler/xla/service/hlo_rematerialization_test.cc b/tensorflow/compiler/xla/service/hlo_rematerialization_test.cc index 0a4f2776891cfc932b4fc0627daaa9b5408f420a..2358969f38ee66e3eb024215cba4c62da3d6a32f 100644 --- a/tensorflow/compiler/xla/service/hlo_rematerialization_test.cc +++ b/tensorflow/compiler/xla/service/hlo_rematerialization_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_matchers.h" #include "tensorflow/compiler/xla/service/hlo_opcode.h" #include "tensorflow/compiler/xla/service/hlo_ordering.h" #include "tensorflow/compiler/xla/shape_util.h" @@ -30,12 +31,16 @@ limitations under the License. namespace xla { namespace { -class HloOrderingTest : public HloTestBase { +namespace op = xla::testing::opcode_matchers; + +using ::testing::_; + +class HloRematerializationTest : public HloTestBase { protected: // Creates and returns a computation which can benefit from // rematerialization. The computation looks like: // - // F32[1] %param = {...} + // F32[] %param = {...} // F32[1024] %bcast = broadcast(%param) // F32[1024] %negate = negate(%bcast) // F32[2048] %concat_1 = concat({%negate, %negate}) @@ -52,7 +57,7 @@ class HloOrderingTest : public HloTestBase { const string& suffix = "") { auto builder = HloComputation::Builder(TestName() + suffix); auto param = builder.AddInstruction( - HloInstruction::CreateParameter(0, vec1_shape_, "param")); + HloInstruction::CreateParameter(0, scalar_shape_, "param")); auto bcast = builder.AddInstruction( HloInstruction::CreateBroadcast(vec1024_shape_, param, {})); auto negate = builder.AddInstruction( @@ -62,7 +67,8 @@ class HloOrderingTest : public HloTestBase { /*dimension=*/0)); auto slice_1 = builder.AddInstruction(HloInstruction::CreateSlice( vec1_shape_, concat_1, /*start_indices=*/{0}, - /*limit_indices=*/{1})); + /*limit_indices=*/{1}, + /*strides=*/{1})); auto concat_2 = builder.AddInstruction(HloInstruction::CreateConcatenate( ShapeUtil::MakeShape(xla::F32, {1025}), {bcast, slice_1}, /*dimension=*/0)); @@ -70,14 +76,15 @@ class HloOrderingTest : public HloTestBase { // which is necessary to use this computation in a while. builder.AddInstruction(HloInstruction::CreateSlice(vec1_shape_, concat_2, /*start_indices=*/{0}, - /*limit_indices=*/{1})); + /*limit_indices=*/{1}, + /*strides=*/{1})); return builder.Build(); } // Creates and returns a computation which includes a while and can benefit // from rematerialization. The computation looks like: // - // F32[1] %param = {...} + // F32[] %param = {...} // F32[1024] %bcast = broadcast(%param) // F32[1] %slice_1 = slice(%bcast, {0:1}) // F32[1] %while = while(%slice_1, while_body, while_cond) @@ -93,12 +100,13 @@ class HloOrderingTest : public HloTestBase { const string& suffix = "") { auto builder = HloComputation::Builder(TestName() + suffix); auto param = builder.AddInstruction( - HloInstruction::CreateParameter(0, vec1_shape_, "param")); + HloInstruction::CreateParameter(0, scalar_shape_, "param")); auto bcast = builder.AddInstruction( HloInstruction::CreateBroadcast(vec1024_shape_, param, {})); auto slice_1 = builder.AddInstruction( HloInstruction::CreateSlice(vec1_shape_, bcast, /*start_indices=*/{0}, - /*limit_indices=*/{1})); + /*limit_indices=*/{1}, + /*strides=*/{1})); auto while_inst = builder.AddInstruction(HloInstruction::CreateWhile( vec1_shape_, while_cond, while_body, slice_1)); auto concat = builder.AddInstruction(HloInstruction::CreateConcatenate( @@ -106,7 +114,8 @@ class HloOrderingTest : public HloTestBase { /*dimension=*/0)); builder.AddInstruction(HloInstruction::CreateSlice(vec1_shape_, concat, /*start_indices=*/{0}, - /*limit_indices=*/{1})); + /*limit_indices=*/{1}, + /*strides=*/{1})); return builder.Build(); } @@ -117,7 +126,7 @@ class HloOrderingTest : public HloTestBase { builder.AddInstruction( HloInstruction::CreateParameter(0, vec1_shape_, "param")); builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(true))); + HloInstruction::CreateConstant(Literal::CreateR0(true))); return builder.Build(); } @@ -127,33 +136,33 @@ class HloOrderingTest : public HloTestBase { } // Various shapes used in the canned computations. + const Shape scalar_shape_ = ShapeUtil::MakeShape(xla::F32, {}); const Shape vec1_shape_ = ShapeUtil::MakeShape(xla::F32, {1}); const Shape vec1024_shape_ = ShapeUtil::MakeShape(xla::F32, {1024}); }; // Test rematerialization of a single computation produced by // MakeRematerializableComputation. -TEST_F(HloOrderingTest, SingleComputation) { - HloModule module(TestName()); +TEST_F(HloRematerializationTest, SingleComputation) { + auto module = CreateNewModule(); HloComputation* computation = - module.AddEntryComputation(MakeRematerializableComputation()); + module->AddEntryComputation(MakeRematerializableComputation()); // Find and save the original broadcast instruction which should be // rematerialized. const HloInstruction* slice = computation->root_instruction(); - ASSERT_EQ(HloOpcode::kSlice, slice->opcode()); + ASSERT_THAT(slice, op::Slice(op::Concatenate(op::Broadcast(_), _))); const HloInstruction* concat = slice->operand(0); - ASSERT_EQ(HloOpcode::kConcatenate, concat->opcode()); const HloInstruction* bcast = concat->operand(0); - ASSERT_EQ(HloOpcode::kBroadcast, bcast->opcode()); SequentialHloOrdering::HloModuleSequence sequence; // Computation requires 16KB without rematerialization, but uses only 12KB // with rematerialization so pick a memory limit between these values (14KB). - TF_ASSIGN_OR_ASSERT_OK( - bool changed, HloRematerialization::RematerializeAndSchedule( - ByteSizeOf, - /*memory_limit_bytes=*/14 * 1024, &module, &sequence)); + TF_ASSERT_OK_AND_ASSIGN( + bool changed, + HloRematerialization::RematerializeAndSchedule( + ByteSizeOf, + /*memory_limit_bytes=*/14 * 1024, module.get(), &sequence)); EXPECT_TRUE(changed); // Root should not have changed. @@ -161,8 +170,7 @@ TEST_F(HloOrderingTest, SingleComputation) { // The broadcast should have been rematerialized. const HloInstruction* remat_bcast = concat->operand(0); - EXPECT_EQ(HloOpcode::kBroadcast, remat_bcast->opcode()); - EXPECT_NE(bcast, remat_bcast); + EXPECT_THAT(remat_bcast, op::Broadcast(::testing::Ne(bcast))); // The rematerialized broadcast should be immediate before the concat in the // sequence. @@ -175,18 +183,19 @@ TEST_F(HloOrderingTest, SingleComputation) { // Test rematerialization of a single computation produced by // MakeRematerializableComputation but with a sufficiently high memory limit // such that no instructions are rematerialized. -TEST_F(HloOrderingTest, SingleComputationNoRematerialization) { - HloModule module(TestName()); +TEST_F(HloRematerializationTest, SingleComputationNoRematerialization) { + auto module = CreateNewModule(); HloComputation* computation = - module.AddEntryComputation(MakeRematerializableComputation()); + module->AddEntryComputation(MakeRematerializableComputation()); EXPECT_EQ(computation->instruction_count(), 7); SequentialHloOrdering::HloModuleSequence sequence; - TF_ASSIGN_OR_ASSERT_OK( - bool changed, HloRematerialization::RematerializeAndSchedule( - ByteSizeOf, - /*memory_limit_bytes=*/20 * 1024, &module, &sequence)); + TF_ASSERT_OK_AND_ASSIGN( + bool changed, + HloRematerialization::RematerializeAndSchedule( + ByteSizeOf, + /*memory_limit_bytes=*/20 * 1024, module.get(), &sequence)); // No instructions should have been materialized. EXPECT_FALSE(changed); @@ -199,21 +208,21 @@ TEST_F(HloOrderingTest, SingleComputationNoRematerialization) { // only one computation needs to have an instruction rematerialized. The entry // computation should be the one chosen because rematerialization in the while // will presumably be more expensive. -TEST_F(HloOrderingTest, RematerializeAroundWhile) { - HloModule module(TestName()); +TEST_F(HloRematerializationTest, RematerializeAroundWhile) { + auto module = CreateNewModule(); auto cond_builder = HloComputation::Builder(TestName() + ".cond"); cond_builder.AddInstruction( HloInstruction::CreateParameter(0, vec1_shape_, "param")); cond_builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(true))); + HloInstruction::CreateConstant(Literal::CreateR0(true))); HloComputation* while_cond = - module.AddEmbeddedComputation(cond_builder.Build()); + module->AddEmbeddedComputation(cond_builder.Build()); - HloComputation* body_computation = module.AddEmbeddedComputation( + HloComputation* body_computation = module->AddEmbeddedComputation( MakeRematerializableComputation(/*suffix=*/".body")); HloComputation* entry_computation = - module.AddEntryComputation(MakeRematerializableWhileComputation( + module->AddEntryComputation(MakeRematerializableWhileComputation( while_cond, /*while_body=*/body_computation)); EXPECT_EQ(entry_computation->instruction_count(), 6); @@ -223,10 +232,11 @@ TEST_F(HloOrderingTest, RematerializeAroundWhile) { // while so the peak memory use of the module is 18KB. Set the memory limit a // bit lower (17KB) to force rematerialization of the entry computation. SequentialHloOrdering::HloModuleSequence sequence; - TF_ASSIGN_OR_ASSERT_OK( - bool changed, HloRematerialization::RematerializeAndSchedule( - ByteSizeOf, - /*memory_limit_bytes=*/17 * 1024, &module, &sequence)); + TF_ASSERT_OK_AND_ASSIGN( + bool changed, + HloRematerialization::RematerializeAndSchedule( + ByteSizeOf, + /*memory_limit_bytes=*/17 * 1024, module.get(), &sequence)); EXPECT_TRUE(changed); // Only the entry computation should have a rematerialized instruction added. @@ -237,31 +247,32 @@ TEST_F(HloOrderingTest, RematerializeAroundWhile) { // Test rematerialization of a computation which calls another computation via a // while. Both the entry computation and while body computation should have // computations rematerialized. -TEST_F(HloOrderingTest, RematerializeEntryAndWhileBody) { - HloModule module(TestName()); +TEST_F(HloRematerializationTest, RematerializeEntryAndWhileBody) { + auto module = CreateNewModule(); auto cond_builder = HloComputation::Builder(TestName() + ".cond"); cond_builder.AddInstruction( HloInstruction::CreateParameter(0, vec1_shape_, "param")); cond_builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(true))); + HloInstruction::CreateConstant(Literal::CreateR0(true))); HloComputation* while_cond = - module.AddEmbeddedComputation(cond_builder.Build()); + module->AddEmbeddedComputation(cond_builder.Build()); - HloComputation* body_computation = module.AddEmbeddedComputation( + HloComputation* body_computation = module->AddEmbeddedComputation( MakeRematerializableComputation(/*suffix=*/".body")); HloComputation* entry_computation = - module.AddEntryComputation(MakeRematerializableWhileComputation( + module->AddEntryComputation(MakeRematerializableWhileComputation( while_cond, /*while_body=*/body_computation)); EXPECT_EQ(entry_computation->instruction_count(), 6); EXPECT_EQ(body_computation->instruction_count(), 7); SequentialHloOrdering::HloModuleSequence sequence; - TF_ASSIGN_OR_ASSERT_OK( - bool changed, HloRematerialization::RematerializeAndSchedule( - ByteSizeOf, - /*memory_limit_bytes=*/15 * 1024, &module, &sequence)); + TF_ASSERT_OK_AND_ASSIGN( + bool changed, + HloRematerialization::RematerializeAndSchedule( + ByteSizeOf, + /*memory_limit_bytes=*/15 * 1024, module.get(), &sequence)); EXPECT_TRUE(changed); // Both computations should have a rematerialized instruction added. @@ -271,25 +282,25 @@ TEST_F(HloOrderingTest, RematerializeEntryAndWhileBody) { // Test rematerialization of a doubly nested computation. All computations // should have an instruction rematerialized. -TEST_F(HloOrderingTest, RematerializeNestedComputations) { - HloModule module(TestName()); +TEST_F(HloRematerializationTest, RematerializeNestedComputations) { + auto module = CreateNewModule(); auto cond_builder = HloComputation::Builder(TestName() + ".cond"); cond_builder.AddInstruction( HloInstruction::CreateParameter(0, vec1_shape_, "param")); cond_builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(true))); + HloInstruction::CreateConstant(Literal::CreateR0(true))); HloComputation* while_cond = - module.AddEmbeddedComputation(cond_builder.Build()); + module->AddEmbeddedComputation(cond_builder.Build()); - HloComputation* inner_computation = module.AddEmbeddedComputation( + HloComputation* inner_computation = module->AddEmbeddedComputation( MakeRematerializableComputation(/*suffix=*/".inner")); HloComputation* middle_computation = - module.AddEmbeddedComputation(MakeRematerializableWhileComputation( + module->AddEmbeddedComputation(MakeRematerializableWhileComputation( while_cond, /*while_body=*/inner_computation, /*suffix=*/".middle")); HloComputation* entry_computation = - module.AddEntryComputation(MakeRematerializableWhileComputation( + module->AddEntryComputation(MakeRematerializableWhileComputation( while_cond, /*while_body=*/middle_computation)); EXPECT_EQ(entry_computation->instruction_count(), 6); @@ -299,10 +310,11 @@ TEST_F(HloOrderingTest, RematerializeNestedComputations) { // If all computations are maximally rematerialized then peak memory usage is // ~12K so pick something slightly larger. SequentialHloOrdering::HloModuleSequence sequence; - TF_ASSIGN_OR_ASSERT_OK( - bool changed, HloRematerialization::RematerializeAndSchedule( - ByteSizeOf, - /*memory_limit_bytes=*/13 * 1024, &module, &sequence)); + TF_ASSERT_OK_AND_ASSIGN( + bool changed, + HloRematerialization::RematerializeAndSchedule( + ByteSizeOf, + /*memory_limit_bytes=*/13 * 1024, module.get(), &sequence)); EXPECT_TRUE(changed); // All computations should have a rematerialized instruction added. @@ -311,6 +323,209 @@ TEST_F(HloOrderingTest, RematerializeNestedComputations) { EXPECT_EQ(inner_computation->instruction_count(), 8); } +TEST_F(HloRematerializationTest, InstructionRematerializedMultipleTimes) { + // Test that a single instruction is rematerialized several times. Module: + // + // Entry computation: + // F32[] %param = {...} + // F32[1024] %bcast = broadcast(%param) + // F32[1024] %add_1 = add(%bcast, bcast) + // F32[1024] %call_1 = call(Subcomputation, {%add_1}) + // F32[1024] %add_2 = add(%bcast, call_1) + // F32[1024] %call_2 = call(SubComputation, {%add_2}) + // F32[1024] %add_3 = add(%bcast, call_2) + // F32[1024] %call_3 = call(Subcomputation, {%add_3}) + // F32[1024] %add_4 = add(%bcast, call_3) + // + // Subcomputation: + // F32[1024] %param = {...} + // F32[2048] %concat = concat({%param, %param}) + // F32[1024] %slice = slice(%concat) + // + // 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(); + + HloComputation* subcomputation = nullptr; + { + auto builder = HloComputation::Builder(TestName() + ".subcomputation"); + auto param = builder.AddInstruction( + HloInstruction::CreateParameter(0, vec1024_shape_, "param")); + auto concat = builder.AddInstruction(HloInstruction::CreateConcatenate( + ShapeUtil::MakeShape(xla::F32, {2048}), {param, param}, + /*dimension=*/0)); + builder.AddInstruction(HloInstruction::CreateSlice( + vec1024_shape_, concat, /*start_indices=*/{0}, + /*limit_indices=*/{1024}, /*strides=*/{1})); + subcomputation = module->AddEmbeddedComputation(builder.Build()); + } + + auto builder = HloComputation::Builder(TestName()); + auto param = builder.AddInstruction( + HloInstruction::CreateParameter(0, scalar_shape_, "param")); + auto bcast = builder.AddInstruction( + HloInstruction::CreateBroadcast(vec1024_shape_, param, {})); + auto add_1 = builder.AddInstruction(HloInstruction::CreateBinary( + vec1024_shape_, HloOpcode::kAdd, bcast, bcast)); + auto call_1 = builder.AddInstruction( + HloInstruction::CreateCall(vec1024_shape_, {add_1}, subcomputation)); + auto add_2 = builder.AddInstruction(HloInstruction::CreateBinary( + vec1024_shape_, HloOpcode::kAdd, bcast, call_1)); + auto call_2 = builder.AddInstruction( + HloInstruction::CreateCall(vec1024_shape_, {add_2}, subcomputation)); + auto add_3 = builder.AddInstruction(HloInstruction::CreateBinary( + vec1024_shape_, HloOpcode::kAdd, bcast, call_2)); + auto call_3 = builder.AddInstruction( + HloInstruction::CreateCall(vec1024_shape_, {add_3}, subcomputation)); + auto add_4 = builder.AddInstruction(HloInstruction::CreateBinary( + vec1024_shape_, HloOpcode::kAdd, bcast, call_3)); + HloComputation* entry_computation = + module->AddEntryComputation(builder.Build()); + + auto count_broadcasts = [](const HloComputation* computation) { + int64 bcast_count = 0; + for (auto& instruction : computation->instructions()) { + if (instruction->opcode() == HloOpcode::kBroadcast) { + bcast_count++; + } + } + return bcast_count; + }; + + // Before rematerialization there should be a single broadcast instruction in + // the graph. + EXPECT_EQ(count_broadcasts(entry_computation), 1); + EXPECT_EQ(entry_computation->instruction_count(), 9); + + EXPECT_EQ(add_2->operand(0), bcast); + EXPECT_EQ(add_3->operand(0), bcast); + EXPECT_EQ(add_4->operand(0), bcast); + + SequentialHloOrdering::HloModuleSequence sequence; + // Pick a memory limit some where between 24KB (initial peak memory including + // parameter and output) and 20KB (peak memory possible with + // rematerialization). + TF_ASSERT_OK_AND_ASSIGN( + bool changed, + HloRematerialization::RematerializeAndSchedule( + ByteSizeOf, + /*memory_limit_bytes=*/22 * 1024, module.get(), &sequence)); + EXPECT_TRUE(changed); + + // The broadcast should have been rematerialized 3 times. + EXPECT_EQ(count_broadcasts(entry_computation), 4); + EXPECT_EQ(entry_computation->instruction_count(), 12); + + // The operands of add_2, add_3, and add_4 should all be rematerialized + // broadcasts. + EXPECT_NE(add_2->operand(0), bcast); + EXPECT_THAT(add_2->operand(0), op::Broadcast(param)); + EXPECT_NE(add_3->operand(0), bcast); + EXPECT_THAT(add_3->operand(0), op::Broadcast(param)); + EXPECT_NE(add_4->operand(0), bcast); + EXPECT_THAT(add_4->operand(0), op::Broadcast(param)); +} + +class IndirectUseTest : public HloRematerializationTest, + public ::testing::WithParamInterface {}; + +TEST_P(IndirectUseTest, IndirectUseNotRematerialized) { + // Test that an rematerializable instruction is not rematerialized if it has + // an indirect use. Test is parameterized on whether the value has an indirect + // use, and the instruction should be rematerialized iff the value has no + // indirect use. Module: + // + // Entry computation: + // F32[] %param = {...} + // F32[1024] %bcast = broadcast(%param) + // F32[1024] %add_1 = add(%bcast, bcast) + // F32[1024] %call = call(Subcomputation, {%add_1}) + // F32[1024] %add_2 = add(%bcast, call) + // {F32[1024], F32[1024]} %tuple = tuple(%bcast, %add_2) + // F32[1024] %gte = GetTupleElememt(%tuple, 0) + // F32[1024] %negate = negate(%gte) + // + // Subcomputation: + // F32[1024] %param = {...} + // F32[2048] %concat = concat({%param, %param}) + // F32[1024] %slice = slice(%concat) + // + // The value %bcast is live across the call and rematerialization of %bcast + // across that point would reduce peak memory use by 4KB. However, %bcast is + // used indirectly in the %negate so rematerialization should not happen. + // + // This test is parameterized on whether the broadcast has an indirect use or + // not. The indirect use is controlled by the index of the GetTupleElement + // instruction. If the element is 0, then the %negate operand aliases %bcast + // (ie %bcast is used indirectly by %negate), otherwise the %negate operand + // aliases %add_2. + const bool indirectly_used = GetParam(); + auto module = CreateNewModule(); + + HloComputation* subcomputation = nullptr; + { + auto builder = HloComputation::Builder(TestName() + ".subcomputation"); + auto param = builder.AddInstruction( + HloInstruction::CreateParameter(0, vec1024_shape_, "param")); + auto concat = builder.AddInstruction(HloInstruction::CreateConcatenate( + ShapeUtil::MakeShape(xla::F32, {2048}), {param, param}, + /*dimension=*/0)); + builder.AddInstruction(HloInstruction::CreateSlice( + vec1024_shape_, concat, /*start_indices=*/{0}, + /*limit_indices=*/{1024}, /*strides=*/{1})); + subcomputation = module->AddEmbeddedComputation(builder.Build()); + } + + auto builder = HloComputation::Builder(TestName()); + auto param = builder.AddInstruction( + HloInstruction::CreateParameter(0, scalar_shape_, "param")); + auto bcast = builder.AddInstruction( + HloInstruction::CreateBroadcast(vec1024_shape_, param, {})); + auto add_1 = builder.AddInstruction(HloInstruction::CreateBinary( + vec1024_shape_, HloOpcode::kAdd, bcast, bcast)); + auto call_1 = builder.AddInstruction( + HloInstruction::CreateCall(vec1024_shape_, {add_1}, subcomputation)); + auto add_2 = builder.AddInstruction(HloInstruction::CreateBinary( + vec1024_shape_, HloOpcode::kAdd, bcast, call_1)); + auto tuple = + builder.AddInstruction(HloInstruction::CreateTuple({bcast, add_2})); + auto gte = builder.AddInstruction(HloInstruction::CreateGetTupleElement( + vec1024_shape_, tuple, indirectly_used ? 0 : 1)); + builder.AddInstruction( + HloInstruction::CreateUnary(vec1024_shape_, HloOpcode::kNegate, gte)); + HloComputation* entry_computation = + module->AddEntryComputation(builder.Build()); + + EXPECT_EQ(entry_computation->instruction_count(), 8); + + SequentialHloOrdering::HloModuleSequence sequence; + // Pick a memory limit some where between 24KB (initial peak memory including + // parameter and output) and 20KB (peak memory possible with + // rematerialization). + TF_ASSERT_OK_AND_ASSIGN( + bool changed, + HloRematerialization::RematerializeAndSchedule( + ByteSizeOf, + /*memory_limit_bytes=*/22 * 1024, module.get(), &sequence)); + // Rematerialization should only occur if the rematerializable instruction has + // no indirect uses. + if (indirectly_used) { + EXPECT_FALSE(changed); + EXPECT_EQ(entry_computation->instruction_count(), 8); + } else { + EXPECT_TRUE(changed); + EXPECT_EQ(entry_computation->instruction_count(), 9); + } +} + +INSTANTIATE_TEST_CASE_P(IndirectUseTestInstantiation, IndirectUseTest, + ::testing::Values(true, false)); + } // namespace } // namespace xla + +int main(int argc, char** argv) { + return xla::ParseDebugOptionsFlagsAndRunTests(argc, argv); +} diff --git a/tensorflow/compiler/xla/service/hlo_scheduling.cc b/tensorflow/compiler/xla/service/hlo_scheduling.cc new file mode 100644 index 0000000000000000000000000000000000000000..25be448c8d186514e5d5d04382f4733fee3af68b --- /dev/null +++ b/tensorflow/compiler/xla/service/hlo_scheduling.cc @@ -0,0 +1,435 @@ +/* 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/compiler/xla/service/hlo_scheduling.h" + +#include +#include + +#include "tensorflow/compiler/xla/service/heap_simulator.h" +#include "tensorflow/compiler/xla/service/hlo_computation.h" +#include "tensorflow/compiler/xla/service/tuple_points_to_analysis.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/types.h" +#include "tensorflow/compiler/xla/util.h" +#include "tensorflow/core/lib/core/errors.h" +#include "tensorflow/core/lib/strings/str_util.h" +#include "tensorflow/core/lib/strings/stringprintf.h" +#include "tensorflow/core/platform/logging.h" + +namespace xla { + +StatusOr MinimumMemoryForSequence( + const SequentialHloOrdering::HloModuleSequence& module_sequence, + const LogicalBuffer::SizeFunction& size_function) { + if (module_sequence.empty()) { + return 0; + } + + const HloModule* module = module_sequence.begin()->first->parent(); + TF_ASSIGN_OR_RETURN(std::unique_ptr points_to_analysis, + TuplePointsToAnalysis::Run(module)); + + // The absolute minimum memory required for a given sequence of instructions + // is determined by the sequence of Alloc and Free calls on a simulated heap, + // ignoring fragmentation. We run the heap simulation on the whole module, + // rather than summing each computation, since it gives us a better lower + // bound, by minimizing the liveness of sub-computations. + TF_ASSIGN_OR_RETURN( + HeapSimulator::Result result, + HeapSimulator::Run(MakeUnique(), *module, + module_sequence, *points_to_analysis, size_function)); + return result.heap_size; +} + +namespace { + +// Class implementing a list scheduler of HLO instructions which produces a +// sequence which minimizes memory usage. +class ListScheduler { + public: + // Construct and return a memory-minimizing sequence of HLO instructions + // containing the given HLO computation. + static StatusOr> Run( + const HloComputation& computation, + const TuplePointsToAnalysis& points_to_analysis, + const LogicalBuffer::SizeFunction& size_function) { + ListScheduler scheduler(computation, points_to_analysis, size_function); + return scheduler.CreateSchedule(); + } + + // Returns whether the memory used by the given HLO should be ignored by the + // scheduling heuristic. + static bool IgnoreInstruction(const HloInstruction& instruction) { + return instruction.opcode() == HloOpcode::kParameter || + instruction.opcode() == HloOpcode::kConstant; + } + + private: + // The scheduling priority of an instruction is first the number of bytes + // freed by scheduling the instruction, and second (tie-breaker) by the number + // of users. This is represented as a std::pair containing these two values + // (first element is the bytes freed). std::pair provides the necessary + // comparison operators. + using Priority = std::pair; + + ListScheduler(const HloComputation& computation, + const TuplePointsToAnalysis& points_to_analysis, + const LogicalBuffer::SizeFunction& size_function) + : computation_(computation), + points_to_analysis_(points_to_analysis), + size_function_(size_function) { + // Create a map containing the LogicalBuffer uses for each HLO + // 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()) { + std::unordered_set instr_uses; + for (auto* operand : instruction->operands()) { + for (const LogicalBuffer* buffer : + points_to_analysis.GetBuffersDefinedByInstruction(operand)) { + instr_uses.insert(buffer); + } + } + buffer_uses_[instruction.get()] = std::vector( + instr_uses.begin(), instr_uses.end()); + } + + // Create map containing the number of unscheduled uses (hlo instructions) + // of each logical buffer. + for (auto& instruction : computation.instructions()) { + for (auto* buffer : points_to_analysis.GetBuffersDefinedByInstruction( + instruction.get())) { + unscheduled_use_count_[buffer] = 0; + } + } + for (auto& instruction : computation.instructions()) { + for (const LogicalBuffer* buffer : buffer_uses_.at(instruction.get())) { + ++unscheduled_use_count_[buffer]; + } + } + + // 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()) + .CreateFlattenedSet()) { + ++unscheduled_use_count_[live_out_buffer]; + } + } + + // Returns whether the memory used by the given buffer should be ignored by + // the scheduling heuristic. + static bool IgnoreBuffer(const LogicalBuffer& buffer) { + return IgnoreInstruction(*buffer.instruction()); + } + + // An entry in the worklist used by CreateSchedule. Corresponds to one + // HloInstruction, plus some cached metadata, saved for the purposes of making + // BytesFreedIfScheduled fast. + struct ReadyListEntry { + const HloInstruction* instruction; + + // The total size of all buffers defined by this instruction. + int64 bytes_defined; + + // For each buffer B used by this instruction, we keep a pair (B, U), where + // U is the number of uses of B that have not yet been scheduled. This pair + // is a pointer into the unscheduled_use_count_ map, so it gets updated for + // free when we update counts in the map. + std::vector*> + used_buffer_unscheduled_use_counts; + }; + + // Creates a ReadyListEntry for the given instruction. + ReadyListEntry MakeReadyListEntry(const HloInstruction* instruction) { + ReadyListEntry entry; + entry.instruction = instruction; + + entry.bytes_defined = 0; + for (auto* buffer : + points_to_analysis_.GetBuffersDefinedByInstruction(instruction)) { + if (!IgnoreBuffer(*buffer)) { + entry.bytes_defined += size_function_(*buffer); + } + } + + for (auto* buffer : buffer_uses_.at(instruction)) { + if (IgnoreBuffer(*buffer)) { + continue; + } + auto unscheduled_use_count_it = unscheduled_use_count_.find(buffer); + CHECK(unscheduled_use_count_it != unscheduled_use_count_.end()); + entry.used_buffer_unscheduled_use_counts.push_back( + &*unscheduled_use_count_it); + } + return entry; + } + + // Returns the number of bytes freed if the HLO instruction is scheduled. + int64 BytesFreedIfScheduled(const ReadyListEntry& entry) { + int64 freed_bytes = 0; + for (const auto& kv : entry.used_buffer_unscheduled_use_counts) { + auto buffer = kv->first; + auto use_count = kv->second; + if (use_count == 1) { + freed_bytes += size_function_(*buffer); + } + } + return freed_bytes - entry.bytes_defined; + } + + // Constructs the scheduling priority of the given instruction. + Priority GetPriority(const ReadyListEntry& entry) { + return {BytesFreedIfScheduled(entry), entry.instruction->user_count()}; + } + + std::vector CreateSchedule() { + std::vector schedule; + + // Populate the ready list with instructions which have no operands or + // control predecessors. + std::unordered_map unscheduled_pred_count; + 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()) { + unscheduled_pred_count[user]++; + } + for (const HloInstruction* succ : instruction->control_successors()) { + unscheduled_pred_count[succ]++; + } + } + + std::list ready_list; + for (auto& instruction : computation_.instructions()) { + // Instruction with no operands or control predecessors will + // not be in the map. + if (unscheduled_pred_count.count(instruction.get()) == 0) { + ready_list.push_back(MakeReadyListEntry(instruction.get())); + } + } + + while (!ready_list.empty()) { + // Select the highest priority HLO instruction from the ready list. + auto best_it = ready_list.begin(); + Priority best_priority = GetPriority(*best_it); + for (auto ready_it = std::next(ready_list.begin()); + ready_it != ready_list.end(); ++ready_it) { + Priority priority = GetPriority(*ready_it); + if (priority > best_priority) { + best_it = ready_it; + best_priority = priority; + } + } + + // Remove the selected instruction from the ready list and add it to the + // schedule. + const HloInstruction* best = best_it->instruction; + ready_list.erase(best_it); + schedule.push_back(best); + scheduled_instructions_.insert(best); + + // Update the unscheduled uses of the logical buffers. + for (const LogicalBuffer* buffer : buffer_uses_.at(best)) { + CHECK_GT(unscheduled_use_count_.at(buffer), 0); + --unscheduled_use_count_[buffer]; + } + + // Add new instructions to ready list. + auto update_pred_count = [&](HloInstruction* inst) { + int64 pred_count = --unscheduled_pred_count.at(inst); + CHECK_GE(pred_count, 0); + if (pred_count == 0) { + ready_list.push_back(MakeReadyListEntry(inst)); + } + }; + // TODO(b/34466113): Replace this and above with successors() or + // predecessors() when these methods are added to HloInstruction. + for (HloInstruction* user : best->users()) { + update_pred_count(user); + } + for (HloInstruction* succ : best->control_successors()) { + update_pred_count(succ); + } + } + CHECK_EQ(schedule.size(), computation_.instructions().size()); + CHECK_EQ(scheduled_instructions_.size(), + computation_.instructions().size()); + + return schedule; + } + + const HloComputation& computation_; + const TuplePointsToAnalysis& points_to_analysis_; + const LogicalBuffer::SizeFunction& size_function_; + + // A map containing the LogicalBuffers that each instruction uses. + std::unordered_map> + buffer_uses_; + + // A map containing the count of unscheduled HLOs which using a particular + // LogicalBuffer. We rely on iterator stability in this map. + std::unordered_map unscheduled_use_count_; + + // Set of instructions which have been scheduled. + std::unordered_set scheduled_instructions_; +}; + +int64 SumLogicalBufferSizes( + const TuplePointsToAnalysis::BufferDefinitionVector& buffers, + const LogicalBuffer::SizeFunction& size_function) { + int64 size = 0; + for (const LogicalBuffer* buffer : buffers) { + size += size_function(*buffer); + } + return size; +} + +StatusOr> RunDFSMemoryScheduler( + const HloComputation& computation, + const TuplePointsToAnalysis& points_to_analysis, + const LogicalBuffer::SizeFunction& size_function) { + // This ordering is based on DFS post-order, with a heuristic to decide which + // operand to visit first. The heuristic is based on 'extra_users', which is + // simply users-1 for each instruction. By subtracting 1, we're saying that + // instructions with no users or a single user don't count; instructions with + // lots of fan-out will be visited earlier. + tensorflow::gtl::FlatMap extra_users; + tensorflow::gtl::FlatMap total_sizes; + for (const HloInstruction* hlo : computation.MakeInstructionPostOrder()) { + if (ListScheduler::IgnoreInstruction(*hlo)) { + extra_users[hlo] = 0; + total_sizes[hlo] = 0; + continue; + } + extra_users[hlo] = hlo->users().empty() ? 0 : hlo->users().size() - 1; + total_sizes[hlo] = SumLogicalBufferSizes( + points_to_analysis.GetBuffersDefinedByInstruction(hlo), size_function); + tensorflow::gtl::FlatSet unique_operands( + hlo->operands().begin(), hlo->operands().end()); + for (const HloInstruction* operand : unique_operands) { + extra_users[hlo] += extra_users[operand]; + total_sizes[hlo] += total_sizes[operand]; + } + } + CHECK_EQ(extra_users.size(), computation.instructions().size()); + CHECK_EQ(total_sizes.size(), computation.instructions().size()); + + // Construct a total order based on DFS post-order, visiting operands in + // decreasing cumulative extra user order, and next by cumulative size, with a + // tiebreaker by name for determinism. + std::vector sequence; + FunctionVisitor visitor([&sequence](HloInstruction* hlo) { + sequence.push_back(hlo); + return Status::OK(); + }); + TF_RETURN_IF_ERROR(computation.AcceptWithOperandOrder( + &visitor, [&extra_users, &total_sizes](const HloInstruction* a, + const HloInstruction* b) { + if (extra_users[a] != extra_users[b]) { + return extra_users[a] > extra_users[b]; + } + if (total_sizes[a] != total_sizes[b]) { + return total_sizes[a] > total_sizes[b]; + } + return a->name() < b->name(); + })); + CHECK_EQ(sequence.size(), computation.instructions().size()); + return sequence; +} + +StatusOr MinimumMemoryForComputation( + const HloComputation& computation, + const std::vector& sequence, + const TuplePointsToAnalysis& points_to_analysis, + const LogicalBuffer::SizeFunction& size_function) { + TF_ASSIGN_OR_RETURN( + HeapSimulator::Result result, + HeapSimulator::Run(MakeUnique(), computation, + sequence, points_to_analysis, size_function)); + return result.heap_size; +} + +StatusOr> CreateMemoryMinimizingSequence( + const HloComputation& computation, + const TuplePointsToAnalysis& points_to_analysis, + const LogicalBuffer::SizeFunction& size_function) { + // We try both a list-scheduler based ordering and a DFS based ordering, and + // choose whichever returns a lower min-memory, not accounting for + // fragmentation. + // + // Note that this is just a heuristic. One obvious inaccuracy is that the + // memory required for sub-computations might be different when considered + // within the caller's context. But it's good enough for now. + TF_ASSIGN_OR_RETURN( + std::vector list_sequence, + ListScheduler::Run(computation, points_to_analysis, size_function)); + TF_ASSIGN_OR_RETURN( + const int64 list_memory, + MinimumMemoryForComputation(computation, list_sequence, + points_to_analysis, size_function)); + VLOG(2) << "Min-memory list sequence: " << list_memory << " bytes"; + + TF_ASSIGN_OR_RETURN( + std::vector dfs_sequence, + RunDFSMemoryScheduler(computation, points_to_analysis, size_function)); + TF_ASSIGN_OR_RETURN( + const int64 dfs_memory, + MinimumMemoryForComputation(computation, dfs_sequence, points_to_analysis, + size_function)); + VLOG(2) << "Min-memory dfs sequence: " << dfs_memory << " bytes"; + + if (list_memory <= dfs_memory) { + VLOG(2) << "Chose min-memory list sequence: " << list_memory << " bytes"; + return list_sequence; + } else { + VLOG(2) << "Chose min-memory dfs sequence: " << dfs_memory << " bytes"; + return dfs_sequence; + } +} + +} // namespace + +StatusOr +CreateMemoryMinimizingSequence( + const HloModule& module, const LogicalBuffer::SizeFunction& size_function) { + SequentialHloOrdering::HloModuleSequence sequence; + TF_ASSIGN_OR_RETURN(std::unique_ptr points_to_analysis, + TuplePointsToAnalysis::Run(&module)); + for (const auto& computation : module.computations()) { + if (computation->IsFusionComputation()) { + continue; + } + TF_ASSIGN_OR_RETURN(sequence[computation.get()], + CreateMemoryMinimizingSequence( + *computation, *points_to_analysis, size_function)); + } + return sequence; +} + +StatusOr> CreateMemoryMinimizingSequence( + const HloComputation& computation, + const LogicalBuffer::SizeFunction& size_function) { + CHECK(!computation.IsFusionComputation()); + TF_ASSIGN_OR_RETURN(std::unique_ptr points_to_analysis, + TuplePointsToAnalysis::Run(computation.parent())); + return CreateMemoryMinimizingSequence(computation, *points_to_analysis, + size_function); +} + +} // namespace xla diff --git a/tensorflow/compiler/xla/service/hlo_scheduling.h b/tensorflow/compiler/xla/service/hlo_scheduling.h new file mode 100644 index 0000000000000000000000000000000000000000..ec92a56b962152b15981f868369683144aa7c76a --- /dev/null +++ b/tensorflow/compiler/xla/service/hlo_scheduling.h @@ -0,0 +1,50 @@ +/* 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_COMPILER_XLA_SERVICE_HLO_SCHEDULING_H_ +#define TENSORFLOW_COMPILER_XLA_SERVICE_HLO_SCHEDULING_H_ + +#include + +#include "tensorflow/compiler/xla/service/hlo_instruction.h" +#include "tensorflow/compiler/xla/service/hlo_module.h" +#include "tensorflow/compiler/xla/service/hlo_ordering.h" +#include "tensorflow/compiler/xla/service/logical_buffer.h" +#include "tensorflow/compiler/xla/statusor.h" +#include "tensorflow/compiler/xla/types.h" + +namespace xla { + +// Returns the minimum memory required to compute the given module sequence, +// assuming no fragmentation. +StatusOr MinimumMemoryForSequence( + const SequentialHloOrdering::HloModuleSequence& module_sequence, + const LogicalBuffer::SizeFunction& size_function); + +// Returns an HloModuleSequence which seeks to minimize the memory required for +// the computation. size_function is the function returning the number of bytes +// required for a LogicalBuffer. +StatusOr +CreateMemoryMinimizingSequence( + const HloModule& module, const LogicalBuffer::SizeFunction& size_function); + +// Overload of above that computes the sequence for a single computation. +StatusOr> CreateMemoryMinimizingSequence( + const HloComputation& computation, + const LogicalBuffer::SizeFunction& size_function); + +} // namespace xla + +#endif // TENSORFLOW_COMPILER_XLA_SERVICE_HLO_SCHEDULING_H_ diff --git a/tensorflow/compiler/xla/service/hlo_scheduling_test.cc b/tensorflow/compiler/xla/service/hlo_scheduling_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..d09d22ee40638c5beed3f4eaf3723be0f6b6bf96 --- /dev/null +++ b/tensorflow/compiler/xla/service/hlo_scheduling_test.cc @@ -0,0 +1,97 @@ +/* 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/hlo_scheduling.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_opcode.h" +#include "tensorflow/compiler/xla/service/hlo_ordering.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/compiler/xla/xla_data.pb.h" + +namespace xla { +namespace { + +class MinimumMemoryForSequenceTest : public HloTestBase {}; + +TEST_F(MinimumMemoryForSequenceTest, MultiComputation) { + auto module = CreateNewModule(); + const Shape scalar_shape = ShapeUtil::MakeShape(xla::F32, {}); + const Shape tuple_shape = + ShapeUtil::MakeTupleShape({scalar_shape, scalar_shape}); + + auto cond_builder = HloComputation::Builder("WhileCond"); + // Tuple param: 24 bytes (each elem has 8 byte pointer, 4 byte element) + HloInstruction* cond_param = cond_builder.AddInstruction( + HloInstruction::CreateParameter(0, tuple_shape, "cond_param")); + HloInstruction* cond_iter = cond_builder.AddInstruction( + HloInstruction::CreateGetTupleElement(scalar_shape, cond_param, 0)); + HloInstruction* cond_data = cond_builder.AddInstruction( + HloInstruction::CreateGetTupleElement(scalar_shape, cond_param, 1)); + // Free cond_param[] (16 bytes), Alloc PRED[] (1 byte) + HloInstruction* cond_lt = cond_builder.AddInstruction( + HloInstruction::CreateBinary(ShapeUtil::MakeShape(PRED, {}), + HloOpcode::kLt, cond_iter, cond_data)); + HloComputation* cond_computation = + module->AddEmbeddedComputation(cond_builder.Build()); + + auto body_builder = HloComputation::Builder("WhileBody"); + // Tuple param: 24 bytes (each elem has 8 byte pointer, 4 byte element) + HloInstruction* body_param = body_builder.AddInstruction( + HloInstruction::CreateParameter(0, tuple_shape, "body_param")); + HloComputation* body_computation = + module->AddEmbeddedComputation(body_builder.Build()); + + auto builder = HloComputation::Builder(TestName()); + // Entry params: 8 bytes (4 bytes per param), TOTAL=8 + HloInstruction* iter = builder.AddInstruction( + HloInstruction::CreateParameter(0, scalar_shape, "param_iter")); + HloInstruction* data = builder.AddInstruction( + HloInstruction::CreateParameter(1, scalar_shape, "param_data")); + // Tuple: 16 bytes (8 bytes per pointer), TOTAL=24 + HloInstruction* tuple = + builder.AddInstruction(HloInstruction::CreateTuple({iter, data})); + // While: 8 bytes (4 bytes per element), TOTAL=32 + // Both cond and body use a max of 24 bytes, TOTAL=56 + HloInstruction* while_op = builder.AddInstruction(HloInstruction::CreateWhile( + tuple_shape, cond_computation, body_computation, tuple)); + HloComputation* entry_computation = + module->AddEntryComputation(builder.Build()); + + auto size_fn = [](const LogicalBuffer& buffer) { + return ShapeUtil::ByteSizeOf(buffer.shape(), /*pointer_size=*/8); + }; + + SequentialHloOrdering::HloModuleSequence module_sequence; + module_sequence[cond_computation] = {cond_param, cond_iter, cond_data, + cond_lt}; + module_sequence[body_computation] = {body_param}; + module_sequence[entry_computation] = {iter, data, tuple, while_op}; + EXPECT_EQ(56, + MinimumMemoryForSequence(module_sequence, size_fn).ValueOrDie()); +} + +} // namespace +} // namespace xla + +int main(int argc, char** argv) { + return xla::ParseDebugOptionsFlagsAndRunTests(argc, argv); +} diff --git a/tensorflow/compiler/xla/service/hlo_subcomputation_unification_test.cc b/tensorflow/compiler/xla/service/hlo_subcomputation_unification_test.cc index 14800b53420a18c89dc478ee41d8d7c258f728fc..e3d287d4c91708577b712261842b6ae231fb188b 100644 --- a/tensorflow/compiler/xla/service/hlo_subcomputation_unification_test.cc +++ b/tensorflow/compiler/xla/service/hlo_subcomputation_unification_test.cc @@ -66,16 +66,16 @@ class HloSubcomputationUnificationTest : public HloTestBase { }; TEST_F(HloSubcomputationUnificationTest, UnifyIdentities) { - auto hlo_module = MakeUnique("test_module"); + auto module = CreateNewModule(); auto builder = HloComputation::Builder(TestName()); auto callee1 = - hlo_module->AddEmbeddedComputation(CreateR0S32IdentityComputation()); + module->AddEmbeddedComputation(CreateR0S32IdentityComputation()); auto callee2 = - hlo_module->AddEmbeddedComputation(CreateR0S32IdentityComputation()); + module->AddEmbeddedComputation(CreateR0S32IdentityComputation()); auto constant = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(5))); + HloInstruction::CreateConstant(Literal::CreateR0(5))); auto x = builder.AddInstruction( HloInstruction::CreateCall(r0s32_, {constant}, callee1)); auto y = builder.AddInstruction( @@ -83,37 +83,38 @@ TEST_F(HloSubcomputationUnificationTest, UnifyIdentities) { builder.AddInstruction( HloInstruction::CreateBinary(r0s32_, HloOpcode::kAdd, x, y)); - hlo_module->AddEntryComputation(builder.Build()); + module->AddEntryComputation(builder.Build()); - EXPECT_EQ(3, hlo_module->computations().size()); + EXPECT_EQ(3, module->computations().size()); EXPECT_NE(x->to_apply(), y->to_apply()); if (VLOG_IS_ON(1)) { - hlo_graph_dumper::DumpGraph(*hlo_module->entry_computation(), - "before unification", false, false, nullptr); + hlo_graph_dumper::DumpGraph(*module->entry_computation(), + "before unification", + module->config().debug_options()); } - EXPECT_TRUE( - HloSubcomputationUnification().Run(hlo_module.get()).ValueOrDie()); + EXPECT_TRUE(HloSubcomputationUnification().Run(module.get()).ValueOrDie()); if (VLOG_IS_ON(1)) { - hlo_graph_dumper::DumpGraph(*hlo_module->entry_computation(), - "after unification", false, false, nullptr); + hlo_graph_dumper::DumpGraph(*module->entry_computation(), + "after unification", + module->config().debug_options()); } - EXPECT_EQ(2, hlo_module->computations().size()); + EXPECT_EQ(2, module->computations().size()); EXPECT_EQ(x->to_apply(), y->to_apply()); } TEST_F(HloSubcomputationUnificationTest, UnifyAdditions) { - auto hlo_module = MakeUnique("test_module"); + auto module = CreateNewModule(); auto builder = HloComputation::Builder(TestName()); auto callee1 = - hlo_module->AddEmbeddedComputation(CreateR0S32AdditionComputation()); + module->AddEmbeddedComputation(CreateR0S32AdditionComputation()); auto callee2 = - hlo_module->AddEmbeddedComputation(CreateR0S32AdditionComputation()); + module->AddEmbeddedComputation(CreateR0S32AdditionComputation()); auto constant1 = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(5))); + HloInstruction::CreateConstant(Literal::CreateR0(5))); auto constant2 = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(3))); + HloInstruction::CreateConstant(Literal::CreateR0(3))); auto x = builder.AddInstruction( HloInstruction::CreateCall(r0s32_, {constant1, constant2}, callee1)); auto y = builder.AddInstruction( @@ -121,33 +122,34 @@ TEST_F(HloSubcomputationUnificationTest, UnifyAdditions) { builder.AddInstruction( HloInstruction::CreateBinary(r0s32_, HloOpcode::kAdd, x, y)); - hlo_module->AddEntryComputation(builder.Build()); + module->AddEntryComputation(builder.Build()); - EXPECT_EQ(3, hlo_module->computations().size()); + EXPECT_EQ(3, module->computations().size()); EXPECT_NE(x->to_apply(), y->to_apply()); if (VLOG_IS_ON(1)) { - hlo_graph_dumper::DumpGraph(*hlo_module->entry_computation(), - "before unification", false, false, nullptr); + hlo_graph_dumper::DumpGraph(*module->entry_computation(), + "before unification", + module->config().debug_options()); } - EXPECT_TRUE( - HloSubcomputationUnification().Run(hlo_module.get()).ValueOrDie()); + EXPECT_TRUE(HloSubcomputationUnification().Run(module.get()).ValueOrDie()); if (VLOG_IS_ON(1)) { - hlo_graph_dumper::DumpGraph(*hlo_module->entry_computation(), - "after unification", false, false, nullptr); + hlo_graph_dumper::DumpGraph(*module->entry_computation(), + "after unification", + module->config().debug_options()); } - EXPECT_EQ(2, hlo_module->computations().size()); + EXPECT_EQ(2, module->computations().size()); EXPECT_EQ(x->to_apply(), y->to_apply()); } // Do not unify subcomputations with different parameter shapes. TEST_F(HloSubcomputationUnificationTest, DifferentParameterShapes) { - auto hlo_module = MakeUnique("test_module"); + auto module = CreateNewModule(); auto builder = HloComputation::Builder(TestName()); - auto callee1 = hlo_module->AddEmbeddedComputation( - CreateR1S32AdditionComputation(r1s32_5_)); - auto callee2 = hlo_module->AddEmbeddedComputation( - CreateR1S32AdditionComputation(r1s32_3_)); + auto callee1 = + module->AddEmbeddedComputation(CreateR1S32AdditionComputation(r1s32_5_)); + auto callee2 = + module->AddEmbeddedComputation(CreateR1S32AdditionComputation(r1s32_3_)); auto param1 = builder.AddInstruction( HloInstruction::CreateParameter(0, r1s32_5_, "param1")); @@ -160,28 +162,29 @@ TEST_F(HloSubcomputationUnificationTest, DifferentParameterShapes) { builder.AddInstruction(HloInstruction::CreateConcatenate( ShapeUtil::MakeShape(S32, {8}), {x, y}, 0)); - hlo_module->AddEntryComputation(builder.Build()); + module->AddEntryComputation(builder.Build()); - EXPECT_EQ(3, hlo_module->computations().size()); + EXPECT_EQ(3, module->computations().size()); EXPECT_NE(x->to_apply(), y->to_apply()); if (VLOG_IS_ON(1)) { - hlo_graph_dumper::DumpGraph(*hlo_module->entry_computation(), - "before unification", false, false, nullptr); + hlo_graph_dumper::DumpGraph(*module->entry_computation(), + "before unification", + module->config().debug_options()); } - EXPECT_FALSE( - HloSubcomputationUnification().Run(hlo_module.get()).ValueOrDie()); + EXPECT_FALSE(HloSubcomputationUnification().Run(module.get()).ValueOrDie()); if (VLOG_IS_ON(1)) { - hlo_graph_dumper::DumpGraph(*hlo_module->entry_computation(), - "after unification", false, false, nullptr); + hlo_graph_dumper::DumpGraph(*module->entry_computation(), + "after unification", + module->config().debug_options()); } - EXPECT_EQ(3, hlo_module->computations().size()); + EXPECT_EQ(3, module->computations().size()); EXPECT_NE(x->to_apply(), y->to_apply()); } // Regression test for b/31466798. Checks that entry_computation is still valid // after unification. TEST_F(HloSubcomputationUnificationTest, TwoIdenticalComputations) { - HloModule module(TestName()); + auto module = CreateNewModule(); for (int i = 0; i < 2; ++i) { HloComputation::Builder builder("pow"); auto x = @@ -191,15 +194,19 @@ TEST_F(HloSubcomputationUnificationTest, TwoIdenticalComputations) { builder.AddInstruction( HloInstruction::CreateBinary(r0f32_, HloOpcode::kPower, x, y)); if (i == 0) { - module.AddEmbeddedComputation(builder.Build()); + module->AddEmbeddedComputation(builder.Build()); } else { - module.AddEntryComputation(builder.Build()); + module->AddEntryComputation(builder.Build()); } } - EXPECT_TRUE(HloSubcomputationUnification().Run(&module).ValueOrDie()); - EXPECT_EQ(1, module.computations().size()); - EXPECT_EQ(module.computations().front().get(), module.entry_computation()); + EXPECT_TRUE(HloSubcomputationUnification().Run(module.get()).ValueOrDie()); + EXPECT_EQ(1, module->computations().size()); + EXPECT_EQ(module->computations().front().get(), module->entry_computation()); } } // namespace xla + +int main(int argc, char** argv) { + return xla::ParseDebugOptionsFlagsAndRunTests(argc, argv); +} diff --git a/tensorflow/compiler/xla/service/hlo_tfgraph_builder.cc b/tensorflow/compiler/xla/service/hlo_tfgraph_builder.cc new file mode 100644 index 0000000000000000000000000000000000000000..5a4c93b59a6810b962e3c8f54b2964dffa8ecd6d --- /dev/null +++ b/tensorflow/compiler/xla/service/hlo_tfgraph_builder.cc @@ -0,0 +1,214 @@ +/* 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/hlo_tfgraph_builder.h" +#include "tensorflow/compiler/xla/layout_util.h" +#include "tensorflow/compiler/xla/literal_util.h" +#include "tensorflow/compiler/xla/service/hlo_opcode.h" +#include "tensorflow/compiler/xla/shape_util.h" +#include "tensorflow/core/framework/attr_value.pb.h" +#include "tensorflow/core/framework/op.h" +#include "tensorflow/core/framework/tensor_shape.pb.h" +#include "tensorflow/core/lib/strings/str_util.h" +#include "tensorflow/core/lib/strings/strcat.h" + +using ::tensorflow::GraphDef; +using ::tensorflow::NodeDef; +using ::tensorflow::TensorShapeProto; +using ::tensorflow::strings::StrAppend; +using ::tensorflow::strings::StrCat; +using ::tensorflow::str_util::Join; + +namespace xla { +namespace hlo_graph_dumper { +namespace { + +string GetOpDefName(const HloInstruction* instruction) { + string name = StrCat("hlo-", HloOpcodeString(instruction->opcode())); + tensorflow::str_util::TitlecaseString(&name, "-"); + name.erase(std::remove(name.begin(), name.end(), '-'), name.end()); + + if (instruction->opcode() == HloOpcode::kFusion) { + string fusion_name = ToString(instruction->fusion_kind()); + StrAppend(&name, tensorflow::StringPiece(fusion_name).substr(1)); + } + return name; +} + +TensorShapeProto GetTensorShape(const HloInstruction* instruction) { + TensorShapeProto tensor_shape; + const Shape& shape = instruction->shape(); + for (auto dim : shape.dimensions()) { + tensor_shape.add_dim()->set_size(dim); + } + return tensor_shape; +} + +} // namespace + +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(); + }; + std::replace_if(name->begin(), name->end(), pred, '_'); +} + +Status HloTfGraphBuilder::AddComputation(const HloComputation& computation) { + VLOG(2) << "Adding computation " << computation.name(); + for (auto embedded : computation.MakeEmbeddedComputationsList()) { + for (auto& instruction : embedded->instructions()) { + TF_RETURN_IF_ERROR(AddInstruction(instruction.get())); + } + } + for (auto& instruction : computation.instructions()) { + TF_RETURN_IF_ERROR(AddInstruction(instruction.get())); + } + return Status::OK(); +} + +const GraphDef& HloTfGraphBuilder::GetGraphDef() const { return graph_def_; } + +const string& HloTfGraphBuilder::GetNodeNameForInstruction( + const HloInstruction* instruction) { + if (ContainsKey(instruction_to_node_name_, instruction)) { + return instruction_to_node_name_[instruction]; + } + string node_name; + // If an instruction is fused, put it in the subgraph of the fusion; + // otherwise, put it in the computation subgraph. + const HloComputation* computation = instruction->parent(); + if (computation->IsFusionComputation()) { + node_name = GetNodeNameForInstruction(computation->FusionInstruction()); + } else { + node_name = computation->name(); + if (!instruction->metadata().op_name().empty()) { + // Always make computations contain TF ops but not the other way around. + StrAppend(&node_name, "/", instruction->metadata().op_name()); + } + } + string instruction_name = instruction->name(); + if (instruction->opcode() == HloOpcode::kParameter) { + StrAppend(&instruction_name, ".", instruction->parameter_number()); + } + StrAppend(&node_name, "/", instruction_name); + CleanNodeName(&node_name); + auto ret = + instruction_to_node_name_.insert(std::make_pair(instruction, node_name)); + CHECK(ret.second); + return ret.first->second; +} + +void HloTfGraphBuilder::SetNodeAttrs(const HloInstruction* instruction, + NodeDef* node_def) const { + auto& attrs = *node_def->mutable_attr(); + + // Set the number of arguments for instructions that have variadic operands. + if (HloOpcodeIsVariadic(instruction->opcode())) { + tensorflow::AttrValue attr_value; + attr_value.set_i(instruction->operands().size()); + attrs["arg_num"] = attr_value; + } + + // Set the node type. + attrs["type"].set_s( + xla::PrimitiveType_Name(instruction->shape().element_type())); + + // Set the framework op (e.g. Tensorflow op) that generated this XLA op. + attrs["tf_op_type"].set_s(instruction->metadata().op_type()); + attrs["tf_op_name"].set_s(instruction->metadata().op_name()); + + // Set the shape of the output tensor. "_output_shapes" is a special attribute + // name used by Tensorboard for shapes of output tensors. + tensorflow::AttrValue shapes; + *shapes.mutable_list()->add_shape() = GetTensorShape(instruction); + attrs["_output_shapes"] = shapes; + + // Set the layout. + if (LayoutUtil::HasLayout(instruction->shape())) { + string layout_string; + if (ShapeUtil::IsTuple(instruction->shape())) { + // For tuples, emit the full shape because the layout of a tuple is not + // represented in a single Layout field. + layout_string = ShapeUtil::HumanStringWithLayout(instruction->shape()); + } else { + layout_string = StrCat( + "{", Join(instruction->shape().layout().minor_to_major(), ","), "}"); + } + attrs["layout"].set_s(layout_string); + } + + // Set op-specific attributes. + switch (instruction->opcode()) { + case HloOpcode::kConcatenate: + case HloOpcode::kBroadcast: + case HloOpcode::kReduce: + case HloOpcode::kReverse: + case HloOpcode::kTranspose: + for (auto dim : instruction->dimensions()) { + attrs["dims"].mutable_list()->add_i(dim); + } + break; + case HloOpcode::kGetTupleElement: + attrs["index"].set_i(instruction->tuple_index()); + break; + case HloOpcode::kRng: + attrs["dist"].set_s( + RandomDistribution_Name(instruction->random_distribution())); + break; + case HloOpcode::kConstant: + if (ShapeUtil::IsScalar(instruction->shape())) { + attrs["value"].set_s(instruction->literal().GetAsString({})); + } + break; + case HloOpcode::kCustomCall: + attrs["custom_call_target"].set_s(instruction->custom_call_target()); + break; + default: + break; + } +} + +Status HloTfGraphBuilder::AddInstruction(const HloInstruction* instruction) { + if (!visited_instructions_.insert(instruction).second) { + // Skip instructions that have already been added. + return Status::OK(); + } + + NodeDef* node_def = graph_def_.add_node(); + node_def->set_name(GetNodeNameForInstruction(instruction)); + node_def->set_op(GetOpDefName(instruction)); + SetNodeAttrs(instruction, node_def); + if (instruction->opcode() == HloOpcode::kFusion) { + for (auto& fused_instruction : instruction->fused_instructions()) { + TF_RETURN_IF_ERROR(AddInstruction(fused_instruction.get())); + } + } + // Add all edges including control edges. + for (unsigned i = 0; i < instruction->operands().size(); ++i) { + *node_def->add_input() = GetNodeNameForInstruction(instruction->operand(i)); + } + // Called computations are control dependencies. + for (const auto* called_computation : instruction->called_computations()) { + *node_def->add_input() = StrCat( + "^", GetNodeNameForInstruction(called_computation->root_instruction())); + } + return Status::OK(); +} + +} // namespace hlo_graph_dumper +} // namespace xla diff --git a/tensorflow/compiler/xla/service/hlo_tfgraph_builder.h b/tensorflow/compiler/xla/service/hlo_tfgraph_builder.h new file mode 100644 index 0000000000000000000000000000000000000000..b2c578af912ac0b777d1bc72a198504735a6b845 --- /dev/null +++ b/tensorflow/compiler/xla/service/hlo_tfgraph_builder.h @@ -0,0 +1,58 @@ +/* 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 THIRD_PARTY_TENSORFLOW_COMPILER_XLA_SERVICE_HLO_TFGRAPH_BUILDER_H_ +#define THIRD_PARTY_TENSORFLOW_COMPILER_XLA_SERVICE_HLO_TFGRAPH_BUILDER_H_ + +#include "tensorflow/compiler/xla/service/hlo_computation.h" +#include "tensorflow/core/framework/graph.pb.h" +#include "tensorflow/core/framework/node_def.pb.h" + +namespace xla { +namespace hlo_graph_dumper { + +// This constructs a tensorflow graph for HLO computations. +class HloTfGraphBuilder { + public: + // Adds a computation to the graph. + Status AddComputation(const HloComputation& computation); + + const tensorflow::GraphDef& GetGraphDef() const; + + private: + // Gets the node name of an instruction. The node name is hierarchical. For + // example, if an instruction is fused, it will be put in a subgraph of the + // fusion instruction. + const string& GetNodeNameForInstruction(const HloInstruction* instruction); + + void SetNodeAttrs(const HloInstruction* instruction, + tensorflow::NodeDef* node_def) const; + + Status AddInstruction(const HloInstruction* instruction); + + tensorflow::GraphDef graph_def_; + // This records instructions that have been visited. + std::unordered_set visited_instructions_; + // A cache that maps instruction to the node name. + std::unordered_map instruction_to_node_name_; +}; + +// Cleans the node name to make it a valid name in a tensorflow graph. +void CleanNodeName(string* name); + +} // namespace hlo_graph_dumper +} // namespace xla + +#endif // THIRD_PARTY_TENSORFLOW_COMPILER_XLA_SERVICE_HLO_TFGRAPH_BUILDER_H_ diff --git a/tensorflow/compiler/xla/service/hlo_tfgraph_builder_test.cc b/tensorflow/compiler/xla/service/hlo_tfgraph_builder_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..8e9d93e367e51cb69f0a38ae7aa8d9539e78ad8a --- /dev/null +++ b/tensorflow/compiler/xla/service/hlo_tfgraph_builder_test.cc @@ -0,0 +1,188 @@ +/* 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/hlo_tfgraph_builder.h" +#include "tensorflow/compiler/xla/client/computation_builder.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" + +namespace xla { +namespace hlo_graph_dumper { +namespace { + +using ::tensorflow::GraphDef; + +class HloTfGraphBuilderTest : public HloTestBase { + protected: + HloTfGraphBuilderTest() {} + HloTfGraphBuilder generator_; + + // Create a computation which takes a scalar and returns its negation. + std::unique_ptr CreateNegateComputation() { + auto builder = HloComputation::Builder("Negate"); + auto param = builder.AddInstruction( + HloInstruction::CreateParameter(0, r0f32_, "param0")); + builder.AddInstruction( + HloInstruction::CreateUnary(r0f32_, HloOpcode::kNegate, param)); + return builder.Build(); + } + + // Creates a computation which calls map with the given computation. + std::unique_ptr CreateMapComputation( + HloComputation *map_computation) { + auto builder = HloComputation::Builder("Map"); + auto param = builder.AddInstruction( + HloInstruction::CreateParameter(0, r0f32_, "param0")); + builder.AddInstruction( + HloInstruction::CreateMap(r0f32_, {param}, map_computation)); + return builder.Build(); + } + Shape r0f32_ = ShapeUtil::MakeShape(PrimitiveType::F32, {}); +}; + +static const tensorflow::AttrValue &GetNodeAttr(const tensorflow::NodeDef &node, + const string &attr_name) { + auto attr = node.attr().find(attr_name); + CHECK(attr != node.attr().end()); + return attr->second; +} + +TEST_F(HloTfGraphBuilderTest, CheckConcatenateDimsAndShapes) { + auto builder = HloComputation::Builder("Concatenate"); + Shape shape = ShapeUtil::MakeShape(PrimitiveType::F32, {2, 2}); + auto param_1 = builder.AddInstruction( + HloInstruction::CreateParameter(0, shape, "param0")); + auto param_2 = builder.AddInstruction( + HloInstruction::CreateParameter(1, shape, "param1")); + builder.AddInstruction(HloInstruction::CreateConcatenate( + ShapeUtil::MakeShape(PrimitiveType::F32, {2, 4}), {param_1, param_2}, 1)); + TF_CHECK_OK(generator_.AddComputation(*builder.Build())); + GraphDef graph_def = generator_.GetGraphDef(); + EXPECT_EQ(graph_def.node_size(), 3); + const auto &node = graph_def.node(2); + EXPECT_EQ(node.name(), "Concatenate/concatenate"); + + // Check dimensions. + auto dims_value = GetNodeAttr(node, "dims"); + EXPECT_EQ(dims_value.list().i_size(), 1); + EXPECT_EQ(dims_value.list().i(0), 1); + + // Check shapes. + auto shape_value = GetNodeAttr(node, "_output_shapes"); + EXPECT_EQ(shape_value.list().shape_size(), 1); + EXPECT_EQ(shape_value.list().shape(0).dim_size(), 2); + EXPECT_EQ(shape_value.list().shape(0).dim(0).size(), 2); + EXPECT_EQ(shape_value.list().shape(0).dim(1).size(), 4); +} + +TEST_F(HloTfGraphBuilderTest, CheckScalarValue) { + auto builder = HloComputation::Builder("Const"); + HloInstruction *instruction = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(123))); + OpMetadata metadata; + metadata.set_op_name("x"); + metadata.set_op_type("y"); + instruction->set_metadata(metadata); + TF_CHECK_OK(generator_.AddComputation(*builder.Build())); + GraphDef graph_def = generator_.GetGraphDef(); + EXPECT_EQ(graph_def.node_size(), 1); + const auto &node = graph_def.node(0); + EXPECT_EQ(GetNodeAttr(node, "value").s(), "123"); + EXPECT_EQ(GetNodeAttr(node, "type").s(), "S32"); + EXPECT_EQ(GetNodeAttr(node, "tf_op_name").s(), "x"); + EXPECT_EQ(GetNodeAttr(node, "tf_op_type").s(), "y"); +} + +TEST_F(HloTfGraphBuilderTest, SimpleNegateComputation) { + auto negate_computation = CreateNegateComputation(); + TF_CHECK_OK(generator_.AddComputation(*negate_computation)); + GraphDef graph_def = generator_.GetGraphDef(); + EXPECT_EQ(graph_def.node_size(), 2); + EXPECT_EQ(graph_def.node(0).name(), "Negate/param0.0"); + EXPECT_EQ(graph_def.node(0).op(), "HloParameter"); + EXPECT_EQ(graph_def.node(1).name(), "Negate/negate"); + EXPECT_EQ(graph_def.node(1).op(), "HloNegate"); + EXPECT_EQ(graph_def.node(1).input_size(), 1); + EXPECT_EQ(graph_def.node(1).input(0), "Negate/param0.0"); +} + +TEST_F(HloTfGraphBuilderTest, GreaterThanOrEqualTo) { + auto builder = HloComputation::Builder("GE"); + auto param_1 = builder.AddInstruction( + HloInstruction::CreateParameter(0, r0f32_, "param0")); + auto param_2 = builder.AddInstruction( + HloInstruction::CreateParameter(1, r0f32_, "param1")); + builder.AddInstruction( + HloInstruction::CreateBinary(r0f32_, HloOpcode::kGe, param_1, param_2)); + TF_CHECK_OK(generator_.AddComputation(*builder.Build())); + GraphDef graph_def = generator_.GetGraphDef(); + EXPECT_EQ(graph_def.node_size(), 3); + EXPECT_EQ(graph_def.node(0).name(), "GE/param0.0"); + EXPECT_EQ(graph_def.node(1).name(), "GE/param1.1"); + EXPECT_EQ(graph_def.node(2).input_size(), 2); + EXPECT_EQ(graph_def.node(2).name(), "GE/greater-than-or-equal-to"); + EXPECT_EQ(graph_def.node(2).op(), "HloGreaterThanOrEqualTo"); +} + +TEST_F(HloTfGraphBuilderTest, IncorparateTfOpsStructure) { + auto builder = HloComputation::Builder("GE"); + auto param_1 = builder.AddInstruction( + HloInstruction::CreateParameter(0, r0f32_, "param0")); + auto param_2 = builder.AddInstruction( + HloInstruction::CreateParameter(1, r0f32_, "param1")); + auto ge = builder.AddInstruction( + HloInstruction::CreateBinary(r0f32_, HloOpcode::kGe, param_1, param_2)); + OpMetadata metadata; + metadata.set_op_name("x/y"); + metadata.set_op_type("Y"); + ge->set_metadata(metadata); + TF_CHECK_OK(generator_.AddComputation(*builder.Build())); + GraphDef graph_def = generator_.GetGraphDef(); + EXPECT_EQ(graph_def.node_size(), 3); + EXPECT_EQ(graph_def.node(0).name(), "GE/param0.0"); + EXPECT_EQ(graph_def.node(1).name(), "GE/param1.1"); + EXPECT_EQ(graph_def.node(2).input_size(), 2); + EXPECT_EQ(graph_def.node(2).name(), "GE/x/y/greater-than-or-equal-to"); + EXPECT_EQ(graph_def.node(2).op(), "HloGreaterThanOrEqualTo"); +} + +TEST_F(HloTfGraphBuilderTest, EmbeddedComputationsDiamond) { + // Create computations with a diamond-shaped callgraph. + auto negate_computation = CreateNegateComputation(); + auto map1_computation = CreateMapComputation(negate_computation.get()); + auto map2_computation = CreateMapComputation(negate_computation.get()); + + auto builder = HloComputation::Builder(TestName()); + auto param = builder.AddInstruction( + HloInstruction::CreateParameter(0, r0f32_, "param0")); + auto map1 = builder.AddInstruction( + HloInstruction::CreateMap(r0f32_, {param}, map1_computation.get())); + auto map2 = builder.AddInstruction( + HloInstruction::CreateMap(r0f32_, {param}, map2_computation.get())); + builder.AddInstruction( + HloInstruction::CreateBinary(r0f32_, HloOpcode::kAdd, map1, map2)); + auto computation = builder.Build(); + TF_CHECK_OK(generator_.AddComputation(*computation)); + EXPECT_GT(generator_.GetGraphDef().node_size(), 0); +} + +} // namespace +} // namespace hlo_graph_dumper +} // namespace xla + +int main(int argc, char **argv) { + return xla::ParseDebugOptionsFlagsAndRunTests(argc, argv); +} diff --git a/tensorflow/compiler/xla/service/hlo_value.cc b/tensorflow/compiler/xla/service/hlo_value.cc new file mode 100644 index 0000000000000000000000000000000000000000..e6cf0d37b8a0f42dc04cfaad067a4741bc803705 --- /dev/null +++ b/tensorflow/compiler/xla/service/hlo_value.cc @@ -0,0 +1,340 @@ +/* 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/hlo_value.h" + +#include +#include + +#include "tensorflow/compiler/xla/map_util.h" +#include "tensorflow/compiler/xla/ptr_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/service/hlo_opcode.h" +#include "tensorflow/compiler/xla/shape_util.h" +#include "tensorflow/compiler/xla/status.h" +#include "tensorflow/compiler/xla/types.h" +#include "tensorflow/compiler/xla/util.h" +#include "tensorflow/core/lib/core/errors.h" +#include "tensorflow/core/lib/strings/str_util.h" +#include "tensorflow/core/lib/strings/strcat.h" +#include "tensorflow/core/platform/logging.h" + +namespace xla { + +using ::tensorflow::str_util::Join; +using ::tensorflow::strings::StrAppend; +using ::tensorflow::strings::StrCat; + +const Shape& HloPosition::shape() const { + return ShapeUtil::GetSubshape(instruction->shape(), index); +} + +string HloPosition::ToString() const { + string index_str = + ShapeUtil::IsTuple(instruction->shape()) ? (" " + index.ToString()) : ""; + return StrCat(instruction->name(), index_str); +} + +std::ostream& operator<<(std::ostream& out, const HloPosition& position) { + out << position.ToString(); + return out; +} + +string HloUse::ToString() const { + string index_str = + ShapeUtil::IsTuple(instruction->operand(operand_number)->shape()) + ? (" " + operand_index.ToString()) + : ""; + return StrCat(instruction->name(), ", operand ", operand_number, index_str); +} + +std::ostream& operator<<(std::ostream& out, const HloUse& use) { + out << use.ToString(); + return out; +} + +HloValue::HloValue(HloValue::Id id, HloInstruction* instruction, + const ShapeIndex& index, bool is_phi) + : id_(id), is_phi_(is_phi) { + // The defining position is always the first element in the positions_ vector. + AddPosition(instruction, index); +} + +bool HloValue::operator==(const HloValue& other) const { + bool equal = defining_instruction() == other.defining_instruction() && + defining_index() == other.defining_index(); + // If the values are equal they most both be phi (or non phi). + CHECK(!(equal && is_phi() != other.is_phi())); + return equal; +} + +bool HloValue::operator!=(const HloValue& other) const { + return !(*this == other); +} + +string HloValue::ToShortString() const { + string index_str = ShapeUtil::IsTuple(defining_instruction()->shape()) + ? defining_index().ToString() + : ""; + return StrCat(id_, " ", is_phi_ ? "PHI " : "", defining_instruction()->name(), + index_str); +} + +string HloValue::ToString(int indent) const { + string indentation(indent, ' '); + string out = StrCat(indentation, ToShortString(), ", positions:\n"); + for (const HloPosition& position : positions()) { + StrAppend(&out, indentation, " ", position.ToString(), "\n"); + } + StrAppend(&out, indentation, " uses:\n"); + for (const HloUse& use : uses()) { + StrAppend(&out, indentation, " ", use.ToString(), "\n"); + } + return out; +} + +namespace { + +// Returns true if the instruction 'user' may use the value at the given +// ShapeIndex in the given operand. Generally, instruction which pass through +// values transparently without reading the value are not considered to use the +// value. +bool MayUseOperandValue(int64 operand_number, const ShapeIndex& index, + const HloInstruction* user) { + switch (user->opcode()) { + case HloOpcode::kGetTupleElement: + case HloOpcode::kCopy: + // These instructions only access the top-level values of their + // operand. Non-top-level (nested) values are passed through + // transparently. + CHECK_EQ(operand_number, 0); + return index.empty(); + case HloOpcode::kSelect: + // Select does not use any nested elements of its selected-from operands + // (operand 1 and 2) + CHECK_GE(operand_number, 0); + CHECK_LE(operand_number, 2); + return operand_number == 0 || index.empty(); + + case HloOpcode::kCall: + case HloOpcode::kTuple: + // These instructions always pass through their operands transparently. + return false; + + case HloOpcode::kWhile: + // Though the while instructions passes through its operands, we return + // true because in SSA form there may be a Phi at the parameter of the + // while which is considered a use of its incoming value because the Phi + // input values are not passed through into the body computation. Because + // this function is used in both SSA and non-SSA forms of the analysis + // conservatively return true. + return true; + + default: + return true; + } +} + +} // namespace + +void HloValue::AddPosition(HloInstruction* instruction, + const ShapeIndex& index) { + HloPosition new_position{instruction, index}; + + // The new position must not already exist in positions_. + for (const HloPosition& position : positions_) { + DCHECK_NE(position, new_position); + } + + positions_.push_back(std::move(new_position)); + + // Update uses. + for (HloInstruction* user : instruction->users()) { + for (int64 operand_number : user->OperandIndices(instruction)) { + if (MayUseOperandValue(operand_number, index, user)) { + HloUse new_use{user, operand_number, index}; + + // The new use must not already exist in uses_. + for (const HloUse& use : uses_) { + DCHECK_NE(use, new_use); + } + + uses_.push_back(std::move(new_use)); + } + } + } + + // Update liveout status of this HloValue. + const HloModule& module = *instruction->parent()->parent(); + if (instruction == module.entry_computation()->root_instruction()) { + live_out_of_module_ = true; + } + + if (instruction == instruction->parent()->root_instruction()) { + live_out_of_computation_ = true; + } +} + +void HloValue::RemovePosition(HloInstruction* instruction, + const ShapeIndex& index) { + // The defining position cannot be removed. + CHECK(!(instruction == defining_instruction() && index == defining_index())); + + int64 size_before = positions_.size(); + positions_.erase( + std::remove_if(positions_.begin(), positions_.end(), + [instruction, &index](const HloPosition& position) { + return position.instruction == instruction && + position.index == index; + }), + positions_.end()); + // Only a single position should have been removed. + CHECK_EQ(positions_.size(), size_before - 1); + + // Update uses which referred to this position. + uses_.erase(std::remove_if(uses_.begin(), uses_.end(), + [instruction, &index](const HloUse& use) { + return use.instruction->operand( + use.operand_number) == instruction && + use.operand_index == index; + }), + uses_.end()); + + // Returns whether this value is contained in the given instruction's output. + auto is_contained_in = [this](const HloInstruction* instruction) { + for (const HloPosition& position : positions()) { + if (position.instruction == instruction) { + return true; + } + } + return false; + }; + + const HloModule& module = *instruction->parent()->parent(); + if (instruction == module.entry_computation()->root_instruction()) { + // Value has been removed from a position in the entry root instruction. + live_out_of_module_ = + is_contained_in(module.entry_computation()->root_instruction()); + } + if (instruction == defining_instruction()->parent()->root_instruction()) { + // Value has been removed from the root of the computation the value has + // been defined in. + live_out_of_computation_ = + is_contained_in(defining_instruction()->parent()->root_instruction()); + } +} + +void HloValue::RecomputeUses() { + uses_.clear(); + for (const HloPosition& position : positions()) { + for (HloInstruction* user : position.instruction->users()) { + for (int64 operand_number : user->OperandIndices(position.instruction)) { + if (MayUseOperandValue(operand_number, position.index, user)) { + uses_.push_back(HloUse{user, operand_number, position.index}); + } + } + } + } +} + +std::ostream& operator<<(std::ostream& out, const HloValue& value) { + out << value.ToShortString(); + return out; +} + +void HloValueSet::SortAndUniquifyValues() { + std::sort(values_.begin(), values_.end(), HloValue::IdLessThan); + values_.erase(std::unique(values_.begin(), values_.end(), HloValue::IdEqual), + values_.end()); +} + +string HloValueSet::ToString() const { + return StrCat("HloValueSet: ", + Join(values_, ", ", [](string* result, const HloValue* value) { + result->append(value->ToShortString()); + })); +} + +bool HloValueSet::AssignUnionOf( + tensorflow::gtl::ArraySlice inputs) { + HloValueSet union_set; + for (const HloValueSet* input : inputs) { + for (const HloValue* value : input->values()) { + union_set.values_.push_back(value); + } + } + union_set.SortAndUniquifyValues(); + if (*this != union_set) { + *this = union_set; + return true; + } + return false; +} + +bool HloValueSet::AddValue(const HloValue* value) { + auto it = std::lower_bound(values_.begin(), values_.end(), value, + HloValue::IdLessThan); + if (it == values_.end() || (*it)->id() != value->id()) { + values_.insert(it, value); + return true; + } + return false; // already exists +} + +std::ostream& operator<<(std::ostream& out, const HloValueSet& value_set) { + out << value_set.ToString(); + return out; +} + +bool InstructionValueSet::AssignUnionOf( + tensorflow::gtl::ArraySlice inputs) { + CHECK_GT(inputs.size(), 0); + for (int i = 1; i < inputs.size(); ++i) { + DCHECK(ShapeUtil::Compatible(inputs[0]->shape(), inputs[i]->shape())); + } + bool changed = false; + for (auto& pair : *this) { + const ShapeIndex& index = pair.first; + HloValueSet& value_set = pair.second; + + std::vector input_value_sets; + for (const InstructionValueSet* input : inputs) { + input_value_sets.push_back(&input->element(index)); + } + changed |= value_set.AssignUnionOf(input_value_sets); + } + + return changed; +} + +std::ostream& operator<<(std::ostream& out, + const InstructionValueSet& instruction_value_set) { + out << instruction_value_set.ToString(); + return out; +} + +string InstructionValueSet::ToString() const { + string out = + StrCat("InstructionValueSet(", ShapeUtil::HumanString(shape()), ")\n"); + ForEachElement([this, &out](const ShapeIndex& index, + const HloValueSet& value_set) { + StrAppend(&out, " ", index.ToString(), " : ", value_set.ToString(), "\n"); + }); + return out; +} + +} // namespace xla diff --git a/tensorflow/compiler/xla/service/hlo_value.h b/tensorflow/compiler/xla/service/hlo_value.h new file mode 100644 index 0000000000000000000000000000000000000000..6872bc76a82253b916e826aa1afabc3d309c1d12 --- /dev/null +++ b/tensorflow/compiler/xla/service/hlo_value.h @@ -0,0 +1,283 @@ +/* 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_VALUE_H_ +#define TENSORFLOW_COMPILER_XLA_SERVICE_HLO_VALUE_H_ + +#include +#include +#include + +#include "tensorflow/compiler/xla/service/hlo_instruction.h" +#include "tensorflow/compiler/xla/shape_tree.h" +#include "tensorflow/compiler/xla/types.h" +#include "tensorflow/compiler/xla/xla_data.pb.h" +#include "tensorflow/core/lib/gtl/array_slice.h" +#include "tensorflow/core/platform/macros.h" + +namespace xla { + +// Abstraction which identifies a specific point in the XLA graph. An +// HloPosition specifies a ShapeIndex within the output of a specific +// instruction. +struct HloPosition { + HloInstruction* instruction; + ShapeIndex index; + + // Returns the shape at this position. + const Shape& shape() const; + + string ToString() const; + + bool operator==(const HloPosition& other) const { + return instruction == other.instruction && index == other.index; + } + bool operator!=(const HloPosition& other) const { return !(*this == other); } + + // Stable less-than operator using instruction id and index. + bool operator<(const HloPosition& other) const { + return instruction->unique_id() < other.instruction->unique_id() || + (instruction->unique_id() == other.instruction->unique_id() && + index < other.index); + } +}; + +std::ostream& operator<<(std::ostream& out, const HloPosition& position); + +// Defines a single use of an HLO value. +struct HloUse { + // Instruction at which the value is used. + HloInstruction* instruction; + + // The operand number in which the value is appears. + int64 operand_number; + + // The shape index within the operand in which the value appears. + ShapeIndex operand_index; + + string ToString() const; + + bool operator==(const HloUse& other) const { + return instruction == other.instruction && + operand_number == other.operand_number && + operand_index == other.operand_index; + } + + bool operator!=(const HloUse& other) const { return !(*this == other); } +}; + +std::ostream& operator<<(std::ostream& out, const HloUse& use); + +// Class describing a value used by the dataflow analysis. XLA arrays are +// trivially a single HloValue. Tuples are made up of more than one HloValue: an +// HloValue for the pointer vector, and an HloValue for each child element. +// +// Every HloValue is defined by a particular instruction and most instructions +// define only a single HloValue. Instructions which define a single HloValue +// include array-shaped instructions such as Add but also includes Tuple-shaped +// instructions such as Tuple. The Tuple instruction defines a single HloValue +// which is a vector of pointers to the values containing the Tuple +// instruction's operands. Though the result of the Tuple instruction includes +// multiple values only the top-level HloValue (the vector of pointers) is +// defined by the Tuple instruction. The values containing the tuple elements +// are defined by earlier instructions, usually the operands of the Tuple +// instruction. +// +// Instructions which construct both the tuple *and* the tuple elements define +// more than one HloValue. This includes (at least) tuple-shaped Constant, +// Parameter, Infeed and While instructions. These tuple-shaped instructions do +// not assemble a tuple from existing HloValues like the Tuple instruction does, +// but rather define all the HloValues in the tuple. +class HloValue { + public: + using Id = int64; + + // Predicate comparing HloValues by increasing id, useful for std::sort. + static bool IdLessThan(const HloValue* a, const HloValue* b) { + return a->id() < b->id(); + } + + // Predicate comparing HloValues by equal id, useful for std::unique. + static bool IdEqual(const HloValue* a, const HloValue* b) { + return a->id() == b->id(); + } + + // Construct an HloValue defined by 'instruction' at shape index 'index'. If + // is_phi is true, then this value is a phi value, for example, at the + // parameter of a while body computation. Phi values are only used in the SSA + // dataflow analysis (HloDataflowAnalysis::ssa_form_ is true). + HloValue(Id id, HloInstruction* instruction, const ShapeIndex& index, + bool is_phi = false); + + // Return a unique identifier for this HloValue. This value is used for stable + // sorting and iteration + Id id() const { return id_; } + + // Returns whether this value is a phi value. + bool is_phi() const { return is_phi_; } + + // Return the position where this value is defined. + const HloPosition& defining_position() const { return positions_[0]; } + + // Return the instruction which defines this HloValue. + HloInstruction* defining_instruction() const { + return defining_position().instruction; + } + + // Return the shape index at which this HloValue is defined in the output of + // its defining instruction. + const ShapeIndex& defining_index() const { return defining_position().index; } + + // Return the shape of this HloValue. + const Shape& shape() const { return defining_position().shape(); } + + // Add or remove a position at which the HloValue appears. The definition + // position can not be removed. The uses of the HloValue are updated. + void AddPosition(HloInstruction* instruction, const ShapeIndex& index); + void RemovePosition(HloInstruction* instruction, const ShapeIndex& index); + + // Remove all positions except the defining position. Updates uses. + void ClearPositions(); + + // Return all positions of the HloValue in the module. + const std::vector& positions() const { return positions_; } + + // Return all uses of the HloValue. + const std::vector& uses() const { return uses_; } + + void RecomputeUses(); + + // Get whether this HloValue is live out of the module. + bool live_out_of_module() const { return live_out_of_module_; } + + // Get whether this HloValue is live out of the computation it is defined in. + bool live_out_of_computation() const { return live_out_of_computation_; } + + bool operator==(const HloValue& other) const; + bool operator!=(const HloValue& other) const; + + // Return a single-line string representation of the value. + string ToShortString() const; + + string ToString(int indent = 0) const; + + private: + // Unique identifier for this HloValue. Used for stable sorting and iteration. + const Id id_; + + // Whether this instruction is a phi value. + const bool is_phi_; + + // The set of positions of this HloValue. The first element is always the + // position of the definition. + std::vector positions_; + + // The set of uses of this HloValue. + std::vector uses_; + + // 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); + +// A class representing the possible set of HloValues at a particular point +// (shape index in the output of an instruction) in the XLA graph. This set +// contains the set of reaching HloValue definitions. For a simple array-shaped +// instruction like Add, the HloValueSet of the top-level of the instruction's +// output trivially contains only the HloValue defined by the instruction. For +// instructions which have non-trivial dataflow such as Tuple or Select, the +// HloValueSets of the instruction's output contains one or more HloValues +// defined by the instruction's operands or defined further up in the XLA graph. +class HloValueSet { + public: + HloValueSet() = default; + + explicit HloValueSet(tensorflow::gtl::ArraySlice values) + : values_(values.begin(), values.end()) { + SortAndUniquifyValues(); + } + + // Sets this value set to the union of the given value sets. Returns whether + // this value set changed. + bool AssignUnionOf(tensorflow::gtl::ArraySlice inputs); + + // Return the vector of HloValues in the set. Values in the vector are unique + // and stably sorted by value id. + const std::vector& values() const { return values_; } + + // Adds the value to the set. Returns true iff the value was added and didn't + // already exist in the set. + bool AddValue(const HloValue* value); + + // Clear all values from the set. + void Clear() { values_.clear(); } + + // Return the unique HLO value in the set. CHECKs if the set does not contain + // exactly one value. + const HloValue& GetUniqueValue() const { + CHECK_EQ(values_.size(), 1); + return *values_[0]; + } + + bool operator==(const HloValueSet& other) const { + if (values_.size() != other.values_.size()) return false; + for (size_t i = 0; i < values_.size(); ++i) { + if (values_[i]->id() != other.values_[i]->id()) { + return false; + } + } + return true; + } + bool operator!=(const HloValueSet& other) const { return !(*this == other); } + + string ToString() const; + + private: + // Sorts value_ and removes duplicates. This should be called after adding any + // elements to values_. + void SortAndUniquifyValues(); + + // HloValues sorted by HloValue::Id. + std::vector values_; +}; + +std::ostream& operator<<(std::ostream& out, const HloValueSet& hlo_value); + +// A class collecting the HloValues which might be contained in the output of +// an HLO instruction. For array-shaped instructions, an InstructionValueSet +// trivially holds a single HloValueSet. Tuple-shaped InstructionValueSets +// hold multiple HloValueSets. +class InstructionValueSet : public ShapeTree { + public: + InstructionValueSet(const Shape& shape) : ShapeTree(shape) {} + + // Sets this value set to the union of the given value sets. Returns whether + // this value set changed. + bool AssignUnionOf( + tensorflow::gtl::ArraySlice inputs); + + string ToString() const; +}; + +std::ostream& operator<<(std::ostream& out, + const InstructionValueSet& instruction_value_set); + +} // namespace xla + +#endif // TENSORFLOW_COMPILER_XLA_SERVICE_HLO_VALUE_H_ diff --git a/tensorflow/compiler/xla/service/hlo_verifier.cc b/tensorflow/compiler/xla/service/hlo_verifier.cc index 035b570ed3419503ad2325c5fdb46118b5076187..c44be716cdf995cda8ac3a768568b594e7e5328b 100644 --- a/tensorflow/compiler/xla/service/hlo_verifier.cc +++ b/tensorflow/compiler/xla/service/hlo_verifier.cc @@ -14,22 +14,322 @@ limitations under the License. ==============================================================================*/ #include "tensorflow/compiler/xla/service/hlo_verifier.h" +#include "tensorflow/compiler/xla/service/shape_inference.h" +#include "tensorflow/compiler/xla/status_macros.h" +#include "tensorflow/core/lib/core/errors.h" +#include "tensorflow/core/lib/gtl/flatmap.h" namespace xla { +namespace { + +// Visitor which verifies that the output shape is correctly set. Verifies +// against the inferred shape for the instruction. +// TODO(b/26024837): Check output shape for all instruction types. +class ShapeVerifier : public DfsHloVisitor { + public: + explicit ShapeVerifier( + const std::function& shape_size_fn) + : shape_size_fn_(shape_size_fn) {} + + Status HandleElementwiseUnary(HloInstruction* hlo) override { + return CheckUnaryShape(hlo); + } + + Status HandleElementwiseBinary(HloInstruction* hlo) override { + return CheckBinaryShape(hlo); + } + + Status HandleClamp(HloInstruction* clamp, HloInstruction* min, + HloInstruction* arg, HloInstruction* max) override { + return CheckTernaryShape(clamp); + } + + Status HandleSelect(HloInstruction* select, HloInstruction* pred, + HloInstruction* on_true, + HloInstruction* on_false) override { + return CheckTernaryShape(select); + } + + Status HandleConcatenate( + HloInstruction* concatenate, + tensorflow::gtl::ArraySlice operands) override { + return tensorflow::Status::OK(); + } + + Status HandleConvert(HloInstruction* convert) override { + return tensorflow::Status::OK(); + } + + Status HandleCopy(HloInstruction* copy) override { + return CheckUnaryShape(copy); + } + + Status HandleDot(HloInstruction* dot, HloInstruction* lhs, + HloInstruction* rhs) override { + return CheckBinaryShape(dot); + } + + Status HandleConvolution(HloInstruction* convolution, HloInstruction* lhs, + HloInstruction* rhs, const Window& window) override { + return tensorflow::Status::OK(); + } + + Status HandleCrossReplicaSum(HloInstruction* crs) override { + return tensorflow::Status::OK(); + } + + Status HandleReducePrecision(HloInstruction* reduce_precision) override { + return tensorflow::Status::OK(); + } + + Status HandleInfeed(HloInstruction* infeed) override { + return tensorflow::Status::OK(); + } + + Status HandleOutfeed(HloInstruction* outfeed) override { + return tensorflow::Status::OK(); + } + + Status HandleRng(HloInstruction* random, + RandomDistribution distribution) override { + return tensorflow::Status::OK(); + } + + Status HandleReverse(HloInstruction* reverse, + HloInstruction* operand) override { + return tensorflow::Status::OK(); + } + + Status HandleSort(HloInstruction* sort, HloInstruction* operand) override { + return tensorflow::Status::OK(); + } + + Status HandleConstant(HloInstruction* constant, + const Literal& literal) override { + return tensorflow::Status::OK(); + } + + Status HandleGetTupleElement(HloInstruction* get_tuple_element, + HloInstruction* operand) override { + return tensorflow::Status::OK(); + } + + Status HandleReduce(HloInstruction* reduce, HloInstruction* arg, + HloInstruction* init_value, + tensorflow::gtl::ArraySlice dimensions, + HloComputation* function) override { + return tensorflow::Status::OK(); + } + + Status HandleBitcast(HloInstruction* bitcast) override { + // Bitcasts can be any shape, as long as the size matches the operand size. + TF_RET_CHECK(shape_size_fn_(bitcast->shape()) == + shape_size_fn_(bitcast->operand(0)->shape())); + return tensorflow::Status::OK(); + } + + Status HandleBroadcast(HloInstruction* broadcast) override { + return tensorflow::Status::OK(); + } + + Status HandleReshape(HloInstruction* reshape) override { + return tensorflow::Status::OK(); + } + + Status HandleTranspose(HloInstruction* transpose) override { + return tensorflow::Status::OK(); + } + + Status HandleParameter(HloInstruction* parameter) override { + return tensorflow::Status::OK(); + } + + Status HandleFusion(HloInstruction* fusion) override { + return tensorflow::Status::OK(); + } + + Status HandleCall(HloInstruction* call) override { + return tensorflow::Status::OK(); + } + + Status HandleCustomCall(HloInstruction* custom_call, + tensorflow::gtl::ArraySlice operands, + tensorflow::StringPiece custom_call_target) override { + return tensorflow::Status::OK(); + } + + Status HandleSlice(HloInstruction* slice, HloInstruction* operand) override { + return tensorflow::Status::OK(); + } + + Status HandleDynamicSlice(HloInstruction* dynamic_slice, + HloInstruction* operand, + HloInstruction* start_indices) override { + return tensorflow::Status::OK(); + } + + Status HandleDynamicUpdateSlice(HloInstruction* dynamic_update_slice, + HloInstruction* operand, + HloInstruction* update, + HloInstruction* start_indices) override { + return tensorflow::Status::OK(); + } + + Status HandleTuple( + HloInstruction* tuple, + tensorflow::gtl::ArraySlice operands) override { + return CheckVariadicShape(tuple); + } + + Status HandleMap( + HloInstruction* map, + tensorflow::gtl::ArraySlice operands, + HloComputation* function, + tensorflow::gtl::ArraySlice static_operands) override { + return tensorflow::Status::OK(); + } + + Status HandleReduceWindow(HloInstruction* reduce_window, + HloInstruction* operand, const Window& window, + HloComputation* function) override { + return tensorflow::Status::OK(); + } + + Status HandleSelectAndScatter(HloInstruction* instruction) override { + return tensorflow::Status::OK(); + } + + Status HandleWhile(HloInstruction* xla_while) override { + return tensorflow::Status::OK(); + } + + Status HandlePad(HloInstruction* pad) override { + return tensorflow::Status::OK(); + } + + Status HandleSend(HloInstruction* send) override { + return tensorflow::Status::OK(); + } + + Status HandleRecv(HloInstruction* recv) override { + return tensorflow::Status::OK(); + } + + Status HandleBatchNormTraining(HloInstruction* batchNormTraining) override { + return tensorflow::Status::OK(); + } + + Status HandleBatchNormInference(HloInstruction* batchNormInference) override { + return tensorflow::Status::OK(); + } + + Status HandleBatchNormGrad(HloInstruction* batchNormGrad) override { + return tensorflow::Status::OK(); + } + + Status FinishVisit(HloInstruction* root) override { + return tensorflow::Status::OK(); + } + + private: + // Check the instruction's shape against the given expected shape and return + // an appropriate error if there is a mismatch. + Status CheckShape(const HloInstruction* instruction, + const Shape& expected_shape) { + if (!ShapeUtil::Compatible(instruction->shape(), expected_shape)) { + return InvalidArgument( + "Expected instruction to have shape compatible with %s, actual " + "shape is %s:\n%s", + ShapeUtil::HumanString(expected_shape).c_str(), + ShapeUtil::HumanString(instruction->shape()).c_str(), + instruction->ToString().c_str()); + } + return tensorflow::Status::OK(); + } + + // Check a unary (binary, etc) instruction's shape against the inferred shape. + Status CheckUnaryShape(const HloInstruction* instruction) { + TF_ASSIGN_OR_RETURN(const Shape expected, + ShapeInference::InferUnaryOpShape( + instruction->opcode(), instruction->operand(0))); + return CheckShape(instruction, expected); + } + Status CheckBinaryShape(const HloInstruction* instruction) { + TF_ASSIGN_OR_RETURN(const Shape expected, + ShapeInference::InferBinaryOpShape( + instruction->opcode(), instruction->operand(0), + instruction->operand(1))); + return CheckShape(instruction, expected); + } + Status CheckTernaryShape(const HloInstruction* instruction) { + TF_ASSIGN_OR_RETURN(const Shape expected, + ShapeInference::InferTernaryOpShape( + instruction->opcode(), instruction->operand(0), + instruction->operand(1), instruction->operand(2))); + return CheckShape(instruction, expected); + } + Status CheckVariadicShape(const HloInstruction* instruction) { + TF_ASSIGN_OR_RETURN(const Shape expected, + ShapeInference::InferVariadicOpShape( + instruction->opcode(), instruction->operands())); + return CheckShape(instruction, expected); + } + + // Returns the size of a Shape in bytes. + const std::function shape_size_fn_; +}; + +string ComputationsToString( + tensorflow::gtl::ArraySlice computations) { + return tensorflow::str_util::Join( + computations, ",", [](string* s, const HloComputation* computation) { + s->append(computation->name()); + }); +} + +} // namespace + StatusOr HloVerifier::Run(HloModule* module) { + tensorflow::gtl::FlatMap instructions; + ShapeVerifier shape_verifier(shape_size_fn_); + for (auto& computation : module->computations()) { for (const auto& instruction : computation->instructions()) { TF_RET_CHECK(instruction->parent() == computation.get()); if (instruction->opcode() == HloOpcode::kFusion) { + TF_RET_CHECK( + ContainersEqual(instruction->called_computations(), + {instruction->fused_instructions_computation()})) + << "Fusion HLO calls computations other than the " + "fused_instructions_computation: " + << instruction->ToString() + << " instruction->fused_instructions_computation(): " + << instruction->fused_instructions_computation()->ToString() + << " instruction->called_computations(): " + << ComputationsToString(instruction->called_computations()); + for (const auto& fused : instruction->fused_instructions()) { - TF_RET_CHECK(fused->parent() == computation.get()) + TF_RET_CHECK(fused->parent() == + instruction->fused_instructions_computation()) << "Fused HLO was missing a parent: " << fused->ToString() << " parent: " << fused->parent() << " computation: " << computation.get(); } } + + auto previous = instructions.find(instruction->name()); + TF_RET_CHECK(previous == instructions.end()) + << "HLO has name that is not unique within module:\n" + << instruction->ToString() + << " in computation: " << computation->name() + << "\nPrevious HLO with same name:\n" + << previous->second->ToString() + << " in computation: " << previous->second->parent()->name(); + instructions[instruction->name()] = instruction.get(); } + + TF_RETURN_IF_ERROR(computation->Accept(&shape_verifier)); } return false; diff --git a/tensorflow/compiler/xla/service/hlo_verifier.h b/tensorflow/compiler/xla/service/hlo_verifier.h index 5159420b3fbea5d3d01950fa379e8ba39437ab85..bc6800dae54609b83a0e4ad92b0b4985ea750df6 100644 --- a/tensorflow/compiler/xla/service/hlo_verifier.h +++ b/tensorflow/compiler/xla/service/hlo_verifier.h @@ -24,12 +24,18 @@ namespace xla { // the module. class HloVerifier : public HloPassInterface { public: + explicit HloVerifier(const std::function& shape_size_fn) + : shape_size_fn_(shape_size_fn) {} ~HloVerifier() override = default; tensorflow::StringPiece name() const override { return "verifier"; } // Note: always returns false (no instructions are ever modified by this // pass). StatusOr Run(HloModule* module) override; + + private: + // Returns the size of a Shape in bytes. + const std::function shape_size_fn_; }; } // namespace xla diff --git a/tensorflow/compiler/xla/service/human_readable_profile_builder.cc b/tensorflow/compiler/xla/service/human_readable_profile_builder.cc new file mode 100644 index 0000000000000000000000000000000000000000..d620f45d27eba706fbd7fc30d3b27b0d963475d4 --- /dev/null +++ b/tensorflow/compiler/xla/service/human_readable_profile_builder.cc @@ -0,0 +1,131 @@ +/* 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/human_readable_profile_builder.h" +#include "tensorflow/compiler/xla/metric_table_report.h" +#include "tensorflow/compiler/xla/util.h" +#include "tensorflow/core/lib/strings/numbers.h" +#include "tensorflow/core/lib/strings/strcat.h" +#include "tensorflow/core/lib/strings/stringprintf.h" + +namespace xla { + +using tensorflow::strings::Appendf; +using tensorflow::strings::HumanReadableElapsedTime; +using tensorflow::strings::HumanReadableNumBytes; +using tensorflow::strings::StrAppend; + +string HumanReadableProfileBuilder::ToString() const { + string s; + + Appendf(&s, "Execution profile for %s: (%s @ f_nom)\n", + computation_name_.c_str(), + HumanReadableElapsedTime(CyclesToSeconds(total_cycles_)).c_str()); + + auto append_op = [&](const OpInfo& op) { + string bytes_per_sec; + string bytes_per_cycle; + if (op.cycles <= 0 || op.bytes_accessed < 0) { + bytes_per_sec = ""; + bytes_per_cycle = ""; + } else { + bytes_per_sec = + HumanReadableNumBytes(op.bytes_accessed / CyclesToSeconds(op.cycles)); + bytes_per_cycle = HumanReadableNumBytes(op.bytes_accessed / op.cycles); + } + + double cycles_percent = 0; + if (total_cycles_ > 0) { + cycles_percent = op.cycles / static_cast(total_cycles_) * 100; + } + + double nsecs = op.cycles / clock_rate_ghz_; + Appendf(&s, + "%15lld cycles (%6.2f%%) :: %12.1f usec (%12.1f optimal) :: %18s " + ":: %18s :: %12s/s :: %12s/cycle :: %s\n", + op.cycles, cycles_percent, CyclesToMicroseconds(op.cycles), + op.optimal_seconds * 1e6, + op.flop_count <= 0 + ? "" + : HumanReadableNumFlops(op.flop_count, nsecs).c_str(), + op.transcendental_count <= 0 ? "" + : HumanReadableNumTranscendentalOps( + op.transcendental_count, nsecs) + .c_str(), + bytes_per_sec.c_str(), bytes_per_cycle.c_str(), op.name.c_str()); + }; + + float optimal_seconds_sum = 0.0; + for (const auto& op : op_infos_) { + optimal_seconds_sum += op.optimal_seconds; + } + + append_op({"[total]", "[total]", /*category=*/"", total_cycles_, -1, -1, -1, + optimal_seconds_sum}); + + // Sort ops in decreasing order of cycles. + 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; }); + for (const auto& op : sorted_ops) { + append_op(op); + } + + if (total_cycles_ <= 0) { + StrAppend(&s, "****** 0 total cycles ******\n"); + } else { + // Only show an optimal discrepancy table if at least one value was + // specified. Estimates are non-negative, so if the sum is greater than + // zero, then at least one summand was greater than zero. + if (optimal_seconds_sum > 0) { + MetricTableReport table; + table.SetMetricName("microseconds above estimated optimum"); + table.SetEntryName("ops"); + table.SetShowCategoryTable(); + float total_discrepancy_in_microseconds = 0.0f; + for (const auto& op : sorted_ops) { + MetricTableReport::Entry entry; + entry.text = op.name; + entry.short_text = op.short_name; + entry.category_text = op.category; + entry.metric = + CyclesToMicroseconds(op.cycles) - op.optimal_seconds * 1e6; + total_discrepancy_in_microseconds += entry.metric; + table.AddEntry(std::move(entry)); + } + StrAppend(&s, table.MakeReport(total_discrepancy_in_microseconds)); + } + + { + MetricTableReport table; + table.SetMetricName("microseconds"); + table.SetEntryName("ops"); + table.SetShowCategoryTable(); + for (const auto& op : sorted_ops) { + MetricTableReport::Entry entry; + entry.text = op.name; + entry.short_text = op.short_name; + entry.category_text = op.category; + entry.metric = CyclesToMicroseconds(op.cycles); + table.AddEntry(std::move(entry)); + } + StrAppend(&s, table.MakeReport(CyclesToMicroseconds(total_cycles_))); + } + } + return s; +} + +} // namespace xla diff --git a/tensorflow/compiler/xla/service/human_readable_profile_builder.h b/tensorflow/compiler/xla/service/human_readable_profile_builder.h new file mode 100644 index 0000000000000000000000000000000000000000..fc24acd2713f4cd8af2816ffdf085e84a4920cbc --- /dev/null +++ b/tensorflow/compiler/xla/service/human_readable_profile_builder.h @@ -0,0 +1,85 @@ +/* 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_HUMAN_READABLE_PROFILE_BUILDER_H_ +#define TENSORFLOW_COMPILER_XLA_SERVICE_HUMAN_READABLE_PROFILE_BUILDER_H_ + +#include + +#include "tensorflow/compiler/xla/types.h" +#include "tensorflow/core/lib/core/stringpiece.h" +#include "tensorflow/core/platform/logging.h" +#include "tensorflow/core/platform/types.h" + +namespace xla { + +// HumanReadableProfileBuilder helps you create a textual profile of a +// computation, suitable for consumption by humans. +class HumanReadableProfileBuilder { + public: + explicit HumanReadableProfileBuilder(tensorflow::StringPiece computation_name, + int64 total_cycles, + double clock_rate_ghz) + : computation_name_(computation_name.ToString()), + total_cycles_(total_cycles), + clock_rate_ghz_(clock_rate_ghz) { + CHECK_GE(clock_rate_ghz, 1e-9); + } + + int64 total_cycles() const { return total_cycles_; } + + // Adds an operation to the profile. If you don't know the number of + // floating-point ops or bytes touched by the op, pass -1 for that param. + void AddOp(tensorflow::StringPiece op_name, + tensorflow::StringPiece short_name, + tensorflow::StringPiece category, int64 cycles, int64 flop_count, + int64 transcendental_count, int64 bytes_accessed, + float optimal_seconds) { + op_infos_.push_back( + {op_name.ToString(), short_name.ToString(), category.ToString(), cycles, + flop_count, transcendental_count, bytes_accessed, optimal_seconds}); + } + + // Gets the human-readable profile. + string ToString() const; + + private: + struct OpInfo { + string name; + string short_name; + string category; + int64 cycles; + int64 flop_count; + int64 transcendental_count; + int64 bytes_accessed; + float optimal_seconds; + }; + + double CyclesToSeconds(int64 cycles) const { + return cycles / clock_rate_ghz_ / 1e9; + } + double CyclesToMicroseconds(int64 cycles) const { + return cycles / clock_rate_ghz_ / 1000.0; + } + + string computation_name_; + int64 total_cycles_; + double clock_rate_ghz_; + std::vector op_infos_; +}; + +} // namespace xla + +#endif // TENSORFLOW_COMPILER_XLA_SERVICE_HUMAN_READABLE_PROFILE_BUILDER_H_ diff --git a/tensorflow/compiler/xla/service/inliner_test.cc b/tensorflow/compiler/xla/service/inliner_test.cc index 0054edcf6ab3b5134abbc43a8b326d56919364bc..84bfbb30c30d84a6a233a60fb420b43c3fe3454c 100644 --- a/tensorflow/compiler/xla/service/inliner_test.cc +++ b/tensorflow/compiler/xla/service/inliner_test.cc @@ -22,13 +22,16 @@ limitations under the License. #include "tensorflow/compiler/xla/ptr_util.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/shape_util.h" -#include "tensorflow/compiler/xla/test_helpers.h" +#include "tensorflow/compiler/xla/test.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" +namespace op = xla::testing::opcode_matchers; + namespace xla { namespace { @@ -48,26 +51,26 @@ TEST_F(InlinerTest, MapMax) { auto max_f32 = max_builder.Build(); auto builder = HloComputation::Builder("MapMaxFunction"); - auto lhs = builder.AddInstruction(HloInstruction::CreateConstant( - LiteralUtil::CreateR1({1, 2, 3, 4}))); - auto rhs = builder.AddInstruction(HloInstruction::CreateConstant( - LiteralUtil::CreateR1({4, 3, 2, 1}))); + auto lhs = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR1({1, 2, 3, 4}))); + auto rhs = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR1({4, 3, 2, 1}))); builder.AddInstruction( HloInstruction::CreateMap(lhs->shape(), {lhs, rhs}, max_f32.get())); auto computation = builder.Build(); - auto hlo_module = MakeUnique("test_module"); + auto hlo_module = CreateNewModule(); hlo_module->AddEmbeddedComputation(std::move(max_f32)); hlo_module->AddEntryComputation(std::move(computation)); - HloInstruction* root = hlo_module->entry_computation()->root_instruction(); + Inliner inliner; EXPECT_TRUE(inliner.Run(hlo_module.get()).ValueOrDie()); - root = hlo_module->entry_computation()->root_instruction(); - EXPECT_EQ(root->opcode(), HloOpcode::kMaximum); + EXPECT_THAT(hlo_module->entry_computation()->root_instruction(), + op::Maximum(lhs, rhs)); // Verify execution on CPU. auto result = ExecuteAndTransfer(std::move(hlo_module), {}); - auto expected = LiteralUtil::CreateR1({4, 3, 3, 4}); + auto expected = Literal::CreateR1({4, 3, 3, 4}); LiteralTestUtil::ExpectEqual(*result, *expected); } @@ -80,30 +83,34 @@ TEST_F(InlinerTest, MapConstant) { HloInstruction::CreateParameter(0, r0f32, "x")); (void)param1; const2_builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(2.0f))); + HloInstruction::CreateConstant(Literal::CreateR0(2.0f))); auto const2_f32 = const2_builder.Build(); auto builder = HloComputation::Builder("MapConstFunction"); auto lhs = builder.AddInstruction(HloInstruction::CreateConstant( - LiteralUtil::CreateR2({{1, 2, 3, 4}, {5, 6, 7, 8}}))); + Literal::CreateR2({{1, 2, 3, 4}, {5, 6, 7, 8}}))); builder.AddInstruction( HloInstruction::CreateMap(lhs->shape(), {lhs}, const2_f32.get())); auto computation = builder.Build(); - auto hlo_module = MakeUnique("test_module"); + auto hlo_module = CreateNewModule(); hlo_module->AddEmbeddedComputation(std::move(const2_f32)); hlo_module->AddEntryComputation(std::move(computation)); HloInstruction* root = hlo_module->entry_computation()->root_instruction(); Inliner inliner; EXPECT_TRUE(inliner.Run(hlo_module.get()).ValueOrDie()); root = hlo_module->entry_computation()->root_instruction(); - EXPECT_EQ(root->opcode(), HloOpcode::kBroadcast); + EXPECT_THAT(root, op::Broadcast(op::Constant())); // Verify execution on CPU. auto result = ExecuteAndTransfer(std::move(hlo_module), {}); - auto expected = LiteralUtil::CreateR2({{2, 2, 2, 2}, {2, 2, 2, 2}}); + auto expected = Literal::CreateR2({{2, 2, 2, 2}, {2, 2, 2, 2}}); LiteralTestUtil::ExpectEqual(*result, *expected); } } // namespace } // namespace xla + +int main(int argc, char** argv) { + return xla::ParseDebugOptionsFlagsAndRunTests(argc, argv); +} diff --git a/tensorflow/compiler/xla/service/instruction_fusion.cc b/tensorflow/compiler/xla/service/instruction_fusion.cc index c162945bcae33f4e94b8fbd3a7e48bacce802925..265be54116c406de433da5f07bdec724a1f9f580 100644 --- a/tensorflow/compiler/xla/service/instruction_fusion.cc +++ b/tensorflow/compiler/xla/service/instruction_fusion.cc @@ -28,7 +28,6 @@ limitations under the License. #include "tensorflow/core/platform/logging.h" namespace xla { - /*static*/ bool InstructionFusion::IsExpensive( const HloInstruction& instruction) { switch (instruction.opcode()) { @@ -43,6 +42,7 @@ namespace xla { case HloOpcode::kConstant: case HloOpcode::kConvert: case HloOpcode::kCopy: + case HloOpcode::kCos: case HloOpcode::kDynamicSlice: case HloOpcode::kDynamicUpdateSlice: case HloOpcode::kEq: @@ -64,10 +64,12 @@ namespace xla { case HloOpcode::kNegate: case HloOpcode::kOutfeed: case HloOpcode::kPad: + case HloOpcode::kReducePrecision: case HloOpcode::kReshape: case HloOpcode::kReverse: case HloOpcode::kSelect: case HloOpcode::kSign: + case HloOpcode::kSin: case HloOpcode::kSlice: case HloOpcode::kSubtract: case HloOpcode::kTranspose: @@ -75,6 +77,9 @@ namespace xla { return false; // Expensive instructions. + case HloOpcode::kBatchNormTraining: + case HloOpcode::kBatchNormInference: + case HloOpcode::kBatchNormGrad: case HloOpcode::kCall: case HloOpcode::kConvolution: case HloOpcode::kCrossReplicaSum: @@ -106,19 +111,109 @@ namespace xla { return false; } -namespace { -// Returns true if fusing producer into consumer would cause producer to be -// duplicated. This is the case if producer has uses other than consumer. -bool FusionWouldDuplicate(const HloInstruction& producer, - const HloInstruction& consumer) { - return !(producer.users().size() == 1 && consumer.IsUserOf(&producer)); +// An "effectively unary" operation is one that has one "large" +// input with the others being negligible in terms of memory usage. +// We use "has a smaller true rank than the output" as a heuristic +// for "negligible" memory usage. +bool InstructionFusion::EffectivelyUnary(HloInstruction* hlo) { + int64 output_rank = 0; + ShapeUtil::ForEachSubshape( + hlo->shape(), + [&output_rank](const Shape& subshape, const ShapeIndex& shape_index) { + if (ShapeUtil::IsArray(subshape)) { + 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) { + return false; + } + if (operand->opcode() == HloOpcode::kConstant && + ShapeUtil::IsEffectiveScalar(operand->shape())) { + return false; + } + return ShapeUtil::TrueRank(operand->shape()) >= + output_rank; + }) <= 1; +} + +bool InstructionFusion::CanFuseOnAllPaths( + const HloReachabilityMap& reachability_map, HloInstruction* producer, + HloInstruction* consumer, DoNotFuseSet* do_not_fuse) { + auto could_fuse_on_all_paths = [&] { + // First check to see if we have already marked this producer as infeasible + // to fuse into consumer. + if (do_not_fuse->count(producer) > 0) { + return false; + } + // Make sure it is possible for producer and consumer to exist in a fusion + // node. + if (!producer->IsFusable() || !consumer->IsFusable()) { + return false; + } + // We do an upward walk of the graph from consumer towards all paths which + // lead to producer to find any unfusable paths. + for (int64 i = 0, e = consumer->operand_count(); i < e; ++i) { + auto* consumer_operand = consumer->mutable_operand(i); + if (consumer_operand == producer) { + // This is the base case: our upward crawl ends but we need to make sure + // that fusion from consumer can happen. + if (!ShouldFuse(consumer, i)) { + return false; + } + } else if (reachability_map.IsReachable(producer, consumer_operand)) { + // The reachability map told us that consumer_operand is a node on the + // path to producer. We need to further investigate from + // consumer_operand. + + // First check if we have already ruled out fusing producer into + // consumer_operand. + if (do_not_fuse->count(consumer_operand) > 0) { + return false; + } + // Make sure it is possible for consumer_operand to exist in a fusion + // node. + if (!consumer_operand->IsFusable()) { + return false; + } + // The producer is reachable from consumer_operand which means we need + // to be able to fuse consumer_operand into consumer in order for + // producer to be fusable into consumer on all paths. + if (!ShouldFuse(consumer, i)) { + return false; + } + // Perform the recursive step: make sure producer can be fused into + // consumer_operand on all paths. + if (!CanFuseOnAllPaths(reachability_map, producer, consumer_operand, + do_not_fuse)) { + return false; + } + } + } + return true; + }; + if (could_fuse_on_all_paths()) { + return true; + } + // We couldn't fuse on all paths, record this result. + do_not_fuse->insert(producer); + return false; } -} // namespace StatusOr InstructionFusion::Run(HloModule* module) { bool changed = false; + module_ = module; + std::vector computations; for (auto& computation : module->computations()) { - computation_ = computation.get(); + if (computation->IsFusionComputation()) { + continue; + } + computations.push_back(computation.get()); + } + for (auto& computation : computations) { + CHECK(!computation->IsFusionComputation()); + computation_ = computation; // We want to be able to remove arbitrary instructions from the post order // and also compare positions of instructions in the post order. To make @@ -130,11 +225,43 @@ StatusOr InstructionFusion::Run(HloModule* module) { computation_->MakeInstructionPostOrder(); std::vector post_order(post_order_list.begin(), post_order_list.end()); + tensorflow::gtl::FlatMap post_order_index; for (size_t i = 0; i < post_order.size(); ++i) { InsertOrDie(&post_order_index, post_order[i], i); } + DoNotFuseSet do_not_fuse; + auto reachability = computation->ComputeReachability(); + + auto cheap_to_duplicate = [this](HloInstruction* producer) { + if (producer->opcode() == HloOpcode::kBroadcast) { + return true; + } + if (producer->opcode() == HloOpcode::kConstant && + ShapeUtil::IsEffectiveScalar(producer->shape())) { + return true; + } + if (EffectivelyUnary(producer)) { + return true; + } + return false; + }; + + for (HloInstruction* consumer : post_order) { + for (HloInstruction* producer : consumer->operands()) { + if (cheap_to_duplicate(producer)) { + continue; + } + if (CanFuseOnAllPaths(*reachability, producer, consumer, + &do_not_fuse)) { + CHECK_EQ(do_not_fuse.count(producer), 0); + } else { + CHECK_GT(do_not_fuse.count(producer), 0); + } + } + } + // Instruction fusion effectively fuses edges in the computation graph // (producer instruction -> consumer instruction) so we iterate over all // edges. When we fuse an edge, we create a copy of the producer inside the @@ -216,29 +343,37 @@ StatusOr InstructionFusion::Run(HloModule* module) { for (int64 i : sorted_operand_numbers) { HloInstruction* operand = instruction->mutable_operand(i); - if (operand->IsFusable() && ShouldFuse(instruction, i)) { - HloInstruction* fusion_instruction = Fuse(operand, instruction); - - // Fusing an instruction into a fusion instruction can change the - // operand set of the fusion instruction. For simplicity just push the - // instruction to the top of the post_order and reconsider it for - // further fusion in the next iteration of the outer loop. - post_order.push_back(fusion_instruction); - InsertOrDie(&post_order_index, fusion_instruction, - post_order.size() - 1); - changed = true; - - if (operand->user_count() == 0) { - // Operand is now dead. Remove from post order by setting it's - // location to nullptr. - post_order[FindOrDie(post_order_index, operand)] = nullptr; - post_order_index.erase(operand); - - // Remove from computation. - TF_RETURN_IF_ERROR(computation_->RemoveInstruction(operand)); - } - break; + + if (!operand->IsFusable()) { + continue; + } + if (!ShouldFuse(instruction, i)) { + continue; + } + if (do_not_fuse.count(operand) > 0) { + continue; + } + HloInstruction* fusion_instruction = Fuse(operand, instruction); + + // Fusing an instruction into a fusion instruction can change the + // operand set of the fusion instruction. For simplicity just push the + // instruction to the top of the post_order and reconsider it for + // further fusion in the next iteration of the outer loop. + post_order.push_back(fusion_instruction); + InsertOrDie(&post_order_index, fusion_instruction, + post_order.size() - 1); + changed = true; + + if (operand->user_count() == 0) { + // Operand is now dead. Remove from post order by setting it's + // location to nullptr. + post_order[FindOrDie(post_order_index, operand)] = nullptr; + post_order_index.erase(operand); + + // Remove from computation. + TF_RETURN_IF_ERROR(computation_->RemoveInstruction(operand)); } + break; } } } @@ -249,18 +384,21 @@ HloInstruction* InstructionFusion::Fuse(HloInstruction* producer, HloInstruction* consumer) { HloInstruction* fusion_instruction; - VLOG(2) << "Fusing " << producer << " into " << consumer; - + VLOG(2) << "Fusing " << producer->ToString() << " into " + << consumer->ToString(); + auto kind = ChooseKind(producer, consumer); if (consumer->opcode() == HloOpcode::kFusion) { fusion_instruction = consumer; + if (kind != fusion_instruction->fusion_kind()) { + fusion_instruction->set_fusion_kind(kind); + } } else { - fusion_instruction = - computation_->AddInstruction(HloInstruction::CreateFusion( - consumer->shape(), ChooseKind(producer, consumer), consumer)); + fusion_instruction = computation_->AddInstruction( + HloInstruction::CreateFusion(consumer->shape(), kind, consumer)); TF_CHECK_OK(computation_->ReplaceInstruction(consumer, fusion_instruction)); } - fusion_instruction->FuseInstruction(producer); + fusion_instruction->FuseInstruction(producer); return fusion_instruction; } @@ -275,13 +413,15 @@ bool InstructionFusion::ShouldFuse(HloInstruction* consumer, if (consumer->opcode() == HloOpcode::kFusion && consumer->fusion_kind() != HloInstruction::FusionKind::kLoop && - consumer->fusion_kind() != HloInstruction::FusionKind::kInput) { + consumer->fusion_kind() != HloInstruction::FusionKind::kInput && + consumer->fusion_kind() != HloInstruction::FusionKind::kOutput) { return false; } - // Cost condition: not fuse (expensive producers) and (consumers who reuse - // operand elements). - if (consumer->ReusesOperandElements(operand_index) && + // Cost condition: not fuse (simple, expensive producers) and (consumers who + // reuse operand elements). + if (producer->opcode() != HloOpcode::kFusion && + consumer->ReusesOperandElements(operand_index) && is_expensive_(*producer)) { return false; } diff --git a/tensorflow/compiler/xla/service/instruction_fusion.h b/tensorflow/compiler/xla/service/instruction_fusion.h index a9f3723f2dfcc1b3b697d34eb9510f5857a443f0..0eb8d03489d29b701f09c2912bc906778f02a99b 100644 --- a/tensorflow/compiler/xla/service/instruction_fusion.h +++ b/tensorflow/compiler/xla/service/instruction_fusion.h @@ -66,11 +66,36 @@ class InstructionFusion : public HloPassInterface { virtual HloInstruction::FusionKind ChooseKind(const HloInstruction* producer, const HloInstruction* consumer); + // Fuses producer into consumer. + virtual HloInstruction* Fuse(HloInstruction* producer, + HloInstruction* consumer); + + // An "effectively unary" operation is one that has one "large" + // input with the others being negligible in terms of memory usage. + // We use "has a smaller true rank than the output" as a heuristic + // for "negligible" memory usage. + bool EffectivelyUnary(HloInstruction* hlo); + + // Returns true if fusing producer into consumer would cause producer to be + // duplicated. This is the case if producer has uses other than consumer. + bool FusionWouldDuplicate(const HloInstruction& producer, + const HloInstruction& consumer) { + return !(producer.users().size() == 1 && consumer.IsUserOf(&producer)); + } + // Current HloComputation instance the loop fuser is traversing. HloComputation* computation_; + HloModule* module_; private: - HloInstruction* Fuse(HloInstruction* producer, HloInstruction* consumer); + // The set of producers whose consumers we cannot fuse into. + using DoNotFuseSet = std::unordered_set; + + // Whether or not we can fuse consumer into original_producer on all paths + // from the producer to the consumer where nodes are HLOs and edges are uses. + bool CanFuseOnAllPaths(const HloReachabilityMap& reachability_map, + HloInstruction* producer, HloInstruction* consumer, + DoNotFuseSet* do_not_fuse); // Used to determine if an HLO is expensive. Expensive operations will not be // duplicated. diff --git a/tensorflow/compiler/xla/service/instruction_fusion_test.cc b/tensorflow/compiler/xla/service/instruction_fusion_test.cc index a4c269f0ebd40b2a1ab46619fec24e76ffd73ff0..b3e0007dcc2d43028b49cc48477a0a69153b13c8 100644 --- a/tensorflow/compiler/xla/service/instruction_fusion_test.cc +++ b/tensorflow/compiler/xla/service/instruction_fusion_test.cc @@ -15,8 +15,11 @@ limitations under the License. #include "tensorflow/compiler/xla/service/instruction_fusion.h" +#include "tensorflow/compiler/xla/service/hlo_matchers.h" #include "tensorflow/compiler/xla/tests/hlo_test_base.h" +namespace op = xla::testing::opcode_matchers; + namespace xla { using InstructionFusionTest = HloTestBase; @@ -25,14 +28,14 @@ TEST_F(InstructionFusionTest, CostlyProducerAndOperandElementReusingConsumerNotFused) { HloComputation::Builder builder(TestName()); HloInstruction* const0 = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(5))); + HloInstruction::CreateConstant(Literal::CreateR0(5))); HloInstruction* exp1 = builder.AddInstruction(HloInstruction::CreateUnary( ShapeUtil::MakeShape(S32, {}), HloOpcode::kExp, const0)); HloInstruction* broadcast2 = builder.AddInstruction(HloInstruction::CreateBroadcast( ShapeUtil::MakeShape(S32, {1}), exp1, {0})); - auto module = MakeUnique(TestName()); + auto module = CreateNewModule(); auto computation = module->AddEntryComputation(builder.Build()); EXPECT_EQ(broadcast2, computation->root_instruction()); EXPECT_TRUE( @@ -46,61 +49,61 @@ TEST_F(InstructionFusionTest, NonCostlyProducerAndOperandElementReusingConsumerFused) { HloComputation::Builder builder(TestName()); HloInstruction* const0 = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(5))); + HloInstruction::CreateConstant(Literal::CreateR0(5))); HloInstruction* negate1 = builder.AddInstruction(HloInstruction::CreateUnary( ShapeUtil::MakeShape(S32, {}), HloOpcode::kNegate, const0)); HloInstruction* broadcast2 = builder.AddInstruction(HloInstruction::CreateBroadcast( ShapeUtil::MakeShape(S32, {1}), negate1, {0})); - auto module = MakeUnique(TestName()); + auto module = CreateNewModule(); auto computation = module->AddEntryComputation(builder.Build()); EXPECT_EQ(broadcast2, computation->root_instruction()); EXPECT_TRUE( InstructionFusion(InstructionFusion::IsExpensive, /*may_duplicate=*/true) .Run(module.get()) .ValueOrDie()); - EXPECT_EQ(HloOpcode::kFusion, computation->root_instruction()->opcode()); + EXPECT_THAT(computation->root_instruction(), op::Fusion()); } TEST_F(InstructionFusionTest, CostlyProducerAndNonOperandElementReusingConsumerFused_Reshape) { HloComputation::Builder builder(TestName()); HloInstruction* const0 = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(5))); + HloInstruction::CreateConstant(Literal::CreateR0(5))); HloInstruction* exp1 = builder.AddInstruction(HloInstruction::CreateUnary( ShapeUtil::MakeShape(S32, {}), HloOpcode::kExp, const0)); HloInstruction* reshape2 = builder.AddInstruction( HloInstruction::CreateReshape(ShapeUtil::MakeShape(S32, {}), exp1)); - auto module = MakeUnique(TestName()); + auto module = CreateNewModule(); auto computation = module->AddEntryComputation(builder.Build()); EXPECT_EQ(reshape2, computation->root_instruction()); EXPECT_TRUE( InstructionFusion(InstructionFusion::IsExpensive, /*may_duplicate=*/true) .Run(module.get()) .ValueOrDie()); - EXPECT_EQ(HloOpcode::kFusion, computation->root_instruction()->opcode()); + EXPECT_THAT(computation->root_instruction(), op::Fusion()); } TEST_F(InstructionFusionTest, CostlyProducerAndNonOperandElementReusingConsumerFused_Transpose) { HloComputation::Builder builder(TestName()); HloInstruction* const0 = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(5))); + HloInstruction::CreateConstant(Literal::CreateR0(5))); HloInstruction* exp1 = builder.AddInstruction(HloInstruction::CreateUnary( ShapeUtil::MakeShape(S32, {}), HloOpcode::kExp, const0)); HloInstruction* transpose2 = builder.AddInstruction( HloInstruction::CreateTranspose(ShapeUtil::MakeShape(S32, {}), exp1, {})); - auto module = MakeUnique(TestName()); + auto module = CreateNewModule(); auto computation = module->AddEntryComputation(builder.Build()); EXPECT_EQ(transpose2, computation->root_instruction()); EXPECT_TRUE( InstructionFusion(InstructionFusion::IsExpensive, /*may_duplicate=*/true) .Run(module.get()) .ValueOrDie()); - EXPECT_EQ(HloOpcode::kFusion, computation->root_instruction()->opcode()); + EXPECT_THAT(computation->root_instruction(), op::Fusion()); } TEST_F(InstructionFusionTest, PotentialBitcastReshapeOfParameterUnfused) { @@ -110,7 +113,7 @@ TEST_F(InstructionFusionTest, PotentialBitcastReshapeOfParameterUnfused) { auto reshape1 = builder.AddInstruction( HloInstruction::CreateReshape(ShapeUtil::MakeShape(S32, {1, 1}), param0)); - auto module = MakeUnique(TestName()); + auto module = CreateNewModule(); auto computation = module->AddEntryComputation(builder.Build()); EXPECT_EQ(reshape1, computation->root_instruction()); EXPECT_FALSE( @@ -126,7 +129,7 @@ TEST_F(InstructionFusionTest, PotentialBitcastSimpleReshapeOfParameterUnfused) { auto reshape1 = builder.AddInstruction( HloInstruction::CreateReshape(ShapeUtil::MakeShape(S32, {1, 1}), param0)); - auto module = MakeUnique(TestName()); + auto module = CreateNewModule(); auto computation = module->AddEntryComputation(builder.Build()); EXPECT_EQ(reshape1, computation->root_instruction()); EXPECT_FALSE( @@ -142,7 +145,7 @@ TEST_F(InstructionFusionTest, PotentialBitcastTransposeOfParameterUnfused) { auto transpose1 = builder.AddInstruction(HloInstruction::CreateTranspose( ShapeUtil::MakeShape(S32, {}), param0, {})); - auto module = MakeUnique(TestName()); + auto module = CreateNewModule(); auto computation = module->AddEntryComputation(builder.Build()); EXPECT_EQ(transpose1, computation->root_instruction()); EXPECT_FALSE( @@ -151,4 +154,73 @@ TEST_F(InstructionFusionTest, PotentialBitcastTransposeOfParameterUnfused) { .ValueOrDie()); } +TEST_F(InstructionFusionTest, AvoidDuplicationIfNotAllFusable) { + HloComputation::Builder builder(TestName()); + auto shape = ShapeUtil::MakeShape(F32, {16, 16}); + auto param0 = + builder.AddInstruction(HloInstruction::CreateParameter(0, shape, "0")); + auto param1 = + builder.AddInstruction(HloInstruction::CreateParameter(1, shape, "1")); + HloInstruction* binary1 = builder.AddInstruction( + HloInstruction::CreateBinary(shape, HloOpcode::kAdd, param0, param1)); + builder.AddInstruction(HloInstruction::CreateSend(binary1, 0)); + HloInstruction* unary = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kAbs, binary1)); + + auto module = CreateNewModule(); + auto computation = module->AddEntryComputation(builder.Build()); + EXPECT_EQ(unary, computation->root_instruction()); + EXPECT_FALSE( + InstructionFusion(InstructionFusion::IsExpensive, /*may_duplicate=*/true) + .Run(module.get()) + .ValueOrDie()); +} + +TEST_F(InstructionFusionTest, AllowUnaryDuplication) { + HloComputation::Builder builder(TestName()); + auto shape = ShapeUtil::MakeShape(F32, {16, 16}); + auto param0 = + builder.AddInstruction(HloInstruction::CreateParameter(0, shape, "0")); + HloInstruction* unary1 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kFloor, param0)); + builder.AddInstruction(HloInstruction::CreateSend(unary1, 0)); + HloInstruction* unary2 = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kAbs, unary1)); + + auto module = CreateNewModule(); + auto computation = module->AddEntryComputation(builder.Build()); + EXPECT_EQ(unary2, computation->root_instruction()); + EXPECT_TRUE( + InstructionFusion(InstructionFusion::IsExpensive, /*may_duplicate=*/true) + .Run(module.get()) + .ValueOrDie()); +} + +TEST_F(InstructionFusionTest, AllowEffectiveUnaryDuplication) { + auto shape = ShapeUtil::MakeShape(F32, {16, 16}); + auto small_shape = ShapeUtil::MakeShape(F32, {16}); + HloComputation::Builder builder(TestName()); + auto param0 = builder.AddInstruction( + HloInstruction::CreateParameter(0, small_shape, "0")); + auto param1 = + builder.AddInstruction(HloInstruction::CreateParameter(1, shape, "1")); + HloInstruction* binary1 = builder.AddInstruction( + HloInstruction::CreateBinary(shape, HloOpcode::kAdd, param0, param1)); + builder.AddInstruction(HloInstruction::CreateSend(binary1, 0)); + HloInstruction* unary = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kAbs, binary1)); + + auto module = CreateNewModule(); + auto computation = module->AddEntryComputation(builder.Build()); + EXPECT_EQ(unary, computation->root_instruction()); + EXPECT_TRUE( + InstructionFusion(InstructionFusion::IsExpensive, /*may_duplicate=*/true) + .Run(module.get()) + .ValueOrDie()); +} + } // namespace xla + +int main(int argc, char** argv) { + return xla::ParseDebugOptionsFlagsAndRunTests(argc, argv); +} diff --git a/tensorflow/compiler/xla/service/layout_assignment.cc b/tensorflow/compiler/xla/service/layout_assignment.cc index 5e7bd4a7ce8a1152973979d4a8fdb790a7fbd219..47ca070edab93177483092caecd0997fc7e5f518 100644 --- a/tensorflow/compiler/xla/service/layout_assignment.cc +++ b/tensorflow/compiler/xla/service/layout_assignment.cc @@ -60,8 +60,9 @@ std::ostream& operator<<(std::ostream& out, } BufferLayoutConstraint::BufferLayoutConstraint(const Layout& layout, - const LogicalBuffer& buffer) - : layout_(layout), buffer_(&buffer) { + const LogicalBuffer& buffer, + bool mandatory) + : LayoutConstraint(mandatory), layout_(layout), buffer_(&buffer) { CHECK(LayoutUtil::ValidateLayoutForShape(layout, buffer.shape()).ok()); } @@ -73,8 +74,9 @@ string BufferLayoutConstraint::ToString() const { OperandLayoutConstraint::OperandLayoutConstraint( const ShapeLayout& shape_layout, const HloInstruction* instruction, - int64 operand_no) - : shape_layout_(shape_layout), + int64 operand_no, bool mandatory) + : LayoutConstraint(mandatory), + shape_layout_(shape_layout), instruction_(instruction), operand_no_(operand_no) { CHECK(shape_layout_.LayoutIsSet()); @@ -99,9 +101,13 @@ LayoutConstraints::LayoutConstraints( const HloComputation* computation) : points_to_analysis_(points_to_analysis), computation_(computation) { // Gather all array-shaped logical buffers into unconstrained_buffer_ids. - for (auto& buffer : points_to_analysis_.logical_buffers()) { - if (buffer->IsArray()) { - unconstrained_buffer_ids_.insert(buffer->id()); + for (LogicalBuffer::Id id = 0; id < points_to_analysis_.num_logical_buffers(); + id++) { + auto& buffer = points_to_analysis_.logical_buffer(id); + // The points to analysis is computed per module, restrict constraints to + // array buffers in this computation. + if (buffer.IsArray() && buffer.instruction()->parent() == computation) { + unconstrained_buffer_ids_.insert(buffer.id()); } } } @@ -124,7 +130,8 @@ bool LayoutConstraints::OperandBufferForwarded( } Status LayoutConstraints::SetBufferLayout(const Layout& layout, - const LogicalBuffer& buffer) { + const LogicalBuffer& buffer, + bool mandatory) { VLOG(3) << "SetBufferLayout : " << buffer << " : " << LayoutUtil::HumanString(layout); @@ -139,26 +146,38 @@ Status LayoutConstraints::SetBufferLayout(const Layout& layout, TF_RETURN_IF_ERROR( LayoutUtil::ValidateLayoutForShape(layout, buffer.shape())); - const Layout* curr_layout = BufferLayout(buffer); - if (curr_layout != nullptr) { - if (!LayoutUtil::Equal(*curr_layout, layout)) { + const BufferLayoutConstraint* curr_constraint = + GetBufferLayoutConstraint(buffer); + if (curr_constraint != nullptr) { + if (LayoutUtil::Equal(curr_constraint->layout(), layout)) { + // New constraint matches existing constraint. Nothing to do. + return Status::OK(); + } + if (curr_constraint->mandatory()) { return FailedPrecondition( "Buffer %s already has the layout constraint %s, cannot add " "incompatible constraint %s", buffer.ToString().c_str(), - LayoutUtil::HumanString(*curr_layout).c_str(), + LayoutUtil::HumanString(curr_constraint->layout()).c_str(), LayoutUtil::HumanString(layout).c_str()); } - // New constraint matches existing constraint. Nothing to do. - return Status::OK(); } - auto new_constraint_it = buffer_constraints_.insert( - {&buffer, BufferLayoutConstraint(layout, buffer)}); - added_constraints_.push_back(&new_constraint_it.first->second); + auto iter = buffer_constraints_.find(&buffer); + bool overwrite = iter != buffer_constraints_.end(); + if (!overwrite) { + iter = buffer_constraints_ + .insert(std::make_pair( + &buffer, BufferLayoutConstraint(layout, buffer, mandatory))) + .first; + } else { + iter->second = BufferLayoutConstraint(layout, buffer, /*mandatory=*/true); + } + added_constraints_.push_back(&iter->second); // Remove buffer from the set of unconstrained buffers. - TF_RET_CHECK(unconstrained_buffer_ids_.count(buffer.id()) == 1); + TF_RET_CHECK(unconstrained_buffer_ids_.count(buffer.id()) == + static_cast(!overwrite)); unconstrained_buffer_ids_.erase(buffer.id()); return Status::OK(); @@ -166,23 +185,27 @@ Status LayoutConstraints::SetBufferLayout(const Layout& layout, Status LayoutConstraints::SetOperandLayout(const Shape& shape_with_layout, const HloInstruction* instruction, - int64 operand_no) { + int64 operand_no, bool mandatory) { VLOG(3) << "SetOperandLayout : " << instruction->name() << ", operand " << operand_no << " : " << ShapeUtil::HumanStringWithLayout(shape_with_layout); - const ShapeLayout* curr_shape_layout = OperandLayout(instruction, operand_no); + const OperandLayoutConstraint* curr_shape_layout = + GetOperandLayoutConstraint(instruction, operand_no); if (curr_shape_layout != nullptr) { - if (!curr_shape_layout->MatchesLayoutInShape(shape_with_layout)) { + if (curr_shape_layout->shape_layout().MatchesLayoutInShape( + shape_with_layout)) { + // New constraint matches existing constraint. Nothing to do. + return Status::OK(); + } + if (curr_shape_layout->mandatory()) { return FailedPrecondition( "Operand %lld of instruction %s already has a layout constraint " "%s, cannot add incompatible constraint %s", operand_no, instruction->name().c_str(), - curr_shape_layout->ToString().c_str(), + curr_shape_layout->shape_layout().ToString().c_str(), ShapeUtil::HumanStringWithLayout(shape_with_layout).c_str()); } - // New constraint matches existing constraint. Nothing to do. - return Status::OK(); } // If any buffers in the operand occur in the output of the instruction, then @@ -196,22 +219,31 @@ Status LayoutConstraints::SetOperandLayout(const Shape& shape_with_layout, } auto key = std::make_pair(instruction, operand_no); - auto new_constraint_it = operand_constraints_.insert( - {key, OperandLayoutConstraint(ShapeLayout(shape_with_layout), instruction, - operand_no)}); - added_constraints_.push_back(&new_constraint_it.first->second); + auto iter = operand_constraints_.find(key); + if (iter == operand_constraints_.end()) { + auto pair = std::make_pair( + key, OperandLayoutConstraint(ShapeLayout(shape_with_layout), + instruction, operand_no, mandatory)); + iter = operand_constraints_.insert(pair).first; + } else { + iter->second = + OperandLayoutConstraint(ShapeLayout(shape_with_layout), instruction, + operand_no, /*mandatory=*/true); + } + added_constraints_.push_back(&iter->second); return Status::OK(); } Status LayoutConstraints::SetArrayOperandLayout( - const Layout& layout, const HloInstruction* instruction, int64 operand_no) { + const Layout& layout, const HloInstruction* instruction, int64 operand_no, + bool mandatory) { const HloInstruction* operand = instruction->operand(operand_no); TF_RET_CHECK(ShapeUtil::IsArray(operand->shape())); Shape shape(operand->shape()); *shape.mutable_layout() = layout; TF_RETURN_IF_ERROR(LayoutUtil::ValidateLayoutInShape(shape)); - return SetOperandLayout(shape, instruction, operand_no); + return SetOperandLayout(shape, instruction, operand_no, mandatory); } Status LayoutConstraints::SetResultLayout(const Shape& shape_with_layout) { @@ -253,7 +285,7 @@ Status LayoutConstraints::SetInstructionLayout( // Create a BufferLayoutConstraint for each array shape in the output of the // instruction. - return ShapeUtil::ForEachSubshape( + return ShapeUtil::ForEachSubshapeWithStatus( shape_with_layout, [this, instruction](const Shape& subshape, const ShapeIndex& index) -> Status { @@ -274,15 +306,31 @@ Status LayoutConstraints::SetInstructionLayout( const Layout* LayoutConstraints::BufferLayout( const LogicalBuffer& buffer) const { + if (const auto* constraint = GetBufferLayoutConstraint(buffer)) { + return &constraint->layout(); + } + return nullptr; +} + +const BufferLayoutConstraint* LayoutConstraints::GetBufferLayoutConstraint( + const LogicalBuffer& buffer) const { auto it = buffer_constraints_.find(&buffer); - return it == buffer_constraints_.end() ? nullptr : &it->second.layout(); + return it == buffer_constraints_.end() ? nullptr : &it->second; } const ShapeLayout* LayoutConstraints::OperandLayout( const HloInstruction* instruction, int64 operand_no) const { + if (const auto* constraint = + GetOperandLayoutConstraint(instruction, operand_no)) { + return &constraint->shape_layout(); + } + return nullptr; +} + +const OperandLayoutConstraint* LayoutConstraints::GetOperandLayoutConstraint( + const HloInstruction* instruction, int64 operand_no) const { auto it = operand_constraints_.find(std::make_pair(instruction, operand_no)); - return it == operand_constraints_.end() ? nullptr - : &it->second.shape_layout(); + return it == operand_constraints_.end() ? nullptr : &it->second; } const ShapeLayout* LayoutConstraints::ResultLayout() const { @@ -338,12 +386,15 @@ Status LayoutAssignment::AddMandatoryConstraints( // instruction. // TODO(b/31425034): Change infeeds to be more like parameters, with // shapes in the ComputationLayout. - shape_with_layout = &instruction->shape(); + DCHECK(!LayoutUtil::IsPadded(instruction->shape())); + TF_RETURN_IF_ERROR(constraints->SetInstructionLayout(instruction->shape(), + instruction.get())); } else if (instruction->opcode() == HloOpcode::kOutfeed) { // Constrain the input to the Outfeed instruction to be the expected // layout of the Outfeed. TF_RETURN_IF_ERROR(constraints->SetOperandLayout( - instruction->outfeed_shape(), instruction.get(), 0)); + instruction->outfeed_shape(), instruction.get(), 0, + /*mandatory=*/true)); } else if (instruction->opcode() == HloOpcode::kParameter) { // Parameter layouts must match the respective layout in // ComputationLayout. @@ -375,7 +426,7 @@ Status LayoutAssignment::AddMandatoryConstraints( for (int64 i = 0; i < instruction->operand_count(); ++i) { TF_RETURN_IF_ERROR(constraints->SetOperandLayout( called_computation_layout.parameter_layout(i).shape(), - instruction.get(), i)); + instruction.get(), i, /*mandatory=*/true)); } } else if (instruction->opcode() == HloOpcode::kWhile) { // Layout of input and output of kWhile instruction must be equal and must @@ -426,7 +477,8 @@ Status LayoutAssignment::AddMandatoryConstraints( TF_RETURN_IF_ERROR(constraints->SetInstructionLayout( body_layout.result_shape(), instruction.get())); TF_RETURN_IF_ERROR(constraints->SetOperandLayout( - body_layout.result_shape(), instruction.get(), 0)); + body_layout.result_shape(), instruction.get(), 0, + /*mandatory=*/true)); } else if (instruction->opcode() == HloOpcode::kCustomCall) { // Add constraints for kCustomCall instruction operands and instructions. // For now we only support row major layouts for all inputs and outputs. @@ -450,7 +502,7 @@ Status LayoutAssignment::AddMandatoryConstraints( Shape row_major_operand_shape(row_major_shape(operand_shape)); TF_RETURN_IF_ERROR(constraints->SetOperandLayout( - row_major_operand_shape, instruction.get(), i)); + row_major_operand_shape, instruction.get(), i, /*mandatory=*/true)); } } } @@ -561,6 +613,9 @@ Status CheckLayouts( TF_ASSIGN_OR_RETURN(auto points_to_analysis, TuplePointsToAnalysis::Run(module)); for (auto& computation : module->computations()) { + if (computation->IsFusionComputation()) { + continue; + } for (auto& instruction : computation->instructions()) { // Verify every instruction has a layout and the layout is valid for the // shape. @@ -572,11 +627,10 @@ Status CheckLayouts( // which could be the source of the subshape value. const PointsToSet& points_to_set = points_to_analysis->GetPointsToSet(instruction.get()); - TF_RETURN_IF_ERROR(points_to_set.ForEachElement( - [&instruction]( - ShapeIndex index, bool is_leaf, - const std::vector& buffers) -> Status { - if (is_leaf) { + TF_RETURN_IF_ERROR(points_to_set.ForEachElementWithStatus( + [&instruction](ShapeIndex index, + const PointsToSet::BufferList& buffers) -> Status { + if (ShapeUtil::IsLeafIndex(instruction->shape(), index)) { const Shape& instruction_subshape = ShapeUtil::GetSubshape(instruction->shape(), index); for (const LogicalBuffer* buffer : buffers) { @@ -659,44 +713,6 @@ LayoutAssignment::LayoutAssignment(ComputationLayout* entry_computation_layout) } } -namespace { - -// Given a pemutation of `{0, 1, ..., n}` `indices`, returns a permutation of -// `{0, 1, ..., n - to_delete.size() + to_insert.size()}` by deleting the -// indices `to_delete` wherever in `indices` they are, and inserting the indices -// `to_insert` arbitrarily at the back. -tensorflow::protobuf::RepeatedField -DeleteAndInsertIndices( - std::vector to_delete, std::vector to_insert, - tensorflow::protobuf::RepeatedField indices) { - std::sort(to_delete.begin(), to_delete.end(), std::greater()); - std::sort(to_insert.begin(), to_insert.end(), std::less()); - for (auto index : to_delete) { - auto i = indices.begin(); - while (i != indices.end()) { - if (*i == index) { - i = indices.erase(i); - } else { - if (*i > index) { - (*i)--; - } - ++i; - } - } - } - for (auto index : to_insert) { - for (auto i = indices.begin(); i != indices.end(); ++i) { - if (*i >= index) { - (*i)++; - } - } - indices.Add(index); - } - return indices; -} - -} // namespace - std::unique_ptr LayoutAssignment::ChooseOperandLayoutFromOutputLayout( const Layout& output_layout, const HloInstruction* instruction, int64 operand_no) { @@ -719,31 +735,46 @@ std::unique_ptr LayoutAssignment::ChooseOperandLayoutFromOutputLayout( } if (instruction->opcode() == HloOpcode::kReshape) { - // Pick the operand layout that makes the reshape a bitcast. If the reshape - // only inserts or deletes degenerate dimensions, we can easily compute the - // desired layout by accordingly inserting and deleting the elements in the - // minor-to-major list. - bool merely_inserts_or_deletes_1_sized_dims; - std::vector inserted_indices, deleted_indices; - std::tie(merely_inserts_or_deletes_1_sized_dims, deleted_indices, - inserted_indices) = - instruction->ReshapeMerelyInsertsOrDeletes1SizedDimensions(); - if (merely_inserts_or_deletes_1_sized_dims) { - Layout operand_layout = LayoutUtil::MakeLayout( - AsInt64Slice(DeleteAndInsertIndices(inserted_indices, deleted_indices, - output_layout.minor_to_major()))); + // Prefer the operand layout that makes the reshape an bitcast. If any + // dimension bound is 1 in the operand shape, there may be several such + // layouts. So if 'output_layout' is the default layout, try if the + // reshape is a bitcast when using the same layout. This may avoid copy + // operations. + if (ShapeUtil::TrueRank(operand->shape()) == 1 && + ShapeUtil::Rank(instruction->shape()) == 1) { + // Don't assign a layout in case of R1 -> effective R1 reshape. + return nullptr; + } + const Shape& output_shape = instruction->shape(); + Shape output_shape_with_layout = ShapeUtil::MakeShapeWithLayout( + output_shape.element_type(), AsInt64Slice(output_shape.dimensions()), + AsInt64Slice(output_layout.minor_to_major())); + Shape operand_shape = operand->shape(); + *operand_shape.mutable_layout() = + LayoutUtil::GetDefaultLayoutForShape(operand_shape); + if (ShapeUtil::ReshapeIsBitcast(operand_shape, output_shape_with_layout)) { + return MakeUnique(operand_shape.layout()); + } + auto aligned_operand_shape = + ShapeUtil::AlignLayouts(output_shape_with_layout, operand_shape); + if (aligned_operand_shape) { + auto operand_layout = aligned_operand_shape.value().layout(); TF_CHECK_OK( - LayoutUtil::ValidateLayoutForShape(operand_layout, operand->shape())); + LayoutUtil::ValidateLayoutForShape(operand_layout, operand_shape)); return MakeUnique(operand_layout); } } if (instruction->opcode() == HloOpcode::kTranspose) { // Pick the operand layout that makes the transpose a bitcast. - std::vector perm = - ComposePermutations(instruction->dimensions(), - AsInt64Slice(output_layout.minor_to_major())); - Layout operand_layout = LayoutUtil::MakeLayout(perm); + int64 rank = ShapeUtil::Rank(instruction->shape()); + std::vector new_minor_to_major(rank); + for (int64 i = 0; i < rank; ++i) { + int64 output_dim = output_layout.minor_to_major(i); + int64 operand_dim = instruction->dimensions(output_dim); + new_minor_to_major[i] = operand_dim; + } + Layout operand_layout = LayoutUtil::MakeLayout(new_minor_to_major); TF_CHECK_OK( LayoutUtil::ValidateLayoutForShape(operand_layout, operand->shape())); return MakeUnique(operand_layout); @@ -768,31 +799,47 @@ std::unique_ptr LayoutAssignment::ChooseOutputLayoutFromOperandLayout( } if (user->opcode() == HloOpcode::kReshape) { - // Pick the user layout that makes the reshape a bitcast. - bool merely_inserts_or_deletes_1_sized_dims; - std::vector inserted_indices, deleted_indices; - std::tie(merely_inserts_or_deletes_1_sized_dims, deleted_indices, - inserted_indices) = - user->ReshapeMerelyInsertsOrDeletes1SizedDimensions(); - if (merely_inserts_or_deletes_1_sized_dims) { - Layout user_layout = LayoutUtil::MakeLayout(AsInt64Slice( - DeleteAndInsertIndices(deleted_indices, inserted_indices, - operand_layout.minor_to_major()))); + // Prefer the user layout that makes the reshape an bitcast. If any + // dimension bound is 1 in the user shape, there may be several such + // layouts. So if 'operand_layout' is the default layout, try if the + // reshape is a bitcast when using the same layout. This may avoid copy + // operations. + if (ShapeUtil::Rank(operand->shape()) == 1 && + ShapeUtil::TrueRank(user->shape()) == 1) { + // Don't assign a layout in case of R1 -> effective R1 reshape. + return nullptr; + } + Shape operand_shape_with_layout = ShapeUtil::MakeShapeWithLayout( + operand->shape().element_type(), + AsInt64Slice(operand->shape().dimensions()), + AsInt64Slice(operand_layout.minor_to_major())); + Shape output_shape = user->shape(); + *output_shape.mutable_layout() = + LayoutUtil::GetDefaultLayoutForShape(output_shape); + if (ShapeUtil::ReshapeIsBitcast(output_shape, operand_shape_with_layout)) { + return MakeUnique(output_shape.layout()); + } + auto aligned_user_shape = + ShapeUtil::AlignLayouts(operand_shape_with_layout, output_shape); + if (aligned_user_shape) { + auto user_layout = aligned_user_shape.value().layout(); TF_CHECK_OK( - LayoutUtil::ValidateLayoutForShape(user_layout, user->shape())); + LayoutUtil::ValidateLayoutForShape(user_layout, output_shape)); return MakeUnique(user_layout); } } if (user->opcode() == HloOpcode::kTranspose) { - // Pick the user layout that makes the reshape a bitcast. - // To become a bitcast, the layouts need to satisfy - // collapsing_order * output_layout = input_layout - // so output_layout = inverse(collapsing_order) * input_layout - std::vector perm = - Permute(InversePermutation(user->dimensions()), - AsInt64Slice(operand_layout.minor_to_major())); - Layout user_layout = LayoutUtil::MakeLayout(perm); + // Pick the user layout that makes the transpose a bitcast. + int64 rank = ShapeUtil::Rank(user->shape()); + std::vector new_minor_to_major(rank); + auto inverse_dimensions = InversePermutation(user->dimensions()); + for (int64 i = 0; i < rank; ++i) { + int64 operand_dim = operand_layout.minor_to_major(i); + int64 user_dim = inverse_dimensions[operand_dim]; + new_minor_to_major[i] = user_dim; + } + Layout user_layout = LayoutUtil::MakeLayout(new_minor_to_major); TF_CHECK_OK(LayoutUtil::ValidateLayoutForShape(user_layout, user->shape())); return MakeUnique(user_layout); } @@ -883,17 +930,17 @@ Status LayoutAssignment::PropagateUseConstraintToDefs( // match the given layout. const PointsToSet& points_to_set = constraints->points_to_analysis().GetPointsToSet(instruction); - return points_to_set.ForEachElement( + return points_to_set.ForEachElementWithStatus( [this, &shape_layout, constraints]( - const ShapeIndex& index, bool is_leaf, - const std::vector& buffers) -> Status { - if (is_leaf) { + const ShapeIndex& index, + const PointsToSet::BufferList& buffers) -> Status { + if (ShapeUtil::IsLeafIndex(shape_layout.shape(), index)) { for (const LogicalBuffer* buffer : buffers) { if (constraints->BufferLayout(*buffer) == nullptr && ShapeUtil::IsArray(buffer->shape())) { TF_RETURN_IF_ERROR(constraints->SetBufferLayout( ShapeUtil::GetSubshape(shape_layout.shape(), index).layout(), - *buffer)); + *buffer, /*mandatory=*/true)); } } } @@ -936,7 +983,8 @@ Status LayoutAssignment::PropagateOperandConstraint( operand_constraint.shape_layout().layout(), user, operand_constraint.operand_no()); if (layout != nullptr) { - TF_RETURN_IF_ERROR(constraints->SetBufferLayout(*layout, *buffer)); + TF_RETURN_IF_ERROR( + constraints->SetBufferLayout(*layout, *buffer, /*mandatory=*/false)); } } return Status::OK(); @@ -966,11 +1014,19 @@ Status LayoutAssignment::PropagateBufferConstraint( instruction, operand_no); if (operand_layout != nullptr) { TF_RETURN_IF_ERROR(constraints->SetArrayOperandLayout( - *operand_layout, instruction, operand_no)); + *operand_layout, instruction, operand_no, /*mandatory=*/true)); } } } } + return PropagateBufferConstraintToUses(buffer_constraint, constraints); +} + +Status LayoutAssignment::PropagateBufferConstraintToUses( + const BufferLayoutConstraint& buffer_constraint, + LayoutConstraints* constraints) { + const LogicalBuffer& buffer = buffer_constraint.buffer(); + TF_RET_CHECK(buffer.IsArray()); // Propagate the layout to all array uses of the logical buffer. This skips // uses of the buffer where the buffer is the element of a tuple. @@ -983,7 +1039,7 @@ Status LayoutAssignment::PropagateBufferConstraint( if (constraints->OperandLayout(user, operand_no) == nullptr && !constraints->OperandBufferForwarded(user, operand_no)) { TF_RETURN_IF_ERROR(constraints->SetArrayOperandLayout( - buffer_constraint.layout(), user, operand_no)); + buffer_constraint.layout(), user, operand_no, /*mandatory=*/false)); } } @@ -1019,7 +1075,7 @@ StatusOr InferArrayLayout( TF_RET_CHECK( !points_to_analysis.InstructionDefinesBufferAtIndex(instruction, index)); - const std::vector& source_buffers = + const auto& source_buffers = points_to_analysis.GetPointsToSet(instruction).element(index); TF_RET_CHECK(!source_buffers.empty()); @@ -1040,7 +1096,7 @@ StatusOr InferArrayLayout( *first_buffer_layout)) { // The points-to set is ambiguous for this index and the different source // buffers have different layouts. This case is possible in valid XLA - // computations because we do not propagate BufferLayoutConstaints to all + // computations because we do not propagate BufferLayoutConstraints to all // LogicalBuffers which may alias the constrained LogicalBuffer at some // point in the computation. return FailedPrecondition( @@ -1203,7 +1259,7 @@ Status LayoutAssignment::AssignLayouts(const LayoutConstraints& constraints, // Any remaining layouts in the output of the instruction must be // inferrable using points-to analysis. - TF_RETURN_IF_ERROR(ShapeUtil::ForEachMutableSubshape( + TF_RETURN_IF_ERROR(ShapeUtil::ForEachMutableSubshapeWithStatus( instruction->mutable_shape(), [instruction, &constraints](Shape* subshape, const ShapeIndex& index) { if (subshape->has_layout() || !ShapeUtil::IsArray(*subshape)) { @@ -1223,6 +1279,9 @@ Status LayoutAssignment::AssignLayouts(const LayoutConstraints& constraints, TF_RETURN_IF_ERROR(SetFusionLayouts(instruction)); } + // Execute extra verification step once the layout has been finalized. + TF_RETURN_IF_ERROR(Verify(instruction)); + // Verify all layouts in the shape have been set. TF_RET_CHECK(LayoutUtil::HasLayout(instruction->shape())); } @@ -1243,18 +1302,17 @@ Status LayoutAssignment::AssignLayouts(const LayoutConstraints& constraints, } Status LayoutAssignment::RunOnComputation( - const ComputationLayout& computation_layout, HloComputation* computation) { + const ComputationLayout& computation_layout, + const TuplePointsToAnalysis& points_to_analysis, + HloComputation* computation) { DCHECK(computation_layout.LayoutIsSet()); InsertOrDie(&computation_layouts_, computation, computation_layout); VLOG(2) << "LayoutAssignment::RunOnComputation(" << computation->name() << ")"; VLOG(2) << " ComputationLayout = " << computation_layout.ToString(); - TF_ASSIGN_OR_RETURN(auto points_to_analysis, - TuplePointsToAnalysis::Run(computation->parent())); - - // Construct LayoutConstaints with all layout constraints of the computation. - LayoutConstraints constraints(*points_to_analysis, computation); + // Construct LayoutConstraints with all layout constraints of the computation. + LayoutConstraints constraints(points_to_analysis, computation); // Add constraints required for correctness on all backends (eg, entry // parameter layout constraints). @@ -1275,10 +1333,11 @@ Status LayoutAssignment::RunOnComputation( // Arbitrarily pick the first unconstrained buffer and give it the default // layout. By construction unconstrained_buffers() has a stable sort based // on LogicalBuffer::Id. - const LogicalBuffer& buffer = points_to_analysis->GetBuffer( + const LogicalBuffer& buffer = points_to_analysis.GetBuffer( *constraints.unconstrained_buffer_ids().begin()); TF_RETURN_IF_ERROR(constraints.SetBufferLayout( - LayoutUtil::GetDefaultLayoutForShape(buffer.shape()), buffer)); + LayoutUtil::GetDefaultLayoutForShape(buffer.shape()), buffer, + /*mandatory=*/false)); TF_RETURN_IF_ERROR(PropagateConstraints(&constraints)); @@ -1300,23 +1359,29 @@ StatusOr LayoutAssignment::Run(HloModule* module) { if (VLOG_IS_ON(10)) { hlo_graph_dumper::DumpGraph(*module->entry_computation(), "before layout assignment", - /*show_addresses=*/false, - /*show_layouts=*/true); + module->config().debug_options()); } + TF_ASSIGN_OR_RETURN(auto points_to_analysis, + TuplePointsToAnalysis::Run(module)); + // Assign layouts to computations in an order such that a callee computation // is handled before its caller computation. This ensures that the layout of // all callers of a computation will agree. for (auto* computation : module->MakeComputationPostOrder()) { if (computation == module->entry_computation()) { TF_RETURN_IF_ERROR(RunOnComputation(*entry_computation_layout_, + *points_to_analysis, module->entry_computation())); + } else if (computation->IsFusionComputation()) { + continue; } else { ComputationLayout computation_layout(computation->ComputeProgramShape()); // Setting all embedded computations to the default layout is potentially // suboptimal. computation_layout.SetToDefaultLayout(); - TF_RETURN_IF_ERROR(RunOnComputation(computation_layout, computation)); + TF_RETURN_IF_ERROR(RunOnComputation(computation_layout, + *points_to_analysis, computation)); } } @@ -1327,8 +1392,7 @@ StatusOr LayoutAssignment::Run(HloModule* module) { if (VLOG_IS_ON(10)) { hlo_graph_dumper::DumpGraph(*module->entry_computation(), "after layout assignment", - /*show_addresses=*/false, - /*show_layouts=*/true); + module->config().debug_options()); } // All layouts are reset then reassigned by this pass. diff --git a/tensorflow/compiler/xla/service/layout_assignment.h b/tensorflow/compiler/xla/service/layout_assignment.h index 61dc7b120752d57cf09423f38546441de2fc8dd9..118d68dc476c23c6945ec38e061617e7e29f357c 100644 --- a/tensorflow/compiler/xla/service/layout_assignment.h +++ b/tensorflow/compiler/xla/service/layout_assignment.h @@ -46,10 +46,16 @@ namespace xla { // gathered together in LayoutConstraints object. class LayoutConstraint { public: - LayoutConstraint() = default; + LayoutConstraint(bool mandatory) : mandatory_(mandatory) {} virtual ~LayoutConstraint() = default; virtual string ToString() const = 0; + + // True if this constraint cannot be overwritten by a different constraint. + bool mandatory() const { return mandatory_; } + + private: + bool mandatory_; }; std::ostream& operator<<(std::ostream& out, const LayoutConstraint& constraint); @@ -58,7 +64,8 @@ std::ostream& operator<<(std::ostream& out, const LayoutConstraint& constraint); // array produced by a particular instruction. class BufferLayoutConstraint : public LayoutConstraint { public: - BufferLayoutConstraint(const Layout& layout, const LogicalBuffer& buffer); + BufferLayoutConstraint(const Layout& layout, const LogicalBuffer& buffer, + bool mandatory); const LogicalBuffer& buffer() const { return *buffer_; } const Layout& layout() const { return layout_; } @@ -66,7 +73,7 @@ class BufferLayoutConstraint : public LayoutConstraint { string ToString() const override; private: - const Layout layout_; + Layout layout_; const LogicalBuffer* buffer_; }; @@ -78,7 +85,8 @@ class BufferLayoutConstraint : public LayoutConstraint { class OperandLayoutConstraint : public LayoutConstraint { public: OperandLayoutConstraint(const ShapeLayout& shape_layout, - const HloInstruction* instruction, int64 operand_no); + const HloInstruction* instruction, int64 operand_no, + bool mandatory); const ShapeLayout& shape_layout() const { return shape_layout_; } const HloInstruction* instruction() const { return instruction_; } @@ -90,7 +98,7 @@ class OperandLayoutConstraint : public LayoutConstraint { string ToString() const override; private: - const ShapeLayout shape_layout_; + ShapeLayout shape_layout_; const HloInstruction* instruction_; int64 operand_no_; }; @@ -99,7 +107,7 @@ class OperandLayoutConstraint : public LayoutConstraint { class ResultLayoutConstraint : public LayoutConstraint { public: explicit ResultLayoutConstraint(const ShapeLayout& shape_layout) - : shape_layout_(shape_layout) {} + : LayoutConstraint(/*mandatory=*/true), shape_layout_(shape_layout) {} const ShapeLayout& shape_layout() const { return shape_layout_; } string ToString() const override; @@ -124,8 +132,7 @@ class LayoutConstraints { // Return a vector containing the constraints which have been added to the // LayoutConstraints object since the construction of the object or since the // last time ConsumeAddedConstraints() has been called. This is used to - // identify - // newly added constraints when propagating layouts. + // identify newly added constraints when propagating layouts. std::vector ConsumeAddedConstraints() { std::vector ret_vec(std::move(added_constraints_)); added_constraints_.clear(); @@ -137,23 +144,29 @@ class LayoutConstraints { // instruction, or the layout of the result of the computation, respectively, // if it has been constrained. Otherwise return nullptr. const Layout* BufferLayout(const LogicalBuffer& buffer) const; + const BufferLayoutConstraint* GetBufferLayoutConstraint( + const LogicalBuffer& buffer) const; const ShapeLayout* OperandLayout(const HloInstruction* instruction, int64 operand_no) const; + const OperandLayoutConstraint* GetOperandLayoutConstraint( + const HloInstruction* instruction, int64 operand_no) const; const ShapeLayout* ResultLayout() const; // Add a constraint on the layout of a LogicalBuffer, the layout of the // operand of the instruction, or the layout of the result of the computation, // respectively. - Status SetBufferLayout(const Layout& layout, const LogicalBuffer& buffer); + Status SetBufferLayout(const Layout& layout, const LogicalBuffer& buffer, + bool mandatory = true); Status SetOperandLayout(const Shape& shape_with_layout, - const HloInstruction* instruction, int64 operand_no); + const HloInstruction* instruction, int64 operand_no, + bool mandatory = true); Status SetResultLayout(const Shape& shape_with_layout); // Convenience wrapper around SetOperandLayout for setting the layout of a // operand using a Layout object. The operand must be array-shaped. Status SetArrayOperandLayout(const Layout& layout, const HloInstruction* instruction, - int64 operand_no); + int64 operand_no, bool mandatory = true); // Convenience wrapper around SetBufferLayout. Sets the layouts of all buffers // created by the instruction to the layouts in the given shape. The @@ -233,6 +246,39 @@ class LayoutAssignment : public HloPassInterface { const ResultLayoutConstraint& layout_constraint, LayoutConstraints* constraints); + // Called after layouts of an instruction have been finalized to allow + // subclasses to check for platform specific assumptions. + virtual Status Verify(const HloInstruction* instruction) { + return Status::OK(); + } + + // Propagates a buffer layout constraint into the operands that use it. + Status PropagateBufferConstraintToUses( + const BufferLayoutConstraint& layout_constraint, + LayoutConstraints* constraints); + + // Propagates a layout constraint on the use of the result of the given + // instruction to the definitions of the LogicalBuffers which make up the + // result. + Status PropagateUseConstraintToDefs(const ShapeLayout& shape_layout, + const HloInstruction* instruction, + LayoutConstraints* constraints); + + // Chooses a layout of operand `operand_no` of `instruction` that minimizes + // the cost of `instruction`. `output_layout` is the layout of `instruction`. + // Returns null if it can't decide the best layout. + // Precondition: `instruction` and the operand are array-shaped. + std::unique_ptr ChooseOperandLayoutFromOutputLayout( + const Layout& output_layout, const HloInstruction* instruction, + int64 operand_no); + // Given the layout of `user`'s `operand_no`-th operand, chooses a layout of + // `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( + const Layout& operand_layout, const HloInstruction* user, + int64 operand_no); + private: // Adds constraints which must be satisfied for correctness on all // backends. Called once prior to propagating constraints. @@ -253,6 +299,7 @@ class LayoutAssignment : public HloPassInterface { // added, then propagated until all LogicalBuffers in the computation are // constrained. Status RunOnComputation(const ComputationLayout& computation_layout, + const TuplePointsToAnalysis& points_to_analysis, HloComputation* computation); // Assign layouts to the instructions of a computation which satisfy the given @@ -267,30 +314,9 @@ class LayoutAssignment : public HloPassInterface { // required for correctness. Status PropagateConstraints(LayoutConstraints* constraints); - // Propagates a layout constraint on the use of the result of the given - // instruction to the definitions of the LogicalBuffers which make up the - // result. - Status PropagateUseConstraintToDefs(const ShapeLayout& shape_layout, - const HloInstruction* instruction, - LayoutConstraints* constraints); - - // Chooses a layout of operand `operand_no` of `instruction` that minimizes - // the cost of `instruction`. `output_layout` is the layout of `instruction`. - // Returns null if it can't decide the best layout. - // Precondition: `instruction` and the operand are array-shaped. - std::unique_ptr ChooseOperandLayoutFromOutputLayout( - const Layout& output_layout, const HloInstruction* instruction, - int64 operand_no); - // Given the layout of `user`'s `operand_no`-th operand, chooses a layout of - // `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( - const Layout& operand_layout, const HloInstruction* user, - int64 operand_no); - ComputationLayout* entry_computation_layout_; + protected: // Map containing the layouts of all computations assigned so // far. Computations are handled in a topological sort where computations are // handled before their caller instructions so the layouts of caller diff --git a/tensorflow/compiler/xla/service/layout_assignment_test.cc b/tensorflow/compiler/xla/service/layout_assignment_test.cc index 6361907b0e4ad8e21baec88b975f88fc65e42b38..f69c043f32b4e688a543d277164eb91b364b51dc 100644 --- a/tensorflow/compiler/xla/service/layout_assignment_test.cc +++ b/tensorflow/compiler/xla/service/layout_assignment_test.cc @@ -26,10 +26,12 @@ 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/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_test_base.h" #include "tensorflow/compiler/xla/tests/test_utils.h" @@ -38,9 +40,13 @@ limitations under the License. #include "tensorflow/core/lib/core/status.h" #include "tensorflow/core/lib/gtl/array_slice.h" +namespace op = xla::testing::opcode_matchers; + namespace xla { namespace { +using ::testing::ElementsAre; + class LayoutAssignmentTest : public HloTestBase { protected: void AssignLayouts(HloModule* module, @@ -63,8 +69,8 @@ TEST_F(LayoutAssignmentTest, ComputationLayout) { HloInstruction::CreateParameter(1, ashape, "param1")); auto add = builder.AddInstruction( HloInstruction::CreateBinary(ashape, HloOpcode::kAdd, param0, param1)); - HloModule module(TestName()); - HloComputation* computation = module.AddEntryComputation(builder.Build()); + auto module = CreateNewModule(); + HloComputation* computation = module->AddEntryComputation(builder.Build()); Layout layout = LayoutUtil::MakeLayout(minor_to_major); Shape shape(ashape); @@ -75,7 +81,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(module.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())); @@ -93,8 +99,8 @@ TEST_F(LayoutAssignmentTest, ComputationLayoutMixedLayout) { HloInstruction::CreateParameter(1, ashape, "param1")); builder.AddInstruction( HloInstruction::CreateBinary(ashape, HloOpcode::kAdd, param0, param1)); - HloModule module(TestName()); - HloComputation* computation = module.AddEntryComputation(builder.Build()); + auto module = CreateNewModule(); + HloComputation* computation = module->AddEntryComputation(builder.Build()); Layout col_major_layout = LayoutUtil::MakeLayout({1, 0}); Shape col_major_shape(ashape); @@ -111,7 +117,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(module.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( @@ -142,8 +148,8 @@ TEST_F(LayoutAssignmentTest, FusionInstruction) { auto negate2 = builder.AddInstruction( HloInstruction::CreateUnary(ashape, HloOpcode::kNegate, negate1)); - HloModule module(TestName()); - HloComputation* computation = module.AddEntryComputation(builder.Build()); + auto module = CreateNewModule(); + HloComputation* computation = module->AddEntryComputation(builder.Build()); auto fusion = computation->CreateFusionInstruction( {negate2, negate1, add}, HloInstruction::FusionKind::kLoop); @@ -156,7 +162,7 @@ TEST_F(LayoutAssignmentTest, FusionInstruction) { ComputationLayout computation_layout(computation->ComputeProgramShape()); *computation_layout.mutable_result_layout() = shape_layout; - AssignLayouts(&module, &computation_layout); + AssignLayouts(module.get(), &computation_layout); EXPECT_TRUE(LayoutUtil::Equal( layout, fusion->fused_parameter(0)->shape().layout())); @@ -191,13 +197,13 @@ TEST_F(LayoutAssignmentTest, TupleLayout) { auto negate = builder.AddInstruction(HloInstruction::CreateUnary( constant0->shape(), HloOpcode::kNegate, get_element0)); - HloModule module(TestName()); - module.AddEntryComputation(builder.Build()); + auto module = CreateNewModule(); + module->AddEntryComputation(builder.Build()); ComputationLayout computation_layout( - module.entry_computation()->ComputeProgramShape()); + module->entry_computation()->ComputeProgramShape()); - AssignLayouts(&module, &computation_layout); + AssignLayouts(module.get(), &computation_layout); EXPECT_FALSE( LayoutUtil::LayoutsInShapesEqual(constant0->shape(), constant1->shape())); @@ -224,22 +230,22 @@ TEST_F(LayoutAssignmentTest, TupleSelect) { HloInstruction::CreateTuple({constant0, constant1})); auto pred = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(true))); + HloInstruction::CreateConstant(Literal::CreateR0(true))); auto select = builder.AddInstruction(HloInstruction::CreateTernary( tuple0->shape(), HloOpcode::kSelect, pred, tuple0, tuple1)); - HloModule module(TestName()); - module.AddEntryComputation(builder.Build()); + auto module = CreateNewModule(); + module->AddEntryComputation(builder.Build()); ComputationLayout computation_layout( - module.entry_computation()->ComputeProgramShape()); + module->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(module.get(), &computation_layout); EXPECT_TRUE(LayoutUtil::LayoutsInShapesEqual(result_shape, select->shape())); } @@ -258,17 +264,17 @@ TEST_F(LayoutAssignmentTest, ConflictingLayoutTuple) { // tuple and assigning the layouts of the copied arrays as needed. auto builder = HloComputation::Builder(TestName()); auto constant = builder.AddInstruction(HloInstruction::CreateConstant( - LiteralUtil::CreateR2({{1.0, 2.0}, {3.0, 4.0}}))); + Literal::CreateR2({{1.0, 2.0}, {3.0, 4.0}}))); auto inner_tuple = builder.AddInstruction(HloInstruction::CreateTuple({constant})); auto nested_tuple = builder.AddInstruction( HloInstruction::CreateTuple({inner_tuple, inner_tuple})); - HloModule module(TestName()); - module.AddEntryComputation(builder.Build()); + auto module = CreateNewModule(); + module->AddEntryComputation(builder.Build()); ComputationLayout computation_layout( - module.entry_computation()->ComputeProgramShape()); + module->entry_computation()->ComputeProgramShape()); Shape result_shape = nested_tuple->shape(); *ShapeUtil::GetMutableSubshape(&result_shape, /*index=*/{0, 0}) = ShapeUtil::MakeShapeWithLayout(F32, {2, 2}, {1, 0}); @@ -278,7 +284,7 @@ TEST_F(LayoutAssignmentTest, ConflictingLayoutTuple) { result_shape)); LayoutAssignment layout_assignment(&computation_layout); - AssignLayouts(&module, &computation_layout); + AssignLayouts(module.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 @@ -294,9 +300,9 @@ TEST_F(LayoutAssignmentTest, ConflictingLayoutTuple) { EXPECT_TRUE( AlgebraicSimplifier(/*is_layout_sensitive=*/true, [](const Shape&, const Shape&) { return false; }) - .Run(&module) + .Run(module.get()) .ValueOrDie()); - HloInstruction* root = module.entry_computation()->root_instruction(); + HloInstruction* root = module->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}), @@ -304,18 +310,16 @@ TEST_F(LayoutAssignmentTest, ConflictingLayoutTuple) { EXPECT_TRUE(ShapeUtil::Equal(ShapeUtil::GetSubshape(result_shape, {1}), root->operand(1)->shape())); - // Verify some of the structure of the HLO graph. - EXPECT_EQ(constant, root->operand(0)->operand(0)); - EXPECT_EQ(HloOpcode::kCopy, root->operand(1)->operand(0)->opcode()); - EXPECT_EQ(HloOpcode::kConstant, - root->operand(1)->operand(0)->operand(0)->opcode()); + // Verify the structure of the HLO graph. + EXPECT_THAT(root, + op::Tuple(op::Tuple(constant), op::Tuple(op::Copy(constant)))); } TEST_F(LayoutAssignmentTest, ElementwiseAndReshape) { // param -> log -> reshape -> tanh auto builder = HloComputation::Builder(TestName()); Shape ashape = ShapeUtil::MakeShape(F32, {1, 2, 3, 1}); - Shape bshape = ShapeUtil::MakeShape(F32, {2, 1, 3}); + Shape bshape = ShapeUtil::MakeShape(F32, {3, 1, 2}); auto param = builder.AddInstruction( HloInstruction::CreateParameter(0, ashape, "param")); auto log = builder.AddInstruction( @@ -325,28 +329,29 @@ TEST_F(LayoutAssignmentTest, ElementwiseAndReshape) { auto tanh = builder.AddInstruction( HloInstruction::CreateUnary(bshape, HloOpcode::kTanh, reshape)); - HloModule module(TestName()); - HloComputation* computation = module.AddEntryComputation(builder.Build(tanh)); + auto module = CreateNewModule(); + HloComputation* computation = + module->AddEntryComputation(builder.Build(tanh)); Shape ashape_with_layout(ashape); Shape bshape_with_layout(bshape); - *ashape_with_layout.mutable_layout() = LayoutUtil::MakeLayout({0, 1, 2, 3}); - *bshape_with_layout.mutable_layout() = LayoutUtil::MakeLayout({0, 1, 2}); + *ashape_with_layout.mutable_layout() = LayoutUtil::MakeLayout({0, 2, 1, 3}); + *bshape_with_layout.mutable_layout() = LayoutUtil::MakeLayout({2, 1, 0}); ComputationLayout computation_layout(computation->ComputeProgramShape()); *computation_layout.mutable_parameter_layout(0) = ShapeLayout(ashape_with_layout); *computation_layout.mutable_result_layout() = ShapeLayout(bshape_with_layout); - AssignLayouts(&module, &computation_layout); + AssignLayouts(module.get(), &computation_layout); auto log_minor_to_major = AsInt64Slice(log->shape().layout().minor_to_major()); - EXPECT_LT(PositionInContainer(log_minor_to_major, 1), + EXPECT_GT(PositionInContainer(log_minor_to_major, 1), PositionInContainer(log_minor_to_major, 2)); auto reshape_minor_to_major = AsInt64Slice(reshape->shape().layout().minor_to_major()); - EXPECT_LT(PositionInContainer(reshape_minor_to_major, 0), + EXPECT_GT(PositionInContainer(reshape_minor_to_major, 0), PositionInContainer(reshape_minor_to_major, 2)); } @@ -366,8 +371,8 @@ TEST_F(LayoutAssignmentTest, ElementwiseAndTranspose) { HloInstruction::CreateTranspose(bshape, log, {1, 0})); auto tanh = builder.AddInstruction( HloInstruction::CreateUnary(bshape, HloOpcode::kTanh, transpose)); - HloModule module(TestName()); - auto computation = module.AddEntryComputation(builder.Build(tanh)); + auto module = CreateNewModule(); + auto computation = module->AddEntryComputation(builder.Build(tanh)); Shape ashape_with_layout(ashape); Shape bshape_with_layout(bshape); @@ -378,7 +383,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(module.get(), &computation_layout); EXPECT_TRUE( LayoutUtil::Equal(ashape_with_layout.layout(), log->shape().layout())); @@ -402,9 +407,9 @@ TEST_F(LayoutAssignmentTest, BroadcastAndTranspose) { HloInstruction::CreateBroadcast(bshape, param, {1, 2})); auto transpose = builder.AddInstruction( HloInstruction::CreateTranspose(cshape, broadcast, {2, 1, 0})); - HloModule module(TestName()); + auto module = CreateNewModule(); HloComputation* computation = - module.AddEntryComputation(builder.Build(transpose)); + module->AddEntryComputation(builder.Build(transpose)); Shape input_shape_with_layout(ashape); Shape output_shape_with_layout(cshape); @@ -417,10 +422,10 @@ TEST_F(LayoutAssignmentTest, BroadcastAndTranspose) { ShapeLayout(input_shape_with_layout); *computation_layout.mutable_result_layout() = ShapeLayout(output_shape_with_layout); - AssignLayouts(&module, &computation_layout); + AssignLayouts(module.get(), &computation_layout); - EXPECT_TRUE(ContainersEqual(broadcast->shape().layout().minor_to_major(), - tensorflow::gtl::ArraySlice{0, 1, 2})); + EXPECT_THAT(broadcast->shape().layout().minor_to_major(), + ElementsAre(0, 1, 2)); } TEST_F(LayoutAssignmentTest, ReshapeOperandHasMultipleUsers) { @@ -451,9 +456,9 @@ TEST_F(LayoutAssignmentTest, ReshapeOperandHasMultipleUsers) { HloInstruction::CreateBroadcast(f32_234, tanh, {2})); auto tuple = builder.AddInstruction( HloInstruction::CreateTuple({transpose, broadcast2})); - HloModule module(TestName()); + auto module = CreateNewModule(); HloComputation* computation = - module.AddEntryComputation(builder.Build(tuple)); + module->AddEntryComputation(builder.Build(tuple)); ComputationLayout computation_layout(computation->ComputeProgramShape()); Shape param_shape_with_layout(f32_4); @@ -470,17 +475,121 @@ TEST_F(LayoutAssignmentTest, ReshapeOperandHasMultipleUsers) { *computation_layout.mutable_result_layout() = ShapeLayout(ShapeUtil::MakeTupleShape( {transpose_shape_with_layout, broadcast2_shape_with_layout})); - AssignLayouts(&module, &computation_layout); - - EXPECT_TRUE(ContainersEqual(broadcast->shape().layout().minor_to_major(), - tensorflow::gtl::ArraySlice{0, 1})); - EXPECT_TRUE(ContainersEqual(transpose->shape().layout().minor_to_major(), - tensorflow::gtl::ArraySlice{1, 0})); - EXPECT_TRUE(ContainersEqual(tanh->shape().layout().minor_to_major(), - tensorflow::gtl::ArraySlice{0, 1})); + AssignLayouts(module.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)); + EXPECT_THAT(tanh->shape().layout().minor_to_major(), ElementsAre(0, 1)); } -// Add test which fails due to copy tuple. +class OperandsMustBeTheSameLayoutAssignment : public LayoutAssignment { + public: + explicit OperandsMustBeTheSameLayoutAssignment( + ComputationLayout* entry_computation_layout) + : LayoutAssignment(entry_computation_layout) {} + + protected: + Status PropagateBufferConstraint( + const BufferLayoutConstraint& buffer_constraint, + LayoutConstraints* constraints) override { + const LogicalBuffer& buffer = buffer_constraint.buffer(); + const HloInstruction* instruction = buffer.instruction(); + + // Force the operands' layout to the output layout. + 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())) { + continue; + } + TF_RETURN_IF_ERROR(constraints->SetArrayOperandLayout( + buffer_constraint.layout(), instruction, operand_no, + /*mandatory=*/true)); + } + return PropagateBufferConstraintToUses(buffer_constraint, constraints); + } +}; + +TEST_F(LayoutAssignmentTest, MakeOperandsTheSame) { + // param0 -> concatenate -> reshape + // param1 -^ + auto builder = HloComputation::Builder(TestName()); + Shape ashape = ShapeUtil::MakeShape(F32, {50, 1}); + Shape bshape = ShapeUtil::MakeShape(F32, {50, 2}); + Shape cshape = ShapeUtil::MakeShape(F32, {100}); + auto param0 = builder.AddInstruction( + HloInstruction::CreateParameter(0, ashape, "param")); + auto param1 = builder.AddInstruction( + HloInstruction::CreateParameter(1, ashape, "param")); + auto concatenate = builder.AddInstruction( + HloInstruction::CreateConcatenate(bshape, {param0, param1}, 1)); + auto reshape = builder.AddInstruction( + HloInstruction::CreateReshape(cshape, concatenate)); + auto module = CreateNewModule(); + HloComputation* computation = + module->AddEntryComputation(builder.Build(reshape)); + + Shape param0_shape_with_layout(ashape); + Shape param1_shape_with_layout(ashape); + *param0_shape_with_layout.mutable_layout() = LayoutUtil::MakeLayout({0, 1}); + *param1_shape_with_layout.mutable_layout() = LayoutUtil::MakeLayout({1, 0}); + ComputationLayout computation_layout(computation->ComputeProgramShape()); + *computation_layout.mutable_parameter_layout(0) = + ShapeLayout(param0_shape_with_layout); + *computation_layout.mutable_parameter_layout(1) = + ShapeLayout(param1_shape_with_layout); + OperandsMustBeTheSameLayoutAssignment layout_assignment(&computation_layout); + EXPECT_IS_OK(layout_assignment.Run(module.get()).status()); + + EXPECT_EQ(HloOpcode::kCopy, concatenate->operand(0)->opcode()); + EXPECT_THAT(concatenate->operand(0)->shape().layout().minor_to_major(), + ElementsAre(1, 0)); + EXPECT_THAT(concatenate->operand(1)->shape().layout().minor_to_major(), + ElementsAre(1, 0)); + EXPECT_THAT(concatenate->shape().layout().minor_to_major(), + ElementsAre(1, 0)); +} + +// Test layout assignment of a transpose into a bitcast based on its operand. +TEST_F(LayoutAssignmentTest, TransposeToBitcastFromOperand) { + auto builder = HloComputation::Builder(TestName()); + Shape input_shape_with_layout = + ShapeUtil::MakeShapeWithLayout(F32, {3, 5, 6, 7}, {2, 0, 3, 1}); + auto param = builder.AddInstruction( + 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(); + HloComputation* computation = + module->AddEntryComputation(builder.Build(transpose)); + ComputationLayout computation_layout(computation->ComputeProgramShape()); + AssignLayouts(module.get(), &computation_layout); + EXPECT_TRUE(ShapeUtil::TransposeIsBitcast(transpose->operand(0)->shape(), + transpose->shape(), {2, 3, 0, 1})); +} +// Test layout assignment of a transpose into a bitcast based on its user. +TEST_F(LayoutAssignmentTest, TransposeToBitcastToUser) { + auto builder = HloComputation::Builder(TestName()); + Shape input_shape = ShapeUtil::MakeShape(F32, {3, 5, 6, 7}); + auto constant = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(1.0f))); + auto broadcast = builder.AddInstruction( + 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(); + HloComputation* computation = + module->AddEntryComputation(builder.Build(transpose)); + ComputationLayout computation_layout(computation->ComputeProgramShape()); + AssignLayouts(module.get(), &computation_layout); + EXPECT_TRUE(ShapeUtil::TransposeIsBitcast(transpose->operand(0)->shape(), + transpose->shape(), {2, 3, 0, 1})); +} } // namespace } // namespace xla + +int main(int argc, char** argv) { + return xla::ParseDebugOptionsFlagsAndRunTests(argc, argv); +} diff --git a/tensorflow/compiler/xla/service/liveness_util.cc b/tensorflow/compiler/xla/service/liveness_util.cc index caaf56a5516fcf9f21d8754feec04db23381809e..317271dfdd6f0f58519dd106bcf7ac8108642365 100644 --- a/tensorflow/compiler/xla/service/liveness_util.cc +++ b/tensorflow/compiler/xla/service/liveness_util.cc @@ -28,8 +28,9 @@ limitations under the License. namespace xla { -bool DoesNotUseOperandBuffer(HloInstruction* operand, const ShapeIndex& index, - HloInstruction* user, +bool DoesNotUseOperandBuffer(const HloInstruction* operand, + const ShapeIndex& index, + const HloInstruction* user, const TuplePointsToAnalysis& points_to_analysis) { CHECK(user->IsUserOf(operand)) << "user: " << user->ToString() << " operand: " << operand->ToString(); @@ -79,7 +80,7 @@ std::vector> GetAllUsesOfInstructionAtIndex( HloInstruction* instruction, const ShapeIndex& index, const TuplePointsToAnalysis& points_to_analysis) { std::vector> uses; - const std::vector& points_to = + const PointsToSet::BufferList& points_to = points_to_analysis.GetPointsToSet(instruction).element(index); for (const LogicalBuffer* buffer : points_to) { for (const BufferAlias& alias : @@ -98,53 +99,114 @@ std::vector> GetAllUsesOfInstructionAtIndex( return uses; } +// Returns true if there is exactly one use of 'operand' at 'operand_index' +// in 'fusion.fused_instructions', where the singleton use is the fused +// root at operand index 'use_operand_index'. Returns false otherwise. +// +// REQUIRES: 'fusion' opcode is a kFusion instruction. +bool HasUniqueFusedUseOfOperandAt( + HloInstruction* operand, const ShapeIndex& operand_index, + HloInstruction* fusion, const int64 use_operand_index, + const TuplePointsToAnalysis& points_to_analysis) { + CHECK_EQ(HloOpcode::kFusion, fusion->opcode()); + // Check that 'operand' is unique in the operand list of 'fusion'. + if (fusion->OperandIndices(operand).size() > 1) { + return false; + } + // 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) { + return fusion->operand(fused_param->parameter_number()) == operand; + }); + if (fused_param_it == fused_params.end()) { + return false; + } + auto* fused_param = *fused_param_it; + // Get all uses of 'operand' at 'index' from 'fusion.fused_instructions'. + auto fused_param_uses = GetAllUsesOfInstructionAtIndex( + fused_param, operand_index, points_to_analysis); + // Return true iff there is exactly one use of 'operand' at 'index', and + // this singleton use is the fused root (at index in 'use_operand_indices'). + return fused_param_uses.size() == 1 && + fused_param_uses[0].first == fusion->fused_expression_root() && + fused_param_uses[0].second == use_operand_index; +} + } // namespace // User and operand can share buffers iff both instructions emit the same shape -// and layout, and 'user' meets one of the following two qualifications: -// *) Is element-wise. -// *) Is a loop fusion instruction where the only use of 'operand' at 'index' -// in the set 'user.fused_instructions' is a DynamicUpdateSlice fused root -// at operand 0. -// *) Use of 'operand' is DynamicUpdateSlice at operand index 0. +// and layout, and 'user' meets one of the following qualifications: +// +// (1) Is element-wise. Or... +// (2) Is a loop fusion instruction where the only use of 'operand' at 'index' +// in the set 'user.fused_instructions' is a DynamicUpdateSlice fused root +// at operand 0. Or... +// (3) Is a kDot -> kAdd (or fused kTransposeDot -> kAdd) output fusion +// instruction where the only use of 'operand' at 'index' in the set +// 'user.fused_instructions' is a kAdd fused root at operand 0 or 1. Or... +// (4) The 'user' of 'operand' is DynamicUpdateSlice or While at operand index +// 0. +// +// (2) and (3) can only be determined if points-to analysis is available. bool CanShareOperandBufferWithUser( HloInstruction* operand, const ShapeIndex& operand_index, HloInstruction* user, const ShapeIndex& user_index, - const TuplePointsToAnalysis& points_to_analysis) { + const TuplePointsToAnalysis* points_to_analysis) { CHECK(user->IsUserOf(operand)) << "user: " << user->ToString() << " operand: " << operand->ToString(); - Shape operand_subshape = + const Shape& operand_subshape = ShapeUtil::GetSubshape(operand->shape(), operand_index); - Shape user_subshape = ShapeUtil::GetSubshape(user->shape(), user_index); + const Shape& user_subshape = + ShapeUtil::GetSubshape(user->shape(), user_index); // Check that operand and user emit the same shape and layout. if (!ShapeUtil::Equal(operand_subshape, user_subshape)) { return false; } - // Check if 'user' is a loop fusion instruction with a kDynamicUpdateSlice - // fused root instruction. - if (user->opcode() == HloOpcode::kFusion && - user->fusion_kind() == HloInstruction::FusionKind::kLoop && - user->fused_expression_root()->opcode() == - HloOpcode::kDynamicUpdateSlice) { - for (auto& fused_param : user->fused_parameters()) { - // Find fusion parameter associated with 'operand'. - if (user->operand(fused_param->parameter_number()) != operand) { - continue; - } - // Get all uses of 'operand' at 'index' from 'user.fused_instructions'. - auto fused_param_uses = GetAllUsesOfInstructionAtIndex( - fused_param, operand_index, points_to_analysis); - // Return true iff there is exactly one use of 'operand' at 'index', and - // this singleton use is the fused root at operand index 0. - if (fused_param_uses.size() == 1 && - fused_param_uses[0].first == user->fused_expression_root() && - fused_param_uses[0].second == 0) { - return true; + if (points_to_analysis != nullptr && user->opcode() == HloOpcode::kFusion) { + if (user->fusion_kind() == HloInstruction::FusionKind::kLoop && + user->fused_expression_root()->opcode() == + HloOpcode::kDynamicUpdateSlice) { + // Loop fusion with kDynamicUpdateSlice fused root. + // + // Returns true iff there is exactly one use of 'operand' at shape index + // 'operand_index', and this singleton use is the fused root at operand + // index 0. + return HasUniqueFusedUseOfOperandAt(operand, operand_index, user, 0, + *points_to_analysis); + } else if (user->fusion_kind() == HloInstruction::FusionKind::kOutput && + user->fused_expression_root()->opcode() == HloOpcode::kAdd) { + // Output fusion with kAdd fused root. + + // Check if one operand of kAdd fused root is either kDot, or nested + // kFusion of kind kTransposeDot. + 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::kDot || + (operand->opcode() == HloOpcode::kFusion && + operand->fusion_kind() == + HloInstruction::FusionKind::kTransposeDot); + }); + if (add_operand_it == add->operands().end()) { + return false; } - break; + auto* matched_add_operand = *add_operand_it; + // Calculate operand index of 'add' operand which was not matched above. + const int64 other_add_operand_index = + matched_add_operand == add->operand(0) ? 1 : 0; + // Returns true iff there is exactly one use of 'operand' at shape index + // 'operand_index', and this singleton use is the fused root (at operand + // index 'other_add_operand_index'). + return HasUniqueFusedUseOfOperandAt(operand, operand_index, user, + other_add_operand_index, + *points_to_analysis); } - return false; - } else if (user->opcode() == HloOpcode::kDynamicUpdateSlice) { + } + if (user->opcode() == HloOpcode::kDynamicUpdateSlice || + user->opcode() == HloOpcode::kWhile) { // We eliminated other users in BufferLiveness::live_range_strictly_before, // so here we just need to check that the use is at operand index 0. std::vector operand_indices = user->OperandIndices(operand); diff --git a/tensorflow/compiler/xla/service/liveness_util.h b/tensorflow/compiler/xla/service/liveness_util.h index 410a7b1b519e117f21c01938cb8e4a5b1c358ad2..c7799e5ab5d0c0d0477c09fa7e6a36c67312a72b 100644 --- a/tensorflow/compiler/xla/service/liveness_util.h +++ b/tensorflow/compiler/xla/service/liveness_util.h @@ -18,9 +18,6 @@ limitations under the License. #ifndef TENSORFLOW_COMPILER_XLA_SERVICE_LIVENESS_UTIL_H_ #define TENSORFLOW_COMPILER_XLA_SERVICE_LIVENESS_UTIL_H_ -#include -#include - #include "tensorflow/compiler/xla/service/hlo_instruction.h" #include "tensorflow/compiler/xla/service/tuple_points_to_analysis.h" #include "tensorflow/compiler/xla/shape_util.h" @@ -32,19 +29,21 @@ namespace xla { // 'operand'. Returns false otherwise. // // REQUIRES: 'operand' is an operand of 'user'. -bool DoesNotUseOperandBuffer(HloInstruction* operand, const ShapeIndex& index, - HloInstruction* user, +bool DoesNotUseOperandBuffer(const HloInstruction* operand, + const ShapeIndex& index, + const HloInstruction* user, const TuplePointsToAnalysis& points_to_analysis); // Returns true if 'user' (at 'user_index') can share a buffer with its operand -// 'operand' (at 'operand_index'). -// Returns false otherwise. +// 'operand' (at 'operand_index'). Returns false otherwise. Optionally takes a +// points-to analysis argument. Without the analysis, the result is more +// conservative (returns false more often). // // REQUIRES: 'operand' is an operand of 'user'. bool CanShareOperandBufferWithUser( HloInstruction* operand, const ShapeIndex& operand_index, HloInstruction* user, const ShapeIndex& user_index, - const TuplePointsToAnalysis& points_to_analysis); + const TuplePointsToAnalysis* points_to_analysis = nullptr); } // namespace xla diff --git a/tensorflow/compiler/xla/service/liveness_util_test.cc b/tensorflow/compiler/xla/service/liveness_util_test.cc index 2ff71d6f3c8eff58b83783fc867d5874c6c700a3..6a4fde87614750d21cf9572e7f447bba924379c4 100644 --- a/tensorflow/compiler/xla/service/liveness_util_test.cc +++ b/tensorflow/compiler/xla/service/liveness_util_test.cc @@ -27,16 +27,14 @@ namespace { class PointsToAnalysisTestBase : public HloTestBase { protected: void BuildModule(std::unique_ptr computation) { - module_ = MakeUnique(TestName()); + module_ = CreateNewModule(); computation_ = module_->AddEntryComputation(std::move(computation)); } void RunAnalysis() { CHECK_NOTNULL(module_.get()); points_to_analysis_ = - TuplePointsToAnalysis::Run(module_.get(), - /*include_loop_fusion_instructions=*/true) - .ConsumeValueOrDie(); + TuplePointsToAnalysis::Run(module_.get()).ConsumeValueOrDie(); } void BuildModuleAndRunAnalysis(std::unique_ptr computation) { @@ -87,9 +85,9 @@ TEST_F(DoesNotUseOperandBufferTest, FusedDynamicUpdateSlice) { // Create a DynamicUpdateSlice instruction of tuple element 1. auto starts = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR1({2}))); + HloInstruction::CreateConstant(Literal::CreateR1({2}))); auto update = builder.AddInstruction(HloInstruction::CreateConstant( - LiteralUtil::CreateR1({2.f, 2.f, 2.f}))); + Literal::CreateR1({2.f, 2.f, 2.f}))); auto dynamic_update_slice = builder.AddInstruction(HloInstruction::CreateDynamicUpdateSlice( data_shape, gte1, update, starts)); @@ -124,10 +122,10 @@ TEST_F(CanShareOperandBufferWithUserTest, ElementWiseSameShape) { BuildModuleAndRunAnalysis(builder.Build()); - EXPECT_TRUE( - CanShareOperandBufferWithUser(param, {}, exp, {}, *points_to_analysis_)); - EXPECT_TRUE( - CanShareOperandBufferWithUser(exp, {}, log, {}, *points_to_analysis_)); + EXPECT_TRUE(CanShareOperandBufferWithUser(param, {}, exp, {}, + points_to_analysis_.get())); + EXPECT_TRUE(CanShareOperandBufferWithUser(exp, {}, log, {}, + points_to_analysis_.get())); } TEST_F(CanShareOperandBufferWithUserTest, ElementWiseDifferentShape) { @@ -145,9 +143,28 @@ TEST_F(CanShareOperandBufferWithUserTest, ElementWiseDifferentShape) { BuildModuleAndRunAnalysis(builder.Build()); EXPECT_FALSE(CanShareOperandBufferWithUser(param0, {}, result, {}, - *points_to_analysis_)); + points_to_analysis_.get())); EXPECT_FALSE(CanShareOperandBufferWithUser(param1, {}, result, {}, - *points_to_analysis_)); + points_to_analysis_.get())); +} + +TEST_F(CanShareOperandBufferWithUserTest, CopyShares) { + auto builder = HloComputation::Builder(TestName()); + + Shape shape = ShapeUtil::MakeShape(F32, {8}); + auto param = builder.AddInstruction( + HloInstruction::CreateParameter(0, shape, "param")); + auto exp = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kExp, param)); + auto copy = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kCopy, exp)); + + BuildModuleAndRunAnalysis(builder.Build()); + + EXPECT_TRUE(CanShareOperandBufferWithUser(param, {}, exp, {}, + points_to_analysis_.get())); + EXPECT_TRUE(CanShareOperandBufferWithUser(exp, {}, copy, {}, + points_to_analysis_.get())); } TEST_F(CanShareOperandBufferWithUserTest, FusedDynamicUpdateSlice) { @@ -163,9 +180,9 @@ TEST_F(CanShareOperandBufferWithUserTest, FusedDynamicUpdateSlice) { // Create a DynamicUpdateSlice instruction of tuple element 1. auto starts = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR1({2}))); + HloInstruction::CreateConstant(Literal::CreateR1({2}))); auto update = builder.AddInstruction(HloInstruction::CreateConstant( - LiteralUtil::CreateR1({2.f, 2.f, 2.f}))); + Literal::CreateR1({2.f, 2.f, 2.f}))); auto dynamic_update_slice = builder.AddInstruction(HloInstruction::CreateDynamicUpdateSlice( data_shape, gte1, update, starts)); @@ -180,10 +197,176 @@ TEST_F(CanShareOperandBufferWithUserTest, FusedDynamicUpdateSlice) { // The fusion instruction can share with tuple element 1. EXPECT_FALSE(CanShareOperandBufferWithUser(tuple, {0}, fusion, {}, - *points_to_analysis_)); + points_to_analysis_.get())); EXPECT_TRUE(CanShareOperandBufferWithUser(tuple, {1}, fusion, {}, - *points_to_analysis_)); + points_to_analysis_.get())); +} + +TEST_F(CanShareOperandBufferWithUserTest, DynamicUpdateSliceCanShare) { + auto builder = HloComputation::Builder(TestName()); + + Shape data_shape = ShapeUtil::MakeShape(F32, {8}); + Shape update_shape = ShapeUtil::MakeShape(F32, {4}); + Shape starts_shape = ShapeUtil::MakeShape(S32, {1}); + auto data = builder.AddInstruction( + 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 dus = builder.AddInstruction(HloInstruction::CreateDynamicUpdateSlice( + data_shape, data, update, starts)); + + BuildModuleAndRunAnalysis(builder.Build()); + + // The DynamicUpdateSlice instruction can share with the data operand, but not + // with update or starts. + EXPECT_TRUE(CanShareOperandBufferWithUser(data, {}, dus, {}, + points_to_analysis_.get())); + EXPECT_FALSE(CanShareOperandBufferWithUser(update, {}, dus, {}, + points_to_analysis_.get())); + EXPECT_FALSE(CanShareOperandBufferWithUser(starts, {}, dus, {}, + points_to_analysis_.get())); +} + +TEST_F(CanShareOperandBufferWithUserTest, FusedDotAdd) { + auto builder = HloComputation::Builder(TestName()); + Shape data_shape = ShapeUtil::MakeShape(F32, {2, 2}); + + auto a = builder.AddInstruction(HloInstruction::CreateConstant( + Literal::CreateR2({{1.0, 0.0}, {0.0, 1.0}}))); + auto b = builder.AddInstruction(HloInstruction::CreateConstant( + Literal::CreateR2({{2.0, 2.0}, {2.0, 2.0}}))); + + auto dot = builder.AddInstruction( + HloInstruction::CreateBinary(data_shape, HloOpcode::kDot, a, b)); + + auto one = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(1.0))); + auto add_operand = builder.AddInstruction( + HloInstruction::CreateBroadcast(data_shape, one, {1})); + + auto add = builder.AddInstruction(HloInstruction::CreateBinary( + data_shape, HloOpcode::kAdd, dot, add_operand)); + + BuildModule(builder.Build()); + auto fusion = computation_->CreateFusionInstruction( + {add, dot}, HloInstruction::FusionKind::kOutput); + RunAnalysis(); + + // Output fused dot add should be able to share buffer with 'add_operand'. + EXPECT_TRUE(CanShareOperandBufferWithUser(add_operand, {}, fusion, {}, + points_to_analysis_.get())); +} + +TEST_F(CanShareOperandBufferWithUserTest, FusedTransposeDotAdd) { + auto builder = HloComputation::Builder(TestName()); + Shape data_shape = ShapeUtil::MakeShape(F32, {2, 2}); + + auto a = builder.AddInstruction(HloInstruction::CreateConstant( + Literal::CreateR2({{1.0, 0.0}, {0.0, 1.0}}))); + auto b = builder.AddInstruction(HloInstruction::CreateConstant( + Literal::CreateR2({{2.0, 2.0}, {2.0, 2.0}}))); + auto b_t = builder.AddInstruction( + HloInstruction::CreateTranspose(data_shape, b, {1, 0})); + + auto dot = builder.AddInstruction( + HloInstruction::CreateBinary(data_shape, HloOpcode::kDot, a, b_t)); + + auto one = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(1.0))); + auto add_operand = builder.AddInstruction( + HloInstruction::CreateBroadcast(data_shape, one, {1})); + + auto add = builder.AddInstruction(HloInstruction::CreateBinary( + data_shape, HloOpcode::kAdd, dot, add_operand)); + + BuildModule(builder.Build()); + + auto nested_fusion = computation_->CreateFusionInstruction( + {dot, b_t}, HloInstruction::FusionKind::kTransposeDot); + + auto fusion = computation_->CreateFusionInstruction( + {add, nested_fusion}, HloInstruction::FusionKind::kOutput); + RunAnalysis(); + + // Output fused transpose-dot-add should be share buffer with 'add_operand'. + EXPECT_TRUE(CanShareOperandBufferWithUser(add_operand, {}, fusion, {}, + points_to_analysis_.get())); +} + +TEST_F(CanShareOperandBufferWithUserTest, OutputFusionCantAliasOperandBuffer) { + auto builder = HloComputation::Builder(TestName()); + Shape data_shape = ShapeUtil::MakeShape(F32, {2, 2}); + + auto one = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(1.0))); + auto operand = builder.AddInstruction( + HloInstruction::CreateBroadcast(data_shape, one, {1})); + + auto reverse = builder.AddInstruction( + HloInstruction::CreateReverse(data_shape, operand, {0, 1})); + + auto two = builder.AddInstruction(HloInstruction::CreateConstant( + Literal::CreateR2({{2.0, 2.0}, {2.0, 2.0}}))); + + auto add = builder.AddInstruction( + HloInstruction::CreateBinary(data_shape, HloOpcode::kAdd, reverse, two)); + + BuildModule(builder.Build()); + auto fusion = computation_->CreateFusionInstruction( + {add, two, reverse}, HloInstruction::FusionKind::kOutput); + RunAnalysis(); + + // Output fused operand->reverse->add cannot alias operand buffer 'operand'. + EXPECT_FALSE(CanShareOperandBufferWithUser(operand, {}, fusion, {}, + points_to_analysis_.get())); +} + +TEST_F(CanShareOperandBufferWithUserTest, WhileCanShare) { + Shape data_shape = ShapeUtil::MakeShape(F32, {8}); + + auto make_cond = [this, &data_shape]() { + auto builder = HloComputation::Builder(TestName() + ".Cond"); + auto data = builder.AddInstruction( + HloInstruction::CreateParameter(0, data_shape, "data")); + builder.AddInstruction(HloInstruction::CreateBinary( + ShapeUtil::MakeShape(PRED, {}), HloOpcode::kEq, data, data)); + return builder.Build(); + }; + + auto make_body = [this, &data_shape]() { + auto builder = HloComputation::Builder(TestName() + ".Body"); + auto data = builder.AddInstruction( + HloInstruction::CreateParameter(0, data_shape, "data")); + builder.AddInstruction( + HloInstruction::CreateBinary(data_shape, HloOpcode::kAdd, data, data)); + return builder.Build(); + }; + + module_ = CreateNewModule(); + HloComputation* cond_computation = + module_->AddEmbeddedComputation(make_cond()); + HloComputation* body_computation = + module_->AddEmbeddedComputation(make_body()); + + auto builder = HloComputation::Builder(TestName()); + auto data = builder.AddInstruction( + HloInstruction::CreateParameter(0, data_shape, "data")); + auto whil = builder.AddInstruction(HloInstruction::CreateWhile( + data_shape, cond_computation, body_computation, data)); + computation_ = module_->AddEntryComputation(builder.Build()); + + RunAnalysis(); + + // The While instruction can share with the data operand. + EXPECT_TRUE(CanShareOperandBufferWithUser(data, {}, whil, {}, + points_to_analysis_.get())); } } // namespace } // namespace xla + +int main(int argc, char** argv) { + return xla::ParseDebugOptionsFlagsAndRunTests(argc, argv); +} diff --git a/tensorflow/compiler/xla/service/llvm_compiler.h b/tensorflow/compiler/xla/service/llvm_compiler.h new file mode 100644 index 0000000000000000000000000000000000000000..b2e72871c10192c84349b117797c7bd7e6ee251a --- /dev/null +++ b/tensorflow/compiler/xla/service/llvm_compiler.h @@ -0,0 +1,67 @@ +/* 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_LLVM_COMPILER_H_ +#define TENSORFLOW_COMPILER_XLA_SERVICE_LLVM_COMPILER_H_ + +#include "llvm/IR/Module.h" +#include "tensorflow/compiler/xla/service/compiler.h" + +namespace xla { + +// Interface for an LLVM-based compiler. This provides the ability to register +// hooks to inspect the LLVM IR during compilation, both before and after +// optimizations are applied. +// +// Hooks get called once per HLO module being compiled. The following should not +// be relied on: +// * The order in which hooks get called. +// * Whether or not a hook gets called if a compilation exits with a non-OK +// status. +class LLVMCompiler : public Compiler { + public: + ~LLVMCompiler() override {} + + // A callback of this type can be run before and/or after IR-level + // optimization to e.g. dump out the generated IR to disk or gather some + // statistics. + using ModuleHook = std::function; + + void SetPreOptimizationHook(ModuleHook hook) { + CHECK(!user_pre_optimization_hook_) + << "Pre-optimization hook is already set"; + CHECK(hook) << "hook cannot be null"; + user_pre_optimization_hook_ = hook; + } + + void RemovePreOptimizationHook() { user_pre_optimization_hook_ = nullptr; } + + void SetPostOptimizationHook(ModuleHook hook) { + CHECK(!user_post_optimization_hook_) + << "Post-optimization hook is already set"; + CHECK(hook) << "hook cannot be null"; + user_post_optimization_hook_ = hook; + } + + void RemovePostOptimizationHook() { user_post_optimization_hook_ = nullptr; } + + protected: + ModuleHook user_pre_optimization_hook_; + ModuleHook user_post_optimization_hook_; +}; + +} // namespace xla + +#endif // TENSORFLOW_COMPILER_XLA_SERVICE_LLVM_COMPILER_H_ diff --git a/tensorflow/compiler/xla/service/llvm_ir/BUILD b/tensorflow/compiler/xla/service/llvm_ir/BUILD index 1edfec4dae5d50a4e5bd2f8079d322f856819d00..61945bd128e68b59bd0a1156882c5b29d6be2a27 100644 --- a/tensorflow/compiler/xla/service/llvm_ir/BUILD +++ b/tensorflow/compiler/xla/service/llvm_ir/BUILD @@ -29,7 +29,6 @@ cc_library( ":ir_array", ":llvm_util", "//tensorflow/compiler/xla:types", - "//tensorflow/compiler/xla/legacy_flags:alias_analysis_flags", "//tensorflow/compiler/xla/service:buffer_assignment", "//tensorflow/compiler/xla/service:hlo", "//tensorflow/compiler/xla/service:logical_buffer", @@ -47,9 +46,6 @@ cc_library( "//tensorflow/compiler/xla:shape_util", "//tensorflow/compiler/xla:types", "//tensorflow/compiler/xla:xla_data_proto", - "//tensorflow/compiler/xla/legacy_flags:llvm_util_flags", - "//tensorflow/compiler/xla/service:buffer_assignment", - "//tensorflow/compiler/xla/service:hlo_module_config", "//tensorflow/core:lib", "@llvm//:core", "@llvm//:support", @@ -80,6 +76,7 @@ cc_library( deps = [ ":ir_array", ":llvm_util", + "//tensorflow/compiler/xla:shape_util", "//tensorflow/compiler/xla:types", "//tensorflow/compiler/xla:xla_data_proto", "//tensorflow/core:lib", @@ -94,6 +91,7 @@ cc_library( deps = [ ":ir_array", ":llvm_loop", + ":ops", "//tensorflow/compiler/xla:shape_util", "//tensorflow/compiler/xla:status_macros", "//tensorflow/compiler/xla:statusor", diff --git a/tensorflow/compiler/xla/service/llvm_ir/README.md b/tensorflow/compiler/xla/service/llvm_ir/README.md index 9fe7152477fefd0de528aa41f6062902cea34835..9e4cdd45dca7199c3f053da7b1e1f6f43c0944ac 100644 --- a/tensorflow/compiler/xla/service/llvm_ir/README.md +++ b/tensorflow/compiler/xla/service/llvm_ir/README.md @@ -1,2 +1,2 @@ -Common utilites and abstractions for handling and emitting LLVM IR for XLA +Common utilities and abstractions for handling and emitting LLVM IR for XLA backends. diff --git a/tensorflow/compiler/xla/service/llvm_ir/alias_analysis.cc b/tensorflow/compiler/xla/service/llvm_ir/alias_analysis.cc index fc337a89246e7c7c9fe8ff5a5d92a0b54904beb1..5e28e37600c18a351e8647d48119f073277f56e1 100644 --- a/tensorflow/compiler/xla/service/llvm_ir/alias_analysis.cc +++ b/tensorflow/compiler/xla/service/llvm_ir/alias_analysis.cc @@ -17,8 +17,7 @@ limitations under the License. #include -#include "external/llvm/include/llvm/IR/MDBuilder.h" -#include "tensorflow/compiler/xla/legacy_flags/alias_analysis_flags.h" +#include "llvm/IR/MDBuilder.h" #include "tensorflow/compiler/xla/service/llvm_ir/llvm_util.h" #include "tensorflow/compiler/xla/service/logical_buffer.h" #include "tensorflow/compiler/xla/types.h" @@ -29,7 +28,8 @@ namespace llvm_ir { // Sentry allocation used to represent parameters of the entry computation in // alias_scope_metadata_ and noalias_metadata_. static const BufferAllocation* kParameterAllocation = new BufferAllocation( - /*index=*/-1, /*size=*/0, /*is_thread_local=*/false, /*is_reusable=*/false); + /*index=*/-1, /*size=*/0, /*is_thread_local=*/false, /*is_reusable=*/false, + LogicalBuffer::Color(0)); void AliasAnalysis::AddAliasingInformationToIrArray(const HloInstruction& hlo, llvm_ir::IrArray* array) { @@ -50,28 +50,41 @@ void AliasAnalysis::AddAliasingInformationToIrArray(const HloInstruction& hlo, buffer_slice = *slices.begin(); } - llvm::MDNode*& alias_scope_md = alias_scope_metadata_[buffer_slice]; - if (alias_scope_md == nullptr) { - alias_scope_md = - GetAliasScopeMetadataForBuffer(buffer_slice, GetAliasDomain()); + if (module_.config().debug_options().xla_llvm_enable_alias_scope_metadata()) { + llvm::MDNode*& alias_scope_md = alias_scope_metadata_[buffer_slice]; + if (alias_scope_md == nullptr) { + alias_scope_md = + GetAliasScopeMetadataForBuffer(buffer_slice, GetAliasDomain()); + } + if (alias_scope_md != nullptr) { + array->AddAliasScopeMetadata(alias_scope_md); + } } - array->AddAliasScopeMetadata(alias_scope_md); - llvm::MDNode*& noalias_md = noalias_metadata_[buffer_slice]; - if (noalias_md == nullptr) { - noalias_md = GetNoaliasMetadataForBuffer(buffer_slice, GetAliasDomain(), - assignment_, hlo); + if (module_.config().debug_options().xla_llvm_enable_noalias_metadata()) { + llvm::MDNode*& noalias_md = noalias_metadata_[buffer_slice]; + if (noalias_md == nullptr) { + noalias_md = GetNoaliasMetadataForBuffer(buffer_slice, GetAliasDomain(), + assignment_, hlo); + } + if (noalias_md != nullptr) { + array->AddNoaliasMetadata(noalias_md); + } } - array->AddNoaliasMetadata(noalias_md); - // Parameters of the entry computation are never stored to, loading from a - // parameter pointer should always return the same result within a loop. - 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()) { - array->AddInvariantLoad(llvm::MDNode::get(*context_, /*MDs=*/{})); + if (module_.config() + .debug_options() + .xla_llvm_enable_invariant_load_metadata()) { + // Parameters of the entry computation are never stored to, loading from a + // parameter pointer should always return the same result within a loop. + 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()) { + array->AddInvariantLoad(llvm::MDNode::get(*context_, /*MDs=*/{})); + } } } } @@ -86,12 +99,6 @@ llvm::MDNode* AliasAnalysis::GetAliasDomain() { llvm::MDNode* AliasAnalysis::GetAliasScopeMetadataForBuffer( const BufferAllocation::Slice& buffer_slice, llvm::MDNode* domain) { - legacy_flags::AliasAnalysisFlags* flags = - legacy_flags::GetAliasAnalysisFlags(); - if (!flags->xla_emit_alias_scope) { - return nullptr; - } - // While we could synthesize an alias.scope, doing so is not more profitable // than LLVM's default behavior. if (buffer_slice.allocation() == kParameterAllocation) { @@ -108,12 +115,6 @@ llvm::MDNode* AliasAnalysis::GetAliasScopeMetadataForBuffer( llvm::MDNode* AliasAnalysis::GetNoaliasMetadataForBuffer( const BufferAllocation::Slice& buffer_slice, llvm::MDNode* domain, const BufferAssignment& assignment, const HloInstruction& hlo) { - legacy_flags::AliasAnalysisFlags* flags = - legacy_flags::GetAliasAnalysisFlags(); - if (!flags->xla_emit_alias_scope) { - return nullptr; - } - // We want to construct a list of buffers which: // // 1. Do not alias the given buffer. diff --git a/tensorflow/compiler/xla/service/llvm_ir/alias_analysis.h b/tensorflow/compiler/xla/service/llvm_ir/alias_analysis.h index 9eb1cbaa341bce3b53d7b71c8d3de93768ec4a8d..5244ac61e56307857aca659854647bd6c3e991d7 100644 --- a/tensorflow/compiler/xla/service/llvm_ir/alias_analysis.h +++ b/tensorflow/compiler/xla/service/llvm_ir/alias_analysis.h @@ -16,7 +16,7 @@ limitations under the License. #ifndef TENSORFLOW_COMPILER_XLA_SERVICE_LLVM_IR_ALIAS_ANALYSIS_H_ #define TENSORFLOW_COMPILER_XLA_SERVICE_LLVM_IR_ALIAS_ANALYSIS_H_ -#include "external/llvm/include/llvm/IR/Module.h" +#include "llvm/IR/Module.h" #include "tensorflow/compiler/xla/service/buffer_assignment.h" #include "tensorflow/compiler/xla/service/hlo_instruction.h" #include "tensorflow/compiler/xla/service/llvm_ir/ir_array.h" 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 b259d348708c227a3e580fd352422e457284129d..c914fc2df7af081d236dd11ff7b383f2a072fbd9 100644 --- a/tensorflow/compiler/xla/service/llvm_ir/fused_ir_emitter.cc +++ b/tensorflow/compiler/xla/service/llvm_ir/fused_ir_emitter.cc @@ -17,8 +17,8 @@ limitations under the License. #include -#include "external/llvm/include/llvm/IR/BasicBlock.h" -#include "external/llvm/include/llvm/IR/Value.h" +#include "llvm/IR/BasicBlock.h" +#include "llvm/IR/Value.h" #include "tensorflow/compiler/xla/service/elemental_ir_emitter.h" #include "tensorflow/compiler/xla/service/llvm_ir/ir_array.h" #include "tensorflow/compiler/xla/service/llvm_ir/llvm_util.h" @@ -128,6 +128,27 @@ Status FusedIrEmitter::HandleParameter(HloInstruction* parameter) { return Status::OK(); } +Status FusedIrEmitter::HandleTuple( + HloInstruction* tuple, + tensorflow::gtl::ArraySlice operands) { + std::vector operand_elemental_ir_types; + for (HloInstruction* operand : operands) { + operand_elemental_ir_types.push_back(llvm_ir::PrimitiveTypeToIrType( + operand->shape().element_type(), ir_builder_)); + } + generators_[tuple] = + [=](const IrArray::Index& index) -> StatusOr { + llvm::Value* ret = llvm::UndefValue::get(llvm::StructType::get( + ir_builder_->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)); + ret = ir_builder_->CreateInsertValue(ret, val_i, i); + } + return ret; + }; + return Status::OK(); +} + Status FusedIrEmitter::FinishVisit(HloInstruction* root) { fused_root_ = root; return tensorflow::Status::OK(); 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 79007b7099a32973cada7a9986ff95c5e4aabec6..a24e104067f19e45ab2566beedbb8217913bad12 100644 --- a/tensorflow/compiler/xla/service/llvm_ir/fused_ir_emitter.h +++ b/tensorflow/compiler/xla/service/llvm_ir/fused_ir_emitter.h @@ -19,8 +19,8 @@ limitations under the License. #include #include -#include "external/llvm/include/llvm/IR/IRBuilder.h" -#include "external/llvm/include/llvm/IR/Value.h" +#include "llvm/IR/IRBuilder.h" +#include "llvm/IR/Value.h" #include "tensorflow/compiler/xla/service/dfs_hlo_visitor_with_default.h" #include "tensorflow/compiler/xla/service/elemental_ir_emitter.h" #include "tensorflow/compiler/xla/service/hlo_instruction.h" @@ -54,6 +54,11 @@ class FusedIrEmitter : public DfsHloVisitorWithDefault { Status HandleParameter(HloInstruction* parameter) override; + // Emits the ir value for each element in the tuple. + Status HandleTuple( + HloInstruction* tuple, + tensorflow::gtl::ArraySlice operands) override; + Status FinishVisit(HloInstruction* root) override; // Returns the generator function for the root of the fused computation. diff --git a/tensorflow/compiler/xla/service/llvm_ir/ir_array.cc b/tensorflow/compiler/xla/service/llvm_ir/ir_array.cc index 38728d2e1f34d03d7eec9fb42e35f99397e26d8a..e36c791c1a52f4e0699cc15ef913fbd2bdcca557 100644 --- a/tensorflow/compiler/xla/service/llvm_ir/ir_array.cc +++ b/tensorflow/compiler/xla/service/llvm_ir/ir_array.cc @@ -15,8 +15,8 @@ limitations under the License. #include "tensorflow/compiler/xla/service/llvm_ir/ir_array.h" -#include "external/llvm/include/llvm/IR/Constants.h" -#include "external/llvm/include/llvm/IR/Instructions.h" +#include "llvm/IR/Constants.h" +#include "llvm/IR/Instructions.h" #include "tensorflow/compiler/xla/layout_util.h" #include "tensorflow/compiler/xla/service/llvm_ir/llvm_util.h" #include "tensorflow/compiler/xla/shape_util.h" @@ -85,7 +85,7 @@ IrArray::IrArray(llvm::Value* base_ptr, const Shape& shape) ++depth; } - if (ShapeUtil::Rank(*shape_) == 0) { + if (!ShapeUtil::IsArray(*shape_) || ShapeUtil::IsScalar(*shape_)) { DCHECK(depth == 1 || depth == 0) << depth; } else { DCHECK_EQ(depth, ShapeUtil::Rank(*shape_)) << shape.ShortDebugString(); @@ -153,6 +153,28 @@ IrArray::Index IrArray::Index::SourceIndexOfReshape( return Index(source_multidim_index); } +IrArray::Index IrArray::Index::SourceIndexOfSlice( + const Shape& shape, tensorflow::gtl::ArraySlice starts, + tensorflow::gtl::ArraySlice strides, + llvm::IRBuilder<>* builder) const { + Index source_index(multidim_.size()); + for (int i = 0; i < multidim_.size(); ++i) { + int64 stride = strides[i]; + auto type = multidim_[i]->getType(); + + if (stride != 1) { + source_index[i] = builder->CreateAdd( + builder->CreateMul(multidim_[i], + llvm::ConstantInt::get(type, stride)), + llvm::ConstantInt::get(type, starts[i])); + } else { + source_index[i] = builder->CreateAdd( + multidim_[i], llvm::ConstantInt::get(type, starts[i])); + } + } + return source_index; +} + IrArray::Index IrArray::Index::SourceIndexOfTranspose( const Shape& shape, const Shape& operand_shape, tensorflow::gtl::ArraySlice dimension_mapping, @@ -228,6 +250,18 @@ llvm::Value* IrArray::EmitArrayElementAddress( llvm_ir::AsStringRef(name)); } +void IrArray::AnnotateLoadStoreInstructionWithMetadata( + llvm::Instruction* instruction) const { + CHECK(llvm::isa(instruction) || + llvm::isa(instruction)); + + for (const auto& kind_md_pair : metadata_) { + CHECK(kind_md_pair.first != llvm::LLVMContext::MD_invariant_load || + llvm::isa(instruction)); + instruction->setMetadata(kind_md_pair.first, kind_md_pair.second); + } +} + llvm::Value* IrArray::EmitReadArrayElement(const Index& index, llvm::IRBuilder<>* ir_builder, tensorflow::StringPiece name) const { @@ -236,11 +270,7 @@ llvm::Value* IrArray::EmitReadArrayElement(const Index& index, llvm::LoadInst* load = ir_builder->CreateLoad(element_address); llvm_ir::SetTbaaForInstruction(load, GetShape(), /*is_pointer_to=*/false); - for (const std::pair& kind_md_pair : metadata_) { - int kind = kind_md_pair.first; - llvm::MDNode* md = kind_md_pair.second; - load->setMetadata(kind, md); - } + AnnotateLoadStoreInstructionWithMetadata(load); return load; } @@ -250,12 +280,7 @@ void IrArray::EmitWriteArrayElement(const Index& index, llvm::Value* value, llvm::StoreInst* store = ir_builder->CreateStore(value, element_address); llvm_ir::SetTbaaForInstruction(store, GetShape(), /*is_pointer_to=*/false); - for (const std::pair& kind_md_pair : metadata_) { - int kind = kind_md_pair.first; - CHECK_NE(kind, llvm::LLVMContext::MD_invariant_load); - llvm::MDNode* md = kind_md_pair.second; - store->setMetadata(kind, md); - } + AnnotateLoadStoreInstructionWithMetadata(store); } IrArray IrArray::CastToShape(const Shape& new_shape, diff --git a/tensorflow/compiler/xla/service/llvm_ir/ir_array.h b/tensorflow/compiler/xla/service/llvm_ir/ir_array.h index 91cb3a679fd67fffb29f8a935cc3c65e9442136b..1ed7e99a829f5b0daa709913554d2300503ca33e 100644 --- a/tensorflow/compiler/xla/service/llvm_ir/ir_array.h +++ b/tensorflow/compiler/xla/service/llvm_ir/ir_array.h @@ -19,8 +19,8 @@ limitations under the License. #include #include -#include "external/llvm/include/llvm/IR/IRBuilder.h" -#include "external/llvm/include/llvm/IR/Value.h" +#include "llvm/IR/IRBuilder.h" +#include "llvm/IR/Value.h" #include "tensorflow/compiler/xla/map_util.h" #include "tensorflow/compiler/xla/types.h" #include "tensorflow/compiler/xla/xla_data.pb.h" @@ -83,7 +83,7 @@ class IrArray { Index(tensorflow::gtl::ArraySlice multidim, const Shape& shape, llvm::IRBuilder<>* ir_builder); - // Consturcts an index from both a multi-dimensional index and a linear + // Constructs an index from both a multi-dimensional index and a linear // index. "shape" has the same meaning as that in the constructor that takes // only a linear index. Index(tensorflow::gtl::ArraySlice multidim, @@ -108,6 +108,8 @@ class IrArray { const_iterator begin() const { return multidim().begin(); } const_iterator end() const { return multidim().end(); } + llvm::Value* back() const { return multidim().back(); } + bool LinearValidOnShape(const Shape& a) const; // Given that "this" is the target index of a reshape from `operand_shape` @@ -115,6 +117,16 @@ class IrArray { Index SourceIndexOfReshape(const Shape& shape, const Shape& operand_shape, llvm::IRBuilder<>* builder) const; + // Returns the index into the source operand from which a slice operation + // selects a value to be placed into index "this". The slice is described + // by starting indices `starts` and stride values `strides`. + // + // Precondition: "this" is an index into a slice whose shape is `shape`. + Index SourceIndexOfSlice(const Shape& shape, + tensorflow::gtl::ArraySlice starts, + tensorflow::gtl::ArraySlice strides, + llvm::IRBuilder<>* builder) const; + // Given that "this" is the target index of a transpose from `operand_shape` // to `shape` with the given dimension mapping, returns the source index. Index SourceIndexOfTranspose( @@ -183,6 +195,10 @@ class IrArray { llvm::IRBuilder<>* ir_builder, tensorflow::StringPiece name = "") const; + // Attach metadata this IrArray instance knows about to "instruction". + void AnnotateLoadStoreInstructionWithMetadata( + llvm::Instruction* instruction) const; + // Emit IR to read an array element at the given index. Returns the read // result (effectively, a Value loaded from memory). This method seamlessly // handles scalar shapes by broadcasting their value to all indices (index is @@ -204,17 +220,22 @@ class IrArray { llvm::IRBuilder<>* ir_builder) const; void AddAliasScopeMetadata(llvm::MDNode* alias_scope) { + CHECK_NE(alias_scope, nullptr); AddMetadata(llvm::LLVMContext::MD_alias_scope, alias_scope); } void AddNoaliasMetadata(llvm::MDNode* noalias) { + CHECK_NE(noalias, nullptr); AddMetadata(llvm::LLVMContext::MD_noalias, noalias); } void AddInvariantLoad(llvm::MDNode* invariant_load) { + CHECK_NE(invariant_load, nullptr); AddMetadata(llvm::LLVMContext::MD_invariant_load, invariant_load); } + const std::map& metadata() const { return metadata_; } + // Bumps the "which_dimension" value within the provided index by the provided // addend. static Index BumpIndex(const Index& index, int64 which_dimension, diff --git a/tensorflow/compiler/xla/service/llvm_ir/llvm_loop.cc b/tensorflow/compiler/xla/service/llvm_ir/llvm_loop.cc index 4ccded61e7376b63df19ecf324d0f693b71f8e9a..c78aa6d38036c7f76d4e7752774935aab6d066a6 100644 --- a/tensorflow/compiler/xla/service/llvm_ir/llvm_loop.cc +++ b/tensorflow/compiler/xla/service/llvm_ir/llvm_loop.cc @@ -18,10 +18,11 @@ limitations under the License. #include #include -#include "external/llvm/include/llvm/IR/Constants.h" -#include "external/llvm/include/llvm/IR/Function.h" -#include "external/llvm/include/llvm/IR/Instructions.h" +#include "llvm/IR/Constants.h" +#include "llvm/IR/Function.h" +#include "llvm/IR/Instructions.h" #include "tensorflow/compiler/xla/service/llvm_ir/llvm_util.h" +#include "tensorflow/compiler/xla/shape_util.h" #include "tensorflow/compiler/xla/types.h" #include "tensorflow/compiler/xla/xla_data.pb.h" #include "tensorflow/core/lib/strings/strcat.h" @@ -143,12 +144,19 @@ llvm::BasicBlock* ForLoop::CreateBasicBlockWithSuffix( std::unique_ptr ForLoopNest::AddLoop(tensorflow::StringPiece suffix, llvm::Value* start_index, llvm::Value* end_index) { + return AddLoop(suffix, start_index, end_index, ir_builder_->getInt64(1)); +} + +std::unique_ptr ForLoopNest::AddLoop(tensorflow::StringPiece suffix, + llvm::Value* start_index, + llvm::Value* end_index, + llvm::Value* stride) { if (inner_loop_body_bb_ != nullptr) { // Create this loop inside the previous one. ir_builder_->SetInsertPoint(&*inner_loop_body_bb_->getFirstInsertionPt()); } - std::unique_ptr loop = ForLoop::EmitForLoop( - suffix, start_index, end_index, ir_builder_->getInt64(1), ir_builder_); + std::unique_ptr loop = + ForLoop::EmitForLoop(suffix, start_index, end_index, stride, ir_builder_); if (outer_loop_preheader_bb_ == nullptr) { outer_loop_preheader_bb_ = loop->GetPreheaderBasicBlock(); @@ -171,6 +179,15 @@ std::unique_ptr ForLoopNest::AddLoop(int64 start_index, ir_builder_->getInt64(end_index)); } +std::unique_ptr ForLoopNest::AddLoop(int64 start_index, + int64 end_index, int64 stride, + tensorflow::StringPiece suffix) { + CHECK_LE(start_index, end_index); + return AddLoop(suffix, ir_builder_->getInt64(start_index), + ir_builder_->getInt64(end_index), + ir_builder_->getInt64(stride)); +} + IrArray::Index ForLoopNest::AddLoopsForShape(const Shape& shape, tensorflow::StringPiece suffix) { std::vector dimensions(ShapeUtil::Rank(shape)); diff --git a/tensorflow/compiler/xla/service/llvm_ir/llvm_loop.h b/tensorflow/compiler/xla/service/llvm_ir/llvm_loop.h index 60ac0444bcde002db6fd6bfa2630c9b78157e491..832e998d6ae659108c216c8586a159c769ec2c7d 100644 --- a/tensorflow/compiler/xla/service/llvm_ir/llvm_loop.h +++ b/tensorflow/compiler/xla/service/llvm_ir/llvm_loop.h @@ -19,9 +19,9 @@ limitations under the License. #include #include -#include "external/llvm/include/llvm/IR/BasicBlock.h" -#include "external/llvm/include/llvm/IR/IRBuilder.h" -#include "external/llvm/include/llvm/IR/Value.h" +#include "llvm/IR/BasicBlock.h" +#include "llvm/IR/IRBuilder.h" +#include "llvm/IR/Value.h" #include "tensorflow/compiler/xla/service/llvm_ir/ir_array.h" #include "tensorflow/compiler/xla/types.h" #include "tensorflow/compiler/xla/xla_data.pb.h" @@ -167,12 +167,22 @@ class ForLoopNest { // Adds a loop to the nest. If no loop has been added yet then emit a loop at // the current insert point of the given builder. If one or more loops have // been added then emit loop inside the body of the last added loop. + std::unique_ptr AddLoop(tensorflow::StringPiece suffix, + llvm::Value* start_index, + llvm::Value* end_index, llvm::Value* stride); + + // Like the above, except that it defaults to a stride of one. std::unique_ptr AddLoop(tensorflow::StringPiece suffix, llvm::Value* start_index, llvm::Value* end_index); // A convenient wrapper of the other flavor of AddLoop. The given start and // end index are constant. + std::unique_ptr AddLoop(int64 start_index, int64 end_index, + int64 stride, + tensorflow::StringPiece suffix); + + // Like the above, except that it defaults to a stride of one. std::unique_ptr AddLoop(int64 start_index, int64 end_index, tensorflow::StringPiece suffix); diff --git a/tensorflow/compiler/xla/service/llvm_ir/llvm_util.cc b/tensorflow/compiler/xla/service/llvm_ir/llvm_util.cc index 1e5b0a1a19d526a658d9df6a498a3de2572ba964..0ae75c5b3c61208db9713239d118d1d1cbc3e367 100644 --- a/tensorflow/compiler/xla/service/llvm_ir/llvm_util.cc +++ b/tensorflow/compiler/xla/service/llvm_ir/llvm_util.cc @@ -18,15 +18,15 @@ limitations under the License. #include #include -#include "external/llvm/include/llvm/IR/MDBuilder.h" -#include "external/llvm/include/llvm/IR/Operator.h" -#include "external/llvm/include/llvm/Target/TargetOptions.h" +#include "llvm/IR/MDBuilder.h" +#include "llvm/IR/Operator.h" +#include "llvm/Target/TargetOptions.h" #include "tensorflow/compiler/xla/layout_util.h" -#include "tensorflow/compiler/xla/legacy_flags/llvm_util_flags.h" #include "tensorflow/compiler/xla/literal_util.h" #include "tensorflow/compiler/xla/shape_util.h" #include "tensorflow/compiler/xla/types.h" #include "tensorflow/core/lib/core/casts.h" +#include "tensorflow/core/lib/core/errors.h" #include "tensorflow/core/lib/strings/strcat.h" #include "tensorflow/core/platform/logging.h" #include "tensorflow/core/platform/types.h" @@ -137,6 +137,24 @@ llvm::Type* ShapeToIrType(const Shape& shape, llvm::IRBuilder<>* ir_builder) { return result_type; } +StatusOr EncodeSelfDescribingShapeConstant( + const Shape& shape, int32* shape_size, llvm::IRBuilder<>* ir_builder) { + string encoded_shape = shape.SerializeAsString(); + if (encoded_shape.size() > std::numeric_limits::max()) { + return InternalError("Encoded shape size exceeded int32 size limit."); + } + *shape_size = static_cast(encoded_shape.size()); + return ir_builder->CreateGlobalStringPtr(llvm_ir::AsStringRef(encoded_shape)); +} + +StatusOr DecodeSelfDescribingShapeConstant(const void* shape_ptr, + int32 size_bytes) { + Shape shape; + TF_RET_CHECK(shape.ParseFromArray(shape_ptr, size_bytes)); + TF_RETURN_IF_ERROR(ShapeUtil::ValidateShape(shape)); + return shape; +} + namespace { // Recursively construct a multidimensional LLVM constant which represents the @@ -163,36 +181,36 @@ llvm::Constant* LiteralToConstant(const Literal& literal, int64 dimension_index, llvm::Constant* value; switch (shape.element_type()) { case PRED: - value = llvm::ConstantInt::get( - ir_element_type, LiteralUtil::Get(literal, *multi_index)); + value = llvm::ConstantInt::get(ir_element_type, + literal.Get(*multi_index)); break; case U8: - value = llvm::ConstantInt::get( - ir_element_type, LiteralUtil::Get(literal, *multi_index)); + value = llvm::ConstantInt::get(ir_element_type, + literal.Get(*multi_index)); break; case S32: - value = llvm::ConstantInt::get( - ir_element_type, LiteralUtil::Get(literal, *multi_index)); + value = llvm::ConstantInt::get(ir_element_type, + literal.Get(*multi_index)); break; case U32: - value = llvm::ConstantInt::get( - ir_element_type, LiteralUtil::Get(literal, *multi_index)); + value = llvm::ConstantInt::get(ir_element_type, + literal.Get(*multi_index)); break; case S64: - value = llvm::ConstantInt::get( - ir_element_type, LiteralUtil::Get(literal, *multi_index)); + value = llvm::ConstantInt::get(ir_element_type, + literal.Get(*multi_index)); break; case U64: - value = llvm::ConstantInt::get( - ir_element_type, LiteralUtil::Get(literal, *multi_index)); + value = llvm::ConstantInt::get(ir_element_type, + literal.Get(*multi_index)); break; case F32: - value = llvm::ConstantFP::get( - ir_element_type, LiteralUtil::Get(literal, *multi_index)); + value = llvm::ConstantFP::get(ir_element_type, + literal.Get(*multi_index)); break; case F64: - value = llvm::ConstantFP::get( - ir_element_type, LiteralUtil::Get(literal, *multi_index)); + value = llvm::ConstantFP::get(ir_element_type, + literal.Get(*multi_index)); break; default: LOG(FATAL) << "unsupported type " << shape.element_type(); @@ -357,31 +375,9 @@ void EmitLogging(const char* tag, llvm::Value* value, void SetTbaaForInstruction(llvm::Instruction* instruction, Shape shape, bool is_pointer_to) { - legacy_flags::LlvmUtilFlags* flags = legacy_flags::GetLlvmUtilFlags(); - if (!flags->xla_emit_tbaa) { - return; - } - - llvm::MDBuilder metadata_builder(instruction->getContext()); - llvm::MDNode* root = metadata_builder.createTBAARoot("XLA TBAA"); - string type_name; - if (is_pointer_to) { - type_name += "pointer-to "; - } - // Scalars do not have layout which makes it permissible to omit an explicit - // layout. To make sure that equivalent scalar shapes have the same TBAA, - // remove the (meaningless) explicit layout if one is present. - if (ShapeUtil::Rank(shape) == 0) { - LayoutUtil::ClearLayout(&shape); - } else { - CHECK(shape.has_layout()); - } - type_name += shape.ShortDebugString(); - llvm::MDNode* tbaa_node = - metadata_builder.createTBAANode(llvm_ir::AsStringRef(type_name), root); - instruction->setMetadata(llvm::LLVMContext::MD_tbaa, - metadata_builder.createTBAAStructTagNode( - tbaa_node, tbaa_node, /*Offset=*/0)); + // TODO(b/62903316): TBAA metadata causes LLVM to miscompile generated code, + // most likely because the generated metadata is incorrect. Disable TBAA + // metadata while we resolve this. } void SetAlignmentMetadataForLoad(llvm::LoadInst* load, uint64_t alignment) { @@ -449,24 +445,71 @@ int64 ByteSizeOf(const Shape& shape, const llvm::DataLayout& data_layout) { return ShapeUtil::ByteSizeOf(shape, pointer_size); } -llvm::FastMathFlags GetFastMathFlags(const HloModuleConfig& config) { +llvm::FastMathFlags GetFastMathFlags(bool fast_math_enabled) { llvm::FastMathFlags flags; - if (!config.fast_math_disabled()) { + if (fast_math_enabled) { // UnsafeAlgebra implies NoInfs, NoNaNs, NoSignedZeros, and AllowReciprocal. flags.setUnsafeAlgebra(); } return flags; } -void SetTargetOptions(const HloModuleConfig& config, +void SetTargetOptions(bool fast_math_enabled, llvm::TargetOptions* target_options) { - bool fast = !config.fast_math_disabled(); // In LLVM backend flags, UnsafeFPMath does not explicitly imply // NoInfs, etc. - target_options->UnsafeFPMath = fast; - target_options->NoInfsFPMath = fast; - target_options->NoNaNsFPMath = fast; - target_options->NoSignedZerosFPMath = fast; + target_options->UnsafeFPMath = fast_math_enabled; + target_options->NoInfsFPMath = fast_math_enabled; + target_options->NoNaNsFPMath = fast_math_enabled; + target_options->NoSignedZerosFPMath = fast_math_enabled; +} + +std::map MergeMetadata( + llvm::LLVMContext* context, const std::map& a, + const std::map& b) { + // We should extend this as needed to deal with other kinds of metadata like + // !dereferenceable and !range. + + std::map result; + for (auto kind_md_pair : a) { + if (kind_md_pair.first == llvm::LLVMContext::MD_alias_scope) { + llvm::SmallVector union_of_scopes; + llvm::SmallPtrSet scope_set; + for (const auto& scope_a : kind_md_pair.second->operands()) { + scope_set.insert(llvm::cast(scope_a.get())); + union_of_scopes.push_back(llvm::cast(scope_a.get())); + } + auto it = b.find(kind_md_pair.first); + if (it != b.end()) { + for (const auto& scope_b : it->second->operands()) { + if (!scope_set.count(llvm::cast(scope_b.get()))) { + union_of_scopes.push_back(llvm::cast(scope_b.get())); + } + } + } + result[llvm::LLVMContext::MD_alias_scope] = + llvm::MDNode::get(*context, union_of_scopes); + } else if (kind_md_pair.first == llvm::LLVMContext::MD_noalias) { + llvm::SmallVector intersection_of_scopes; + llvm::SmallPtrSet scope_set; + for (const auto& scope_a : kind_md_pair.second->operands()) { + scope_set.insert(llvm::cast(scope_a.get())); + } + auto it = b.find(kind_md_pair.first); + if (it != b.end()) { + for (const auto& scope_b : it->second->operands()) { + if (scope_set.count(llvm::cast(scope_b))) { + intersection_of_scopes.push_back(llvm::cast(scope_b)); + } + } + } + if (!intersection_of_scopes.empty()) { + result[llvm::LLVMContext::MD_noalias] = + llvm::MDNode::get(*context, intersection_of_scopes); + } + } + } + return result; } } // namespace llvm_ir diff --git a/tensorflow/compiler/xla/service/llvm_ir/llvm_util.h b/tensorflow/compiler/xla/service/llvm_ir/llvm_util.h index 28488ca99912335a4ead43c9c7cd227f85f7db68..6d94603338cba590276db4974910575ffb7cc994 100644 --- a/tensorflow/compiler/xla/service/llvm_ir/llvm_util.h +++ b/tensorflow/compiler/xla/service/llvm_ir/llvm_util.h @@ -20,15 +20,14 @@ limitations under the License. #include #include -#include "external/llvm/include/llvm/ADT/StringRef.h" -#include "external/llvm/include/llvm/IR/BasicBlock.h" -#include "external/llvm/include/llvm/IR/IRBuilder.h" -#include "external/llvm/include/llvm/IR/Instructions.h" -#include "external/llvm/include/llvm/IR/Module.h" -#include "external/llvm/include/llvm/IR/Value.h" -#include "external/llvm/include/llvm/Support/raw_ostream.h" -#include "tensorflow/compiler/xla/service/buffer_assignment.h" -#include "tensorflow/compiler/xla/service/hlo_module_config.h" +#include "llvm/ADT/StringRef.h" +#include "llvm/IR/BasicBlock.h" +#include "llvm/IR/IRBuilder.h" +#include "llvm/IR/Instructions.h" +#include "llvm/IR/Module.h" +#include "llvm/IR/Value.h" +#include "llvm/Support/raw_ostream.h" +#include "tensorflow/compiler/xla/literal_util.h" #include "tensorflow/compiler/xla/types.h" #include "tensorflow/compiler/xla/xla_data.pb.h" #include "tensorflow/core/lib/core/stringpiece.h" @@ -107,6 +106,19 @@ llvm::Type* PrimitiveTypeToIrType(PrimitiveType element_type, // if "shape" is [5 x [10 x f32]], the function returns [5 x [10 x float]]. llvm::Type* ShapeToIrType(const Shape& shape, llvm::IRBuilder<>* ir_builder); +// Returns a value that represents a pointer to a global string constant that +// encodes the shape as a serialized protobuf. +StatusOr EncodeSelfDescribingShapeConstant( + const Shape& shape, int32* shape_size, llvm::IRBuilder<>* ir_builder); + +// Inverses the encoding of a Shape protobuf into an LLVM global variable. +// +// This is intended to be called from the runtime to decode the llvm::Constants +// that are created via ConvertShapeToSelfDescribingConstant and subsequently +// embedded into the program. +StatusOr DecodeSelfDescribingShapeConstant(const void* shape_ptr, + int32 size_bytes); + // Converts a given literal to an IR Constant. Literals have known constant // values at IR emission time. llvm::Constant* ConvertLiteralToIrConstant(const Literal& literal, @@ -130,7 +142,7 @@ llvm::AllocaInst* EmitAllocaAtFunctionEntryWithCount( llvm::Type* type, llvm::Value* element_count, tensorflow::StringPiece name, llvm::IRBuilder<>* ir_builder, int alignment = 0); -// Creates a basic block with the same context and funtion as for the +// Creates a basic block with the same context and function as for the // builder. Inserts at the end of the function if insert_before is // null. llvm::BasicBlock* CreateBasicBlock(llvm::BasicBlock* insert_before, @@ -219,13 +231,21 @@ int64 ByteSizeOf(const Shape& shape, const llvm::DataLayout& data_layout); // Gets an llvm::FastMathFlags that reflects the settings in the given // module config. -llvm::FastMathFlags GetFastMathFlags(const HloModuleConfig& config); +llvm::FastMathFlags GetFastMathFlags(bool fast_math_enabled); // Sets values in the given TargetOptions struct according to the given // compilation options. -void SetTargetOptions(const HloModuleConfig& config, +void SetTargetOptions(bool fast_math_enabled, llvm::TargetOptions* target_options); +// Computes a conservative union of the metadata in "a" and "b". For +// aliasing-related metadata, this means the result can be applied to +// instructions whose aliasing relationship can be described either by "a" *or* +// by "b". +std::map MergeMetadata( + llvm::LLVMContext* context, const std::map& a, + const std::map& b); + } // namespace llvm_ir } // namespace xla diff --git a/tensorflow/compiler/xla/service/llvm_ir/loop_emitter.cc b/tensorflow/compiler/xla/service/llvm_ir/loop_emitter.cc index 9a128b2aa6f2d5e5650624f103c573e671335f7b..8839ec582df844f46f060e26917f15aa297cba3d 100644 --- a/tensorflow/compiler/xla/service/llvm_ir/loop_emitter.cc +++ b/tensorflow/compiler/xla/service/llvm_ir/loop_emitter.cc @@ -19,6 +19,7 @@ limitations under the License. #include #include "tensorflow/compiler/xla/service/llvm_ir/llvm_loop.h" +#include "tensorflow/compiler/xla/service/llvm_ir/ops.h" #include "tensorflow/compiler/xla/shape_util.h" #include "tensorflow/compiler/xla/status_macros.h" #include "tensorflow/compiler/xla/types.h" @@ -51,8 +52,41 @@ LoopEmitter::LoopEmitter(const ElementGenerator& target_element_generator, shape_(target_array.GetShape()), ir_builder_(ir_builder) {} +LoopEmitter::LoopEmitter(const ElementGenerator& target_element_generator, + tensorflow::gtl::ArraySlice target_arrays, + llvm::IRBuilder<>* ir_builder) + : body_emitter_([=](const llvm_ir::IrArray::Index array_index) + -> ::tensorflow::Status { + // Convert target_element_generator to a BodyEmitter. + TF_ASSIGN_OR_RETURN(llvm::Value * target_element, + target_element_generator(array_index)); + if (target_arrays.size() == 1) { + target_arrays[0].EmitWriteArrayElement(array_index, target_element, + ir_builder); + return tensorflow::Status::OK(); + } + + for (int64 i = 0; i < target_arrays.size(); ++i) { + target_arrays[i].EmitWriteArrayElement( + array_index, ir_builder_->CreateExtractValue(target_element, i), + ir_builder); + } + return tensorflow::Status::OK(); + }), + ir_builder_(ir_builder) { + if (target_arrays.size() > 1) { + // The sanity check for multiple outputs. + shape_ = target_arrays[0].GetShape(); + for (int64 i = 1; i < target_arrays.size(); ++i) { + const Shape& element_shape = target_arrays[i].GetShape(); + CHECK(ShapeUtil::SameDimensions(shape_, element_shape)); + } + } else { + shape_ = target_arrays[0].GetShape(); + } +} + IrArray::Index LoopEmitter::EmitIndexAndSetExitBasicBlock() { - CHECK(!ShapeUtil::IsTuple(shape_)); if (ShapeUtil::IsScalar(shape_)) { // No loop needed, so set exit_bb_ to nullptr. exit_bb_ = nullptr; diff --git a/tensorflow/compiler/xla/service/llvm_ir/loop_emitter.h b/tensorflow/compiler/xla/service/llvm_ir/loop_emitter.h index 08171e9e9de294339359f86059f89dcf4939ddea..6ecaa38c53137decc16e9c1e48cbc26980b040d5 100644 --- a/tensorflow/compiler/xla/service/llvm_ir/loop_emitter.h +++ b/tensorflow/compiler/xla/service/llvm_ir/loop_emitter.h @@ -18,9 +18,9 @@ limitations under the License. #include -#include "external/llvm/include/llvm/IR/BasicBlock.h" -#include "external/llvm/include/llvm/IR/IRBuilder.h" -#include "external/llvm/include/llvm/IR/Value.h" +#include "llvm/IR/BasicBlock.h" +#include "llvm/IR/IRBuilder.h" +#include "llvm/IR/Value.h" #include "tensorflow/compiler/xla/service/llvm_ir/ir_array.h" #include "tensorflow/compiler/xla/statusor.h" @@ -47,6 +47,10 @@ class LoopEmitter { // element of the given target array. LoopEmitter(const ElementGenerator& target_element_generator, const IrArray& target_array, llvm::IRBuilder<>* ir_builder); + // Same as previous method except emits multiple targets in an array. + LoopEmitter(const ElementGenerator& target_element_generator, + tensorflow::gtl::ArraySlice target_arrays, + llvm::IRBuilder<>* ir_builder); LoopEmitter(const LoopEmitter&) = delete; LoopEmitter& operator=(const LoopEmitter&) = delete; virtual ~LoopEmitter() = default; diff --git a/tensorflow/compiler/xla/service/llvm_ir/ops.cc b/tensorflow/compiler/xla/service/llvm_ir/ops.cc index e01d25d2502599cb84401e48aae2bbe6ec92b876..ac562e231c8f56184363d6e186c18847d67435ce 100644 --- a/tensorflow/compiler/xla/service/llvm_ir/ops.cc +++ b/tensorflow/compiler/xla/service/llvm_ir/ops.cc @@ -19,7 +19,7 @@ limitations under the License. #include #include -#include "external/llvm/include/llvm/IR/Instructions.h" +#include "llvm/IR/Instructions.h" #include "tensorflow/compiler/xla/service/llvm_ir/llvm_util.h" #include "tensorflow/compiler/xla/shape_util.h" #include "tensorflow/compiler/xla/types.h" diff --git a/tensorflow/compiler/xla/service/llvm_ir/ops.h b/tensorflow/compiler/xla/service/llvm_ir/ops.h index af4063c3401dadedd0b3f42d66f84ea225b8aedd..4e1d9d1080b3a5c8d8a09145f68bcff9d329c929 100644 --- a/tensorflow/compiler/xla/service/llvm_ir/ops.h +++ b/tensorflow/compiler/xla/service/llvm_ir/ops.h @@ -16,8 +16,8 @@ limitations under the License. #ifndef TENSORFLOW_COMPILER_XLA_SERVICE_LLVM_IR_OPS_H_ #define TENSORFLOW_COMPILER_XLA_SERVICE_LLVM_IR_OPS_H_ -#include "external/llvm/include/llvm/IR/IRBuilder.h" -#include "external/llvm/include/llvm/IR/Value.h" +#include "llvm/IR/IRBuilder.h" +#include "llvm/IR/Value.h" #include "tensorflow/compiler/xla/service/llvm_ir/ir_array.h" #include "tensorflow/core/lib/gtl/array_slice.h" #include "tensorflow/core/platform/types.h" diff --git a/tensorflow/compiler/xla/service/local_service.cc b/tensorflow/compiler/xla/service/local_service.cc index 17d7b97b21bd3296711295e0779b0a273c9917e0..1eb4edbe3e1abae7bcf5a65b9cf91bd95e30d945 100644 --- a/tensorflow/compiler/xla/service/local_service.cc +++ b/tensorflow/compiler/xla/service/local_service.cc @@ -19,7 +19,7 @@ limitations under the License. #include #include -#include "tensorflow/compiler/xla/legacy_flags/service_flags.h" +#include "tensorflow/compiler/xla/execution_options_util.h" #include "tensorflow/compiler/xla/ptr_util.h" #include "tensorflow/compiler/xla/service/backend.h" #include "tensorflow/compiler/xla/service/computation_layout.h" @@ -46,13 +46,6 @@ namespace se = ::perftools::gputools; namespace xla { -/* static */ StatusOr> LocalService::NewService( - perftools::gputools::Platform* platform) { - ServiceOptions default_options; - default_options.set_platform(platform); - return NewService(default_options); -} - /* static */ StatusOr> LocalService::NewService( const ServiceOptions& options) { perftools::gputools::Platform* platform = options.platform(); @@ -60,37 +53,20 @@ namespace xla { TF_ASSIGN_OR_RETURN(platform, PlatformUtil::GetDefaultPlatform()); } - TF_ASSIGN_OR_RETURN( - std::unique_ptr backend, - Backend::CreateBackend(platform, options.number_of_replicas())); + BackendOptions backend_options; + backend_options.set_platform(platform).set_intra_op_parallelism_threads( + options.intra_op_parallelism_threads()); + TF_ASSIGN_OR_RETURN(std::unique_ptr backend, + Backend::CreateBackend(backend_options)); - TF_ASSIGN_OR_RETURN(std::unique_ptr compute_constant_backend, - CreateComputeConstantBackend()); - std::unique_ptr service(new LocalService( - std::move(backend), std::move(compute_constant_backend))); + std::unique_ptr service( + new LocalService(options, std::move(backend))); return std::move(service); } -LocalService::LocalService(std::unique_ptr execute_backend, - std::unique_ptr compute_constant_backend) - : Service(std::move(execute_backend), std::move(compute_constant_backend)) { - runs_in_client_process_ = true; -} - -tensorflow::Status LocalService::ResolveArguments( - const tensorflow::gtl::ArraySlice arguments, - int device_ordinal, - std::vector* argument_ptrs) { - TF_ASSIGN_OR_RETURN(std::vector arg_allocations, - ResolveAndValidateArguments( - arguments, execute_backend_.get(), device_ordinal)); - argument_ptrs->resize(arg_allocations.size()); - for (int i = 0; i < arguments.size(); ++i) { - const Allocation& allocation = *arg_allocations[i]; - (*argument_ptrs)[i] = allocation.device_memory(); - } - return tensorflow::Status::OK(); -} +LocalService::LocalService(const ServiceOptions& options, + std::unique_ptr execute_backend) + : Service(options, std::move(execute_backend)) {} namespace { // Returns the space required to allocate a shape. If @@ -102,12 +78,11 @@ int64 RequiredSpace(const Shape& shape, bool allocate_space_for_deep_copy, // TODO(b/33492279) remove once no devices represent result tuples as // contiguous buffers. if (allocate_space_for_deep_copy) { - TF_CHECK_OK(ShapeUtil::ForEachSubshape( + ShapeUtil::ForEachSubshape( shape, [&size, transfer_manager](const Shape& subshape, const ShapeIndex& /*index*/) { size += transfer_manager->GetByteSizeRequirement(subshape); - return tensorflow::Status::OK(); - })); + }); } return size; } @@ -128,70 +103,6 @@ StatusOr LocalService::AllocateBufferOnDevice( allocation_size)); } -StatusOr>> -LocalService::CompileAheadOfTime( - const tensorflow::gtl::ArraySlice - computations, - const AotCompilationOptions& options) { - std::vector> hlo_modules; - std::vector> module_configs; - for (const AheadOfTimeComputationInstance& instance : computations) { - TF_ASSIGN_OR_RETURN(UserComputation * user_computation, - computation_tracker_.Resolve(instance.computation)); - VersionedComputationHandle versioned_handle = - user_computation->GetVersionedHandle(); - - // Dump computation proto state if flag is set. - legacy_flags::ServiceFlags* flags = legacy_flags::GetServiceFlags(); - const string& directory_path = flags->xla_dump_computations_to; - if (!directory_path.empty()) { - TF_ASSIGN_OR_RETURN( - std::unique_ptr session_module, - computation_tracker_.SnapshotComputation(versioned_handle.handle)); - string filename = tensorflow::strings::StrCat( - "computation_", versioned_handle.handle.handle(), "__", - session_module->entry().name(), "__version_", - versioned_handle.version); - TF_RETURN_IF_ERROR(Executable::DumpToDirectory(directory_path, filename, - *session_module)); - } - - TF_ASSIGN_OR_RETURN(std::unique_ptr hlo_module, - computation_tracker_.BuildHloModule( - versioned_handle, - /*include_unreachable_instructions=*/true)); - hlo_modules.push_back(std::move(hlo_module)); - - TF_ASSIGN_OR_RETURN( - std::shared_ptr program_shape, - user_computation->ComputeProgramShape(versioned_handle.version)); - - module_configs.push_back(MakeUnique(*program_shape)); - HloModuleConfig* module_config = module_configs.back().get(); - auto* computation_layout = - module_config->mutable_entry_computation_layout(); - if (flags->xla_hlo_profile) { - module_config->enable_hlo_profiling(true); - } - for (int i = 0; i < instance.argument_layouts.size(); ++i) { - const Shape& argument_layout = *instance.argument_layouts[i]; - if (ShapeUtil::IsTuple(argument_layout)) { - return Unimplemented("tuple arguments not supported yet"); - } - TF_RETURN_IF_ERROR( - computation_layout->mutable_parameter_layout(i)->CopyLayoutFromShape( - argument_layout)); - } - TF_RETURN_IF_ERROR( - computation_layout->mutable_result_layout()->CopyLayoutFromShape( - *instance.result_layout)); - } - - return execute_backend_->compiler()->CompileAheadOfTime( - std::move(hlo_modules), std::move(module_configs), MakeHloDumper(), - options); -} - StatusOr> LocalService::CompileExecutable( const ComputationHandle& computation, const tensorflow::gtl::ArraySlice argument_layouts, @@ -226,31 +137,19 @@ StatusOr> LocalService::CompileExecutable( ValidateResultShapeWithLayout(*result_layout, program_shape->result())); } - // Construct computation layout from the argument layouts. - auto module_config = MakeUnique(*program_shape); - module_config->set_has_hybrid_result(has_hybrid_result); - module_config->set_replica_count(execute_backend_->Replicas().size()); - legacy_flags::ServiceFlags* flags = legacy_flags::GetServiceFlags(); - if (flags->xla_hlo_profile) { - module_config->enable_hlo_profiling(true); - } - auto* computation_layout = module_config->mutable_entry_computation_layout(); - for (int i = 0; i < argument_layouts.size(); ++i) { - const Shape& shape = *argument_layouts[i]; - if (ShapeUtil::IsTuple(shape)) { - return Unimplemented("tuple arguments not supported yet"); - } - TF_RETURN_IF_ERROR( - computation_layout->mutable_parameter_layout(i)->CopyLayoutFromShape( - shape)); - } + ExecutionOptions execution_options = CreateDefaultExecutionOptions(); if (result_layout != nullptr) { - TF_RETURN_IF_ERROR( - computation_layout->mutable_result_layout()->CopyLayoutFromShape( - *result_layout)); + *execution_options.mutable_shape_with_output_layout() = *result_layout; } else { - computation_layout->mutable_result_layout()->SetToDefaultLayout(); + *execution_options.mutable_shape_with_output_layout() = + program_shape->result(); + LayoutUtil::SetToDefaultLayout( + execution_options.mutable_shape_with_output_layout()); } + TF_ASSIGN_OR_RETURN( + std::unique_ptr module_config, + CreateModuleConfig(*program_shape, argument_layouts, &execution_options, + has_hybrid_result)); TF_ASSIGN_OR_RETURN(se::StreamExecutor * executor, execute_backend_->stream_executor(device_ordinal)); @@ -258,7 +157,6 @@ StatusOr> LocalService::CompileExecutable( std::vector argument_buffers( argument_layouts.size()); return BuildExecutable(versioned_handle, std::move(module_config), - /*executable_for_compute_constant=*/false, argument_buffers, execute_backend_.get(), executor); } diff --git a/tensorflow/compiler/xla/service/local_service.h b/tensorflow/compiler/xla/service/local_service.h index df27f0a7a60dca99caf09994f417f1bc45ec15de..c90943f3c0389f7435c7e32ab5f82cb61c91f99c 100644 --- a/tensorflow/compiler/xla/service/local_service.h +++ b/tensorflow/compiler/xla/service/local_service.h @@ -35,22 +35,10 @@ namespace xla { // in the same process as the client. class LocalService : public Service { public: - // Factory for creating a LocalService. The parameter platform is the platform - // that the service should target. If platform is null then the default - // platform is used. - static StatusOr> NewService( - perftools::gputools::Platform* platform); + // Factory for creating a LocalService. static StatusOr> NewService( const ServiceOptions& options); - // For an array of arguments, validate that each is placed on the - // specified device_ordinal, and return the DeviceMemoryBase - // corresponding to each argument. - tensorflow::Status ResolveArguments( - const tensorflow::gtl::ArraySlice arguments, - int device_ordinal, - std::vector* argument_ptrs); - // Return a handle to a buffer large enough to hold shape, allocated // on device_ordinal. If allocate_space_for_deep_copy, the buffer is // large enough to hold all sub-buffers of a tuple shape, otherwise @@ -59,22 +47,6 @@ class LocalService : public Service { const Shape& shape, int device_ordinal, bool allocate_space_for_deep_copy); - // A description of a computation to compile using CompileAheadOfTime. - struct AheadOfTimeComputationInstance { - ComputationHandle computation; - std::vector argument_layouts; - const Shape* result_layout = nullptr; - }; - - // Compiles a list of computations for ahead-of-time execution. This is - // intended for use in static compilation. See - // |LocalClient::CompileAheadOfTime| for additional details. - StatusOr>> - CompileAheadOfTime( - const tensorflow::gtl::ArraySlice - computations, - const AotCompilationOptions& Options); - // Builds an Executable with the given argument layouts and options. If // result_layout is non-null, then the executable is compiled to produce a // result of the given layout. @@ -84,8 +56,8 @@ class LocalService : public Service { const Shape* result_layout, int device_ordinal, bool has_hybrid_result); private: - explicit LocalService(std::unique_ptr backend, - std::unique_ptr compute_constant_backend); + explicit LocalService(const ServiceOptions& options, + std::unique_ptr backend); LocalService(const LocalService&) = delete; void operator=(const LocalService&) = delete; }; diff --git a/tensorflow/compiler/xla/service/logical_buffer.cc b/tensorflow/compiler/xla/service/logical_buffer.cc index c2f9dda66e3d68c774a1c856132ebd8b4f369950..68553bed121917850aaae41c6154f7895ed1add9 100644 --- a/tensorflow/compiler/xla/service/logical_buffer.cc +++ b/tensorflow/compiler/xla/service/logical_buffer.cc @@ -26,10 +26,18 @@ limitations under the License. namespace xla { +LogicalBuffer::LogicalBuffer(HloInstruction* instruction, + const ShapeIndex& index, Id id) + : instruction_(instruction), id_(id), color_(kInvalidColor), index_(index) { + const auto& s = shape(); + is_array_ = ShapeUtil::IsArray(s); + is_tuple_ = ShapeUtil::IsTuple(s); +} + string LogicalBuffer::ToString() const { - return tensorflow::strings::StrCat(instruction_->FullyQualifiedName(), "[", + return tensorflow::strings::StrCat(instruction_->name(), "[", tensorflow::str_util::Join(index_, ","), - "](#", id_, ")"); + "](#", id_, " @", color_.value(), ")"); } std::ostream& operator<<(std::ostream& out, const LogicalBuffer& buffer) { @@ -37,4 +45,26 @@ std::ostream& operator<<(std::ostream& out, const LogicalBuffer& buffer) { return out; } +/*static*/ LogicalBufferProto::Location LogicalBuffer::ToLocationProto( + const HloInstruction& instruction, const ShapeIndex& index) { + LogicalBufferProto::Location proto; + proto.set_computation_name(instruction.parent()->name()); + proto.set_instruction_name(instruction.name()); + for (const int64 index_entry : index) { + proto.add_shape_index(index_entry); + } + return proto; +} + +LogicalBufferProto LogicalBuffer::ToProto(const SizeFunction& size_fn) const { + LogicalBufferProto proto; + proto.set_id(id_); + proto.set_size(size_fn(*this)); + LogicalBufferProto::Location proto_location = + ToLocationProto(*instruction_, index_); + proto.mutable_defined_at()->Swap(&proto_location); + proto.set_color(color_.value()); + return proto; +} + } // namespace xla diff --git a/tensorflow/compiler/xla/service/logical_buffer.h b/tensorflow/compiler/xla/service/logical_buffer.h index d985db4a18e8fade1bec6e59fb3b9b519b5917dc..67b205e289e626f4db16c39a0a9ddf8618678c3a 100644 --- a/tensorflow/compiler/xla/service/logical_buffer.h +++ b/tensorflow/compiler/xla/service/logical_buffer.h @@ -16,15 +16,18 @@ limitations under the License. #ifndef TENSORFLOW_COMPILER_XLA_SERVICE_LOGICAL_BUFFER_H_ #define TENSORFLOW_COMPILER_XLA_SERVICE_LOGICAL_BUFFER_H_ +#include #include #include #include +#include "tensorflow/compiler/xla/service/hlo.pb.h" #include "tensorflow/compiler/xla/service/hlo_instruction.h" #include "tensorflow/compiler/xla/shape_util.h" #include "tensorflow/compiler/xla/types.h" #include "tensorflow/compiler/xla/xla_data.pb.h" #include "tensorflow/core/lib/gtl/array_slice.h" +#include "tensorflow/core/lib/gtl/int_type.h" #include "tensorflow/core/platform/macros.h" #include "tensorflow/core/platform/types.h" @@ -81,6 +84,8 @@ namespace xla { // LogicalBuffer(%tuple_constant, {1, 1}) // Holds value "43" class LogicalBuffer { public: + TF_LIB_GTL_DEFINE_INT_TYPE(Color, int64); + // Id is a unique identifier for the LogicalBuffer to facilitate efficient // collections of LogicalBuffers with stable iteration order. // LogicalBuffers are typically created and accessed through @@ -90,10 +95,9 @@ class LogicalBuffer { // Functions which return the size and alignment of a logical buffer in bytes. using SizeFunction = std::function; - using AlignmentFunction = std::function; + using AlignmentFunction = std::function; - LogicalBuffer(HloInstruction* instruction, const ShapeIndex& index, Id id) - : instruction_(instruction), index_(index), id_(id) {} + LogicalBuffer(HloInstruction* instruction, const ShapeIndex& index, Id id); Id id() const { return id_; } @@ -104,6 +108,22 @@ class LogicalBuffer { // defined. Index used defined as in ShapeUtil::GetSubshape() const ShapeIndex& index() const { return index_; } + // Return the color of the logical buffer. Differently colored buffers can + // not be parts of the same allocation. + Color color() const { + CHECK_NE(color_, kInvalidColor) + << "Should not query the color of a buffer that was never colored"; + return color_; + } + + void set_color(Color color) { + CHECK_NE(color, kInvalidColor) + << "Should not set the color of a buffer to the invalid color"; + color_ = color; + } + + bool has_color() const { return color_ != kInvalidColor; } + // Return the shape of the buffer. This reference points into the shape field // of the instruction defining the buffer. Therefore, the returned shape will // contain the layout of instruction, if any. @@ -116,20 +136,31 @@ class LogicalBuffer { bool IsTopLevel() const { return index_.empty(); } // Whether this buffer contains a tuple. - bool IsTuple() const { return ShapeUtil::IsTuple(shape()); } + bool IsTuple() const { return is_tuple_; } + + // Whether this buffer contains an array. + bool IsArray() const { return is_array_; } // operator< is required for std::set. bool operator<(const LogicalBuffer& other) const { return id_ < other.id_; } - // Whether this buffer contains an array. - bool IsArray() const { return ShapeUtil::IsArray(shape()); } - string ToString() const; + LogicalBufferProto ToProto(const SizeFunction& size_fn) const; + + // Returns the LogicalBufferProto::Location that serializes the given + // instruction and index. + static LogicalBufferProto::Location ToLocationProto( + const HloInstruction& instruction, const ShapeIndex& index); + + const Color kInvalidColor = Color(-1); private: HloInstruction* instruction_; + Id id_ : 62; + bool is_array_ : 1; + bool is_tuple_ : 1; + Color color_; ShapeIndex index_; - Id id_; // Similar to HLO constructs (HloInstruction, etc), pointers are used for // comparison to equality, so disable all copying. diff --git a/tensorflow/compiler/xla/service/logical_buffer_analysis.cc b/tensorflow/compiler/xla/service/logical_buffer_analysis.cc new file mode 100644 index 0000000000000000000000000000000000000000..8041d74baa72905ea1f81e6d67c67fe6430640a1 --- /dev/null +++ b/tensorflow/compiler/xla/service/logical_buffer_analysis.cc @@ -0,0 +1,129 @@ +/* Copyright 2017 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "tensorflow/compiler/xla/service/logical_buffer_analysis.h" + +#include + +#include "tensorflow/compiler/xla/shape_util.h" +#include "tensorflow/core/lib/core/errors.h" +#include "tensorflow/core/platform/logging.h" + +namespace xla { + +/* static */ StatusOr> +LogicalBufferAnalysis::Run(const HloModule* module) { + std::unique_ptr analysis( + new LogicalBufferAnalysis(module)); + TF_RETURN_IF_ERROR(analysis->Analyze()); + return std::move(analysis); +} + +Status LogicalBufferAnalysis::Analyze() { + // Empirically we usually have a few more logical buffers than instructions, + // so reserve 10% more than the number of instructions to avoid frequent + // resizes. + logical_buffers_.clear(); + logical_buffers_.reserve((module_->NumUniqueInstructionIds() * 11) / 10); + + // We filter out fusion computations, and get to them through fusion + // instructions. This is because it's possible to have orphaned (unreachable) + // fusion computations, and we don't want to try to assign buffers to those. + for (auto& computation : module_->computations()) { + if (computation->IsFusionComputation()) { + continue; + } + TF_RETURN_IF_ERROR(computation->Accept(this)); + for (auto& instruction : computation->instructions()) { + if (instruction->opcode() != HloOpcode::kFusion) { + continue; + } + TF_RETURN_IF_ERROR(instruction->fused_expression_root()->Accept(this)); + } + } + return Status::OK(); +} + +LogicalBuffer& LogicalBufferAnalysis::GetBuffer(LogicalBuffer::Id id) const { + CHECK_GE(id, 0); + CHECK_LT(id, logical_buffers_.size()); + return *logical_buffers_[id]; +} + +LogicalBuffer& LogicalBufferAnalysis::GetBuffer(HloInstruction* instruction, + const ShapeIndex& index) const { + return *output_buffers_.at(std::make_pair(instruction, index)); +} + +void LogicalBufferAnalysis::NewLogicalBuffer(HloInstruction* instruction, + const ShapeIndex& index) { + CHECK_EQ(logical_buffers_.size(), next_buffer_id_); + logical_buffers_.emplace_back( + MakeUnique(instruction, index, next_buffer_id_)); + output_buffers_[std::make_pair(instruction, index)] = + logical_buffers_.back().get(); + + ++next_buffer_id_; +} + +Status LogicalBufferAnalysis::DefaultAction(HloInstruction* hlo_instruction) { + // Create a logical buffer for each output of the instruction. + ShapeUtil::ForEachSubshape( + hlo_instruction->shape(), + [this, hlo_instruction](const Shape& shape, const ShapeIndex& index) { + NewLogicalBuffer(hlo_instruction, index); + }); + + return Status::OK(); +} + +Status LogicalBufferAnalysis::HandleGetTupleElement( + HloInstruction* get_tuple_element, HloInstruction* operand) { + // GetTupleElement 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 + NewLogicalBuffer(copy, /*index=*/{}); + return Status::OK(); +} + +Status LogicalBufferAnalysis::HandleBitcast(HloInstruction* bitcast) { + // A kBitcast instruction aliases its operand. That is, the buffer of its + // result *is* the buffer of its operand. + return Status::OK(); +} + +Status LogicalBufferAnalysis::HandleTuple( + HloInstruction* tuple, + tensorflow::gtl::ArraySlice operands) { + // A Tuple instruction only creates the top-level buffer. + NewLogicalBuffer(tuple, /*index=*/{}); + return Status::OK(); +} + +Status LogicalBufferAnalysis::HandleSelect(HloInstruction* select, + HloInstruction* /*pred*/, + HloInstruction* on_true, + HloInstruction* on_false) { + // Select allocates a new buffer and then shallow copies the on_true or + // on_false buffer into this new buffer. + NewLogicalBuffer(select, /*index=*/{}); + return Status::OK(); +} + +} // namespace xla diff --git a/tensorflow/compiler/xla/service/logical_buffer_analysis.h b/tensorflow/compiler/xla/service/logical_buffer_analysis.h new file mode 100644 index 0000000000000000000000000000000000000000..de9fe1b0a4ed3f6f8c466050520a9c4889793c62 --- /dev/null +++ b/tensorflow/compiler/xla/service/logical_buffer_analysis.h @@ -0,0 +1,96 @@ +/* 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 THIRD_PARTY_TENSORFLOW_COMPILER_XLA_SERVICE_LOGICAL_BUFFER_ANALYSIS_H_ +#define THIRD_PARTY_TENSORFLOW_COMPILER_XLA_SERVICE_LOGICAL_BUFFER_ANALYSIS_H_ + +#include "tensorflow/compiler/xla/service/hlo_instruction.h" +#include "tensorflow/compiler/xla/service/hlo_module.h" +#include "tensorflow/compiler/xla/service/logical_buffer.h" +#include "tensorflow/compiler/xla/statusor.h" +#include "tensorflow/core/lib/hash/hash.h" + +namespace xla { +// A class to create all the logical buffers defined by the HLO ops in a module. +class LogicalBufferAnalysis : public DfsHloVisitorWithDefault { + public: + // Runs points-to analysis on 'module'. + static StatusOr> Run( + const HloModule* module); + + // Returns the logical buffer with the given ID. + LogicalBuffer& GetBuffer(LogicalBuffer::Id id) const; + + // Returns the logical buffer that represents the output of a given HLO + // at a given index. + LogicalBuffer& GetBuffer(HloInstruction* instruction, + const ShapeIndex& index) const; + + const std::vector>& logical_buffers() const { + return logical_buffers_; + } + LogicalBuffer::Id num_logical_buffers() const { return next_buffer_id_; } + + private: + explicit LogicalBufferAnalysis(const HloModule* module) : module_(module) {} + Status Analyze(); + + // The module this analysis is performed on. + const HloModule* module_; + + // Create a new logical buffer and return a reference to it. The newly created + // buffer is stored in an internal vector of LogicalBuffers and can be + // accessed with GetBuffer. + void NewLogicalBuffer(HloInstruction* instruction, const ShapeIndex& index); + + Status DefaultAction(HloInstruction* hlo_instruction) override; + Status HandleTuple( + HloInstruction* tuple, + tensorflow::gtl::ArraySlice operands) override; + Status HandleGetTupleElement(HloInstruction* get_tuple_element, + HloInstruction* operand) override; + Status HandleBitcast(HloInstruction* bitcast) override; + Status HandleCopy(HloInstruction* copy) override; + Status HandleSelect(HloInstruction* select, HloInstruction* pred, + HloInstruction* on_true, + HloInstruction* on_false) override; + + // A map from the buffer ID to the logical buffer + std::vector> logical_buffers_; + + struct Hasher { + size_t operator()( + std::pair p) const { + size_t inst_hash = tensorflow::hash()(p.first); + for (auto index = p.second.begin(); index != p.second.end(); ++index) { + inst_hash = tensorflow::Hash64Combine(*index, inst_hash); + } + return inst_hash; + } + }; + + // A map from an hlo + shape index to the logical buffer representing + // the appropriate output. + std::unordered_map, + LogicalBuffer*, Hasher> + output_buffers_; + + // The ID of the next logical buffer created. + LogicalBuffer::Id next_buffer_id_ = 0; +}; + +} // namespace xla + +#endif // THIRD_PARTY_TENSORFLOW_COMPILER_XLA_SERVICE_LOGICAL_BUFFER_ANALYSIS_H_ diff --git a/tensorflow/compiler/xla/service/name_uniquer.cc b/tensorflow/compiler/xla/service/name_uniquer.cc index 4014856b9b243831a087962484128a121680eb1b..069f85af721228c8f5d40cf243eea7f1e5173c62 100644 --- a/tensorflow/compiler/xla/service/name_uniquer.cc +++ b/tensorflow/compiler/xla/service/name_uniquer.cc @@ -29,7 +29,11 @@ string NameUniquer::GetUniqueName(tensorflow::StringPiece prefix) { return root; } else { tensorflow::strings::StrAppend(&root, separator_, *count); + // Increment lookup under old 'root' name. (*count)++; + // Initialize count under new 'root' name. + count = &(generated_names_[root]); + *count = 1; return root; } } diff --git a/tensorflow/compiler/xla/service/platform_util.cc b/tensorflow/compiler/xla/service/platform_util.cc index 116bd3f067635b44b9148d5a1fa802af8b829c12..4f915a0c2eeaca0fe077a907571c8379992185eb 100644 --- a/tensorflow/compiler/xla/service/platform_util.cc +++ b/tensorflow/compiler/xla/service/platform_util.cc @@ -24,6 +24,7 @@ limitations under the License. #include "tensorflow/compiler/xla/statusor.h" #include "tensorflow/compiler/xla/types.h" #include "tensorflow/compiler/xla/util.h" +#include "tensorflow/core/lib/core/threadpool.h" #include "tensorflow/core/lib/strings/str_util.h" #include "tensorflow/core/platform/logging.h" #include "tensorflow/core/platform/stream_executor_no_cuda.h" @@ -140,21 +141,32 @@ PlatformUtil::GetStreamExecutors(se::Platform* platform) { device_count = 1; } std::vector stream_executors(device_count, nullptr); - for (int i = 0; i < device_count; ++i) { - se::StreamExecutorConfig config; - config.ordinal = i; - auto executor_status = platform->GetExecutor(config); - if (executor_status.ok()) { - se::StreamExecutor* executor = executor_status.ValueOrDie(); - if (IsDeviceSupported(executor)) { - stream_executors[i] = executor; - } - } else { - LOG(WARNING) << "unable to create StreamExecutor for " << platform->Name() - << ":" << i << ": " - << executor_status.status().error_message(); + VLOG(1) << "Initializing devices"; + { + tensorflow::thread::ThreadPool thread_pool( + tensorflow::Env::Default(), "device_initialization", device_count); + for (int i = 0; i < device_count; ++i) { + thread_pool.Schedule([platform, i, &stream_executors]() { + VLOG(1) << "Started device init " << i; + se::StreamExecutorConfig config; + config.ordinal = i; + auto executor_status = platform->GetExecutor(config); + if (executor_status.ok()) { + se::StreamExecutor* executor = executor_status.ValueOrDie(); + if (IsDeviceSupported(executor)) { + stream_executors[i] = executor; + } + } else { + LOG(WARNING) << "unable to create StreamExecutor for " + << platform->Name() << ":" << i << ": " + << executor_status.status().error_message(); + } + VLOG(1) << "Finished device init " << i; + }); } + // 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; })) { return InternalError("no supported devices found for platform %s", diff --git a/tensorflow/compiler/xla/service/reduce_precision_insertion.cc b/tensorflow/compiler/xla/service/reduce_precision_insertion.cc new file mode 100644 index 0000000000000000000000000000000000000000..33327dc60fb9e3cd18cfd7fb15c4350c82aaf6a1 --- /dev/null +++ b/tensorflow/compiler/xla/service/reduce_precision_insertion.cc @@ -0,0 +1,317 @@ +/* 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/reduce_precision_insertion.h" + +#include "tensorflow/compiler/xla/service/hlo_module.h" +#include "tensorflow/compiler/xla/shape_util.h" +#include "tensorflow/core/platform/logging.h" + +namespace xla { + +std::vector ReducePrecisionInsertion::instructions_to_modify( + const HloComputation* computation) { + std::vector instruction_list; + + switch (location_) { + case HloReducePrecisionOptions::OP_INPUTS: + case HloReducePrecisionOptions::OP_OUTPUTS: + case HloReducePrecisionOptions::UNFUSED_OP_OUTPUTS: + for (auto& instruction : computation->instructions()) { + VLOG(4) << "Visited instruction: " << instruction->ToString(); + if (instruction_filter_function_(instruction.get())) { + instruction_list.push_back(instruction.get()); + } + } + break; + + case HloReducePrecisionOptions::FUSION_INPUTS_BY_CONTENT: + case HloReducePrecisionOptions::FUSION_OUTPUTS_BY_CONTENT: + for (auto& instruction : computation->instructions()) { + VLOG(4) << "Visited instruction: " << instruction->ToString(); + if (instruction->opcode() != HloOpcode::kFusion) { + continue; + } + for (auto& fused_instruction : + instruction->fused_instructions_computation()->instructions()) { + VLOG(4) << "Checking sub-instruction: " + << fused_instruction->ToString(); + if (instruction_filter_function_(fused_instruction.get())) { + instruction_list.push_back(instruction.get()); + break; + } + } + } + break; + + default: + break; + } + VLOG(1) << "Found " << instruction_list.size() + << " candidate instruction(s) for reduce-precision insertion"; + + return instruction_list; +} + +StatusOr ReducePrecisionInsertion::insert_after( + HloInstruction* instruction) { + // Check that this isn't already an equivalent operation. + if (is_redundant(instruction)) { + VLOG(2) << "Skipped: instruction is already an equivalent" + " reduce-precision instruction:" + << instruction->ToString(); + return false; + } + + // Check that we haven't already inserted an equivalant reduce-precision + // operation after this instruction. (The zero-user case occurs when this is + // the root instruction.) + if (instruction->user_count() > 0) { + bool redundant_followers = true; + for (HloInstruction* user : instruction->users()) { + if (!is_redundant(user)) { + redundant_followers = false; + break; + } + } + if (redundant_followers) { + VLOG(2) << "Skipped: instruction already followed by equivalent" + " reduce-precision instructions"; + return false; + } + } + + HloInstruction* reduced = instruction->parent()->AddInstruction( + HloInstruction::CreateReducePrecision(instruction->shape(), instruction, + exponent_bits_, mantissa_bits_)); + TF_RETURN_IF_ERROR( + instruction->parent()->ReplaceUsesOfInstruction(instruction, reduced)); + return true; +} + +StatusOr ReducePrecisionInsertion::insert_on_inputs( + const std::vector& instructions) { + bool computation_changed = false; + for (auto instruction : instructions) { + VLOG(2) << "Adding reduce-precision operation to inputs of instruction: " + << instruction->ToString(); + for (int64 i = 0; i < instruction->operand_count(); i++) { + HloInstruction* operand = instruction->mutable_operand(i); + VLOG(2) << "Adding to operand " << i << ": " << operand; + + if (!is_valid_shape(operand->shape())) { + VLOG(2) << "Skipped: value is not an F32 vector"; + continue; + } + + if (is_redundant(operand)) { + VLOG(2) << "Skipped: operand is already an equivalent reduce-precision" + " instruction"; + continue; + } + + if (instruction->opcode() == HloOpcode::kFusion && + instruction->fusion_kind() == HloInstruction::FusionKind::kLoop) { + // Insert the reduce-precision operation inside the fusion computation, + // after the corresponding parameter instruction. + TF_ASSIGN_OR_RETURN( + bool instruction_changed, + insert_after(instruction->fused_instructions_computation() + ->parameter_instruction(i))); + computation_changed |= instruction_changed; + } else { + // Look for an existing reduce-precision operation on the operand. (We + // need to be careful not to create a loop, though!) + HloInstruction* reduced = nullptr; + for (auto& user : operand->users()) { + if (user != instruction && + user->opcode() == HloOpcode::kReducePrecision && + user->exponent_bits() == exponent_bits_ && + user->mantissa_bits() == mantissa_bits_) { + reduced = user; + break; + } + } + // If there wasn't an existing reduce-precision operation, create one. + if (!reduced) { + reduced = instruction->parent()->AddInstruction( + HloInstruction::CreateReducePrecision( + operand->shape(), operand, exponent_bits_, mantissa_bits_)); + } + // Insert the reduce-precision operation before the operand. + TF_RETURN_IF_ERROR(instruction->ReplaceOperandWith(i, reduced)); + computation_changed = true; + } + } + } + + return computation_changed; +} + +StatusOr ReducePrecisionInsertion::insert_on_outputs( + const std::vector& instructions) { + bool computation_changed = false; + for (const auto& instruction : instructions) { + VLOG(2) << "Adding reduce-precision operation to output of instruction: " + << instruction->ToString(); + + if (!is_valid_shape(instruction->shape())) { + VLOG(2) << "Skipped: value is not an F32 nonscalar array"; + continue; + } + + if (instruction->opcode() == HloOpcode::kFusion && + instruction->fusion_kind() == HloInstruction::FusionKind::kLoop) { + // Insert the reduce-precision operation as the last operation inside + // the fusion computation. + HloInstruction* fusion_root = instruction->fused_expression_root(); + VLOG(2) << "Inserting new operation after existing fusion root: " + << fusion_root->ToString(); + + TF_ASSIGN_OR_RETURN(bool instruction_changed, insert_after(fusion_root)); + computation_changed |= instruction_changed; + } else { + // Insert the reduce-precision operation after the instruction. + TF_ASSIGN_OR_RETURN(bool instruction_changed, insert_after(instruction)); + computation_changed |= instruction_changed; + } + } + + return computation_changed; +} + +StatusOr ReducePrecisionInsertion::Run(HloModule* module) { + bool changed = false; + VLOG(1) << "Running ReducePrecisionInsertion pass on " << module->name(); + + for (auto& computation : module->computations()) { + if (computation->IsFusionComputation()) { + continue; + } + + StatusOr computation_changed; + switch (location_) { + case HloReducePrecisionOptions::OP_INPUTS: + case HloReducePrecisionOptions::FUSION_INPUTS_BY_CONTENT: + computation_changed = ReducePrecisionInsertion::insert_on_inputs( + instructions_to_modify(computation.get())); + break; + + case HloReducePrecisionOptions::FUSION_OUTPUTS_BY_CONTENT: + case HloReducePrecisionOptions::OP_OUTPUTS: + case HloReducePrecisionOptions::UNFUSED_OP_OUTPUTS: + computation_changed = ReducePrecisionInsertion::insert_on_outputs( + instructions_to_modify(computation.get())); + break; + default: + break; + } + TF_RETURN_IF_ERROR(computation_changed.status()); + + if (computation_changed.ValueOrDie()) { + changed = true; + VLOG(3) << "Computation after reduce-precision insertion:"; + XLA_VLOG_LINES(3, computation->ToString()); + } else { + VLOG(3) << "Computation " << computation->name() << " unchanged"; + } + } + + return changed; +} + +ReducePrecisionInsertion::InstructionFilterFunction +ReducePrecisionInsertion::make_filter_function( + const HloReducePrecisionOptions& reduce_precision_options) { + // Implement the filter function with a lookup table. + std::vector opcode_filter(HloOpcodeCount(), false); + for (const auto& opcode : reduce_precision_options.opcodes_to_suffix()) { + opcode_filter[opcode] = true; + } + if (reduce_precision_options.opname_substrings_to_suffix_size() == 0) { + return [opcode_filter](const HloInstruction* instruction) { + return opcode_filter[static_cast(instruction->opcode())]; + }; + } else { + std::vector opname_substrings; + for (const auto& substring : + reduce_precision_options.opname_substrings_to_suffix()) { + opname_substrings.push_back(substring); + } + return [opcode_filter, + opname_substrings](const HloInstruction* instruction) { + if (!opcode_filter[static_cast(instruction->opcode())]) { + return false; + } + const auto& opname = instruction->metadata().op_name(); + for (const auto& substring : opname_substrings) { + if (opname.find(substring) != string::npos) { + return true; + } + } + return false; + }; + } +} + +HloReducePrecisionOptions ReducePrecisionInsertion::make_options_proto( + const HloReducePrecisionOptions::Location location, const int exponent_bits, + const int mantissa_bits, + const std::function& opcode_filter_function, + const std::vector& opname_substring_list) { + HloReducePrecisionOptions options; + options.set_location(location); + options.set_exponent_bits(exponent_bits); + options.set_mantissa_bits(mantissa_bits); + for (uint32_t opcode = 0; opcode < HloOpcodeCount(); opcode++) { + if (opcode_filter_function(static_cast(opcode))) { + options.add_opcodes_to_suffix(opcode); + } + } + for (auto& string : opname_substring_list) { + options.add_opname_substrings_to_suffix(string); + } + return options; +} + +bool ReducePrecisionInsertion::AddPasses(HloPassPipeline* pipeline, + const DebugOptions& debug_options, + const PassTiming pass_timing) { + bool passes_added = false; + for (const auto& pass_options : + debug_options.hlo_reduce_precision_options()) { + bool add_pass; + switch (pass_options.location()) { + case HloReducePrecisionOptions::OP_INPUTS: + case HloReducePrecisionOptions::OP_OUTPUTS: + add_pass = pass_timing == PassTiming::BEFORE_OPTIMIZATION; + break; + case HloReducePrecisionOptions::UNFUSED_OP_OUTPUTS: + case HloReducePrecisionOptions::FUSION_INPUTS_BY_CONTENT: + case HloReducePrecisionOptions::FUSION_OUTPUTS_BY_CONTENT: + add_pass = pass_timing == PassTiming::AFTER_FUSION; + break; + default: + add_pass = false; + } + if (add_pass) { + pipeline->AddPass(pass_options); + passes_added = true; + } + } + return passes_added; +} + +} // namespace xla diff --git a/tensorflow/compiler/xla/service/reduce_precision_insertion.h b/tensorflow/compiler/xla/service/reduce_precision_insertion.h new file mode 100644 index 0000000000000000000000000000000000000000..afde3cf95c721b59a39b74b4e1ff3f15a335fe97 --- /dev/null +++ b/tensorflow/compiler/xla/service/reduce_precision_insertion.h @@ -0,0 +1,154 @@ +/* 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_REDUCE_PRECISION_INSERTION_H_ +#define TENSORFLOW_COMPILER_XLA_SERVICE_REDUCE_PRECISION_INSERTION_H_ + +#include "tensorflow/compiler/xla/service/buffer_liveness.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_pass_interface.h" +#include "tensorflow/compiler/xla/service/hlo_pass_pipeline.h" +#include "tensorflow/core/lib/gtl/flatmap.h" + +namespace xla { + +// HLO pass which inserts reduce-precision instructions into the HLO graph, for +// purposes of experimenting with the effects of reduced-precision storage of +// intermediate values. +class ReducePrecisionInsertion : public HloPassInterface { + using InstructionFilterFunction = std::function; + + public: + // The exponent_bits and mantissa_bits arguments specify the parameters of + // the instructions to insert. The instructions will be inserted after each + // instruction with an opcode for which the instruction_filter_function + // function returns true and the output type is F32. + explicit ReducePrecisionInsertion( + const int exponent_bits, const int mantissa_bits, + const HloReducePrecisionOptions::Location location, + const InstructionFilterFunction& instruction_filter_function) + : exponent_bits_(exponent_bits), + mantissa_bits_(mantissa_bits), + location_(location), + instruction_filter_function_(instruction_filter_function) {} + + // Version of the constructor that takes an HloReducePrecisionOptions proto + // rather than explicitly-enumerated parameters, for convenience when + // creating passes based on DebugOptions. + explicit ReducePrecisionInsertion( + const HloReducePrecisionOptions& reduce_precision_options) + : exponent_bits_(reduce_precision_options.exponent_bits()), + mantissa_bits_(reduce_precision_options.mantissa_bits()), + location_(reduce_precision_options.location()), + instruction_filter_function_( + make_filter_function(reduce_precision_options)) {} + + ~ReducePrecisionInsertion() override{}; + + tensorflow::StringPiece name() const override { + return "reduce-precision-insertion"; + } + + // Run the pass on the given module. Returns whether the module was changed + // (reduce-precision instructions were inserted). + StatusOr Run(HloModule* module) override; + + // Convert between the (inconvenient) xla.proto HloReducePrecisionOptions + // representation and InstructionFilterFunction functions. + static InstructionFilterFunction make_filter_function( + const HloReducePrecisionOptions& reduce_precision_options); + static HloReducePrecisionOptions make_options_proto( + const HloReducePrecisionOptions::Location location, + const int exponent_bits, const int mantissa_bits, + const std::function& opcode_filter_function, + const std::vector& opname_substring_list = {}); + + // Enumeration to control which passes should be added. + enum class PassTiming { BEFORE_OPTIMIZATION, AFTER_FUSION }; + + // Add ReducePrecisionInsertion passes to an HloPassPipeline based on the list + // of HloReducePrecisionOptions in a DebugOptions proto. Returns true if any + // passes were added. + static bool AddPasses(HloPassPipeline* pipeline, + const DebugOptions& debug_options, + const PassTiming pass_timing); + + private: + // Select the instructions that should have reduce-precision operations + // attached to them. + std::vector instructions_to_modify( + const HloComputation* computation); + + // Insert a reduce-precision operation into the graph on the output of the + // given instruction. + StatusOr insert_after(HloInstruction* instruction); + + // Insert reduce-precision operations into the graph on the inputs of the + // given instructions. (For fusion instructions, the operations will be + // inserted inside the fusion computation, on the outputs of the relevant + // input parameters.) + StatusOr insert_on_inputs( + const std::vector& instructions); + + // Insert reduce-precision operations into the graph on the outputs of the + // given instructions. (For fusion instructions, the operations will be + // inserted inside the fusion computation as a new root.) + StatusOr insert_on_outputs( + const std::vector& instructions); + + // Is this shape valid for inserting a reduce-precision operation? + bool is_valid_shape(const Shape& shape) { + // For now, ReducePrecision is only implemented for F32 arrays, so this + // ignores instructions that produce other data. In particular, this + // currently ignores instructions producing tuples, even if those tuples + // contain F32 arrays inside them. The assumption is that in most cases + // 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); + } + + // Is this instruction one such that following or preceding it with a new + // reduce-precision operation will be redundant? + bool is_redundant(const HloInstruction* instruction) { + return instruction->opcode() == HloOpcode::kReducePrecision && + instruction->exponent_bits() <= exponent_bits_ && + instruction->mantissa_bits() <= mantissa_bits_; + } + + // Parameters for the precision reduction to be added. + const int exponent_bits_; + const int mantissa_bits_; + + // Pass "timing" parameter. This also controls aspects of how the pass + // selects locations to insert instructions. + const HloReducePrecisionOptions::Location location_; + + // User-provided Function to determine whether a given instruction should + // have a reduce-precision instruction inserted in its output stream. + const InstructionFilterFunction instruction_filter_function_; +}; + +} // namespace xla + +#endif // TENSORFLOW_COMPILER_XLA_SERVICE_REDUCE_PRECISION_INSERTION_H_ diff --git a/tensorflow/compiler/xla/service/reduce_precision_insertion_test.cc b/tensorflow/compiler/xla/service/reduce_precision_insertion_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..607abee33d9548269a42d0e4b2d1d7fd7c6ab858 --- /dev/null +++ b/tensorflow/compiler/xla/service/reduce_precision_insertion_test.cc @@ -0,0 +1,559 @@ +/* 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/reduce_precision_insertion.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/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" + +namespace op = xla::testing::opcode_matchers; + +namespace xla { + +using ::testing::UnorderedElementsAre; + +class ReducePrecisionInsertionTest : public HloTestBase { + protected: + bool InsertOps(HloModule* module, + const HloReducePrecisionOptions::Location location, + const std::function& filter) { + ReducePrecisionInsertion op_insertion(5, 10, location, filter); + StatusOr result = op_insertion.Run(module); + EXPECT_IS_OK(result.status()); + return result.ValueOrDie(); + } +}; + +TEST_F(ReducePrecisionInsertionTest, BeforeUnaryInstruction) { + auto builder = HloComputation::Builder(TestName()); + Shape shape = ShapeUtil::MakeShape(F32, {4}); + + // 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 = CreateNewModule(); + 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, BeforeBinaryInstruction) { + auto builder = HloComputation::Builder(TestName()); + Shape shape = ShapeUtil::MakeShape(F32, {4}); + + // Create a simple graph with parameter feeding a binary add function. + + HloInstruction* a = + builder.AddInstruction(HloInstruction::CreateParameter(0, shape, "a")); + HloInstruction* b = + builder.AddInstruction(HloInstruction::CreateParameter(1, shape, "b")); + HloInstruction* c = builder.AddInstruction( + HloInstruction::CreateBinary(shape, HloOpcode::kAdd, a, b)); + + auto module = CreateNewModule(); + auto computation = module->AddEntryComputation(builder.Build()); + + // Confirm expected state before adding ops. + EXPECT_EQ(computation->root_instruction(), c); + EXPECT_EQ(c->operand(0), a); + EXPECT_EQ(c->operand(1), b); + + EXPECT_TRUE(InsertOps(module.get(), HloReducePrecisionOptions::OP_INPUTS, + [](const HloInstruction* instruction) { + return instruction->opcode() == HloOpcode::kAdd; + })); + + // Confirm expected graph after adding ops. + EXPECT_EQ(computation->root_instruction(), c); + EXPECT_THAT(c->operand(0), op::ReducePrecision(a)); + EXPECT_THAT(c->operand(1), op::ReducePrecision(b)); +} + +TEST_F(ReducePrecisionInsertionTest, BeforeZeroInputInstruction) { + auto builder = HloComputation::Builder(TestName()); + Shape shape = ShapeUtil::MakeShape(F32, {4}); + + // 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 = CreateNewModule(); + 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_FALSE(InsertOps(module.get(), HloReducePrecisionOptions::OP_INPUTS, + [](const HloInstruction* instruction) { + return instruction->opcode() == + HloOpcode::kParameter; + })); + + // Confirm that graph has not changed. + EXPECT_EQ(computation->root_instruction(), b); + EXPECT_EQ(b->operand(0), a); +} + +TEST_F(ReducePrecisionInsertionTest, AvoidAddingDuplicateInstructions) { + auto builder = HloComputation::Builder(TestName()); + Shape shape = ShapeUtil::MakeShape(F32, {4}); + + // Create a simple graph with parameter feeding a binary add function. + + HloInstruction* a = + builder.AddInstruction(HloInstruction::CreateParameter(0, shape, "a")); + HloInstruction* b = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kCos, a)); + HloInstruction* c = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kSin, a)); + HloInstruction* d = builder.AddInstruction( + HloInstruction::CreateBinary(shape, HloOpcode::kAdd, b, c)); + + auto module = CreateNewModule(); + auto computation = module->AddEntryComputation(builder.Build()); + + // Confirm expected state before adding ops. + EXPECT_EQ(computation->root_instruction(), d); + EXPECT_EQ(b->operand(0), a); + EXPECT_EQ(c->operand(0), a); + + EXPECT_TRUE(InsertOps(module.get(), HloReducePrecisionOptions::OP_INPUTS, + [](const HloInstruction* instruction) { + return instruction->opcode() == HloOpcode::kCos || + instruction->opcode() == HloOpcode::kSin; + })); + + // Confirm expected graph after adding ops. In particular, we want to confirm + // that the reduced-precision operation added for the input to b is re-used + // for the input to c. + EXPECT_THAT(b->operand(0), op::ReducePrecision(a)); + EXPECT_THAT(c->operand(0), op::ReducePrecision(a)); + EXPECT_EQ(b->operand(0), c->operand(0)); +} + +TEST_F(ReducePrecisionInsertionTest, AfterRootInstruction) { + auto builder = HloComputation::Builder(TestName()); + Shape shape = ShapeUtil::MakeShape(F32, {4}); + + // 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 = CreateNewModule(); + auto computation = module->AddEntryComputation(builder.Build()); + + // Confirm expected state before adding ops. + EXPECT_EQ(computation->root_instruction(), b); + + EXPECT_TRUE(InsertOps(module.get(), HloReducePrecisionOptions::OP_OUTPUTS, + [](const HloInstruction* instruction) { + return instruction->opcode() == HloOpcode::kCos; + })); + + // Confirm expected graph after adding ops. + EXPECT_THAT(computation->root_instruction(), op::ReducePrecision(b)); +} + +TEST_F(ReducePrecisionInsertionTest, AfterNonRootInstruction) { + auto builder = HloComputation::Builder(TestName()); + Shape shape = ShapeUtil::MakeShape(F32, {4}); + + // Create a graph with two parameters feeding into unary cosine functions, + // and the output of those feeds into an add function. Feeding the outputs + // from the suffixed cosine functions into a binary add function allows us to + // confirm that the separate operand streams are not crossed when the new + // instructions are inserted. + HloInstruction* a = + builder.AddInstruction(HloInstruction::CreateParameter(0, shape, "a")); + HloInstruction* a_cos = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kCos, a)); + + HloInstruction* b = + builder.AddInstruction(HloInstruction::CreateParameter(1, shape, "b")); + HloInstruction* b_cos = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kCos, b)); + + HloInstruction* c = builder.AddInstruction( + HloInstruction::CreateBinary(shape, HloOpcode::kAdd, a_cos, b_cos)); + + auto module = CreateNewModule(); + module->AddEntryComputation(builder.Build()); + + // Confirm expected graph before adding ops. + EXPECT_EQ(c->operand(0), a_cos); + EXPECT_EQ(c->operand(1), b_cos); + + EXPECT_TRUE(InsertOps(module.get(), HloReducePrecisionOptions::OP_OUTPUTS, + [](const HloInstruction* instruction) { + return instruction->opcode() == HloOpcode::kCos; + })); + + // Confirm expected graph after adding ops. + EXPECT_THAT(c->operand(0), op::ReducePrecision()); + EXPECT_EQ(c->operand(0)->operand(0), a_cos); + EXPECT_THAT(c->operand(1), op::ReducePrecision()); + EXPECT_EQ(c->operand(1)->operand(0), b_cos); +} + +TEST_F(ReducePrecisionInsertionTest, OutputIsNotFloat) { + auto builder = HloComputation::Builder(TestName()); + Shape shape = ShapeUtil::MakeShape(S32, {4}); + HloInstruction* x = + builder.AddInstruction(HloInstruction::CreateParameter(0, shape, "x")); + HloInstruction* y = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kCos, x)); + + auto module = CreateNewModule(); + auto computation = module->AddEntryComputation(builder.Build()); + + // Confirm expected graph before adding ops. + EXPECT_THAT(x->users(), UnorderedElementsAre(y)); + EXPECT_EQ(computation->root_instruction(), y); + + // Since none of the instructions produce F32 data, this should not change + // the graph. + EXPECT_FALSE( + InsertOps(module.get(), HloReducePrecisionOptions::OP_OUTPUTS, + [](const HloInstruction* instruction) { return true; })); + + // Confirm that graph has not changed. + EXPECT_THAT(x->users(), UnorderedElementsAre(y)); + EXPECT_EQ(computation->root_instruction(), y); +} + +TEST_F(ReducePrecisionInsertionTest, ShouldReduceOutputPrecisionIsFalse) { + auto builder = HloComputation::Builder(TestName()); + Shape shape = ShapeUtil::MakeShape(F32, {4}); + HloInstruction* x = + builder.AddInstruction(HloInstruction::CreateParameter(0, shape, "x")); + HloInstruction* y = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kCos, x)); + + auto module = CreateNewModule(); + auto computation = module->AddEntryComputation(builder.Build()); + + // Confirm expected graph before adding ops. + EXPECT_THAT(x->users(), UnorderedElementsAre(y)); + EXPECT_EQ(computation->root_instruction(), y); + + // Since none of the instructions match the should_reduce_output_precision + // function, this should not change the graph. + EXPECT_FALSE( + InsertOps(module.get(), HloReducePrecisionOptions::OP_OUTPUTS, + [](const HloInstruction* instruction) { return false; })); + + // Confirm that graph has not changed. + EXPECT_THAT(x->users(), UnorderedElementsAre(y)); + EXPECT_EQ(computation->root_instruction(), y); +} + +TEST_F(ReducePrecisionInsertionTest, InsertionIsNotRecursive) { + auto builder = HloComputation::Builder(TestName()); + Shape shape = ShapeUtil::MakeShape(F32, {4}); + HloInstruction* a = + builder.AddInstruction(HloInstruction::CreateParameter(0, shape, "a")); + HloInstruction* b = builder.AddInstruction( + HloInstruction::CreateReducePrecision(shape, a, 8, 23)); + + auto module = CreateNewModule(); + auto computation = module->AddEntryComputation(builder.Build()); + + // Confirm expected state before adding ops. + EXPECT_EQ(computation->root_instruction(), b); + + // This should insert a new ReducePrecision after the existing one, but + // should not then recurse by adding another after the just-inserted one. + EXPECT_TRUE(InsertOps(module.get(), HloReducePrecisionOptions::OP_OUTPUTS, + [](const HloInstruction* instruction) { + return instruction->opcode() == + HloOpcode::kReducePrecision; + })); + + // Confirm expected graph after adding ops. + EXPECT_THAT(computation->root_instruction(), op::ReducePrecision()); + EXPECT_EQ(computation->root_instruction()->operand(0), b); +} + +TEST_F(ReducePrecisionInsertionTest, SkipRedundantReducePrecisionAfter) { + auto builder = HloComputation::Builder(TestName()); + Shape shape = ShapeUtil::MakeShape(F32, {4}); + HloInstruction* x = + builder.AddInstruction(HloInstruction::CreateParameter(0, shape, "x")); + HloInstruction* y = builder.AddInstruction( + HloInstruction::CreateReducePrecision(shape, x, 5, 10)); + + auto module = CreateNewModule(); + auto computation = module->AddEntryComputation(builder.Build()); + + // Confirm expected graph before adding ops. + EXPECT_THAT(x->users(), UnorderedElementsAre(y)); + EXPECT_EQ(computation->root_instruction(), y); + + // Since the new reduce-precision operation would be redundant, this + // should not change the graph. + EXPECT_FALSE(InsertOps(module.get(), HloReducePrecisionOptions::OP_OUTPUTS, + [](const HloInstruction* instruction) { + return instruction->opcode() == + HloOpcode::kParameter; + })); + + // Confirm that graph has not changed. + EXPECT_THAT(x->users(), UnorderedElementsAre(y)); + EXPECT_EQ(computation->root_instruction(), y); +} + +TEST_F(ReducePrecisionInsertionTest, AddNonRedundantReducePrecision) { + auto builder = HloComputation::Builder(TestName()); + Shape shape = ShapeUtil::MakeShape(F32, {4}); + HloInstruction* x = + builder.AddInstruction(HloInstruction::CreateParameter(0, shape, "x")); + HloInstruction* y = builder.AddInstruction( + HloInstruction::CreateReducePrecision(shape, x, 8, 23)); + + auto module = CreateNewModule(); + auto computation = module->AddEntryComputation(builder.Build()); + + // Confirm expected graph before adding ops. + EXPECT_THAT(x->users(), UnorderedElementsAre(y)); + EXPECT_EQ(computation->root_instruction(), y); + + // Since the new reduce-precision operation is not the same as the existing + // one, this should add a new one. + EXPECT_TRUE(InsertOps(module.get(), HloReducePrecisionOptions::OP_OUTPUTS, + [](const HloInstruction* instruction) { + return instruction->opcode() == HloOpcode::kParameter; + })); + + // Confirm that graph is as expected. + EXPECT_EQ(computation->root_instruction(), y); + EXPECT_THAT(y->operand(0), op::ReducePrecision(x)); +} + +TEST_F(ReducePrecisionInsertionTest, IgnoreOpsInsideFusionNode) { + auto builder = HloComputation::Builder(TestName()); + Shape shape = ShapeUtil::MakeShape(F32, {4}); + HloInstruction* x = + builder.AddInstruction(HloInstruction::CreateParameter(0, shape, "x")); + HloInstruction* y = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kCos, x)); + auto module = CreateNewModule(); + auto computation = module->AddEntryComputation(builder.Build()); + + // Manually fuse the kCos operation into a fusion operation. + HloInstruction* z = computation->AddInstruction(HloInstruction::CreateFusion( + shape, HloInstruction::FusionKind::kLoop, y)); + EXPECT_IS_OK(computation->ReplaceUsesOfInstruction(y, z)); + EXPECT_IS_OK(computation->RemoveInstruction(y)); + + // Confirm expected graph before adding reduce-precision ops. + EXPECT_THAT(x->users(), UnorderedElementsAre(z)); + EXPECT_EQ(computation->root_instruction(), z); + HloInstruction* y_fused = z->fused_expression_root(); + EXPECT_EQ(y_fused->opcode(), HloOpcode::kCos); + + // The ReducePrecisionInsertion pass should not see inside the fusion + // operation, so this should not change the graph. + EXPECT_FALSE(InsertOps(module.get(), + HloReducePrecisionOptions::UNFUSED_OP_OUTPUTS, + [](const HloInstruction* instruction) { + return instruction->opcode() == HloOpcode::kCos; + })); + + // Confirm that graph has not changed. + EXPECT_THAT(x->users(), UnorderedElementsAre(z)); + EXPECT_EQ(computation->root_instruction(), z); + EXPECT_EQ(z->fused_expression_root(), y_fused); +} + +TEST_F(ReducePrecisionInsertionTest, OpGetsInsertedInHeadOfFusionNode) { + auto builder = HloComputation::Builder(TestName()); + Shape shape = ShapeUtil::MakeShape(F32, {4}); + HloInstruction* x = + builder.AddInstruction(HloInstruction::CreateParameter(0, shape, "x")); + HloInstruction* y = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kCos, x)); + auto module = CreateNewModule(); + auto computation = module->AddEntryComputation(builder.Build()); + + // Manually fuse the kCos operation into a fusion operation. + HloInstruction* z = computation->AddInstruction(HloInstruction::CreateFusion( + shape, HloInstruction::FusionKind::kLoop, y)); + EXPECT_IS_OK(computation->ReplaceUsesOfInstruction(y, z)); + EXPECT_IS_OK(computation->RemoveInstruction(y)); + + // Confirm expected graph before adding reduce-precision ops. + EXPECT_THAT(x->users(), UnorderedElementsAre(z)); + EXPECT_EQ(computation->root_instruction(), z); + HloInstruction* y_fused = z->fused_expression_root(); + EXPECT_EQ(y_fused->opcode(), HloOpcode::kCos); + z->CheckFusionInstruction(); + + // This should see that the fusion computation contains a kCos operation, + // and insert a new reduce-precision node at its input. + EXPECT_TRUE(InsertOps(module.get(), + HloReducePrecisionOptions::FUSION_INPUTS_BY_CONTENT, + [](const HloInstruction* instruction) { + return instruction->opcode() == HloOpcode::kCos; + })); + + // This should refuse to insert a second reduce-precision operation, as + // it would be redundant with the first. + EXPECT_FALSE(InsertOps(module.get(), + HloReducePrecisionOptions::FUSION_INPUTS_BY_CONTENT, + [](const HloInstruction* instruction) { + return instruction->opcode() == HloOpcode::kCos; + })); + + // Confirm that the top-level computation still only contains the fusion + // instruction, but that the fused computation now has a reduce-precision + // instruction inserted after its parameter instruction. + EXPECT_THAT(x->users(), UnorderedElementsAre(z)); + EXPECT_EQ(computation->root_instruction(), z); + EXPECT_THAT(z->fused_expression_root(), y_fused); + EXPECT_THAT(y_fused->operand(0), op::ReducePrecision(op::Parameter())); + z->CheckFusionInstruction(); +} + +TEST_F(ReducePrecisionInsertionTest, OpGetsInsertedInTailOfFusionNode) { + auto builder = HloComputation::Builder(TestName()); + Shape shape = ShapeUtil::MakeShape(F32, {4}); + HloInstruction* x = + builder.AddInstruction(HloInstruction::CreateParameter(0, shape, "x")); + HloInstruction* y = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kCos, x)); + auto module = CreateNewModule(); + auto computation = module->AddEntryComputation(builder.Build()); + + // Manually fuse the kCos operation into a fusion operation. + HloInstruction* z = computation->AddInstruction(HloInstruction::CreateFusion( + shape, HloInstruction::FusionKind::kLoop, y)); + EXPECT_IS_OK(computation->ReplaceUsesOfInstruction(y, z)); + EXPECT_IS_OK(computation->RemoveInstruction(y)); + z->CheckFusionInstruction(); + + // Confirm expected graph before adding reduce-precision ops. + EXPECT_THAT(x->users(), UnorderedElementsAre(z)); + EXPECT_EQ(computation->root_instruction(), z); + HloInstruction* y_fused = z->fused_expression_root(); + EXPECT_EQ(y_fused->opcode(), HloOpcode::kCos); + + // This should see that the fusion computation contains a kCos operation, + // and insert a new reduce-precision node at its root. + EXPECT_TRUE(InsertOps(module.get(), + HloReducePrecisionOptions::FUSION_OUTPUTS_BY_CONTENT, + [](const HloInstruction* instruction) { + return instruction->opcode() == HloOpcode::kCos; + })); + + // This should refuse to insert a second reduce-precision operation, as + // it would be redundant with the first. + EXPECT_FALSE(InsertOps(module.get(), + HloReducePrecisionOptions::FUSION_OUTPUTS_BY_CONTENT, + [](const HloInstruction* instruction) { + return instruction->opcode() == HloOpcode::kCos; + })); + + // Confirm that the top-level computation still only contains the fusion + // instruction, but that the fused computation now has a reduce-precision + // instruction inserted as its root. + EXPECT_THAT(x->users(), UnorderedElementsAre(z)); + EXPECT_EQ(computation->root_instruction(), z); + EXPECT_THAT(z->fused_expression_root(), op::ReducePrecision(y_fused)); + z->CheckFusionInstruction(); +} + +TEST_F(ReducePrecisionInsertionTest, MakeFilterFunctionNoSubstrings) { + auto builder = HloComputation::Builder(TestName()); + Shape shape = ShapeUtil::MakeShape(F32, {4}); + HloInstruction* a = + builder.AddInstruction(HloInstruction::CreateParameter(0, shape, "a")); + HloInstruction* b = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kCos, a)); + HloInstruction* c = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kSin, a)); + + auto options_proto = ReducePrecisionInsertion::make_options_proto( + HloReducePrecisionOptions::OP_OUTPUTS, 5, 10, + [](const HloOpcode opcode) { return opcode == HloOpcode::kCos; }); + + auto filter_function = + ReducePrecisionInsertion::make_filter_function(options_proto); + + EXPECT_TRUE(filter_function(b)); + EXPECT_FALSE(filter_function(c)); +} + +TEST_F(ReducePrecisionInsertionTest, MakeFilterFunctionWithSubstrings) { + auto builder = HloComputation::Builder(TestName()); + Shape shape = ShapeUtil::MakeShape(F32, {4}); + HloInstruction* a = + builder.AddInstruction(HloInstruction::CreateParameter(0, shape, "a")); + + HloInstruction* b = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kCos, a)); + OpMetadata b_metadata; + b_metadata.set_op_name("FlowTensor/foom"); + b->set_metadata(b_metadata); + + HloInstruction* c = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kCos, a)); + OpMetadata c_metadata; + c_metadata.set_op_name("FlowTensor/barn"); + c->set_metadata(c_metadata); + + auto options_proto = ReducePrecisionInsertion::make_options_proto( + HloReducePrecisionOptions::OP_OUTPUTS, 5, 10, + [](const HloOpcode opcode) { return opcode == HloOpcode::kCos; }, + {"foo", "baz"}); + + auto filter_function = + ReducePrecisionInsertion::make_filter_function(options_proto); + + EXPECT_TRUE(filter_function(b)); + EXPECT_FALSE(filter_function(c)); +} + +} // namespace xla + +int main(int argc, char** argv) { + return xla::ParseDebugOptionsFlagsAndRunTests(argc, argv); +} diff --git a/tensorflow/compiler/xla/service/reshape_mover.cc b/tensorflow/compiler/xla/service/reshape_mover.cc index b72ef95a6a7964aa1f41cd2ceef4cdee76e9f708..a480236cebd9b020436b495df24a25421cebf174 100644 --- a/tensorflow/compiler/xla/service/reshape_mover.cc +++ b/tensorflow/compiler/xla/service/reshape_mover.cc @@ -13,17 +13,8 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ -#include "tensorflow/compiler/xla/service/reshape_mover.h" - -#include -#include "tensorflow/compiler/xla/literal_util.h" -#include "tensorflow/compiler/xla/shape_util.h" -#include "tensorflow/compiler/xla/util.h" - -namespace xla { - -namespace { - +// Implementation note: +// // The general idea behind this pass is that we're converting from this: // %param.A = OldShape // %param.B = OldShape @@ -44,46 +35,76 @@ namespace { // only implicit scalar broadcast is on Pred, not on A or B. Since reshapes or // transposes to a scalar should be cheap, we simply never move them. -// Finds the first non-scalar operand of an instruction that is a reshape or -// transpose and returns the operand if it is found or nullptr if not found. -HloInstruction* FirstNonScalarReshapeOperand(const HloInstruction* hlo) { - for (HloInstruction* operand : hlo->operands()) { - if (!ShapeUtil::IsScalar(operand->shape()) && - (operand->opcode() == HloOpcode::kReshape || - operand->opcode() == HloOpcode::kTranspose)) { - return operand; - } +#include "tensorflow/compiler/xla/service/reshape_mover.h" + +#include +#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/util.h" +#include "tensorflow/core/lib/core/errors.h" + +namespace xla { + +namespace { + +// Checks if an instruction can change its shape simply by adjusting metadata. +// This is the case if it is: +// +// - an instruction does not have any producers like Constants +// or Rng instruction, or is a scalar. +// +// Or +// +// - an reshape/transpose instruction with an operand that can trivially change +// its shape. +bool InstructionCanTriviallyChangeShape(const HloInstruction* instruction) { + // Reshape/Transposes are only trivial if their operand is trivial. + if (instruction->opcode() == HloOpcode::kReshape || + instruction->opcode() == HloOpcode::kTranspose) { + CHECK_EQ(instruction->operand_count(), 1); + return InstructionCanTriviallyChangeShape(instruction->operand(0)); } - return nullptr; -} -// Check if an operand of an instruction can change its shape simply by -// adjusting metadata. This is the case if an operand does not have any -// producers like Constants or Rng instruction, or is a scalar. -bool OperandCanTrivallyChangeShape(const HloInstruction* instruction, - const HloInstruction* operand) { // Scalars can operate with any shape. - if (ShapeUtil::IsScalar(operand->shape())) { + if (ShapeUtil::IsScalar(instruction->shape())) { return true; } // A constant can trivially reshape the literal it holds. - if (operand->opcode() == HloOpcode::kConstant && - ShapeUtil::SameDimensions(operand->shape(), instruction->shape())) { + if (instruction->opcode() == HloOpcode::kConstant) { return true; } // An Rng instruction can be any shape as long as it has one user. Two copies // of the same Rng would be problematic if an Rng of a different shape would // produce random numbers in a different order. - if (operand->opcode() == HloOpcode::kRng && - ShapeUtil::SameDimensions(operand->shape(), instruction->shape()) && - operand->user_count() == 1) { + if (instruction->opcode() == HloOpcode::kRng && + instruction->user_count() == 1) { return true; } return false; } +// Finds the first non-scalar operand of an instruction that is a non-trivial +// reshape or transpose. Returns the operand if it is found or nullptr if not +// found. +HloInstruction* FirstNonScalarAndNonTrivialReshapeOperand( + const HloInstruction* hlo) { + for (HloInstruction* operand : hlo->operands()) { + if (!ShapeUtil::IsScalar(operand->shape()) && + ((operand->opcode() == HloOpcode::kReshape || + operand->opcode() == HloOpcode::kTranspose) && + !InstructionCanTriviallyChangeShape(operand->operand(0)))) { + VLOG(5) << "Found first non-scalar and non-trivial reshape operand of " + << hlo->ToStringNoMetadata() << ":\n\t" + << operand->ToStringNoMetadata(); + return operand; + } + } + return nullptr; +} + // Returns whether `a` and `b` are equivalent for the purposes of this pass. bool AreEquivalentReshapes(const HloInstruction* a, const HloInstruction* b) { if (a->opcode() != b->opcode() || @@ -108,138 +129,201 @@ bool AreEquivalentReshapes(const HloInstruction* a, const HloInstruction* b) { // metadata. bool IsElementwiseOfEquivalentReshapesOrTransposes( const HloInstruction* instruction) { - const std::vector& operands = instruction->operands(); + const auto& operands = instruction->operands(); HloInstruction* first_reshape_operand = - FirstNonScalarReshapeOperand(instruction); - // If there are no reshapes or transposes, then there is nothing to sink below - // the elementwise operation. + FirstNonScalarAndNonTrivialReshapeOperand(instruction); + // If there are no non-trivial reshapes or transposes, then there is nothing + // to sink below the elementwise operation. if (!first_reshape_operand) { return false; } - return (instruction->user_count() > 0 || - instruction == instruction->parent()->root_instruction()) && - instruction->IsElementwise() && !operands.empty() && - // Check whether all operands: - // 1. are all reshapes or transposes that have the same input and - // output shapes as all other reshaped or transposed operands. - // or - // 2. can be any shape like kConstant, kRng, and scalars. - std::all_of( - operands.begin(), operands.end(), - [instruction, - first_reshape_operand](const HloInstruction* operand) { - return AreEquivalentReshapes(first_reshape_operand, operand) || - OperandCanTrivallyChangeShape(instruction, operand); - }); -} + VLOG(3) << "** Checking whether instruction is an elementwise operation of " + "equivalent reshapes/transposes: " + << instruction->ToStringNoMetadata(); + bool result = (instruction->user_count() > 0 || + instruction == instruction->parent()->root_instruction()) && + instruction->IsElementwise() && !operands.empty(); -// Try to sink any reshape or transpose operands of `instruction` across it. We -// do so if `instruction` is elementwise and all operands are equivalent -// reshapes or transposes. -bool TrySinkReshapeOrTranspose(HloComputation* computation, - HloInstruction* instruction) { - if (IsElementwiseOfEquivalentReshapesOrTransposes(instruction)) { - std::vector operands = instruction->operands(); - HloInstruction* old_reshape = FirstNonScalarReshapeOperand(instruction); - CHECK(old_reshape != nullptr); - Shape new_elementwise_shape = old_reshape->operand(0)->shape(); - for (size_t i = 0; i < operands.size(); ++i) { - // All scalar operands remain as-is, even if they're reshape or transpose, - // to simplify handling wrt special scalar broadcast rules for ops like - // Select. Scalar reshapes should be cheap anyways. - if (ShapeUtil::IsScalar(operands[i]->shape())) { + // Check whether all operands: + // 0. Have the same dimensions as the output -- if not, it may be + // implicitly broadcast, which can confound the movement's + // correctness. + // + // And one of the following: + // 1. Are reshapes or transposes that have the same input and + // output shapes as all other reshaped or transposed operands. + // or + // 2. Are one of kConstant, kRng, and scalars that can change shape + // trivially, + if (result) { + for (auto& operand : operands) { + if (!ShapeUtil::SameDimensions(operand->shape(), instruction->shape())) { + VLOG(5) << "Operand shape differs from output shape; may be " + "implicitly broadcast, so preventing " + "movement\n\toperand: " + << operand->ToStringNoMetadata() + << "\n\tinstruction: " << instruction->ToStringNoMetadata(); + result = false; + break; + } + + if (AreEquivalentReshapes(first_reshape_operand, operand)) { + VLOG(5) << "Are equivalent reshapes:\n\tfirst_reshape_operand: " + << first_reshape_operand->ToStringNoMetadata() + << "\n\toperand: " << operand->ToStringNoMetadata(); continue; } - auto element_type = operands[i]->shape().element_type(); - switch (operands[i]->opcode()) { - case HloOpcode::kConstant: { - if (old_reshape->opcode() == HloOpcode::kReshape) { - operands[i] = instruction->parent()->AddInstruction( - HloInstruction::CreateReshape( - ShapeUtil::ChangeElementType(new_elementwise_shape, - element_type), - operands[i])); - } else { - CHECK_EQ(old_reshape->opcode(), HloOpcode::kTranspose); - std::vector inverse_permutation = - InversePermutation(old_reshape->dimensions()); - operands[i] = instruction->parent()->AddInstruction( - HloInstruction::CreateTranspose( - ShapeUtil::ChangeElementType(new_elementwise_shape, - element_type), - operands[i], inverse_permutation)); - } - break; - } - case HloOpcode::kRng: { - CHECK_EQ(operands[i]->user_count(), 1); + + if (InstructionCanTriviallyChangeShape(operand)) { + VLOG(5) << "Operand can trivially change shape: " + << operand->ToStringNoMetadata(); + continue; + } + + // TODO(someone): Look into supporting general ops for the operands as + // well. + VLOG(5) << "Operand is neither equalivant to the first Reshape operand" + "nor can trivially change shape: " + << operand->ToStringNoMetadata(); + result = false; + break; + } + } + + VLOG(3) << "ElementwiseOfEquivalentReshapesOrTransposes result for " + << instruction->ToStringNoMetadata() << ": " << result; + return result; +} + +// Try to sink any reshape or transpose operands of `instruction` across it. We +// do so if `instruction` is elementwise and all operands are either equivalent +// reshapes/transposes or are trivially reshapable. Note that no move is +// performend if there is no nontrivial reshapes/transposes. +StatusOr TrySinkReshapeOrTranspose(HloComputation* computation, + HloInstruction* instruction) { + if (!IsElementwiseOfEquivalentReshapesOrTransposes(instruction)) { + return false; + } + + HloInstruction* old_reshape = + FirstNonScalarAndNonTrivialReshapeOperand(instruction); + TF_RET_CHECK(old_reshape != nullptr); + Shape new_elementwise_shape = old_reshape->operand(0)->shape(); + + VLOG(3) << "** Trying to sink reshape or transpose: " + << instruction->ToStringNoMetadata() + << "\n\told reshape: " << old_reshape->ToStringNoMetadata() + << "\n\tnew elementwise shape: " + << ShapeUtil::HumanString(new_elementwise_shape); + + auto operands = instruction->operands(); + for (size_t i = 0; i < operands.size(); ++i) { + // All scalar operands remain as-is, even if they're reshape or transpose, + // to simplify handling wrt special scalar broadcast rules for ops like + // Select. Scalar reshapes should be cheap anyways. + if (ShapeUtil::IsScalar(operands[i]->shape())) { + continue; + } + PrimitiveType element_type = operands[i]->shape().element_type(); + switch (operands[i]->opcode()) { + case HloOpcode::kConstant: { + if (old_reshape->opcode() == HloOpcode::kReshape) { + VLOG(3) << "Creating reshape for kConstant operand " << i << ": " + << operands[i]->ToStringNoMetadata(); operands[i] = instruction->parent()->AddInstruction( - operands[i]->CloneWithNewOperands( + HloInstruction::CreateReshape( ShapeUtil::ChangeElementType(new_elementwise_shape, element_type), - operands[i]->operands())); - break; + operands[i])); + } else { + TF_RET_CHECK(old_reshape->opcode() == HloOpcode::kTranspose); + std::vector inverse_permutation = + InversePermutation(old_reshape->dimensions()); + operands[i] = instruction->parent()->AddInstruction( + HloInstruction::CreateTranspose( + ShapeUtil::ChangeElementType(new_elementwise_shape, + element_type), + operands[i], inverse_permutation)); } - case HloOpcode::kReshape: - case HloOpcode::kTranspose: - operands[i] = operands[i]->mutable_operand(0); - break; - default: - LOG(FATAL) << "Unexpected opcode while trying to sink reshapes or " - "transposes."; + break; } - } - if (HloOpcode::kFusion == instruction->opcode()) { - // Here we already know `instruction` is elementwise, and no operand is - // implicit broadcast as if it were the operands would not be equivalent - // reshapes, 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_elementwise_shape.dimensions(); - *shape->mutable_layout() = new_elementwise_shape.layout(); + case HloOpcode::kRng: { + CHECK_EQ(operands[i]->user_count(), 1); + operands[i] = instruction->parent()->AddInstruction( + operands[i]->CloneWithNewOperands( + ShapeUtil::ChangeElementType(new_elementwise_shape, + element_type), + operands[i]->operands())); + break; } - } - auto new_elementwise = - computation->AddInstruction(instruction->CloneWithNewOperands( - // `instruction` may change the element type, e.g., from - // operands[0] -> reshape -> convert (`instruction`) - // to - // operands[0] -> convert' -> reshape' - // - // In this case, convert' should have the same element type as - // `convert` and the same dimensions as operands[0]. - ShapeUtil::ChangeElementType(new_elementwise_shape, - instruction->shape().element_type()), - operands)); - std::unique_ptr new_reshape; - switch (old_reshape->opcode()) { case HloOpcode::kReshape: - new_reshape = HloInstruction::CreateReshape(instruction->shape(), - new_elementwise); - break; case HloOpcode::kTranspose: - new_reshape = HloInstruction::CreateTranspose( - instruction->shape(), new_elementwise, old_reshape->dimensions()); + operands[i] = operands[i]->mutable_operand(0); break; default: - LOG(FATAL) << "Bad opcode"; + LOG(FATAL) << "Unexpected opcode while trying to sink reshapes or " + "transposes."; } - TF_CHECK_OK(computation->ReplaceWithNewInstruction(instruction, - std::move(new_reshape))); - return true; } - return false; + if (HloOpcode::kFusion == instruction->opcode()) { + // Here we already know `instruction` is elementwise, and no operand is + // implicit broadcast as if it were the operands would not be equivalent + // reshapes, 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_elementwise_shape.dimensions(); + *shape->mutable_layout() = new_elementwise_shape.layout(); + } + } + HloInstruction* new_elementwise = + computation->AddInstruction(instruction->CloneWithNewOperands( + // `instruction` may change the element type, e.g., from + // operands[0] -> reshape -> convert (`instruction`) + // to + // operands[0] -> convert' -> reshape' + // + // In this case, convert' should have the same element type as + // `convert` and the same dimensions as operands[0]. + ShapeUtil::ChangeElementType(new_elementwise_shape, + instruction->shape().element_type()), + operands)); + + std::unique_ptr new_reshape; + switch (old_reshape->opcode()) { + case HloOpcode::kReshape: + VLOG(3) << "Creating new reshape for new elementwise op: " + << new_elementwise->ToStringNoMetadata(); + new_reshape = + HloInstruction::CreateReshape(instruction->shape(), new_elementwise); + break; + case HloOpcode::kTranspose: + new_reshape = HloInstruction::CreateTranspose( + instruction->shape(), new_elementwise, old_reshape->dimensions()); + break; + default: + LOG(FATAL) << "Bad opcode"; + } + TF_RETURN_IF_ERROR(computation->ReplaceWithNewInstruction( + instruction, std::move(new_reshape))); + return true; } } // namespace StatusOr ReshapeMover::Run(HloModule* module) { bool changed = false; - for (const auto& comp : module->computations()) { + std::vector computations; + for (auto& computation : module->computations()) { + if (computation->IsFusionComputation()) { + continue; + } + computations.push_back(computation.get()); + } + for (const auto& comp : computations) { for (HloInstruction* instruction : comp->MakeInstructionPostOrder()) { - if (TrySinkReshapeOrTranspose(comp.get(), instruction)) { - changed = true; - } + TF_ASSIGN_OR_RETURN(bool did_change, + TrySinkReshapeOrTranspose(comp, instruction)); + changed |= did_change; } } return changed; diff --git a/tensorflow/compiler/xla/service/reshape_mover_test.cc b/tensorflow/compiler/xla/service/reshape_mover_test.cc index 09a673ea809d5799e78f1f6da9983e88321eb6ea..1589d52a256df1914201c866859008c0f1df8a8f 100644 --- a/tensorflow/compiler/xla/service/reshape_mover_test.cc +++ b/tensorflow/compiler/xla/service/reshape_mover_test.cc @@ -20,14 +20,18 @@ limitations under the License. #include "tensorflow/compiler/xla/ptr_util.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/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/types.h" #include "tensorflow/compiler/xla/xla_data.pb.h" #include "tensorflow/core/lib/strings/str_util.h" +namespace op = xla::testing::opcode_matchers; + namespace xla { namespace { using ReshapeMoverTest = HloTestBase; @@ -43,14 +47,58 @@ TEST_F(ReshapeMoverTest, ReshapesWithDifferentInputShapesNotMoved) { builder.AddInstruction(HloInstruction::CreateReshape(root_shape, param0)); auto reshape1 = builder.AddInstruction(HloInstruction::CreateReshape(root_shape, param1)); - auto add = builder.AddInstruction(HloInstruction::CreateBinary( + builder.AddInstruction(HloInstruction::CreateBinary( root_shape, HloOpcode::kAdd, reshape0, reshape1)); - auto module = MakeUnique(TestName()); + auto module = CreateNewModule(); + auto computation = module->AddEntryComputation(builder.Build()); + + EXPECT_THAT(computation->root_instruction(), + op::Add(op::Reshape(param0), op::Reshape(param1))); + + EXPECT_FALSE(ReshapeMover().Run(module.get()).ValueOrDie()); + + EXPECT_THAT(computation->root_instruction(), + op::Add(op::Reshape(param0), op::Reshape(param1))); +} + +// For a graph that looks like: +// +// +- reshape0 - rng0 +// | +// +- const1 +// | +// add +// +// where rng0 has a different shape than reshape0. +// +// 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) { + HloComputation::Builder builder(TestName()); + auto root_shape = ShapeUtil::MakeShape(F32, {8, 7}); + auto rng0 = builder.AddInstruction( + HloInstruction::CreateRng(ShapeUtil::MakeShape(F32, {1, 8, 1, 7, 1}), + RandomDistribution::RNG_UNIFORM, {})); + auto reshape0 = + builder.AddInstruction(HloInstruction::CreateReshape(root_shape, rng0)); + + auto const1 = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateFromShape(root_shape))); + + builder.AddInstruction(HloInstruction::CreateBinary( + root_shape, HloOpcode::kAdd, reshape0, const1)); + + auto module = CreateNewModule(); auto computation = module->AddEntryComputation(builder.Build()); - EXPECT_EQ(add, computation->root_instruction()); + + EXPECT_THAT(computation->root_instruction(), + op::Add(op::Reshape(rng0), const1)); + EXPECT_FALSE(ReshapeMover().Run(module.get()).ValueOrDie()); - EXPECT_EQ(add, computation->root_instruction()); + + EXPECT_THAT(computation->root_instruction(), + op::Add(op::Reshape(rng0), const1)); } TEST_F(ReshapeMoverTest, ScalarReshapesNotMoved) { @@ -64,14 +112,20 @@ TEST_F(ReshapeMoverTest, ScalarReshapesNotMoved) { builder.AddInstruction(HloInstruction::CreateReshape(root_shape, param0)); auto reshape1 = builder.AddInstruction(HloInstruction::CreateReshape(root_shape, param1)); - auto add = builder.AddInstruction(HloInstruction::CreateBinary( + builder.AddInstruction(HloInstruction::CreateBinary( root_shape, HloOpcode::kAdd, reshape0, reshape1)); - auto module = MakeUnique(TestName()); + auto module = CreateNewModule(); auto computation = module->AddEntryComputation(builder.Build()); - EXPECT_EQ(add, computation->root_instruction()); + + EXPECT_THAT(computation->root_instruction(), + op::Add(op::Reshape(param0), op::Reshape(param1))); + EXPECT_FALSE(ReshapeMover().Run(module.get()).ValueOrDie()); - EXPECT_EQ(add, computation->root_instruction()); + + EXPECT_THAT( + computation->root_instruction(), + op::Add(op::Reshape(op::Parameter()), op::Reshape(op::Parameter()))); } TEST_F(ReshapeMoverTest, EquivalentReshapesMoved) { @@ -85,41 +139,202 @@ TEST_F(ReshapeMoverTest, EquivalentReshapesMoved) { builder.AddInstruction(HloInstruction::CreateReshape(root_shape, param0)); auto reshape1 = builder.AddInstruction(HloInstruction::CreateReshape(root_shape, param1)); - auto add = builder.AddInstruction(HloInstruction::CreateBinary( + builder.AddInstruction(HloInstruction::CreateBinary( root_shape, HloOpcode::kAdd, reshape0, reshape1)); - auto module = MakeUnique(TestName()); + auto module = CreateNewModule(); + auto computation = module->AddEntryComputation(builder.Build()); + + EXPECT_THAT(computation->root_instruction(), + op::Add(op::Reshape(param0), op::Reshape(param1))); + EXPECT_TRUE(ReshapeMover().Run(module.get()).ValueOrDie()); + + EXPECT_THAT(computation->root_instruction(), + op::Reshape(op::Add(param0, param1))); + EXPECT_EQ(root_shape.DebugString(), + computation->root_instruction()->shape().DebugString()); +} + +// For a graph that looks like: +// +// +- reshape2 - param2 +// | +// +- reshape1 - param1 +// | +// +- constant0 +// | +// select +// +// Verifies that the reshape1 and reshape2 sink past select: +// +// +- param2 +// | +// +- param1 +// | +// +- reshape3(constant0) +// | +// select +// | +// reshape4 +TEST_F(ReshapeMoverTest, 1ConstantAnd2ReshapesMoved) { + HloComputation::Builder builder(TestName()); + auto root_shape = ShapeUtil::MakeShape(F32, {2, 3}); + auto const0 = builder.AddInstruction(HloInstruction::CreateConstant( + Literal::CreateR2({{true, true, false}, {false, false, true}}))); + + auto param1 = builder.AddInstruction(HloInstruction::CreateParameter( + 0, ShapeUtil::MakeShape(F32, {1, 3, 1, 2}), "param1")); + auto reshape1 = + builder.AddInstruction(HloInstruction::CreateReshape(root_shape, param1)); + + auto param2 = builder.AddInstruction(HloInstruction::CreateParameter( + 1, ShapeUtil::MakeShape(F32, {1, 3, 1, 2}), "param2")); + auto reshape2 = + builder.AddInstruction(HloInstruction::CreateReshape(root_shape, param2)); + + builder.AddInstruction(HloInstruction::CreateTernary( + ShapeUtil::MakeShape(PRED, {2, 3}), HloOpcode::kSelect, const0, reshape1, + reshape2)); + + auto module = CreateNewModule(); auto computation = module->AddEntryComputation(builder.Build()); - EXPECT_EQ(add, computation->root_instruction()); + + EXPECT_THAT(computation->root_instruction(), + op::Select(const0, reshape1, reshape2)); + EXPECT_TRUE(ReshapeMover().Run(module.get()).ValueOrDie()); - auto new_root = computation->root_instruction(); - EXPECT_NE(add, new_root); - EXPECT_EQ(HloOpcode::kReshape, new_root->opcode()); - EXPECT_EQ(root_shape.DebugString(), new_root->shape().DebugString()); + EXPECT_THAT(computation->root_instruction(), + op::Reshape(op::Select(op::Reshape(const0), param1, param2))); + + EXPECT_EQ(const0->shape().DebugString(), + computation->root_instruction()->shape().DebugString()); +} + +// For a graph that looks like: +// +// +- reshape0 - param0 +// | +// +- param1 +// | +// add +// +// Verifies that the reshape0 does not sink below add, because param1 is not +// trivially reshapable nor is a Reshape/Transpose. +TEST_F(ReshapeMoverTest, 1ParameterAnd1ReshapeNotMoved) { + HloComputation::Builder builder(TestName()); + auto root_shape = ShapeUtil::MakeShape(F32, {8, 7}); + auto param0 = builder.AddInstruction(HloInstruction::CreateParameter( + 0, ShapeUtil::MakeShape(F32, {1, 8, 1, 7}), "param0")); + auto reshape0 = + builder.AddInstruction(HloInstruction::CreateReshape(root_shape, param0)); + auto param1 = builder.AddInstruction(HloInstruction::CreateParameter( + 1, ShapeUtil::MakeShape(F32, {1, 8, 1, 7}), "param1")); + builder.AddInstruction(HloInstruction::CreateBinary( + root_shape, HloOpcode::kAdd, reshape0, param1)); + + auto module = CreateNewModule(); + auto computation = module->AddEntryComputation(builder.Build()); + + EXPECT_THAT(computation->root_instruction(), + op::Add(op::Reshape(param0), param1)); + EXPECT_FALSE(ReshapeMover().Run(module.get()).ValueOrDie()); + + EXPECT_THAT(computation->root_instruction(), + op::Add(op::Reshape(param0), param1)); + EXPECT_EQ(root_shape.DebugString(), + computation->root_instruction()->shape().DebugString()); +} + +// For a graph that looks like: +// +// +- pred +// | +// +- reshape0 - const0 +// | +// +- reshape1 - const1 +// | +// select +// +// Verifies that we don't unnecessarily sink reshapes, which are in fact +// trivial reshapes. +TEST_F(ReshapeMoverTest, 2TrivialConstantReshapeNotMoved) { + HloComputation::Builder builder(TestName()); + auto root_shape = ShapeUtil::MakeShape(F32, {2, 3}); + auto const0 = builder.AddInstruction(HloInstruction::CreateConstant( + Literal::CreateR2({{1, 2, 3}, {4, 5, 6}}))); + auto reshape0 = + builder.AddInstruction(HloInstruction::CreateReshape(root_shape, const0)); + + auto const1 = builder.AddInstruction(HloInstruction::CreateConstant( + Literal::CreateR2({{1, 2, 3}, {4, 5, 6}}))); + auto reshape1 = + builder.AddInstruction(HloInstruction::CreateReshape(root_shape, const1)); + + auto pred = builder.AddInstruction(HloInstruction::CreateParameter( + 0, ShapeUtil::MakeShape(PRED, {1, 3, 1, 2}), "pred")); + + builder.AddInstruction(HloInstruction::CreateTernary( + root_shape, HloOpcode::kSelect, pred, reshape0, reshape1)); + + auto module = CreateNewModule(); + auto computation = module->AddEntryComputation(builder.Build()); + + EXPECT_THAT(computation->root_instruction(), + op::Select(pred, op::Reshape(const0), op::Reshape(const1))); + + EXPECT_FALSE(ReshapeMover().Run(module.get()).ValueOrDie()); + + EXPECT_THAT(computation->root_instruction(), + op::Select(pred, op::Reshape(const0), op::Reshape(const1))); + EXPECT_EQ(root_shape.DebugString(), + computation->root_instruction()->shape().DebugString()); } -TEST_F(ReshapeMoverTest, ConstantAndReshapeMoved) { +// For a graph that looks like: +// +// +- reshape0 - param0 +// | +// +- const1 +// | +// add +// +// where there is only 1 non-trivial reshape (reshape0), we sink the reshape +// here for canonicalization benefit: +// +// +- param0 +// | +// +- reshape1 - const1 +// | +// add +// | +// reshape2 +// +// (note that reshape1 here is trivial). +TEST_F(ReshapeMoverTest, 1NonTrivialReshapeMoved) { HloComputation::Builder builder(TestName()); auto root_shape = ShapeUtil::MakeShape(F32, {2, 3}); auto param0 = builder.AddInstruction(HloInstruction::CreateParameter( 0, ShapeUtil::MakeShape(F32, {1, 3, 1, 2}), "param0")); auto const1 = builder.AddInstruction(HloInstruction::CreateConstant( - LiteralUtil::CreateR2({{1, 2, 3}, {4, 5, 6}}))); + Literal::CreateR2({{1, 2, 3}, {4, 5, 6}}))); auto reshape0 = builder.AddInstruction(HloInstruction::CreateReshape(root_shape, param0)); - auto add = builder.AddInstruction(HloInstruction::CreateBinary( + builder.AddInstruction(HloInstruction::CreateBinary( root_shape, HloOpcode::kAdd, reshape0, const1)); - auto module = MakeUnique(TestName()); + auto module = CreateNewModule(); auto computation = module->AddEntryComputation(builder.Build()); - EXPECT_EQ(add, computation->root_instruction()); + + EXPECT_THAT(computation->root_instruction(), + op::Add(op::Reshape(param0), const1)); + EXPECT_TRUE(ReshapeMover().Run(module.get()).ValueOrDie()); - auto new_root = computation->root_instruction(); - EXPECT_NE(add, new_root); - EXPECT_EQ(HloOpcode::kReshape, new_root->opcode()); - EXPECT_EQ(root_shape.DebugString(), new_root->shape().DebugString()); + EXPECT_THAT(computation->root_instruction(), + op::Reshape(op::Add(param0, op::Reshape(const1)))); + EXPECT_EQ(root_shape.DebugString(), + computation->root_instruction()->shape().DebugString()); } TEST_F(ReshapeMoverTest, EquivalentReshapesMovedAcrossFusion) { @@ -136,18 +351,20 @@ TEST_F(ReshapeMoverTest, EquivalentReshapesMovedAcrossFusion) { auto add = builder.AddInstruction(HloInstruction::CreateBinary( root_shape, HloOpcode::kAdd, reshape0, reshape1)); - auto module = MakeUnique(TestName()); - auto computation = module->AddEntryComputation(builder.Build()); - auto fusion = computation->AddInstruction(HloInstruction::CreateFusion( - add->shape(), HloInstruction::FusionKind::kLoop, add)); - TF_CHECK_OK(computation->ReplaceInstruction(add, fusion)); - EXPECT_EQ(fusion, computation->root_instruction()); - EXPECT_TRUE(ReshapeMover().Run(module.get()).ValueOrDie()); + HloModule module(TestName()); + auto computation = module.AddEntryComputation(builder.Build()); + computation->CreateFusionInstruction({add}, + HloInstruction::FusionKind::kLoop); + + EXPECT_THAT(computation->root_instruction(), + op::Fusion(op::Reshape(param0), op::Reshape(param1))); - auto new_root = computation->root_instruction(); - EXPECT_NE(fusion, new_root); - EXPECT_EQ(HloOpcode::kReshape, new_root->opcode()); - EXPECT_EQ(root_shape.DebugString(), new_root->shape().DebugString()); + EXPECT_TRUE(ReshapeMover().Run(&module).ValueOrDie()); + + EXPECT_THAT(computation->root_instruction(), + op::Reshape(op::Fusion(param0, param1))); + EXPECT_EQ(root_shape.DebugString(), + computation->root_instruction()->shape().DebugString()); } TEST_F(ReshapeMoverTest, EquivalentReshapesMovedAcrossSelect) { @@ -166,18 +383,22 @@ TEST_F(ReshapeMoverTest, EquivalentReshapesMovedAcrossSelect) { builder.AddInstruction(HloInstruction::CreateReshape(root_shape, param1)); auto reshape_pred = builder.AddInstruction(HloInstruction::CreateReshape(pred_shape, pred)); - auto select = builder.AddInstruction(HloInstruction::CreateTernary( + builder.AddInstruction(HloInstruction::CreateTernary( root_shape, HloOpcode::kSelect, reshape_pred, reshape0, reshape1)); - auto module = MakeUnique(TestName()); + auto module = CreateNewModule(); auto computation = module->AddEntryComputation(builder.Build()); - EXPECT_EQ(select, computation->root_instruction()); + + EXPECT_THAT( + computation->root_instruction(), + op::Select(op::Reshape(pred), op::Reshape(param0), op::Reshape(param1))); + EXPECT_TRUE(ReshapeMover().Run(module.get()).ValueOrDie()); - auto new_root = computation->root_instruction(); - EXPECT_NE(select, new_root); - EXPECT_EQ(HloOpcode::kReshape, new_root->opcode()); - EXPECT_EQ(root_shape.DebugString(), new_root->shape().DebugString()); + EXPECT_THAT(computation->root_instruction(), + op::Reshape(op::Select(pred, param0, param1))); + EXPECT_EQ(root_shape.DebugString(), + computation->root_instruction()->shape().DebugString()); } TEST_F(ReshapeMoverTest, ScalarReshapeNotMovedAcrossSelect) { @@ -195,13 +416,70 @@ TEST_F(ReshapeMoverTest, ScalarReshapeNotMovedAcrossSelect) { auto select = builder.AddInstruction(HloInstruction::CreateTernary( root_shape, HloOpcode::kSelect, reshape_pred, param0, param1)); - auto module = MakeUnique(TestName()); + auto module = CreateNewModule(); auto computation = module->AddEntryComputation(builder.Build()); - EXPECT_EQ(select, computation->root_instruction()); + EXPECT_THAT(computation->root_instruction(), + op::Select(op::Reshape(pred), param0, param1)); + EXPECT_FALSE(ReshapeMover().Run(module.get()).ValueOrDie()); + + EXPECT_THAT(computation->root_instruction(), + op::Select(op::Reshape(pred), param0, param1)); EXPECT_EQ(select, computation->root_instruction()); } +// Tree looks like: +// +// param0 [1,128,1] +// | +// reshape [128,1] constant [128,1024] +// \ / +// multiply w/implicit broadcast [128,1024] +// +// The reshape mover would like to sink the reshape below the multiply. +// +// Previously we would attempt to insert a reshape of the constant to [1,128,1] +// (which is unsound, because it has a different number of elements) as +// preparation for sinking the reshape. +// +// To eliminate the unsoundness, we outlaw reshape sinking when one of the +// operands is implicitly broadcast in the elementwise consumer. +// +// TODO(b/37799338) However, it would be possible in this case to do a more +// in-depth analysis to get reshape movement to occur: +// +// 1. Note that the broadcast dimension (logical dimension 1) in the operands +// would map back to logical dimension 2 in the param0 node. +// 2. Match rank of the constant to the param0 node (by prepending a trivial 1 +// dimension). +// 3. Reshape to [128,1024] at the root. +// +// But this is not currently done. +TEST_F(ReshapeMoverTest, ImplicitlyBroadcastReshapeIsNotMovedBug37787999) { + HloComputation::Builder builder(TestName()); + auto param0 = builder.AddInstruction(HloInstruction::CreateParameter( + 0, ShapeUtil::MakeShape(F32, {1, 128, 1}), "param0")); + auto reshape = builder.AddInstruction(HloInstruction::CreateReshape( + ShapeUtil::MakeShape(F32, {128, 1}), param0)); + Array2D a(128, 1024); + auto literal = Literal::CreateR2FromArray2D(a); + auto constant = builder.AddInstruction( + HloInstruction::CreateConstant(std::move(literal))); + auto multiply = builder.AddInstruction(HloInstruction::CreateBinary( + constant->shape(), HloOpcode::kMultiply, constant, reshape)); + + auto module = CreateNewModule(); + auto computation = module->AddEntryComputation(builder.Build()); + EXPECT_THAT(computation->root_instruction(), + op::Multiply(op::Constant(), op::Reshape(param0))); + + EXPECT_FALSE(ReshapeMover().Run(module.get()).ValueOrDie()); + + EXPECT_THAT(computation->root_instruction(), + op::Multiply(op::Constant(), op::Reshape(param0))); + EXPECT_EQ(multiply, computation->root_instruction()); +} + // Tree looks like this: // // add1 @@ -236,22 +514,27 @@ TEST_F(ReshapeMoverTest, MultiplePasses) { builder.AddInstruction(HloInstruction::CreateReshape(shape3, param2)); auto reshape3 = builder.AddInstruction(HloInstruction::CreateReshape(shape3, add0)); - auto add1 = builder.AddInstruction(HloInstruction::CreateBinary( - shape3, HloOpcode::kAdd, reshape2, reshape3)); + builder.AddInstruction(HloInstruction::CreateBinary(shape3, HloOpcode::kAdd, + reshape2, reshape3)); - auto module = MakeUnique(TestName()); + auto module = CreateNewModule(); auto computation = module->AddEntryComputation(builder.Build()); - EXPECT_EQ(add1, computation->root_instruction()); + + 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.get()).ValueOrDie()); - EXPECT_EQ(HloOpcode::kReshape, computation->root_instruction()->opcode()); - EXPECT_EQ(HloOpcode::kAdd, - computation->root_instruction()->operand(0)->opcode()); - const auto& add_params = - computation->root_instruction()->operand(0)->operands(); - EXPECT_EQ(2, add_params.size()); - EXPECT_EQ(HloOpcode::kParameter, add_params[0]->opcode()); - EXPECT_EQ(HloOpcode::kReshape, add_params[1]->opcode()); + + EXPECT_THAT( + computation->root_instruction(), + op::Reshape(op::Add(param2, op::Reshape(op::Add(param0, param1))))); } } // namespace } // namespace xla + +int main(int argc, char** argv) { + return xla::ParseDebugOptionsFlagsAndRunTests(argc, argv); +} diff --git a/tensorflow/compiler/xla/service/service.cc b/tensorflow/compiler/xla/service/service.cc index 451bb8c7eadf3e2210788a722d8f75aa3050e30f..d63d33ecb00348dd79553cd93dfce17be6b534f5 100644 --- a/tensorflow/compiler/xla/service/service.cc +++ b/tensorflow/compiler/xla/service/service.cc @@ -20,8 +20,9 @@ limitations under the License. #include #include +#include "tensorflow/compiler/xla/execution_options_util.h" #include "tensorflow/compiler/xla/layout_util.h" -#include "tensorflow/compiler/xla/legacy_flags/service_flags.h" +#include "tensorflow/compiler/xla/legacy_flags/debug_options_flags.h" #include "tensorflow/compiler/xla/ptr_util.h" #include "tensorflow/compiler/xla/service/compiler.h" #include "tensorflow/compiler/xla/service/computation_layout.h" @@ -29,7 +30,7 @@ limitations under the License. #include "tensorflow/compiler/xla/service/executable.h" #include "tensorflow/compiler/xla/service/hlo_computation.h" #include "tensorflow/compiler/xla/service/hlo_cost_analysis.h" -#include "tensorflow/compiler/xla/service/hlo_graph_dumper.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_module_config.h" @@ -47,7 +48,6 @@ limitations under the License. #include "tensorflow/core/platform/env.h" #include "tensorflow/core/platform/logging.h" #include "tensorflow/core/platform/protobuf.h" -#include "tensorflow/core/platform/regexp.h" #include "tensorflow/core/platform/stream_executor_no_cuda.h" #include "tensorflow/core/platform/types.h" @@ -79,8 +79,10 @@ tensorflow::Status RecordArguments( SessionModule* module) { module->clear_arguments(); for (const Allocation* allocation : arg_allocations) { - TF_RETURN_IF_ERROR(LiteralFromAllocation(allocation, allocation->shape(), - module->add_arguments())); + Literal argument; + TF_RETURN_IF_ERROR( + LiteralFromAllocation(allocation, allocation->shape(), &argument)); + *module->add_arguments() = argument.ToProto(); } return tensorflow::Status::OK(); } @@ -89,8 +91,11 @@ tensorflow::Status RecordArguments( tensorflow::Status RecordResult(const Allocation* result_allocation, SessionModule* module) { module->clear_result(); - return LiteralFromAllocation(result_allocation, result_allocation->shape(), - module->mutable_result()); + Literal result; + TF_RETURN_IF_ERROR(LiteralFromAllocation( + result_allocation, result_allocation->shape(), &result)); + *module->mutable_result() = result.ToProto(); + return tensorflow::Status::OK(); } } // namespace @@ -112,6 +117,16 @@ ServiceOptions& ServiceOptions::set_number_of_replicas(int number_of_replicas) { int ServiceOptions::number_of_replicas() const { return number_of_replicas_; } +ServiceOptions& ServiceOptions::set_intra_op_parallelism_threads( + int num_threads) { + intra_op_parallelism_threads_ = num_threads; + return *this; +} + +int ServiceOptions::intra_op_parallelism_threads() const { + return intra_op_parallelism_threads_; +} + /* static */ StatusOr> Service::NewService( perftools::gputools::Platform* platform) { ServiceOptions default_options; @@ -126,74 +141,41 @@ int ServiceOptions::number_of_replicas() const { return number_of_replicas_; } if (platform == nullptr) { TF_ASSIGN_OR_RETURN(platform, PlatformUtil::GetDefaultPlatform()); } - TF_ASSIGN_OR_RETURN( - execute_backend, - Backend::CreateBackend(platform, options.number_of_replicas())); - TF_ASSIGN_OR_RETURN(std::unique_ptr compute_constant_backend, - CreateComputeConstantBackend()); - std::unique_ptr service(new Service( - std::move(execute_backend), std::move(compute_constant_backend))); + BackendOptions backend_options; + backend_options.set_platform(platform); + TF_ASSIGN_OR_RETURN(execute_backend, Backend::CreateBackend(backend_options)); + + std::unique_ptr service( + new Service(options, std::move(execute_backend))); return std::move(service); } -/* static */ StatusOr> -Service::CreateComputeConstantBackend() { - TF_ASSIGN_OR_RETURN(std::vector platforms, - PlatformUtil::GetSupportedPlatforms()); - for (auto* platform : platforms) { - if (platform->id() == se::host::kHostPlatformId) { - return Backend::CreateBackend(platform, /*replica_count=*/1); +Service::Service(const ServiceOptions& options, + std::unique_ptr execute_backend) + : options_(options), execute_backend_(std::move(execute_backend)) { + CHECK(options_.number_of_replicas() > 0); + if (execute_backend_) { + if (execute_backend_->device_count() > 0) { + CHECK_GE(execute_backend_->device_count(), options_.number_of_replicas()) + << "Requested more replicas than there are devices."; } - } - return NotFound("CPU platform not found"); -} - -/* static */ void Service::DumpExecutedHlo(const HloModule& module, - const string& label, - const HloExecutionProfile* profile) { - VLOG(2) << "module name = " << module.name(); - legacy_flags::ServiceFlags* flags = legacy_flags::GetServiceFlags(); - if (!flags->xla_generate_hlo_graph.empty() && - RE2::PartialMatch(module.name(), flags->xla_generate_hlo_graph)) { - hlo_graph_dumper::DumpGraph(*module.entry_computation(), label, - flags->xla_hlo_graph_addresses, - flags->xla_hlo_graph_layout, profile); - } - if (!flags->xla_log_hlo_text.empty() && - RE2::PartialMatch(module.name(), flags->xla_log_hlo_text)) { - LOG(INFO) << "HLO for module " << module.name(); - LOG(INFO) << "Label: " << label; - XLA_LOG_LINES(2, module.ToString()); - } - if (!flags->xla_dump_hlo_text_to.empty()) { - hlo_graph_dumper::DumpText(module, label, flags->xla_dump_hlo_text_to); - } -} - -/* static */ Compiler::HloDumper Service::MakeHloDumper() { - return [](const HloModule& module, const string& label) { - return DumpExecutedHlo(module, label, /*profile=*/nullptr); - }; -} - -Service::Service(std::unique_ptr execute_backend, - std::unique_ptr compute_constant_backend) - : execute_backend_(std::move(execute_backend)), - compute_constant_backend_(std::move(compute_constant_backend)) { - LOG(INFO) << Printf( - "XLA service %p executing computations on platform %s. Devices:", this, - execute_backend_->platform()->Name().c_str()); - 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) << Printf(" StreamExecutor device (%d): %s, %s", i, - description.name().c_str(), - description.platform_version().c_str()); - } else { - LOG(INFO) << Printf(" StreamExecutor device (%d) not supported", i); + LOG(INFO) << Printf( + "XLA service %p executing computations on platform %s. Devices:", this, + execute_backend_->platform()->Name().c_str()); + 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) << Printf(" StreamExecutor device (%d): %s, %s", i, + description.name().c_str(), + description.platform_version().c_str()); + } else { + LOG(INFO) << Printf(" StreamExecutor device (%d) not supported", i); + } } + } else { + VLOG(1) << "XLA compile-only service constructed"; } } @@ -285,52 +267,71 @@ StatusOr> Service::ResolveAndValidateArguments( StatusOr> Service::CreateModuleConfig( const ProgramShape& program_shape, - tensorflow::gtl::ArraySlice arguments, - const ExecutionOptions& execution_options) { - auto module_config = MakeUnique(program_shape); - auto* computation_layout = module_config->mutable_entry_computation_layout(); + tensorflow::gtl::ArraySlice argument_shapes, + const ExecutionOptions* execution_options, bool has_hybrid_result) { + auto config = MakeUnique(program_shape); + auto* computation_layout = config->mutable_entry_computation_layout(); - if (program_shape.parameters_size() != arguments.size()) { + if (program_shape.parameters_size() != argument_shapes.size()) { return InvalidArgument("computation takes %d parameters, but %zu given", - program_shape.parameters_size(), arguments.size()); + program_shape.parameters_size(), + argument_shapes.size()); } - - for (size_t i = 0; i < arguments.size(); ++i) { + for (int i = 0; i < argument_shapes.size(); ++i) { // Verify that shape of arguments matches the shape of the arguments in the // ProgramShape. - if (!ShapeUtil::Compatible(arguments[i]->shape(), + if (!ShapeUtil::Compatible(*argument_shapes[i], program_shape.parameters(i))) { return InvalidArgument( - "computation expects parameter %lu to have shape %s, given shape %s", + "computation expects parameter %d to have shape %s, given shape %s", i, ShapeUtil::HumanString(program_shape.parameters(i)).c_str(), - ShapeUtil::HumanString(arguments[i]->shape()).c_str()); + ShapeUtil::HumanString(*argument_shapes[i]).c_str()); } TF_RETURN_IF_ERROR( computation_layout->mutable_parameter_layout(i)->CopyLayoutFromShape( - arguments[i]->shape())); + *argument_shapes[i])); } - if (!execution_options.has_shape_with_output_layout()) { - computation_layout->mutable_result_layout()->Clear(); - } else { + if (execution_options != nullptr && + execution_options->has_shape_with_output_layout()) { const auto& shape_with_output_layout = - execution_options.shape_with_output_layout(); + execution_options->shape_with_output_layout(); TF_RETURN_IF_ERROR(ValidateResultShapeWithLayout(shape_with_output_layout, program_shape.result())); TF_RETURN_IF_ERROR( computation_layout->mutable_result_layout()->CopyLayoutFromShape( shape_with_output_layout)); + } else { + computation_layout->mutable_result_layout()->Clear(); } - legacy_flags::ServiceFlags* flags = legacy_flags::GetServiceFlags(); - if (flags->xla_hlo_profile) { - module_config->enable_hlo_profiling(true); + config->set_replica_count(options_.number_of_replicas()); + config->set_has_hybrid_result(has_hybrid_result); + if (execution_options != nullptr) { + config->set_seed(execution_options->seed()); + config->set_debug_options(execution_options->debug_options()); + config->enable_hlo_profiling( + execution_options->debug_options().xla_hlo_profile()); + } else { + config->set_debug_options(legacy_flags::GetDebugOptionsFromFlags()); } - module_config->set_replica_count(execute_backend_->Replicas().size()); - module_config->set_fast_math_disabled(execution_options.disable_fast_math()); - module_config->set_seed(execution_options.seed()); + if (execute_backend_ != nullptr && + execute_backend_->eigen_intra_op_thread_pool() != nullptr) { + config->set_intra_op_parallelism_threads( + execute_backend_->eigen_intra_op_thread_pool()->NumThreads()); + } + return std::move(config); +} - return std::move(module_config); +StatusOr> Service::CreateModuleConfig( + const ProgramShape& program_shape, + tensorflow::gtl::ArraySlice arguments, + const ExecutionOptions& execution_options) { + std::vector argument_shapes; + for (const auto* arg : arguments) { + argument_shapes.push_back(&arg->shape()); + } + return CreateModuleConfig(program_shape, argument_shapes, &execution_options); } StatusOr>> Service::BuildExecutables( @@ -342,23 +343,25 @@ StatusOr>> Service::BuildExecutables( // Dump computation proto state if flag is set. std::vector> session_modules; - legacy_flags::ServiceFlags* flags = legacy_flags::GetServiceFlags(); - const string& directory_path = flags->xla_dump_computations_to; - const string& other_directory_path = flags->xla_dump_executions_to; - if ((!directory_path.empty() || !other_directory_path.empty())) { - for (int64 i = 0; i < versioned_handles.size(); ++i) { - TF_ASSIGN_OR_RETURN(std::unique_ptr session_module, - computation_tracker_.SnapshotComputation( - versioned_handles[i].handle)); - if (!directory_path.empty()) { - string filename = Printf("computation_%lld__%s__version_%lld", - versioned_handles[i].handle.handle(), - session_module->entry().name().c_str(), - versioned_handles[i].version); - TF_RETURN_IF_ERROR(Executable::DumpToDirectory(directory_path, filename, - *session_module)); - session_modules.push_back(std::move(session_module)); - } + for (int64 i = 0; i < versioned_handles.size(); ++i) { + const string& directory_path = + module_configs[i]->debug_options().xla_dump_computations_to(); + const string& other_directory_path = + module_configs[i]->debug_options().xla_dump_executions_to(); + if (directory_path.empty() && other_directory_path.empty()) { + continue; + } + TF_ASSIGN_OR_RETURN( + std::unique_ptr session_module, + computation_tracker_.SnapshotComputation(versioned_handles[i].handle)); + if (!directory_path.empty()) { + string filename = Printf("computation_%lld__%s__version_%lld", + versioned_handles[i].handle.handle(), + session_module->entry().name().c_str(), + versioned_handles[i].version); + TF_RETURN_IF_ERROR(Executable::DumpToDirectory(directory_path, filename, + *session_module)); + session_modules.push_back(std::move(session_module)); } } @@ -367,23 +370,24 @@ StatusOr>> Service::BuildExecutables( VLOG(1) << versioned_handle; } + CHECK_EQ(versioned_handles.size(), module_configs.size()); std::vector> modules; - for (const VersionedComputationHandle& versioned_handle : versioned_handles) { + for (int64 i = 0; i < versioned_handles.size(); ++i) { + const VersionedComputationHandle& versioned_handle = versioned_handles[i]; + const HloModuleConfig& config = *module_configs[i]; TF_ASSIGN_OR_RETURN(auto module, computation_tracker_.BuildHloModule( - versioned_handle, + versioned_handle, config, /*include_unreachable_instructions=*/true)); modules.push_back(std::move(module)); } - Compiler::HloDumper hlo_dumper = MakeHloDumper(); - TF_ASSIGN_OR_RETURN(std::vector> executables, - backend->compiler()->Compile( - std::move(modules), std::move(module_configs), - hlo_dumper, std::move(executors))); + TF_ASSIGN_OR_RETURN( + std::vector> executables, + backend->compiler()->Compile(std::move(modules), std::move(executors))); - if (!other_directory_path.empty()) { - for (size_t i = 0; i < versioned_handles.size(); ++i) { + for (size_t i = 0; i < versioned_handles.size(); ++i) { + if (!module_configs[i]->debug_options().xla_dump_executions_to().empty()) { executables[i]->set_session_module(std::move(session_modules[i])); } } @@ -394,7 +398,6 @@ StatusOr>> Service::BuildExecutables( StatusOr> Service::BuildExecutable( const VersionedComputationHandle& versioned_handle, std::unique_ptr module_config, - bool executable_for_compute_constant, const tensorflow::gtl::ArraySlice arguments, Backend* backend, se::StreamExecutor* executor) { @@ -403,11 +406,11 @@ StatusOr> Service::BuildExecutable( // Dump computation proto state if flag is set. std::unique_ptr session_module; - legacy_flags::ServiceFlags* flags = legacy_flags::GetServiceFlags(); - const string& directory_path = flags->xla_dump_computations_to; - const string& other_directory_path = flags->xla_dump_executions_to; - if (!executable_for_compute_constant && - (!directory_path.empty() || !other_directory_path.empty())) { + const string& directory_path = + module_config->debug_options().xla_dump_computations_to(); + const string& other_directory_path = + module_config->debug_options().xla_dump_executions_to(); + if (!directory_path.empty() || !other_directory_path.empty()) { TF_ASSIGN_OR_RETURN( session_module, computation_tracker_.SnapshotComputation(versioned_handle.handle)); @@ -423,20 +426,13 @@ StatusOr> Service::BuildExecutable( TF_ASSIGN_OR_RETURN( std::unique_ptr module, - computation_tracker_.BuildHloModule(versioned_handle, + computation_tracker_.BuildHloModule(versioned_handle, *module_config, /*include_unreachable_instructions=*/ - !executable_for_compute_constant)); - - Compiler::HloDumper hlo_dumper = MakeHloDumper(); - if (executable_for_compute_constant && - !flags->xla_hlo_graph_for_compute_constant) { - hlo_dumper = [](const HloModule&, const string&) {}; - } + true)); TF_ASSIGN_OR_RETURN( std::unique_ptr executable, - backend->compiler()->Compile(std::move(module), std::move(module_config), - hlo_dumper, executor)); + backend->compiler()->Compile(std::move(module), executor)); if (!other_directory_path.empty()) { executable->set_session_module(std::move(session_module)); @@ -472,9 +468,8 @@ StatusOr> Service::BuildAndCacheExecutable( HloModuleConfig original_module_config = *module_config; TF_ASSIGN_OR_RETURN( std::unique_ptr executable_unique_ptr, - BuildExecutable(versioned_handle, std::move(module_config), - /*executable_for_compute_constant=*/false, arguments, - execute_backend_.get(), executor)); + BuildExecutable(versioned_handle, std::move(module_config), arguments, + backend, executor)); if (profile != nullptr) { uint64 end_micros = tensorflow::Env::Default()->NowMicros(); @@ -494,47 +489,55 @@ Service::ExecuteParallelAndRegisterResult( tensorflow::gtl::ArraySlice< std::vector> arguments, - Backend* backend, - tensorflow::gtl::ArraySlice executors, + Backend* backend, tensorflow::gtl::ArraySlice device_handles, tensorflow::gtl::ArraySlice result_tags) { - // TODO(b/33943292): Support for replication when using multiple computations. - TF_RET_CHECK(backend->Replicas().size() == 1); - - // Set up streams. + // Streams where the computation are launched, so we can wait on the streams + // to complete. std::vector::SmartPtr> streams; - for (se::StreamExecutor* executor : executors) { - TF_ASSIGN_OR_RETURN(Pool::SmartPtr stream, - backend->BorrowStream(executor)); - streams.push_back(std::move(stream)); - } - - // Set up run options. - std::vector run_options; - for (const Pool::SmartPtr& stream : streams) { - ExecutableRunOptions options; - options.set_stream(stream.get()); - options.set_allocator(backend->memory_allocator()); - options.set_inter_op_thread_pool(backend->inter_op_thread_pool()); - options.set_intra_op_thread_pool( - backend->eigen_intra_op_thread_pool_device()); - run_options.emplace_back(options, backend->StreamBorrower()); - } - - // Asynchronously launch all executables. + // Global data handles for the computation results, one for each computation. std::vector result_handles; - for (tensorflow::gtl::ArraySlice::size_type i = 0; - i < executables.size(); i++) { - TF_ASSIGN_OR_RETURN( - perftools::gputools::DeviceMemoryBase result, - executables[i]->ExecuteAsyncOnStream(&run_options[i], arguments[i])); - result_handles.push_back(allocation_tracker_.Register( - backend, executors[i]->device_ordinal(), result, - executables[i]->result_shape(), result_tags[i])); + + TF_ASSIGN_OR_RETURN(DeviceAssignment device_assignment, + backend->computation_placer()->AssignDevices( + options_.number_of_replicas(), executables.size())); + + for (int64 i = 0; i < executables.size(); i++) { + // Stream executors for the replicas of the current computation. + TF_ASSIGN_OR_RETURN(auto replicas, Replicas(*backend, device_handles[i])); + for (int64 replica = 0; replica < replicas.size(); ++replica) { + TF_ASSIGN_OR_RETURN(Pool::SmartPtr stream, + backend->BorrowStream(replicas[replica])); + streams.push_back(std::move(stream)); + + // Set up run options. + ExecutableRunOptions options; + options.set_stream(streams.back().get()); + options.set_allocator(backend->memory_allocator()); + options.set_inter_op_thread_pool(backend->inter_op_thread_pool()); + options.set_intra_op_thread_pool( + backend->eigen_intra_op_thread_pool_device()); + options.set_device_assignment(&device_assignment); + ServiceExecutableRunOptions run_options(options, + backend->StreamBorrower()); + + // Asynchronously launch the computation. + TF_ASSIGN_OR_RETURN( + perftools::gputools::DeviceMemoryBase result, + executables[i]->ExecuteAsyncOnStream(&run_options, arguments[i])); + + // All replicas share the same device address for the result allocation, + // so only one of the replicas need to register the result handle. + if (replica == 0) { + result_handles.push_back(allocation_tracker_.Register( + backend, replicas[0]->device_ordinal(), result, + executables[i]->result_shape(), result_tags[i])); + } + } } // Wait for all executions to complete. - for (int64 i = 0; i < result_handles.size(); ++i) { + for (int64 i = 0; i < streams.size(); ++i) { if (!streams[i]->BlockHostUntilDone()) { return InternalError("failed to complete execution for stream %lld", i); } @@ -549,17 +552,23 @@ StatusOr Service::ExecuteAndRegisterResult( arguments, Backend* backend, perftools::gputools::StreamExecutor* executor, const string& result_tag, ExecutionProfile* profile) { - TF_RET_CHECK(!backend->Replicas().empty()); - // Set up streams. std::vector::SmartPtr> streams; - for (se::StreamExecutor* executor : backend->Replicas()) { + TF_ASSIGN_OR_RETURN(auto replicas, + Replicas(*backend, SingleComputationDeviceHandle())); + TF_RET_CHECK(!replicas.empty()); + for (se::StreamExecutor* executor : replicas) { TF_ASSIGN_OR_RETURN(Pool::SmartPtr stream, backend->BorrowStream(executor)); streams.push_back(std::move(stream)); } + TF_ASSIGN_OR_RETURN(DeviceAssignment device_assignment, + backend->computation_placer()->AssignDevices( + options_.number_of_replicas(), + /*computation_count=*/1)); + // Set up run options. std::vector run_options; for (const Pool::SmartPtr& stream : streams) { @@ -569,25 +578,20 @@ StatusOr Service::ExecuteAndRegisterResult( options.set_inter_op_thread_pool(backend->inter_op_thread_pool()); options.set_intra_op_thread_pool( backend->eigen_intra_op_thread_pool_device()); - run_options.emplace_back(options, backend->StreamBorrower()); + options.set_device_assignment(&device_assignment); + run_options.emplace_back(options, backend->StreamBorrower(), + backend->inter_op_thread_pool()); } perftools::gputools::DeviceMemoryBase result; - if (backend->Replicas().size() == 1) { + if (options_.number_of_replicas() == 1) { TF_ASSIGN_OR_RETURN( - result, - ExecuteOnStreamWrapper>( - executable, &run_options[0], profile, execute_backend_.get(), - [&arguments](Executable* executable, - const ServiceExecutableRunOptions* run_options, - HloExecutionProfile* hlo_execution_profile) { - return executable->ExecuteOnStream(run_options, arguments, - hlo_execution_profile); - })); + result, executable->ExecuteOnStreamWrapper( + &run_options[0], profile, arguments)); } else { std::vector< tensorflow::gtl::ArraySlice> - repeated_arguments(backend->Replicas().size(), arguments); + repeated_arguments(options_.number_of_replicas(), arguments); TF_ASSIGN_OR_RETURN(auto results, executable->ExecuteOnStreams( run_options, repeated_arguments)); @@ -615,25 +619,26 @@ tensorflow::Status Service::ExecuteParallel(const ExecuteParallelRequest* arg, std::vector versioned_handles; std::vector> module_configs; std::vector computation_names; + std::vector device_handles; - if (arg->requests_size() > execute_backend_->stream_executors().size()) { + if (arg->requests_size() * options_.number_of_replicas() > + execute_backend_->device_count()) { return FailedPrecondition( "there are not enough stream executors to execute %d computations", arg->requests_size()); } for (int64 i = 0; i < arg->requests_size(); ++i) { - // Get the stream executor on which the computation will run. Select the - // specific device if requested, otherwise select the i'th device from the - // list of available stream executors. - se::StreamExecutor* executor; - if (arg->requests(i).has_device_handle()) { - executor = - execute_backend_ - ->stream_executors()[arg->requests(i).device_handle().handle()]; - } else { - executor = execute_backend_->stream_executors()[i]; + // Get the stream executor for the i'th computation. This stream executor + // is one of the executors to run the replicated computation. + if (!arg->requests(i).has_device_handle()) { + return FailedPrecondition( + "device handles must be given to execute parallel computations"); } + TF_ASSIGN_OR_RETURN( + auto replicas, + Replicas(*execute_backend_, arg->requests(i).device_handle())); + se::StreamExecutor* executor = replicas[0]; CHECK(executor != nullptr); // Resolve the UserComputation object associated with the requested @@ -658,6 +663,7 @@ tensorflow::Status Service::ExecuteParallel(const ExecuteParallelRequest* arg, ResolveAndValidateArguments(request.arguments(), execute_backend_.get(), executor->device_ordinal())); std::vector arguments; + arguments.reserve(arg_allocations.size()); for (const Allocation* allocation : arg_allocations) { arguments.push_back(allocation->device_memory()); } @@ -676,6 +682,7 @@ tensorflow::Status Service::ExecuteParallel(const ExecuteParallelRequest* arg, module_configs.push_back(std::move(module_config)); computation_names.push_back(user_computation->name()); executors.push_back(executor); + device_handles.push_back(arg->requests(i).device_handle()); } // Build the user computations into HloModules and compile to generate the @@ -685,6 +692,7 @@ tensorflow::Status Service::ExecuteParallel(const ExecuteParallelRequest* arg, BuildExecutables(versioned_handles, std::move(module_configs), execute_backend_.get(), executors)); std::vector executable_ptrs; + executable_ptrs.reserve(executables.size()); for (const auto& executable : executables) { executable_ptrs.push_back(executable.get()); } @@ -694,7 +702,7 @@ tensorflow::Status Service::ExecuteParallel(const ExecuteParallelRequest* arg, TF_ASSIGN_OR_RETURN( std::vector outputs, ExecuteParallelAndRegisterResult(executable_ptrs, all_arguments, - execute_backend_.get(), executors, + execute_backend_.get(), device_handles, computation_names)); for (const GlobalDataHandle& output : outputs) { ExecuteResponse response; @@ -708,10 +716,12 @@ tensorflow::Status Service::ExecuteParallel(const ExecuteParallelRequest* arg, tensorflow::Status Service::GetDeviceHandles(const GetDeviceHandlesRequest* arg, GetDeviceHandlesResponse* result) { - const int64 available_device_count = - execute_backend_->stream_executors().size(); - const int64 replicas = execute_backend_->Replicas().size(); - if (available_device_count < arg->device_count() * replicas) { + const int64 available_device_count = execute_backend_->device_count(); + const int64 replica_count = options_.number_of_replicas(); + if (replica_count <= 0) { + return FailedPrecondition("Replica count must be a positive integer"); + } + if (available_device_count < arg->device_count() * replica_count) { return ResourceExhausted( "Requested device count (%lld) exceeds the number of available devices " "on the target (%lld)", @@ -720,8 +730,8 @@ tensorflow::Status Service::GetDeviceHandles(const GetDeviceHandlesRequest* arg, for (int64 i = 0; i < arg->device_count(); ++i) { DeviceHandle device_handle; - device_handle.set_handle( - execute_backend_->stream_executors()[i * replicas]->device_ordinal()); + device_handle.set_handle(i); + device_handle.set_device_count(arg->device_count()); *result->add_device_handles() = device_handle; } @@ -759,6 +769,7 @@ tensorflow::Status Service::Execute(const ExecuteRequest* arg, << module_config->entry_computation_layout().ToString(); std::vector arguments; + arguments.reserve(arg_allocations.size()); for (const Allocation* allocation : arg_allocations) { arguments.push_back(allocation->device_memory()); } @@ -826,6 +837,7 @@ tensorflow::Status Service::ExecuteAsync(const ExecuteAsyncRequest* arg, << module_config->entry_computation_layout().ToString(); std::vector arguments; + arguments.reserve(arg_allocations.size()); for (const Allocation* allocation : arg_allocations) { arguments.push_back(allocation->device_memory()); } @@ -839,11 +851,14 @@ tensorflow::Status Service::ExecuteAsync(const ExecuteAsyncRequest* arg, execute_backend_->default_stream_executor(), &profile)); - TF_RET_CHECK(!execute_backend_->Replicas().empty()); + TF_ASSIGN_OR_RETURN(auto replicas, Replicas(*execute_backend_, + SingleComputationDeviceHandle())); + TF_RET_CHECK(!replicas.empty()); + // Set up streams. std::vector::SmartPtr> streams; - for (se::StreamExecutor* executor : execute_backend_->Replicas()) { + for (se::StreamExecutor* executor : replicas) { TF_ASSIGN_OR_RETURN(Pool::SmartPtr stream, execute_backend_->BorrowStream(executor)); streams.push_back(std::move(stream)); @@ -914,28 +929,31 @@ tensorflow::Status Service::TransferToClient(const TransferToClientRequest* arg, literal_shape = &allocation->shape(); } - return LiteralFromAllocation(allocation, *literal_shape, - result->mutable_literal()); + Literal literal; + auto status = LiteralFromAllocation(allocation, *literal_shape, &literal); + *result->mutable_literal() = literal.ToProto(); + return status; } tensorflow::Status Service::TransferToServer(const TransferToServerRequest* arg, TransferToServerResponse* result) { - const Literal& literal = arg->literal(); + Literal literal = Literal(arg->literal()); const Shape& shape = literal.shape(); - if (ShapeUtil::IsTuple(shape) && execute_backend_->Replicas().size() > 1) { + if (ShapeUtil::IsTuple(shape) && options_.number_of_replicas() > 1) { // TODO(b/32990684): Tuple transfers to host end up allocating further // buffers - implement that correctly. return Unimplemented( "Tuple transfers to the device not supported with replication."); } - se::StreamExecutor* stream_executor; + std::vector replicas; if (arg->has_device_handle()) { - TF_ASSIGN_OR_RETURN(stream_executor, execute_backend_->stream_executor( - arg->device_handle().handle())); + TF_ASSIGN_OR_RETURN(replicas, + Replicas(*execute_backend_, arg->device_handle())); } else { - stream_executor = execute_backend_->default_stream_executor(); + TF_ASSIGN_OR_RETURN( + replicas, Replicas(*execute_backend_, SingleComputationDeviceHandle())); } // Allocate memory on the device, using the stream executor. The size of the @@ -946,14 +964,12 @@ tensorflow::Status Service::TransferToServer(const TransferToServerRequest* arg, TF_ASSIGN_OR_RETURN(se::DeviceMemoryBase allocation, execute_backend_->memory_allocator()->Allocate( - stream_executor->device_ordinal(), allocation_size)); + replicas[0]->device_ordinal(), allocation_size)); *result->mutable_data() = allocation_tracker_.Register( - execute_backend_.get(), stream_executor->device_ordinal(), allocation, - shape, StrCat("TransferToServer literal of size ", allocation_size)); + execute_backend_.get(), replicas[0]->device_ordinal(), allocation, shape, + StrCat("TransferToServer literal of size ", allocation_size)); - TF_ASSIGN_OR_RETURN(auto replicas, execute_backend_->Replicas( - stream_executor->device_ordinal())); for (se::StreamExecutor* executor : replicas) { TF_RETURN_IF_ERROR( execute_backend_->transfer_manager()->TransferLiteralToDevice( @@ -964,7 +980,7 @@ tensorflow::Status Service::TransferToServer(const TransferToServerRequest* arg, tensorflow::Status Service::TransferToInfeed(const TransferToInfeedRequest* arg, TransferToInfeedResponse* result) { - const int64 replica_count = execute_backend_->Replicas().size(); + const int64 replica_count = options_.number_of_replicas(); if (arg->replica_id() < 0 || arg->replica_id() >= replica_count) { return FailedPrecondition( "%s", @@ -976,21 +992,24 @@ tensorflow::Status Service::TransferToInfeed(const TransferToInfeedRequest* arg, se::StreamExecutor* executor; if (arg->has_device_handle()) { - TF_ASSIGN_OR_RETURN(auto replicas, execute_backend_->Replicas( - arg->device_handle().handle())); + TF_ASSIGN_OR_RETURN(auto replicas, + Replicas(*execute_backend_, arg->device_handle())); executor = replicas[arg->replica_id()]; } else { - executor = execute_backend_->Replicas()[arg->replica_id()]; + TF_ASSIGN_OR_RETURN( + auto replicas, + Replicas(*execute_backend_, SingleComputationDeviceHandle())); + executor = replicas[arg->replica_id()]; } return execute_backend_->transfer_manager()->TransferLiteralToInfeed( - executor, arg->literal()); + executor, Literal(arg->literal())); } tensorflow::Status Service::TransferFromOutfeed( const TransferFromOutfeedRequest* arg, TransferFromOutfeedResponse* result) { - const int64 replica_count = execute_backend_->Replicas().size(); + const int64 replica_count = options_.number_of_replicas(); if (arg->replica_id() < 0 || arg->replica_id() >= replica_count) { return FailedPrecondition( "The replica_id=%lld on TransferFromOutfeedRequest not in range [0, " @@ -1000,15 +1019,22 @@ tensorflow::Status Service::TransferFromOutfeed( se::StreamExecutor* executor; if (arg->has_device_handle()) { - TF_ASSIGN_OR_RETURN(auto replicas, execute_backend_->Replicas( - arg->device_handle().handle())); + TF_ASSIGN_OR_RETURN(auto replicas, + Replicas(*execute_backend_, arg->device_handle())); executor = replicas[arg->replica_id()]; } else { - executor = execute_backend_->Replicas()[arg->replica_id()]; + TF_ASSIGN_OR_RETURN( + auto replicas, + Replicas(*execute_backend_, SingleComputationDeviceHandle())); + executor = replicas[arg->replica_id()]; } - return execute_backend_->transfer_manager()->TransferLiteralFromOutfeed( - executor, arg->shape_with_layout(), result->mutable_literal()); + Literal literal; + TF_RETURN_IF_ERROR( + execute_backend_->transfer_manager()->TransferLiteralFromOutfeed( + executor, arg->shape_with_layout(), &literal)); + *result->mutable_literal() = literal.ToProto(); + return tensorflow::Status::OK(); } tensorflow::Status Service::ResetDevice(const ResetDeviceRequest* arg, @@ -1016,70 +1042,6 @@ tensorflow::Status Service::ResetDevice(const ResetDeviceRequest* arg, return execute_backend_->ResetDevices(); } -tensorflow::Status Service::TransferToClientInProcess( - const TransferToClientInProcessRequest* arg, - TransferToClientInProcessResponse* result) { - TF_RETURN_IF_ERROR(CheckRunsInClientProcess("TransferToClientInProcess")); - - TF_ASSIGN_OR_RETURN(const Allocation* allocation, - allocation_tracker_.Resolve(arg->data())); - - void* buffer = reinterpret_cast(arg->buffer()); - int64 size = ShapeUtil::ByteSizeOf(allocation->shape()); - TF_ASSIGN_OR_RETURN( - se::StreamExecutor * executor, - allocation->backend()->stream_executor(allocation->device_ordinal())); - - return allocation->backend()->transfer_manager()->TransferBufferFromDevice( - executor, allocation->device_memory(), size, buffer); -} - -tensorflow::Status Service::TransferToServerInProcess( - const TransferToServerInProcessRequest* arg, - TransferToServerInProcessResponse* result) { - TF_RETURN_IF_ERROR(CheckRunsInClientProcess("TransferToServerInProcess")); - - const Shape& shape = arg->shape(); - - if (ShapeUtil::IsTuple(shape) && execute_backend_->Replicas().size() > 1) { - // TODO(b/32990684): Tuple transfers to host end up allocating further - // buffers - implement that correctly. - return Unimplemented( - "Tuple transfers to the device not supported with replication."); - } - - if (!LayoutUtil::HasLayout(shape)) { - return InvalidArgument("shape must have layout"); - } - - TF_RETURN_IF_ERROR(ShapeUtil::ValidateShape(shape)); - - const void* buffer = reinterpret_cast(arg->buffer()); - - // Allocate memory on the device, using the stream executor. The size of the - // allocation is obtained by examining the shape of the literal passed from - // the client. An allocation handle is returned in the response. - int64 allocation_size = - execute_backend_->transfer_manager()->GetByteSizeRequirement(shape); - se::StreamExecutor* stream_executor = - execute_backend_->default_stream_executor(); - - TF_ASSIGN_OR_RETURN(se::DeviceMemoryBase allocation, - execute_backend_->memory_allocator()->Allocate( - stream_executor->device_ordinal(), allocation_size)); - - *result->mutable_data() = allocation_tracker_.Register( - execute_backend_.get(), stream_executor->device_ordinal(), allocation, - shape, StrCat("TransferToServer literal of size ", allocation_size)); - - for (se::StreamExecutor* executor : execute_backend_->Replicas()) { - TF_RETURN_IF_ERROR( - execute_backend_->transfer_manager()->TransferBufferToDevice( - executor, allocation_size, buffer, &allocation)); - } - return tensorflow::Status::OK(); -} - tensorflow::Status Service::IsConstant(const IsConstantRequest* arg, IsConstantResponse* result) { TF_ASSIGN_OR_RETURN(UserComputation * user_computation, @@ -1113,7 +1075,6 @@ tensorflow::Status Service::ComputeConstant(const ComputeConstantRequest* arg, TF_ASSIGN_OR_RETURN(bool is_constant, user_computation->IsConstant(arg->operand())); - if (!is_constant) { return InvalidArgument("Operand to ComputeConstant depends on parameter."); } @@ -1127,8 +1088,10 @@ tensorflow::Status Service::ComputeConstant(const ComputeConstantRequest* arg, TF_DCHECK_OK(ShapeUtil::ValidateShape(program_shape.result())); - ExecutionOptions execution_options; - execution_options.set_disable_fast_math(true); + ExecutionOptions execution_options = xla::CreateDefaultExecutionOptions(); + execution_options.mutable_debug_options()->set_xla_enable_fast_math(false); + execution_options.mutable_debug_options() + ->set_xla_eliminate_hlo_implicit_broadcast(true); *execution_options.mutable_shape_with_output_layout() = program_shape.result(); @@ -1143,20 +1106,22 @@ tensorflow::Status Service::ComputeConstant(const ComputeConstantRequest* arg, TF_ASSIGN_OR_RETURN(std::unique_ptr module_config, CreateModuleConfig(program_shape, {}, execution_options)); + // Exclude dead parameter instructions for the purpose of computing constants. TF_ASSIGN_OR_RETURN( - std::shared_ptr executable, - BuildExecutable(versioned_handle, std::move(module_config), - /*executable_for_compute_constant=*/true, - /*arguments=*/{}, compute_constant_backend_.get(), - compute_constant_backend_->default_stream_executor())); + std::unique_ptr module, + computation_tracker_.BuildHloModule(versioned_handle, *module_config, + /*include_unreachable_instructions=*/ + false)); + + HloEvaluator evaluator; + TF_ASSIGN_OR_RETURN(auto result_literal, evaluator.Evaluate(*module, {})); + // Since the shape_with_output_layout option in ExecutionOption is + // non-effective to the Evaluator results, explicit relayout here. + if (arg->has_output_layout()) { + result_literal = result_literal->Relayout(arg->output_layout()); + } + *result->mutable_literal() = result_literal->ToProto(); - TF_ASSIGN_OR_RETURN( - *result->mutable_output(), - ExecuteAndRegisterResult( - executable.get(), /*arguments=*/{}, compute_constant_backend_.get(), - compute_constant_backend_->default_stream_executor(), - "constant computed from " + user_computation->name(), - /*profile=*/nullptr)); return tensorflow::Status::OK(); } @@ -1201,15 +1166,18 @@ tensorflow::Status Service::GetComputationStats( VersionedComputationHandle versioned_handle = user_computation->GetVersionedHandle(); - TF_ASSIGN_OR_RETURN(std::unique_ptr module, - computation_tracker_.BuildHloModule(versioned_handle)); + HloModuleConfig config; + config.set_debug_options(arg->debug_options()); + TF_ASSIGN_OR_RETURN( + std::unique_ptr module, + computation_tracker_.BuildHloModule(versioned_handle, config)); - MakeHloDumper()(*module, "computation statistics subject"); + hlo_graph_dumper::MaybeDumpHloModule(*module, + "computation statistics subject"); // Run HLO analysis to get the computation statistics. - HloCostAnalysis analysis([this](const Shape& shape) { - return execute_backend_->compiler()->ShapeSizeBytes(shape); - }); + HloCostAnalysis analysis( + execute_backend_->compiler()->ShapeSizeBytesFunction()); TF_RETURN_IF_ERROR( module->entry_computation()->root_instruction()->Accept(&analysis)); @@ -1221,17 +1189,6 @@ tensorflow::Status Service::GetComputationStats( return tensorflow::Status::OK(); } -tensorflow::Status Service::CheckRunsInClientProcess( - const string& method_name) const { - if (runs_in_client_process_) { - return tensorflow::Status::OK(); - } else { - return FailedPrecondition( - "%s only supported if service runs in the same process as the client", - method_name.c_str()); - } -} - template tensorflow::Status Service::AddInstruction( const RequestT* arg, ResponseT* result, @@ -1250,6 +1207,18 @@ tensorflow::Status Service::Op(const OpRequest* arg, OpResponse* result) { StatusOr handle_status; switch (arg->op_case()) { + case OpRequest::kBatchNormTrainingRequest: + handle_status = computation->AddBatchNormTrainingInstruction( + arg->batch_norm_training_request()); + break; + case OpRequest::kBatchNormInferenceRequest: + handle_status = computation->AddBatchNormInferenceInstruction( + arg->batch_norm_inference_request()); + break; + case OpRequest::kBatchNormGradRequest: + handle_status = computation->AddBatchNormGradInstruction( + arg->batch_norm_grad_request()); + break; case OpRequest::kBinaryOpRequest: handle_status = computation->AddBinaryInstruction(arg->binary_op_request()); @@ -1332,6 +1301,11 @@ tensorflow::Status Service::Op(const OpRequest* arg, OpResponse* result) { computation->AddReduceInstruction(arg->reduce_request(), *to_apply); break; } + case OpRequest::kReducePrecisionRequest: { + handle_status = computation->AddReducePrecisionInstruction( + arg->reduce_precision_request()); + break; + } case OpRequest::kReduceWindowRequest: { TF_ASSIGN_OR_RETURN(UserComputation * to_apply, computation_tracker_.Resolve( @@ -1438,4 +1412,28 @@ tensorflow::Status Service::LoadComputationSnapshot( return tensorflow::Status::OK(); } +DeviceHandle Service::SingleComputationDeviceHandle() const { + DeviceHandle device_handle; + device_handle.set_handle(0); + device_handle.set_device_count(1); + return device_handle; +} + +StatusOr> Service::Replicas( + const Backend& backend, const DeviceHandle& device_handle) const { + std::vector replicas; + for (int replica = 0; replica < options_.number_of_replicas(); ++replica) { + // From the computation placer, find out the device ids of the replicas for + // the given device handle. + TF_ASSIGN_OR_RETURN( + int device_ordinal, + backend.computation_placer()->DeviceId(replica, device_handle.handle(), + options_.number_of_replicas(), + device_handle.device_count())); + TF_ASSIGN_OR_RETURN(auto executor, backend.stream_executor(device_ordinal)); + replicas.push_back(executor); + } + return replicas; +} + } // namespace xla diff --git a/tensorflow/compiler/xla/service/service.h b/tensorflow/compiler/xla/service/service.h index 9600f6989a40c9180d00ccabbeb29cb37a28900a..bb86a53c62e05bb62b93bbac88c2ca251ad0439a 100644 --- a/tensorflow/compiler/xla/service/service.h +++ b/tensorflow/compiler/xla/service/service.h @@ -22,12 +22,11 @@ limitations under the License. #include #include "tensorflow/compiler/xla/executable_run_options.h" -#include "tensorflow/compiler/xla/legacy_flags/service_flags.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/compiler.h" #include "tensorflow/compiler/xla/service/computation_tracker.h" #include "tensorflow/compiler/xla/service/device_memory_allocator.h" #include "tensorflow/compiler/xla/service/executable.h" @@ -58,19 +57,23 @@ class ServiceOptions { perftools::gputools::Platform* platform() const; // Set the number of replicas to use when compiling replicated - // programs. The default is -1 meaning that the value is read from - // the xla_replicas flag. + // programs. ServiceOptions& set_number_of_replicas(int number_of_replicas); int number_of_replicas() const; + // Sets the thread pool size for parallel execution of an individual operator. + ServiceOptions& set_intra_op_parallelism_threads(int num_threads); + int intra_op_parallelism_threads() const; + private: perftools::gputools::Platform* platform_ = nullptr; - int number_of_replicas_ = -1; + int number_of_replicas_ = 1; + int intra_op_parallelism_threads_ = -1; }; -// The XLA service object, which is the same across all -// platforms. It maintains the service state of computations and allocations, -// and delegates target-specific requests to the target-specific infrastructure +// The XLA service object, which is the same across all platforms. It maintains +// the service state of computations and allocations, and delegates +// target-specific requests to the target-specific infrastructure // (target-specific compiler, StreamExecutor). class Service : public ServiceInterface { public: @@ -121,7 +124,7 @@ class Service : public ServiceInterface { // least N * R devices must be available. The devices are assigned based on // the device ordinals such that the first R available devices are assigned to // the first set of replicas, and the next R devices to the second set of - // replicas, etc. Each returned device handles represent the device with the + // replicas, etc. Each returned device handle represents the device with the // replica id 0. tensorflow::Status GetDeviceHandles( const GetDeviceHandlesRequest* arg, @@ -130,6 +133,10 @@ class Service : public ServiceInterface { // Asynchronously executes a computation with provided arguments. Invokes // the provided computation with the provided global data passed as // immutable arguments. Returns a handle to the execution. + // + // (Note: The corresponding function in xla::Client was removed as part of + // b/64116060, in an attempt to simplify our API. We're keeping this around + // for now in case we want to expose this to clients in a different way.) tensorflow::Status ExecuteAsync(const ExecuteAsyncRequest* arg, ExecuteAsyncResponse* result) override; @@ -146,11 +153,6 @@ class Service : public ServiceInterface { const TransferToClientRequest* arg, TransferToClientResponse* result) override; - // Requests that global data be copied into a buffer supplied by the client. - tensorflow::Status TransferToClientInProcess( - const TransferToClientInProcessRequest* arg, - TransferToClientInProcessResponse* result) override; - // Transfers data from a literal provided by the client, into device memory. tensorflow::Status TransferToServer( const TransferToServerRequest* arg, @@ -168,11 +170,6 @@ class Service : public ServiceInterface { const TransferFromOutfeedRequest* arg, TransferFromOutfeedResponse* result) override; - // Transfers data from a buffer provided by the client, into device memory. - tensorflow::Status TransferToServerInProcess( - const TransferToServerInProcessRequest* arg, - TransferToServerInProcessResponse* result) override; - // Resets devices, clearing all existing state on all the devices associated // with this service (including memory allocated on the devices). // @@ -248,13 +245,21 @@ class Service : public ServiceInterface { const Backend& backend() const { return *execute_backend_; } Backend* mutable_backend() { return execute_backend_.get(); } + private: + // A private overload for Service itself, used by other methods within this + // class. + StatusOr> CreateModuleConfig( + const ProgramShape& program_shape, + tensorflow::gtl::ArraySlice arguments, + const ExecutionOptions& execution_options); + protected: friend class LocalExecutable; // The constructor is private. Use the NewService factory to create new // service objects. - Service(std::unique_ptr backend, - std::unique_ptr compute_constant_backend); + Service(const ServiceOptions& options, + std::unique_ptr execute_backend); static StatusOr> CreateComputeConstantBackend(); @@ -265,22 +270,20 @@ class Service : public ServiceInterface { tensorflow::gtl::ArraySlice arguments, const Backend* backend, int device_ordinal); - // Create a Hlo module config foe the given program shape and arguments. + // Create a Hlo module config for the given program shape and arguments. + // execution_options is optional; if not given a default is used. + // has_hybrid_result is used to initialize the same-named field in + // HloModuleConfig -- see that class for documentation. StatusOr> CreateModuleConfig( const ProgramShape& program_shape, - tensorflow::gtl::ArraySlice arguments, - const ExecutionOptions& execution_options); + tensorflow::gtl::ArraySlice argument_shapes, + const ExecutionOptions* execution_options, + bool has_hybrid_result = false); - // Builds an Executable for the given parameters. If - // executable_for_compute_constant is true, then the executable is intended to - // be used for ComputeConstant which means dead parameter instructions are not - // included in the executable.The parameter "profile" can optionally point to - // an ExecutionProfile object which will be filled in with profile data - // relevant to compilation. + // Builds an Executable for the given parameters. StatusOr> BuildExecutable( const VersionedComputationHandle& versioned_handle, std::unique_ptr module_config, - bool executable_for_compute_constant, const tensorflow::gtl::ArraySlice arguments, Backend* backend, perftools::gputools::StreamExecutor* executor); @@ -324,18 +327,9 @@ class Service : public ServiceInterface { std::vector> arguments, Backend* backend, - tensorflow::gtl::ArraySlice - executors, + tensorflow::gtl::ArraySlice device_handles, tensorflow::gtl::ArraySlice result_tags); - // Dumps the executed HLO according to service-associated flags. - static void DumpExecutedHlo(const HloModule& module, const string& label, - const HloExecutionProfile* profile); - - // Returns an HLO dumper for use in the compiler (it refers to flags - // associated with the service). - static Compiler::HloDumper MakeHloDumper(); - // Convenience function for adding a function to a user computation. template tensorflow::Status AddInstruction( @@ -343,32 +337,23 @@ class Service : public ServiceInterface { const std::function(UserComputation*)>& adder); - // If the service is running in the client process - // (runs_in_client_process_ is true) then return - // tensorflow::Status::OK. Otherwise return an appropriate error - // status with the given method name. Used for "InProcess" methods. - tensorflow::Status CheckRunsInClientProcess(const string& method_name) const; - // Convenience function which checks whether the given shape_with_layout // (presumably passed by the client to set the result layout) is valid for the // given computation result shape. tensorflow::Status ValidateResultShapeWithLayout( const Shape& shape_with_layout, const Shape& result_shape) const; - // Convenience wrapper for calling Executable::ExecuteOnStream. Sets up a - // timer for the execution, sets up HLO profiling if enabled, and fills in the - // given ExecutionProfile if non-null. The given execute_func should be a - // function which calls the desired ExecuteOnStream overload with the supplied - // arguments. The ExecuteOnStream overloads return different types so this - // method is templated on return-type of the execute function. - template - static ReturnT ExecuteOnStreamWrapper( - Executable* executable, const ServiceExecutableRunOptions* run_options, - ExecutionProfile* profile, Backend* backend, - std::function - execute_func); + // Returns the stream executors assigned to the replicas represented by the + // given device handle. Each device_handle is a virtual replicated device that + // represents a set of physical devices for the replicas. + StatusOr> Replicas( + const Backend& backend, const DeviceHandle& device_handle) const; + + // Returns the device handle that represents the replicated device for a + // single computation that is not model-parallelized. + DeviceHandle SingleComputationDeviceHandle() const; + + ServiceOptions options_; // Tracks computations built via the API. ComputationTracker computation_tracker_; @@ -390,92 +375,9 @@ class Service : public ServiceInterface { // TODO(b/28616830): Support multiple backends for execution. std::unique_ptr execute_backend_; - // Backend to use when executing ComputeConstant. - std::unique_ptr compute_constant_backend_; - - // Whether the service runs in the same process as the client. - bool runs_in_client_process_ = false; - TF_DISALLOW_COPY_AND_ASSIGN(Service); }; -template -ReturnT Service::ExecuteOnStreamWrapper( - Executable* executable, const ServiceExecutableRunOptions* run_options, - ExecutionProfile* profile, Backend* backend, - std::function - execute_func) { - perftools::gputools::Stream* stream = run_options->stream(); - std::unique_ptr timer; - if (profile != nullptr) { - timer.reset(new perftools::gputools::Timer(stream->parent())); - stream->InitTimer(timer.get()).ThenStartTimer(timer.get()); - } - - VLOG(1) << "enqueueing executable on stream..."; - // If the profiling flag isn't enabled, we pass nullptr as the profile to - // indicate profiling is not requested. - HloExecutionProfile hlo_execution_profile; - legacy_flags::ServiceFlags* flags = legacy_flags::GetServiceFlags(); - HloExecutionProfile* profile_ptr = - flags->xla_hlo_profile && executable->hlo_profiling_enabled() - ? &hlo_execution_profile - : nullptr; - - auto return_value = execute_func(executable, run_options, profile_ptr); - - if (profile != nullptr) { - VLOG(1) << "enqueueing 'stop timer' and blocking host until done..."; - stream->ThenStopTimer(timer.get()).BlockHostUntilDone(); - VLOG(1) << "done with block-host-until-done"; - - // Merge in run time profile information from the executable. - profile->MergeFrom(executable->execution_profile()); - - // Overall execution time (in nanoseconds) from the executor timer. - profile->set_compute_and_transfer_time_ns(timer->Nanoseconds()); - - // TODO(b/28123297): On GPU we end up including transfer time in - // the compute time this way. Instead, we should get the correct - // value by measuring it. Setting the field here at least lets - // benchmarks provide *some* value for GPU computations. - // - // TODO(b/28447609): The value in compute_and_transfer_time_ns is actually - // the compute time without the transfer time, so this way we get the - // correct compute time. We should instead have the correct value for - // compute_and_transfer_time and set compute_time to the compute time. - if (profile->compute_time_ns() == 0) { - profile->set_compute_time_ns(profile->compute_and_transfer_time_ns()); - } - } - - if (profile_ptr != nullptr) { - HloCostAnalysis::ShapeSizeFunction shape_size = - [backend](const Shape& shape) { - return backend->compiler()->ShapeSizeBytes(shape); - }; - std::unordered_set profiled_computations = - profile_ptr->profiled_computations(); - // To ensure we have print the profiles in a stable order, iterate over the - // computations in post order. - std::list all_computations = - executable->module().MakeComputationPostOrder(); - for (xla::HloComputation* computation : all_computations) { - if (profiled_computations.count(computation) > 0) { - string profile_string = profile_ptr->ToString( - *computation, stream->parent()->GetDeviceDescription(), shape_size); - if (!profile_string.empty()) { - XLA_LOG_LINES(tensorflow::INFO, profile_string); - } - } - } - DumpExecutedHlo(executable->module(), "Service::Execute", profile_ptr); - } - - return return_value; -} } // namespace xla #endif // TENSORFLOW_COMPILER_XLA_SERVICE_SERVICE_H_ diff --git a/tensorflow/compiler/xla/service/service_executable_run_options.h b/tensorflow/compiler/xla/service/service_executable_run_options.h index 0d4b214f5f3624971ae68e23f0f4fdba846f9178..017e5ef09ed2f52b862821e9408540d188a1edf5 100644 --- a/tensorflow/compiler/xla/service/service_executable_run_options.h +++ b/tensorflow/compiler/xla/service/service_executable_run_options.h @@ -30,10 +30,12 @@ class ServiceExecutableRunOptions { using StreamBorrower = std::function::SmartPtr>(int)>; - explicit ServiceExecutableRunOptions(ExecutableRunOptions run_options, - StreamBorrower borrow_stream = nullptr) + explicit ServiceExecutableRunOptions( + ExecutableRunOptions run_options, StreamBorrower borrow_stream = nullptr, + tensorflow::thread::ThreadPool* xla_intra_op_thread_pool = nullptr) : run_options_(std::move(run_options)), - borrow_stream_(std::move(borrow_stream)) {} + borrow_stream_(std::move(borrow_stream)), + xla_intra_op_thread_pool_(xla_intra_op_thread_pool) {} // Returns reference or pointer to `ExecutableRunOptions` member. const ExecutableRunOptions& run_options() const { return run_options_; } @@ -53,9 +55,15 @@ 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/session.proto b/tensorflow/compiler/xla/service/session.proto index 4902cb521c2f4f5f51958fe3527d484364b162cd..bb8d1cd2a106ea3e5bb61eee5052bd60c38cd0e2 100644 --- a/tensorflow/compiler/xla/service/session.proto +++ b/tensorflow/compiler/xla/service/session.proto @@ -75,10 +75,10 @@ message SessionModule { repeated SessionComputation embedded_computations = 2; // The arguments passed to the computation. - repeated Literal arguments = 3; + repeated LiteralProto arguments = 3; // The result of the computation. - Literal result = 4; + LiteralProto result = 4; // The name of the platform used to run the computation. string execution_platform = 5; diff --git a/tensorflow/compiler/xla/service/shape_inference.cc b/tensorflow/compiler/xla/service/shape_inference.cc index 9472086e2b49b9a95f9cc09e96a17b7419e067bc..1a24c6c4939af6119e17c400eaad4539e1e5cb1a 100644 --- a/tensorflow/compiler/xla/service/shape_inference.cc +++ b/tensorflow/compiler/xla/service/shape_inference.cc @@ -39,6 +39,112 @@ namespace xla { namespace { +// Return the UnaryOperation proto enum value associated with the given HLO +// opcode. +UnaryOperation OpcodeToUnaryOperation(HloOpcode opcode) { + switch (opcode) { + case HloOpcode::kAbs: + return UNOP_ABS; + case HloOpcode::kCeil: + return UNOP_CEIL; + case HloOpcode::kCos: + return UNOP_COS; + case HloOpcode::kExp: + return UNOP_EXP; + case HloOpcode::kFloor: + return UNOP_FLOOR; + case HloOpcode::kIsFinite: + return UNOP_IS_FINITE; + case HloOpcode::kLog: + return UNOP_LOG; + case HloOpcode::kLogicalNot: + return UNOP_LOGICAL_NOT; + case HloOpcode::kNegate: + return UNOP_NEGATE; + case HloOpcode::kSign: + return UNOP_SIGN; + case HloOpcode::kSin: + return UNOP_SIN; + case HloOpcode::kSort: + return UNOP_SORT; + case HloOpcode::kTanh: + return UNOP_TANH; + default: + LOG(FATAL) << "unhandled opcode " << opcode; + } +} + +// Return the BinaryOperation proto enum value associated with the given HLO +// opcode. +BinaryOperation OpcodeToBinaryOperation(HloOpcode opcode) { + switch (opcode) { + case HloOpcode::kDot: + return BINOP_DOT; + case HloOpcode::kMultiply: + return BINOP_MUL; + case HloOpcode::kAdd: + return BINOP_ADD; + case HloOpcode::kSubtract: + return BINOP_SUB; + case HloOpcode::kIndex: + return BINOP_INDEX; + case HloOpcode::kDivide: + return BINOP_DIV; + case HloOpcode::kEq: + return BINOP_EQ; + case HloOpcode::kGe: + return BINOP_GE; + case HloOpcode::kGt: + return BINOP_GT; + case HloOpcode::kLe: + return BINOP_LE; + case HloOpcode::kLt: + return BINOP_LT; + case HloOpcode::kNe: + return BINOP_NE; + case HloOpcode::kMaximum: + return BINOP_MAX; + case HloOpcode::kMinimum: + return BINOP_MIN; + case HloOpcode::kPower: + return BINOP_POW; + case HloOpcode::kRemainder: + return BINOP_REM; + case HloOpcode::kLogicalOr: + return BINOP_LOGICAL_OR; + case HloOpcode::kLogicalAnd: + return BINOP_LOGICAL_AND; + default: + LOG(FATAL) << "unhandled opcode " << opcode; + } +} + +// Return the TernaryOperation proto enum value associated with the given HLO +// opcode. +TernaryOperation OpcodeToTernaryOperation(HloOpcode opcode) { + switch (opcode) { + case HloOpcode::kClamp: + return TRIOP_CLAMP; + case HloOpcode::kSelect: + return TRIOP_SELECT; + case HloOpcode::kUpdate: + return TRIOP_UPDATE; + default: + LOG(FATAL) << "unhandled opcode " << opcode; + } +} + +// Return the VariadicOperation proto enum value associated with the given HLO +// opcode. +VariadicOperation OpcodeToVariadicOperation(HloOpcode opcode) { + switch (opcode) { + case HloOpcode::kTuple: + return VAROP_TUPLE; + default: + LOG(FATAL) << "unhandled opcode " << opcode; + } +} + // Returns true if no element is present in slice more than once. bool AllUnique(tensorflow::gtl::ArraySlice slice) { return std::set(slice.begin(), slice.end()).size() == slice.size(); @@ -176,14 +282,26 @@ StatusOr InferWindowOutputShape(const Shape& base_shape, } // namespace +/* static */ StatusOr ShapeInference::InferUnaryOpShape( + HloOpcode opcode, const HloInstruction* operand) { + // There is no copy operation at the proto level, so handle copy explicitly. + if (opcode == HloOpcode::kCopy) { + return operand->shape(); + } + + return InferUnaryOpShape(OpcodeToUnaryOperation(opcode), operand->shape()); +} + /* static */ StatusOr ShapeInference::InferUnaryOpShape( UnaryOperation operation, const Shape& arg) { TF_RETURN_IF_ERROR(ExpectNotTupleOrOpaque(arg, "operand of unary operation")); - TF_DCHECK_OK(ShapeUtil::ValidateShape(arg)); + TF_DCHECK_OK(ShapeUtil::ValidateShapeWithOptionalLayout(arg)); switch (operation) { case UNOP_FLOOR: case UNOP_CEIL: + case UNOP_COS: + case UNOP_SIN: case UNOP_EXP: case UNOP_LOG: case UNOP_TANH: @@ -227,7 +345,7 @@ StatusOr InferWindowOutputShape(const Shape& base_shape, /* static */ StatusOr ShapeInference::InferConcatOpShape( tensorflow::gtl::ArraySlice arg_shapes, const int64 dimension) { - if (arg_shapes.size() == 0) { + if (arg_shapes.empty()) { return InvalidArgument("Concatenate expects at least one argument"); } if (dimension < 0 || dimension >= ShapeUtil::Rank(*arg_shapes[0])) { @@ -244,8 +362,11 @@ StatusOr InferWindowOutputShape(const Shape& base_shape, } if (ShapeUtil::Rank(*arg_shape) != ShapeUtil::Rank(*shape)) { return InvalidArgument( - "cannot concatenate arrays with different ranks: %lld vs %lld", - ShapeUtil::Rank(*arg_shape), ShapeUtil::Rank(*shape)); + "Cannot concatenate arrays with different ranks: %lld (%s) vs %lld " + "(%s)", + ShapeUtil::Rank(*arg_shape), + ShapeUtil::HumanString(*arg_shape).c_str(), ShapeUtil::Rank(*shape), + ShapeUtil::HumanString(*shape).c_str()); } if (arg_shape->element_type() != shape->element_type()) { return InvalidArgument( @@ -264,9 +385,9 @@ StatusOr InferWindowOutputShape(const Shape& base_shape, return InvalidArgument( "cannot concatenate arrays that differ in dimensions other than " "the one being concatenated (the other array dimensions must be " - "the same): %s vs %s", + "the same): %s vs %s in dimension %lld", ShapeUtil::HumanString(*arg_shape).c_str(), - ShapeUtil::HumanString(*shape).c_str()); + ShapeUtil::HumanString(*shape).c_str(), dimension); } } } @@ -294,6 +415,30 @@ StatusOr InferWindowOutputShape(const Shape& base_shape, return ShapeUtil::ChangeElementType(operand_shape, new_element_type); } +/* static */ StatusOr ShapeInference::InferReducePrecisionShape( + const Shape& operand_shape, const int exponent_bits, + const int mantissa_bits) { + if (!ShapeUtil::ElementIsFloating(operand_shape)) { + return InvalidArgument( + "expected element type in shape to be floating point for " + "ReducePrecision operation; got %s", + PrimitiveType_Name(operand_shape.element_type()).c_str()); + } + if (exponent_bits < 1) { + // One exponent bit is necessary to distinguish 0 from infinity. Having + // no exponent bits doesn't produce a sensible number, so we require at + // least one. + return InvalidArgument("expected exponent_bits >= 1; got %d", + exponent_bits); + } + if (mantissa_bits < 0) { + // A number with no mantissa bits is still meaningful, however. + return InvalidArgument("expected non-negative mantissa_bits; got %d", + mantissa_bits); + } + return operand_shape; +} + /* static */ StatusOr ShapeInference::InferPadShape( const Shape& operand_shape, const Shape& padding_value_shape, const PaddingConfig& padding_config) { @@ -309,6 +454,10 @@ StatusOr InferWindowOutputShape(const Shape& base_shape, return InvalidArgument( "the rank of the operand and the padding configuration do not match."); } + if (operand_shape.element_type() != padding_value_shape.element_type()) { + return InvalidArgument( + "the element types of the operands to pad do not match"); + } std::vector dimensions(ShapeUtil::Rank(operand_shape)); for (int64 i = 0; i < operand_shape.dimensions_size(); ++i) { dimensions[i] = operand_shape.dimensions(i) + @@ -338,7 +487,7 @@ StatusOr InferWindowOutputShape(const Shape& base_shape, // Check if both element types are the same. if (lhs.element_type() != rhs.element_type()) { - return fail("element types mismatch"); + return fail("element types do not match"); } if (ShapeUtil::Rank(lhs) < 1 || ShapeUtil::Rank(lhs) > 2 || @@ -377,7 +526,7 @@ StatusOr InferWindowOutputShape(const Shape& base_shape, } Shape result = ShapeUtil::MakeShape(lhs.element_type(), dimensions); - TF_DCHECK_OK(ShapeUtil::ValidateShape(result)); + TF_DCHECK_OK(ShapeUtil::ValidateShapeWithOptionalLayout(result)); VLOG(2) << "inferred dot shape: " << ShapeUtil::HumanString(result); return result; } @@ -518,9 +667,11 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( ExpectNotTupleOrOpaque(rhs, "rhs of elementwise binary operation")); if (!ShapeUtil::SameElementType(lhs, rhs)) { - return InvalidArgument("binary op with different element types: %s and %s", - ShapeUtil::HumanString(lhs).c_str(), - ShapeUtil::HumanString(rhs).c_str()); + return InvalidArgument( + "binary op %s with different element types: %s and %s", + BinaryOperation_Name(operation).c_str(), + ShapeUtil::HumanString(lhs).c_str(), + ShapeUtil::HumanString(rhs).c_str()); } if (ShapeUtil::Rank(lhs) == ShapeUtil::Rank(rhs) && @@ -540,7 +691,7 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( 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). + // 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 = @@ -557,6 +708,12 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( } } +/* static */ StatusOr ShapeInference::InferBinaryOpShape( + HloOpcode opcode, const HloInstruction* lhs, const HloInstruction* rhs) { + return InferBinaryOpShape(OpcodeToBinaryOperation(opcode), lhs->shape(), + rhs->shape(), /*broadcast_dimensions=*/{}); +} + /* static */ StatusOr ShapeInference::InferBinaryOpShape( BinaryOperation operation, const Shape& lhs, const Shape& rhs, tensorflow::gtl::ArraySlice broadcast_dimensions) { @@ -565,8 +722,8 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( BinaryOperation_Name(operation).c_str(), ShapeUtil::HumanString(lhs).c_str(), ShapeUtil::HumanString(rhs).c_str(), tensorflow::str_util::Join(broadcast_dimensions, ", ").c_str()); - TF_DCHECK_OK(ShapeUtil::ValidateShape(lhs)); - TF_DCHECK_OK(ShapeUtil::ValidateShape(rhs)); + TF_DCHECK_OK(ShapeUtil::ValidateShapeWithOptionalLayout(lhs)); + TF_DCHECK_OK(ShapeUtil::ValidateShapeWithOptionalLayout(rhs)); TF_RETURN_IF_ERROR(ExpectNotTupleOrOpaque(lhs, "lhs of binary operation")); TF_RETURN_IF_ERROR(ExpectNotTupleOrOpaque(rhs, "rhs of binary operation")); @@ -625,12 +782,19 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( } } +/* static */ StatusOr ShapeInference::InferTernaryOpShape( + HloOpcode opcode, const HloInstruction* lhs, const HloInstruction* rhs, + const HloInstruction* ehs) { + return InferTernaryOpShape(OpcodeToTernaryOperation(opcode), lhs->shape(), + rhs->shape(), ehs->shape()); +} + /* static */ StatusOr ShapeInference::InferTernaryOpShape( TernaryOperation operation, const Shape& lhs, const Shape& rhs, const Shape& ehs) { - TF_DCHECK_OK(ShapeUtil::ValidateShape(lhs)); - TF_DCHECK_OK(ShapeUtil::ValidateShape(rhs)); - TF_DCHECK_OK(ShapeUtil::ValidateShape(ehs)); + TF_DCHECK_OK(ShapeUtil::ValidateShapeWithOptionalLayout(lhs)); + TF_DCHECK_OK(ShapeUtil::ValidateShapeWithOptionalLayout(rhs)); + TF_DCHECK_OK(ShapeUtil::ValidateShapeWithOptionalLayout(ehs)); switch (operation) { case TRIOP_CLAMP: return InferClampShape(lhs, rhs, ehs); @@ -651,9 +815,21 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( } /* static */ StatusOr ShapeInference::InferVariadicOpShape( - VariadicOperation operation, std::vector operand_shapes) { + HloOpcode opcode, + tensorflow::gtl::ArraySlice operands) { + std::vector operand_shapes; + for (const HloInstruction* operand : operands) { + operand_shapes.push_back(&operand->shape()); + } + return InferVariadicOpShape(OpcodeToVariadicOperation(opcode), + operand_shapes); +} + +/* static */ StatusOr ShapeInference::InferVariadicOpShape( + VariadicOperation operation, + tensorflow::gtl::ArraySlice operand_shapes) { for (const Shape* shape : operand_shapes) { - TF_DCHECK_OK(ShapeUtil::ValidateShape(*shape)); + TF_DCHECK_OK(ShapeUtil::ValidateShapeWithOptionalLayout(*shape)); } switch (operation) { case VAROP_TUPLE: { @@ -672,7 +848,7 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( /* static */ StatusOr ShapeInference::InferMapShape( tensorflow::gtl::ArraySlice arg_shapes, const ProgramShape& to_apply) { - if (arg_shapes.size() == 0) { + if (arg_shapes.empty()) { return InvalidArgument("Map expects at least one argument"); } @@ -747,6 +923,408 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( AsInt64Slice(arg_shape->dimensions())); } +/* static */ StatusOr ShapeInference::InferBatchNormTrainingShape( + const Shape& operand_shape, const Shape& offset_shape, + const Shape& scale_shape, int64 feature_index) { + TF_RETURN_IF_ERROR( + ExpectNotTupleOrOpaque(operand_shape, "operand of batch norm training")); + TF_RETURN_IF_ERROR(ExpectNotTupleOrOpaque( + offset_shape, "offset input of batch norm training")); + TF_RETURN_IF_ERROR(ExpectNotTupleOrOpaque( + scale_shape, "scale input of batch norm training")); + + TF_RET_CHECK(ShapeUtil::ValidateShapeWithOptionalLayout(operand_shape) == + tensorflow::Status::OK()); + TF_RET_CHECK(ShapeUtil::ValidateShapeWithOptionalLayout(offset_shape) == + tensorflow::Status::OK()); + TF_RET_CHECK(ShapeUtil::ValidateShapeWithOptionalLayout(scale_shape) == + tensorflow::Status::OK()); + + if (feature_index >= ShapeUtil::Rank(operand_shape)) { + return InvalidArgument( + "Expected feature_index of batch-norm-training to be " + "smaller than the rank of operand_shape; " + "got feature_index %lld, and rank %lld", + feature_index, ShapeUtil::Rank(operand_shape)); + } + + if (feature_index < 0) { + return InvalidArgument( + "Expected feature_index of batch-norm-training to " + "be a non-negative number, got %lld", + feature_index); + } + + if (ShapeUtil::Rank(operand_shape) < 1) { + return InvalidArgument( + "Expected the rank of operand to " + "batch-norm-training to be at least 1; got %lld", + ShapeUtil::Rank(operand_shape)); + } + + if (ShapeUtil::Rank(offset_shape) != 1) { + return InvalidArgument( + "Offset input of batch-norm-training must have" + " rank 1, but has rank %lld.", + ShapeUtil::Rank(offset_shape)); + } + + if (ShapeUtil::Rank(scale_shape) != 1) { + return InvalidArgument( + "Scale input of batch-norm-training must have" + " rank 1, but has rank %lld.", + ShapeUtil::Rank(scale_shape)); + } + + if (!ShapeUtil::ElementIsFloating(operand_shape)) { + return InvalidArgument( + "The operand to batch-norm-training must have a floating point " + "element type, but the shape is %s", + PrimitiveType_Name(operand_shape.element_type()).c_str()); + } + + if (!ShapeUtil::SameElementType(offset_shape, operand_shape)) { + return InvalidArgument( + "The inputs should have the same element type for batch-norm-training, " + "but the shape of offset factor is %s " + "and the shape of operand is %s", + PrimitiveType_Name(offset_shape.element_type()).c_str(), + PrimitiveType_Name(operand_shape.element_type()).c_str()); + } + + if (!ShapeUtil::SameElementType(scale_shape, operand_shape)) { + return InvalidArgument( + "The inputs should have the same element type for batch-norm-training, " + "but the shape of scale factor is %s " + "and the shape of operand is %s", + PrimitiveType_Name(scale_shape.element_type()).c_str(), + PrimitiveType_Name(operand_shape.element_type()).c_str()); + } + + const int64 feature_count = operand_shape.dimensions(feature_index); + Shape output_shape_for_mean_and_var = + ShapeUtil::MakeShape(operand_shape.element_type(), {feature_count}); + + if (ShapeUtil::GetDimension(offset_shape, 0) != feature_count) { + return InvalidArgument( + "The size of offset factor should be the same as feature count," + "but the size of offset factor is %lld " + "and the feature count is %lld", + ShapeUtil::GetDimension(offset_shape, 0), feature_count); + } + + if (ShapeUtil::GetDimension(scale_shape, 0) != feature_count) { + return InvalidArgument( + "The size of scale factor should be the same as feature count," + "but the size of scale factor is %lld " + "and the feature count is %lld", + ShapeUtil::GetDimension(scale_shape, 0), feature_count); + } + + return ShapeUtil::MakeTupleShape({operand_shape, + output_shape_for_mean_and_var, + output_shape_for_mean_and_var}); +} + +/* static */ StatusOr ShapeInference::InferBatchNormInferenceShape( + const Shape& operand_shape, const Shape& offset_shape, + const Shape& scale_shape, const Shape& mean_shape, + const Shape& variance_shape, int64 feature_index) { + TF_RETURN_IF_ERROR( + ExpectNotTupleOrOpaque(operand_shape, "operand of batch norm inference")); + TF_RETURN_IF_ERROR(ExpectNotTupleOrOpaque( + offset_shape, "offset input of batch norm inference")); + TF_RETURN_IF_ERROR(ExpectNotTupleOrOpaque( + scale_shape, "scale input of batch norm inference")); + + TF_RET_CHECK(ShapeUtil::ValidateShapeWithOptionalLayout(operand_shape) == + tensorflow::Status::OK()); + TF_RET_CHECK(ShapeUtil::ValidateShapeWithOptionalLayout(offset_shape) == + tensorflow::Status::OK()); + TF_RET_CHECK(ShapeUtil::ValidateShapeWithOptionalLayout(scale_shape) == + tensorflow::Status::OK()); + TF_RET_CHECK(ShapeUtil::ValidateShapeWithOptionalLayout(mean_shape) == + tensorflow::Status::OK()); + TF_RET_CHECK(ShapeUtil::ValidateShapeWithOptionalLayout(variance_shape) == + tensorflow::Status::OK()); + + if (feature_index >= ShapeUtil::Rank(operand_shape)) { + return InvalidArgument( + "Expected feature_index of batch-norm-inference to be " + "smaller than the rank of operand_shape; " + "got feature_index %lld, and rank %lld", + feature_index, ShapeUtil::Rank(operand_shape)); + } + + if (feature_index < 0) { + return InvalidArgument( + "Expected feature_index of batch-norm-inference to " + "be a non-negative number, got %lld", + feature_index); + } + + if (ShapeUtil::Rank(operand_shape) < 1) { + return InvalidArgument( + "Expected the rank of operand to " + "batch-norm-inference to be at least 1; got %lld", + ShapeUtil::Rank(operand_shape)); + } + + if (ShapeUtil::Rank(offset_shape) != 1) { + return InvalidArgument( + "Offset input of batch-norm-inference must have" + " rank 1, but has rank %lld.", + ShapeUtil::Rank(offset_shape)); + } + + if (ShapeUtil::Rank(scale_shape) != 1) { + return InvalidArgument( + "Scale input of batch-norm-inference must have" + " rank 1, but has rank %lld.", + ShapeUtil::Rank(scale_shape)); + } + + if (!ShapeUtil::ElementIsFloating(operand_shape)) { + return InvalidArgument( + "The operand to batch-norm-inference must have a floating point " + "element type, but the shape is %s", + PrimitiveType_Name(operand_shape.element_type()).c_str()); + } + + if (!ShapeUtil::SameElementType(offset_shape, operand_shape)) { + return InvalidArgument( + "The inputs should have the same element type for " + "batch-norm-inference, " + "but the shape of offset factor is %s " + "and the shape of operand is %s", + PrimitiveType_Name(offset_shape.element_type()).c_str(), + PrimitiveType_Name(operand_shape.element_type()).c_str()); + } + + if (!ShapeUtil::SameElementType(scale_shape, operand_shape)) { + return InvalidArgument( + "The inputs should have the same element type for " + "batch-norm-inference, " + "but the shape of scale factor is %s " + "and the shape of operand is %s", + PrimitiveType_Name(scale_shape.element_type()).c_str(), + PrimitiveType_Name(operand_shape.element_type()).c_str()); + } + + if (!ShapeUtil::SameElementType(mean_shape, operand_shape)) { + return InvalidArgument( + "The inputs should have the same element type for " + "batch-norm-inference, " + "but the shape of mean is %s " + "and the shape of operand is %s", + PrimitiveType_Name(mean_shape.element_type()).c_str(), + PrimitiveType_Name(operand_shape.element_type()).c_str()); + } + + if (!ShapeUtil::SameElementType(variance_shape, operand_shape)) { + return InvalidArgument( + "The inputs should have the same element type for " + "batch-norm-inference, " + "but the shape of variance is %s " + "and the shape of operand is %s", + PrimitiveType_Name(mean_shape.element_type()).c_str(), + PrimitiveType_Name(variance_shape.element_type()).c_str()); + } + + const int64 feature_count = operand_shape.dimensions(feature_index); + Shape output_shape_for_mean_and_var = + ShapeUtil::MakeShape(operand_shape.element_type(), {feature_count}); + + if (ShapeUtil::GetDimension(offset_shape, 0) != feature_count) { + return InvalidArgument( + "The size of offset factor should be the same as feature count," + "but the size of offset factor is %lld " + "and the feature count is %lld", + ShapeUtil::GetDimension(offset_shape, 0), feature_count); + } + + if (ShapeUtil::GetDimension(scale_shape, 0) != feature_count) { + return InvalidArgument( + "The size of scale factor should be the same as feature count," + "but the size of scale factor is %lld " + "and the feature count is %lld", + ShapeUtil::GetDimension(scale_shape, 0), feature_count); + } + + if (ShapeUtil::GetDimension(mean_shape, 0) != feature_count) { + return InvalidArgument( + "The size of mean should be the same as feature count," + "but the size of mean is %lld " + "and the feature count is %lld", + ShapeUtil::GetDimension(mean_shape, 0), feature_count); + } + + if (ShapeUtil::GetDimension(variance_shape, 0) != feature_count) { + return InvalidArgument( + "The size of variance should be the same as feature count," + "but the size of variance is %lld " + "and the feature count is %lld", + ShapeUtil::GetDimension(variance_shape, 0), feature_count); + } + + return operand_shape; +} + +/* static */ StatusOr ShapeInference::InferBatchNormGradShape( + const Shape& operand_shape, const Shape& scale_shape, + const Shape& mean_shape, const Shape& var_shape, + const Shape& output_grad_shape, int64 feature_index) { + TF_RETURN_IF_ERROR( + ExpectNotTupleOrOpaque(operand_shape, "operand of batch norm grad")); + TF_RETURN_IF_ERROR( + ExpectNotTupleOrOpaque(scale_shape, "scale input of batch norm grad")); + TF_RETURN_IF_ERROR( + ExpectNotTupleOrOpaque(mean_shape, "mean input of batch norm grad")); + TF_RETURN_IF_ERROR( + ExpectNotTupleOrOpaque(var_shape, "var input of batch norm grad")); + TF_RETURN_IF_ERROR(ExpectNotTupleOrOpaque( + output_grad_shape, "output_grad input of batch norm grad")); + + TF_RETURN_IF_ERROR(ShapeUtil::ValidateShapeWithOptionalLayout(operand_shape)); + TF_RETURN_IF_ERROR(ShapeUtil::ValidateShapeWithOptionalLayout(mean_shape)); + TF_RETURN_IF_ERROR(ShapeUtil::ValidateShapeWithOptionalLayout(scale_shape)); + TF_RETURN_IF_ERROR(ShapeUtil::ValidateShapeWithOptionalLayout(var_shape)); + TF_RETURN_IF_ERROR( + ShapeUtil::ValidateShapeWithOptionalLayout(output_grad_shape)); + + if (feature_index >= ShapeUtil::Rank(operand_shape)) { + return InvalidArgument( + "Expected feature_index of batch-norm-grad to be " + "smaller than the rank of operand_shape; " + "got feature_index %lld, and rank %lld", + feature_index, ShapeUtil::Rank(operand_shape)); + } + + if (ShapeUtil::Rank(operand_shape) != ShapeUtil::Rank(output_grad_shape)) { + return InvalidArgument( + "Expected operand_shape of batch-norm-grad to have the same rank as" + " output_grad_shape; got rank(oprand_shape) %lld, and" + " rank(output_grad_shape) %lld", + ShapeUtil::Rank(operand_shape), ShapeUtil::Rank(output_grad_shape)); + } + + if (ShapeUtil::Rank(mean_shape) != 1) { + return InvalidArgument( + "Mean input of batch-norm-grad must have" + " rank 1, but has rank %lld.", + ShapeUtil::Rank(mean_shape)); + } + + if (ShapeUtil::Rank(scale_shape) != 1) { + return InvalidArgument( + "Scale input of batch-norm-grad must have" + " rank 1, but has rank %lld.", + ShapeUtil::Rank(scale_shape)); + } + + if (ShapeUtil::Rank(var_shape) != 1) { + return InvalidArgument( + "Var input of batch-norm-grad must have" + " rank 1, but has rank %lld.", + ShapeUtil::Rank(var_shape)); + } + + if (!ShapeUtil::ElementIsFloating(operand_shape)) { + return InvalidArgument( + "The operand to batch-norm-grad must have a floating point " + "element type, but the shape is %s", + PrimitiveType_Name(operand_shape.element_type()).c_str()); + } + + if (!ShapeUtil::ElementIsFloating(output_grad_shape)) { + return InvalidArgument( + "The output_grad to batch-norm-grad must have a floating point " + "element type, but the shape is %s", + PrimitiveType_Name(output_grad_shape.element_type()).c_str()); + } + + if (!ShapeUtil::SameElementType(output_grad_shape, operand_shape)) { + return InvalidArgument( + "The inputs should have the same element type for batch-norm-grad, " + "but the element type of output_grad is %s " + "and the element type of operand is %s", + PrimitiveType_Name(output_grad_shape.element_type()).c_str(), + PrimitiveType_Name(operand_shape.element_type()).c_str()); + } + + if (!ShapeUtil::SameElementType(scale_shape, operand_shape)) { + return InvalidArgument( + "The inputs should have the same element type for batch-norm-grad, " + "but the element type of scale factor is %s " + "and the element type of operand is %s", + PrimitiveType_Name(scale_shape.element_type()).c_str(), + PrimitiveType_Name(operand_shape.element_type()).c_str()); + } + + if (!ShapeUtil::SameElementType(mean_shape, operand_shape)) { + return InvalidArgument( + "The inputs should have the same element type for batch-norm-grad, " + "but the element type of mean is %s " + "and the element type of operand is %s", + PrimitiveType_Name(mean_shape.element_type()).c_str(), + PrimitiveType_Name(operand_shape.element_type()).c_str()); + } + + if (!ShapeUtil::SameElementType(var_shape, operand_shape)) { + return InvalidArgument( + "The inputs should have the same element type for batch-norm-grad, " + "but the element type of mean is %s " + "and the element type of operand is %s", + PrimitiveType_Name(mean_shape.element_type()).c_str(), + PrimitiveType_Name(operand_shape.element_type()).c_str()); + } + + const int64 feature_count = operand_shape.dimensions(feature_index); + + Shape feature_shape = + ShapeUtil::MakeShape(operand_shape.element_type(), {feature_count}); + + if (ShapeUtil::GetDimension(mean_shape, 0) != feature_count) { + return InvalidArgument( + "The size of mean should be the same as feature count," + "but the size of offset factor is %lld " + "and the feature count is %lld", + ShapeUtil::GetDimension(mean_shape, 0), feature_count); + } + + if (ShapeUtil::GetDimension(scale_shape, 0) != feature_count) { + return InvalidArgument( + "The size of scale factor should be the same as feature count," + "but the size of scale factor is %lld " + "and the feature count is %lld", + ShapeUtil::GetDimension(scale_shape, 0), feature_count); + } + + if (ShapeUtil::GetDimension(var_shape, 0) != feature_count) { + return InvalidArgument( + "The size of variance should be the same as feature count," + "but the size of variance is %lld " + "and the feature count is %lld", + ShapeUtil::GetDimension(var_shape, 0), feature_count); + } + + // Verify operand_shape and output_grad_shape have same bounds. + for (int64 i = 0; i < ShapeUtil::Rank(operand_shape); ++i) { + if (ShapeUtil::GetDimension(operand_shape, i) != + ShapeUtil::GetDimension(output_grad_shape, i)) { + return InvalidArgument( + "The bounds of operand shape should be the same as output_grad's," + "but the bound of operand_shape at dimension %lld is %lld " + "and the bound of output_grad_shape is %lld", + i, ShapeUtil::GetDimension(operand_shape, i), + ShapeUtil::GetDimension(output_grad_shape, i)); + } + } + + return ShapeUtil::MakeTupleShape( + {operand_shape, feature_shape, feature_shape}); +} + /* static */ StatusOr ShapeInference::InferConvolveShape( const Shape& lhs, const Shape& rhs, const Window& window, const ConvolutionDimensionNumbers& dnums) { @@ -794,8 +1372,8 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( "lhs: %s", num_dims, ShapeUtil::HumanString(lhs).c_str()); } - TF_DCHECK_OK(ShapeUtil::ValidateShape(lhs)); - TF_DCHECK_OK(ShapeUtil::ValidateShape(rhs)); + TF_DCHECK_OK(ShapeUtil::ValidateShapeWithOptionalLayout(lhs)); + TF_DCHECK_OK(ShapeUtil::ValidateShapeWithOptionalLayout(rhs)); // Verifies that the input and window dimensions are a permutation of // the dimension numbers. @@ -998,7 +1576,8 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( /* static */ StatusOr ShapeInference::InferSliceShape( const Shape& arg, tensorflow::gtl::ArraySlice starts, - tensorflow::gtl::ArraySlice limits) { + tensorflow::gtl::ArraySlice limits, + tensorflow::gtl::ArraySlice strides) { TF_RETURN_IF_ERROR(ExpectNotTupleOrOpaque(arg, "operand of slice")); VLOG(2) << tensorflow::strings::Printf( "slicing shape %s starts={%s} limits={%s}", @@ -1011,6 +1590,11 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( starts.size(), limits.size()); } + if (starts.size() != strides.size()) { + return InvalidArgument("slice start and strides sizes differ: %zu vs %zu", + starts.size(), strides.size()); + } + if (starts.size() != ShapeUtil::Rank(arg)) { return InvalidArgument( "slice index count does not match argument rank: %zu vs %lld", @@ -1021,32 +1605,31 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( for (int64 dimension = 0; dimension < starts.size(); ++dimension) { int64 start_index = starts[dimension]; int64 limit_index = limits[dimension]; + int64 stride = strides[dimension]; if (start_index < 0) { return InvalidArgument("negative start index to slice: %lld", start_index); } - if (limit_index < 0) { - return InvalidArgument("negative limit index to slice: %lld", - limit_index); - } if (limit_index > arg.dimensions(dimension)) { return InvalidArgument( "limit index (%lld) must be less than or equal to dimension " "size (%lld)", limit_index, arg.dimensions(dimension)); } + VLOG(2) << tensorflow::strings::Printf("starts[%lld] = %lld", dimension, + start_index); + VLOG(2) << tensorflow::strings::Printf("limits[%lld] = %lld", dimension, + limit_index); if (start_index > limit_index) { return InvalidArgument( "limit index (%lld) must be greater or equal to " - "start index (%lld) in slice", + "start index (%lld) in slice with positive stride", limit_index, start_index); } - VLOG(2) << tensorflow::strings::Printf("starts[%lld] = %lld", dimension, - start_index); - VLOG(2) << tensorflow::strings::Printf("limits[%lld] = %lld", dimension, - limit_index); - - sizes.push_back(limits[dimension] - starts[dimension]); + if (stride <= 0) { + return InvalidArgument("stride (%lld) must be positive", stride); + } + sizes.push_back((limit_index - start_index + stride - 1) / stride); } return ShapeUtil::MakeShape(arg.element_type(), sizes); @@ -1080,9 +1663,11 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( 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 %lld must match rank %lld of " - "slice input", - start_num_dims, ShapeUtil::Rank(operand_shape)); + "dynamic slice start number of dimensions %lld (%s) must match rank " + "%lld of slice input (%s)", + start_num_dims, ShapeUtil::HumanString(start_indices_shape).c_str(), + ShapeUtil::Rank(operand_shape), + ShapeUtil::HumanString(operand_shape).c_str()); } if (slice_sizes.size() != ShapeUtil::Rank(operand_shape)) { @@ -1094,7 +1679,7 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( for (int64 dim = 0; dim < slice_sizes.size(); ++dim) { const int64 input_dim_size = operand_shape.dimensions(dim); const int64 slice_dim_size = slice_sizes[dim]; - if (slice_dim_size <= 0) { + if (slice_dim_size < 0) { return InvalidArgument("negative size index to dynamic slice: %lld", slice_dim_size); } @@ -1141,9 +1726,11 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( 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 %lld must match " - "rank %lld of slice input", - start_num_dims, ShapeUtil::Rank(operand_shape)); + "dynamic slice start number of dimensions %lld (%s) must match rank " + "%lld of slice input (%s)", + start_num_dims, ShapeUtil::HumanString(start_indices_shape).c_str(), + ShapeUtil::Rank(operand_shape), + ShapeUtil::HumanString(operand_shape).c_str()); } if (ShapeUtil::Rank(update_shape) != ShapeUtil::Rank(operand_shape)) { @@ -1164,9 +1751,9 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( for (int64 dim = 0; dim < ShapeUtil::Rank(operand_shape); ++dim) { const int64 input_dim_size = operand_shape.dimensions(dim); const int64 update_dim_size = update_shape.dimensions(dim); - if (update_dim_size <= 0) { + if (update_dim_size < 0) { return InvalidArgument( - "size index %lld to dynamic update slice must be > 0", + "size index %lld to dynamic update slice must be >= 0", update_dim_size); } if (update_dim_size > input_dim_size) { @@ -1379,10 +1966,17 @@ ShapeInference::InferDegenerateDimensionBroadcastShape( const ProgramShape& to_apply) { // The applied function's arity equals the number of arguments. if (arg_shapes.size() != to_apply.parameters_size()) { + string computation_signature = ShapeUtil::HumanString(to_apply); + string argument_shapes = tensorflow::str_util::Join( + arg_shapes, ", ", [](string* out, const Shape* shape) { + tensorflow::strings::StrAppend(out, ShapeUtil::HumanString(*shape)); + }); return InvalidArgument( "Call applied function arity must match number of arguments; got: " - "arity: %d, arguments: %zu", - to_apply.parameters_size(), arg_shapes.size()); + "arity: %d, arguments: %zu; computation signature: %s; argument " + "shapes: [%s]", + to_apply.parameters_size(), arg_shapes.size(), + computation_signature.c_str(), argument_shapes.c_str()); } // All arguments must be compatible with the program shape. diff --git a/tensorflow/compiler/xla/service/shape_inference.h b/tensorflow/compiler/xla/service/shape_inference.h index c2223423e9223ba8ad995212415f219eea48e2a6..96e3b46c7dece6c945ae9b2a2a0a4eac8a0eb350 100644 --- a/tensorflow/compiler/xla/service/shape_inference.h +++ b/tensorflow/compiler/xla/service/shape_inference.h @@ -21,6 +21,8 @@ limitations under the License. #include +#include "tensorflow/compiler/xla/service/hlo_instruction.h" +#include "tensorflow/compiler/xla/service/hlo_opcode.h" #include "tensorflow/compiler/xla/statusor.h" #include "tensorflow/compiler/xla/types.h" #include "tensorflow/compiler/xla/xla_data.pb.h" @@ -31,32 +33,48 @@ limitations under the License. namespace xla { // For a given operation and input shapes, infers what the resulting shape is -// for the operation. With this functionality, the user does not need to -// specify the expected result type for computations that are built up via the -// API -- the shape that results from an operation is inferred. +// for the operation. With this functionality, the user does not need to specify +// the expected result type for computations that are built up via the API -- +// the shape that results from an operation is inferred. Some methods have +// overloads for inferring shape at the HLO level. +// TODO(b/166374537): Complete HLO level inference overloads and use to +// automatically infer shape in HloInstruction::Create* methods. class ShapeInference { public: // Infers the shape produced by applying the given unary operation to the // given input shape. static StatusOr InferUnaryOpShape(UnaryOperation operation, const Shape& arg); + static StatusOr InferUnaryOpShape(HloOpcode opcode, + const HloInstruction* operand); // Infers the shape produced by applying the given binary operation to the // given input shapes. static StatusOr InferBinaryOpShape( BinaryOperation operation, const Shape& lhs, const Shape& rhs, tensorflow::gtl::ArraySlice broadcast_dimensions); + static StatusOr InferBinaryOpShape(HloOpcode opcode, + const HloInstruction* lhs, + const HloInstruction* rhs); // Infers the shape produced by applying the given ternary operation to the // given input shapes. static StatusOr InferTernaryOpShape(TernaryOperation operation, const Shape& lhs, const Shape& rhs, const Shape& ehs); + static StatusOr InferTernaryOpShape(HloOpcode opcode, + const HloInstruction* lhs, + const HloInstruction* rhs, + const HloInstruction* ehs); // Infers the shape produced by applying the given variadic operation to the // given input operand shapes. static StatusOr InferVariadicOpShape( - VariadicOperation operation, std::vector operand_shapes); + VariadicOperation operation, + tensorflow::gtl::ArraySlice operand_shapes); + static StatusOr InferVariadicOpShape( + HloOpcode opcode, + tensorflow::gtl::ArraySlice operands); // Infers the shape produced by applying the given mapping computation shape // to the given operand shapes. @@ -64,6 +82,28 @@ class ShapeInference { tensorflow::gtl::ArraySlice arg_shapes, const ProgramShape& to_apply); + // Infers the shape produced by InferBatchNormTraining with the given + // operands. + static StatusOr InferBatchNormTrainingShape(const Shape& operand_shape, + const Shape& offset_shape, + const Shape& scale_shape, + int64 feature_index); + + // Infers the shape produced by InferBatchNormInference with the given + // operands. + static StatusOr InferBatchNormInferenceShape( + const Shape& operand_shape, const Shape& offset_shape, + const Shape& scale_shape, const Shape& mean_shape, + const Shape& variance_shape, int64 feature_index); + + // Infers the shape produced by InferBatchNormGrad with the given operands. + static StatusOr InferBatchNormGradShape(const Shape& operand_shape, + const Shape& scale_shape, + const Shape& mean_shape, + const Shape& var_shape, + const Shape& output_grad_shape, + int64 feature_index); + // Infers the shape produced by applying the given convolutional // filter (rhs) to lhs in the way specified by the fields on window. static StatusOr InferConvolveShape( @@ -109,7 +149,8 @@ class ShapeInference { // e.g. slice f32[32x32] 0:16 0:16 -> f32[16x16] static StatusOr InferSliceShape( const Shape& arg, tensorflow::gtl::ArraySlice starts, - tensorflow::gtl::ArraySlice limits); + tensorflow::gtl::ArraySlice limits, + tensorflow::gtl::ArraySlice strides); // Infers the shape produced by a dynamic slice operation of size specified // in 'slice_sizes', with dynamic start indices shape 'start_indices_shape'. @@ -164,6 +205,12 @@ class ShapeInference { static StatusOr InferConvertShape(const Shape& operand_shape, PrimitiveType new_element_type); + // Helper that validates the input data type for a reduce-precision operation, + // and returns the result shape. + static StatusOr InferReducePrecisionShape(const Shape& operand_shape, + const int exponent_bits, + const int mantissa_bits); + // Helper that infers the shape produced by a pad operation based on the // padding configuration. static StatusOr InferPadShape(const Shape& operand_shape, diff --git a/tensorflow/compiler/xla/service/shape_inference_test.cc b/tensorflow/compiler/xla/service/shape_inference_test.cc index 6f968ded568a5ecf631d234f448f66d10f5d3d26..8c731ae2976fd3da275a5c9596a4ac7f738e5fbc 100644 --- a/tensorflow/compiler/xla/service/shape_inference_test.cc +++ b/tensorflow/compiler/xla/service/shape_inference_test.cc @@ -20,12 +20,16 @@ limitations under the License. #include "tensorflow/compiler/xla/shape_util.h" #include "tensorflow/compiler/xla/xla_data.pb.h" +#include "tensorflow/compiler/xla/test.h" #include "tensorflow/compiler/xla/test_helpers.h" #include "tensorflow/compiler/xla/types.h" namespace xla { namespace { +using ::testing::ContainsRegex; +using ::testing::HasSubstr; + class ShapeInferenceTest : public ::testing::Test { protected: // Some handy scalar shapes. @@ -128,23 +132,21 @@ TEST_F(ShapeInferenceTest, SelectBadShapes) { auto inferred_status_error1 = ShapeInference::InferTernaryOpShape( TernaryOperation::TRIOP_SELECT, pred_, matrix_64_48_, matrix_32_64_); ASSERT_FALSE(inferred_status_error1.ok()); - ASSERT_MATCH( - inferred_status_error1.status().error_message(), - testing::ContainsRegex("operands to select must be the same shape")); + ASSERT_THAT(inferred_status_error1.status().error_message(), + HasSubstr("operands to select must be the same shape")); auto inferred_status_error2 = ShapeInference::InferTernaryOpShape( TernaryOperation::TRIOP_SELECT, s32_, matrix_64_48_, matrix_64_48_); ASSERT_FALSE(inferred_status_error2.ok()); - ASSERT_MATCH(inferred_status_error2.status().error_message(), - testing::ContainsRegex("pred operand must have PRED")); + ASSERT_THAT(inferred_status_error2.status().error_message(), + HasSubstr("pred operand must have PRED")); auto inferred_status_error3 = ShapeInference::InferTernaryOpShape( TernaryOperation::TRIOP_SELECT, ShapeUtil::MakeShape(PRED, {64}), matrix_64_48_, matrix_64_48_); ASSERT_FALSE(inferred_status_error3.ok()); - ASSERT_MATCH( - inferred_status_error3.status().error_message(), - testing::ContainsRegex("with non-scalar predicate with dimensionality")); + ASSERT_THAT(inferred_status_error3.status().error_message(), + HasSubstr("with non-scalar predicate with dimensionality")); // Tuples have a TUPLE element type and cannot be the pred of a select. auto inferred_status_error4 = ShapeInference::InferTernaryOpShape( @@ -152,9 +154,8 @@ TEST_F(ShapeInferenceTest, SelectBadShapes) { ShapeUtil::MakeTupleShape({f32_, f32_}), ShapeUtil::MakeTupleShape({f32_, f32_})); ASSERT_FALSE(inferred_status_error4.ok()); - ASSERT_MATCH( - inferred_status_error4.status().error_message(), - testing::ContainsRegex("pred operand must have PRED element type")); + ASSERT_THAT(inferred_status_error4.status().error_message(), + HasSubstr("pred operand must have PRED element type")); } TEST_F(ShapeInferenceTest, ClampAllMatrix) { @@ -298,8 +299,8 @@ TEST_F(SelectAndScatterShapeInferenceTest, SelectAndScatterWrongSourceShape) { operand_shape_, select_program_shape_, window_, source_shape_fail, init_value_shape_, scatter_program_shape_); ASSERT_FALSE(inferred_status_fail.ok()); - ASSERT_MATCH(inferred_status_fail.status().error_message(), - testing::ContainsRegex("source shape does not match")); + ASSERT_THAT(inferred_status_fail.status().error_message(), + HasSubstr("source shape does not match")); } TEST_F(SelectAndScatterShapeInferenceTest, SelectAndScatterWrongSelectShape1) { @@ -309,9 +310,8 @@ TEST_F(SelectAndScatterShapeInferenceTest, SelectAndScatterWrongSelectShape1) { operand_shape_, select_program_shape_fail, window_, source_shape_, init_value_shape_, scatter_program_shape_); ASSERT_FALSE(inferred_status_fail.ok()); - ASSERT_MATCH( - inferred_status_fail.status().error_message(), - testing::ContainsRegex("select function must take 2 parameters")); + ASSERT_THAT(inferred_status_fail.status().error_message(), + HasSubstr("select function must take 2 parameters")); } TEST_F(SelectAndScatterShapeInferenceTest, SelectAndScatterWrongSelectShape2) { @@ -321,8 +321,8 @@ TEST_F(SelectAndScatterShapeInferenceTest, SelectAndScatterWrongSelectShape2) { operand_shape_, select_program_shape_fail, window_, source_shape_, init_value_shape_, scatter_program_shape_); ASSERT_FALSE(inferred_status_fail.ok()); - ASSERT_MATCH(inferred_status_fail.status().error_message(), - testing::ContainsRegex("select function must have rank-0 PRED")); + ASSERT_THAT(inferred_status_fail.status().error_message(), + HasSubstr("select function must have rank-0 PRED")); } TEST_F(SelectAndScatterShapeInferenceTest, SelectAndScatterWrongSelectShape3) { @@ -332,8 +332,8 @@ TEST_F(SelectAndScatterShapeInferenceTest, SelectAndScatterWrongSelectShape3) { operand_shape_, select_program_shape_fail, window_, source_shape_, init_value_shape_, scatter_program_shape_); ASSERT_FALSE(inferred_status_fail.ok()); - ASSERT_MATCH(inferred_status_fail.status().error_message(), - testing::ContainsRegex("select function's first parameter")); + ASSERT_THAT(inferred_status_fail.status().error_message(), + HasSubstr("select function's first parameter")); } TEST_F(SelectAndScatterShapeInferenceTest, SelectAndScatterWrongSelectShape4) { @@ -343,8 +343,8 @@ TEST_F(SelectAndScatterShapeInferenceTest, SelectAndScatterWrongSelectShape4) { operand_shape_, select_program_shape_fail, window_, source_shape_, init_value_shape_, scatter_program_shape_); ASSERT_FALSE(inferred_status_fail.ok()); - ASSERT_MATCH(inferred_status_fail.status().error_message(), - testing::ContainsRegex("select function's second parameter")); + ASSERT_THAT(inferred_status_fail.status().error_message(), + HasSubstr("select function's second parameter")); } TEST_F(ShapeInferenceTest, Convolve) { @@ -498,8 +498,8 @@ TEST_F(ShapeInferenceTest, ConvolveDimensionNumbersOverlapError) { auto inferred_status = ShapeInference::InferConvolveShape(lhs_shape, rhs_shape, window, dnums); ASSERT_FALSE(inferred_status.ok()); - ASSERT_MATCH(inferred_status.status().error_message(), - testing::ContainsRegex("each dimension exactly once")); + ASSERT_THAT(inferred_status.status().error_message(), + HasSubstr("each dimension exactly once")); } TEST_F(ShapeInferenceTest, MapThatChangesElementType) { @@ -536,43 +536,42 @@ TEST_F(ShapeInferenceTest, Map) { auto no_args_error = ShapeInference::InferMapShape( {}, ShapeUtil::MakeProgramShape({f32_, f32_}, f32_)); ASSERT_FALSE(no_args_error.ok()); - ASSERT_MATCH(no_args_error.status().error_message(), - testing::ContainsRegex("expects at least one argument")); + ASSERT_THAT(no_args_error.status().error_message(), + HasSubstr("expects at least one argument")); auto args_diff_shapes_error = ShapeInference::InferMapShape( {&vector_32_, &vector_64_}, ShapeUtil::MakeProgramShape({f32_, f32_}, f32_)); ASSERT_FALSE(args_diff_shapes_error.ok()); - ASSERT_MATCH( - args_diff_shapes_error.status().error_message(), - testing::ContainsRegex("requires all operands to have the same shape")); + ASSERT_THAT(args_diff_shapes_error.status().error_message(), + HasSubstr("requires all operands to have the same shape")); auto arity_error = ShapeInference::InferMapShape( {&vector_32_, &vector_32_}, ShapeUtil::MakeProgramShape({f32_}, f32_)); ASSERT_FALSE(arity_error.ok()); - ASSERT_MATCH(arity_error.status().error_message(), - testing::ContainsRegex("function arity must match")); + ASSERT_THAT(arity_error.status().error_message(), + HasSubstr("function arity must match")); auto output_shape_error = ShapeInference::InferMapShape( {&vector_32_, &vector_32_}, ShapeUtil::MakeProgramShape({f32_, f32_}, vector_32_)); ASSERT_FALSE(output_shape_error.ok()); - ASSERT_MATCH(output_shape_error.status().error_message(), - testing::ContainsRegex("result has to be a scalar")); + ASSERT_THAT(output_shape_error.status().error_message(), + HasSubstr("result has to be a scalar")); auto param_shape_error = ShapeInference::InferMapShape( {&vector_32_, &vector_32_}, ShapeUtil::MakeProgramShape({vector_32_, f32_}, f32_)); ASSERT_FALSE(param_shape_error.ok()); - ASSERT_MATCH(param_shape_error.status().error_message(), - testing::ContainsRegex("parameter has to be a scalar")); + ASSERT_THAT(param_shape_error.status().error_message(), + HasSubstr("parameter has to be a scalar")); auto param_element_type_error = ShapeInference::InferMapShape( {&vector_32_, &vector_32_}, ShapeUtil::MakeProgramShape({f32_, s32_}, f32_)); ASSERT_FALSE(param_element_type_error.ok()); - ASSERT_MATCH(param_element_type_error.status().error_message(), - testing::ContainsRegex("parameter type has to match argument")); + ASSERT_THAT(param_element_type_error.status().error_message(), + HasSubstr("parameter type has to match argument")); Shape arg = ShapeUtil::MakeShape(F32, {20}); ProgramShape to_apply = ShapeUtil::MakeProgramShape({f32_}, f32_); @@ -583,26 +582,26 @@ TEST_F(ShapeInferenceTest, Map) { auto inferred_status_error1 = ShapeInference::InferMapShape( {&arg}, ShapeUtil::MakeProgramShape({f32_, f32_}, f32_)); ASSERT_FALSE(inferred_status_error1.ok()); - ASSERT_MATCH(inferred_status_error1.status().error_message(), - testing::ContainsRegex("arity must match number of arguments")); + ASSERT_THAT(inferred_status_error1.status().error_message(), + HasSubstr("arity must match number of arguments")); auto inferred_status_error2 = ShapeInference::InferMapShape( {&arg}, ShapeUtil::MakeProgramShape({vector_32_}, f32_)); ASSERT_FALSE(inferred_status_error2.ok()); - ASSERT_MATCH(inferred_status_error2.status().error_message(), - testing::ContainsRegex("has to be a scalar")); + ASSERT_THAT(inferred_status_error2.status().error_message(), + HasSubstr("has to be a scalar")); auto inferred_status_error3 = ShapeInference::InferMapShape( {&arg}, ShapeUtil::MakeProgramShape({f32_}, vector_32_)); ASSERT_FALSE(inferred_status_error3.ok()); - ASSERT_MATCH(inferred_status_error3.status().error_message(), - testing::ContainsRegex("has to be a scalar")); + ASSERT_THAT(inferred_status_error3.status().error_message(), + HasSubstr("has to be a scalar")); auto inferred_status_error5 = ShapeInference::InferMapShape( {&arg}, ShapeUtil::MakeProgramShape({s32_}, s32_)); ASSERT_FALSE(inferred_status_error5.ok()); - ASSERT_MATCH(inferred_status_error5.status().error_message(), - testing::ContainsRegex("parameter type has to match argument")); + ASSERT_THAT(inferred_status_error5.status().error_message(), + HasSubstr("parameter type has to match argument")); } TEST_F(ReduceShapeInferenceTest, ReduceVectorToScalar) { @@ -656,8 +655,8 @@ TEST_F(ReduceShapeInferenceTest, ErrorOutOfBoundsDimension) { ShapeUtil::MakeShape(F32, {5, 3}), f32_, /*dimensions_to_reduce=*/{3, 4}, to_apply); EXPECT_FALSE(inferred_status.ok()); - EXPECT_MATCH(inferred_status.status().error_message(), - testing::ContainsRegex("out-of-bounds dimension")); + EXPECT_THAT(inferred_status.status().error_message(), + HasSubstr("out-of-bounds dimension")); } TEST_F(ReduceShapeInferenceTest, ErrorToApplyArity) { @@ -666,8 +665,8 @@ TEST_F(ReduceShapeInferenceTest, ErrorToApplyArity) { ShapeInference::InferReduceShape(ShapeUtil::MakeShape(F32, {5, 3}), f32_, /*dimensions_to_reduce=*/{0}, to_apply); EXPECT_FALSE(inferred_status.ok()); - EXPECT_MATCH(inferred_status.status().error_message(), - testing::ContainsRegex("take 2 parameters")); + EXPECT_THAT(inferred_status.status().error_message(), + HasSubstr("take 2 parameters")); } TEST_F(ReduceShapeInferenceTest, ErrorElementTypeVsApplyType) { @@ -676,23 +675,50 @@ TEST_F(ReduceShapeInferenceTest, ErrorElementTypeVsApplyType) { ShapeInference::InferReduceShape(ShapeUtil::MakeShape(F32, {5, 3}), f32_, /*dimensions_to_reduce=*/{0}, to_apply); EXPECT_FALSE(inferred_status.ok()); - EXPECT_MATCH(inferred_status.status().error_message(), - testing::ContainsRegex("first parameter shape differs")); + EXPECT_THAT(inferred_status.status().error_message(), + HasSubstr("first parameter shape differs")); } TEST_F(ShapeInferenceTest, InferSliceShapeRank2) { Shape matrix_shape = ShapeUtil::MakeShape(F32, {128, 64}); auto inferred_status = - ShapeInference::InferSliceShape(matrix_shape, {32, 0}, {64, 64}); + ShapeInference::InferSliceShape(matrix_shape, {32, 0}, {64, 64}, {1, 1}); ASSERT_IS_OK(inferred_status.status()); Shape inferred = inferred_status.ValueOrDie(); ASSERT_TRUE(ShapeUtil::Equal(ShapeUtil::MakeShape(F32, {32, 64}), inferred)); } +TEST_F(ShapeInferenceTest, InferSliceShapeRank2WithStrides) { + Shape matrix_shape = ShapeUtil::MakeShape(F32, {128, 64}); + auto inferred_status = + ShapeInference::InferSliceShape(matrix_shape, {32, 0}, {64, 64}, {2, 4}); + ASSERT_IS_OK(inferred_status.status()); + Shape inferred = inferred_status.ValueOrDie(); + ASSERT_TRUE(ShapeUtil::Equal(ShapeUtil::MakeShape(F32, {16, 16}), inferred)); +} + +TEST_F(ShapeInferenceTest, InferSliceShapeRank2WithStridesNotIntegral) { + Shape matrix_shape = ShapeUtil::MakeShape(F32, {128, 64}); + auto inferred_status = + ShapeInference::InferSliceShape(matrix_shape, {15, 0}, {20, 13}, {2, 4}); + ASSERT_IS_OK(inferred_status.status()); + Shape inferred = inferred_status.ValueOrDie(); + ASSERT_TRUE(ShapeUtil::Equal(ShapeUtil::MakeShape(F32, {3, 4}), inferred)); +} + +TEST_F(ShapeInferenceTest, InferInvalidStride) { + Shape matrix_shape = ShapeUtil::MakeShape(F32, {128, 64}); + auto inferred_status = + ShapeInference::InferSliceShape(matrix_shape, {127, 0}, {129, 2}, {0, 1}); + ASSERT_FALSE(inferred_status.ok()); + ASSERT_EQ(tensorflow::error::INVALID_ARGUMENT, + inferred_status.status().code()); +} + TEST_F(ShapeInferenceTest, InferOobSliceShapeRank2) { Shape matrix_shape = ShapeUtil::MakeShape(F32, {128, 64}); auto inferred_status = - ShapeInference::InferSliceShape(matrix_shape, {127, 0}, {129, 2}); + ShapeInference::InferSliceShape(matrix_shape, {127, 0}, {129, 2}, {1, 1}); ASSERT_FALSE(inferred_status.ok()); ASSERT_EQ(tensorflow::error::INVALID_ARGUMENT, inferred_status.status().code()); @@ -701,7 +727,7 @@ TEST_F(ShapeInferenceTest, InferOobSliceShapeRank2) { TEST_F(ShapeInferenceTest, InferSliceShapeRank1) { Shape vector_shape = ShapeUtil::MakeShape(F32, {17}); auto inferred_status = - ShapeInference::InferSliceShape(vector_shape, {2}, {4}); + ShapeInference::InferSliceShape(vector_shape, {2}, {4}, {1}); ASSERT_TRUE(inferred_status.ok()); Shape inferred = inferred_status.ValueOrDie(); ASSERT_TRUE(ShapeUtil::Equal(inferred, ShapeUtil::MakeShape(F32, {2}))); @@ -819,8 +845,8 @@ TEST_F(ShapeInferenceTest, ScalarDotVector) { auto inferred_status = ShapeInference::InferBinaryOpShape(BINOP_DOT, f32_, vector_32_, {}); ASSERT_FALSE(inferred_status.ok()); - ASSERT_MATCH(inferred_status.status().error_message(), - testing::ContainsRegex("dot only supports rank")); + ASSERT_THAT(inferred_status.status().error_message(), + HasSubstr("dot only supports rank")); } // 3D 2D: error @@ -828,8 +854,8 @@ TEST_F(ShapeInferenceTest, DotWithRankHigherThanTwo) { auto inferred_status = ShapeInference::InferBinaryOpShape( BINOP_DOT, ShapeUtil::MakeShape(F32, {32, 32, 32}), matrix_32_64_, {}); ASSERT_FALSE(inferred_status.ok()); - ASSERT_MATCH(inferred_status.status().error_message(), - testing::ContainsRegex("dot only supports rank")); + ASSERT_THAT(inferred_status.status().error_message(), + HasSubstr("dot only supports rank")); } // vector vector -> scalar @@ -941,46 +967,43 @@ TEST_F(ShapeInferenceTest, BinOpBroadcastBadDimension) { auto inferred_status_error1 = ShapeInference::InferBinaryOpShape( BinaryOperation::BINOP_ADD, tensor, vec8, {}); ASSERT_FALSE(inferred_status_error1.ok()); - ASSERT_MATCH(inferred_status_error1.status().error_message(), - testing::ContainsRegex("automatic")); + ASSERT_THAT(inferred_status_error1.status().error_message(), + HasSubstr("automatic")); // broadcast_dimension out of bounds for tensor's rank auto inferred_status_error2 = ShapeInference::InferBinaryOpShape( BinaryOperation::BINOP_ADD, tensor, vec8, {3}); ASSERT_FALSE(inferred_status_error2.ok()); - ASSERT_MATCH( - inferred_status_error2.status().error_message(), - testing::ContainsRegex("broadcast dimension number .* too large")); + ASSERT_THAT(inferred_status_error2.status().error_message(), + ContainsRegex("broadcast dimension number .* too large")); // broadcast_dimension doesn't match corresponding dimension auto inferred_status_error3 = ShapeInference::InferBinaryOpShape( BinaryOperation::BINOP_ADD, tensor, vec8, {0}); ASSERT_FALSE(inferred_status_error3.ok()); - ASSERT_MATCH(inferred_status_error3.status().error_message(), - testing::ContainsRegex("broadcast dimension 0 mismatch")); + ASSERT_THAT(inferred_status_error3.status().error_message(), + HasSubstr("broadcast dimension 0 mismatch")); // broadcast_dimensions list too long auto inferred_status_error4 = ShapeInference::InferBinaryOpShape( BinaryOperation::BINOP_ADD, tensor, matrix8_4, {0, 1, 2}); ASSERT_FALSE(inferred_status_error4.ok()); - ASSERT_MATCH( - inferred_status_error4.status().error_message(), - testing::ContainsRegex("size of broadcast_dimensions has to match")); + ASSERT_THAT(inferred_status_error4.status().error_message(), + HasSubstr("size of broadcast_dimensions has to match")); // there's a dimension above the rank of the tensor auto inferred_status_error5 = ShapeInference::InferBinaryOpShape( BinaryOperation::BINOP_ADD, tensor, matrix8_4, {3, 0}); ASSERT_FALSE(inferred_status_error5.ok()); - ASSERT_MATCH( - inferred_status_error5.status().error_message(), - testing::ContainsRegex("broadcast dimension number .* too large")); + ASSERT_THAT(inferred_status_error5.status().error_message(), + ContainsRegex("broadcast dimension number .* too large")); // broadcasting dimensions don't match in this order auto inferred_status_error6 = ShapeInference::InferBinaryOpShape( BinaryOperation::BINOP_ADD, tensor, matrix8_4, {2, 1}); ASSERT_FALSE(inferred_status_error6.ok()); - ASSERT_MATCH(inferred_status_error6.status().error_message(), - testing::ContainsRegex("broadcast dimension 0 mismatch")); + ASSERT_THAT(inferred_status_error6.status().error_message(), + HasSubstr("broadcast dimension 0 mismatch")); // The following two tests make sure that broadcasting dimensions are listed // in a proper (strictly increasing) order, even if the lower-rank array @@ -988,14 +1011,14 @@ TEST_F(ShapeInferenceTest, BinOpBroadcastBadDimension) { auto inferred_status_error7 = ShapeInference::InferBinaryOpShape( BinaryOperation::BINOP_ADD, tensor8_8_8, matrix8_8, {0, 0}); ASSERT_FALSE(inferred_status_error7.ok()); - ASSERT_MATCH(inferred_status_error7.status().error_message(), - testing::ContainsRegex("broadcast dimensions order is wrong")); + ASSERT_THAT(inferred_status_error7.status().error_message(), + HasSubstr("broadcast dimensions order is wrong")); auto inferred_status_error8 = ShapeInference::InferBinaryOpShape( BinaryOperation::BINOP_ADD, tensor8_8_8, matrix8_8, {1, 0}); ASSERT_FALSE(inferred_status_error8.ok()); - ASSERT_MATCH(inferred_status_error8.status().error_message(), - testing::ContainsRegex("broadcast dimensions order is wrong")); + ASSERT_THAT(inferred_status_error8.status().error_message(), + HasSubstr("broadcast dimensions order is wrong")); } // Tests for the while instruction with proper shapes. @@ -1020,30 +1043,30 @@ TEST_F(ShapeInferenceTest, WhileWithBadShapes) { auto inferred_status_error1 = ShapeInference::InferWhileShape(bad_shape_1, body, result_shape); ASSERT_FALSE(inferred_status_error1.ok()); - ASSERT_MATCH(inferred_status_error1.status().error_message(), - testing::ContainsRegex("condition must take 1 arguments")); + ASSERT_THAT(inferred_status_error1.status().error_message(), + HasSubstr("condition must take 1 arguments")); auto bad_shape_2 = ShapeUtil::MakeProgramShape({s32_, result_shape}, result_shape); auto inferred_status_error2 = ShapeInference::InferWhileShape(cond, bad_shape_2, result_shape); ASSERT_FALSE(inferred_status_error2.ok()); - ASSERT_MATCH(inferred_status_error2.status().error_message(), - testing::ContainsRegex("body must take 1 arguments")); + ASSERT_THAT(inferred_status_error2.status().error_message(), + HasSubstr("body must take 1 arguments")); auto bad_shape_3 = ShapeUtil::MakeProgramShape({result_shape}, s32_); auto inferred_status_error3 = ShapeInference::InferWhileShape(bad_shape_3, body, result_shape); ASSERT_FALSE(inferred_status_error3.ok()); - ASSERT_MATCH(inferred_status_error3.status().error_message(), - testing::ContainsRegex("condition must return a boolean")); + ASSERT_THAT(inferred_status_error3.status().error_message(), + HasSubstr("condition must return a boolean")); auto bad_shape_4 = ShapeUtil::MakeProgramShape({result_shape}, vector_32_); auto inferred_status_error4 = ShapeInference::InferWhileShape(cond, bad_shape_4, result_shape); ASSERT_FALSE(inferred_status_error4.ok()); - ASSERT_MATCH(inferred_status_error4.status().error_message(), - testing::ContainsRegex("parameter of condition and body")); + ASSERT_THAT(inferred_status_error4.status().error_message(), + HasSubstr("parameter of condition and body")); } // Tests for the concatenate instruction with proper shapes. @@ -1073,49 +1096,44 @@ TEST_F(ShapeInferenceTest, ConcatenateWithBadShapes) { auto inferred_status_error1 = ShapeInference::InferConcatOpShape({}, /*dimension=*/0); ASSERT_FALSE(inferred_status_error1.ok()); - ASSERT_MATCH( - inferred_status_error1.status().error_message(), - testing::ContainsRegex("Concatenate expects at least one argument")); + ASSERT_THAT(inferred_status_error1.status().error_message(), + HasSubstr("Concatenate expects at least one argument")); auto inferred_status_error2 = ShapeInference::InferConcatOpShape({&vector_32_}, /*dimension=*/-1); ASSERT_FALSE(inferred_status_error2.ok()); - ASSERT_MATCH(inferred_status_error2.status().error_message(), - testing::ContainsRegex( - "dimension to concatenate along out of bounds: -1")); + ASSERT_THAT(inferred_status_error2.status().error_message(), + HasSubstr("dimension to concatenate along out of bounds: -1")); auto inferred_status_error3 = ShapeInference::InferConcatOpShape({&vector_32_}, /*dimension=*/1); ASSERT_FALSE(inferred_status_error3.ok()); - ASSERT_MATCH(inferred_status_error3.status().error_message(), - testing::ContainsRegex( - "dimension to concatenate along out of bounds: 1")); + ASSERT_THAT(inferred_status_error3.status().error_message(), + HasSubstr("dimension to concatenate along out of bounds: 1")); Shape tuple = ShapeUtil::MakeTupleShape({vector_32_}); auto inferred_status_error4 = ShapeInference::InferConcatOpShape( {&vector_32_, &tuple}, /*dimension=*/0); ASSERT_FALSE(inferred_status_error4.ok()); - ASSERT_MATCH( + ASSERT_THAT( inferred_status_error4.status().error_message(), - testing::ContainsRegex( - "Expected non-tuple argument for operand of concatenation.")); + HasSubstr("Expected non-tuple argument for operand of concatenation.")); const Shape vector_s32 = ShapeUtil::MakeShape(S32, {32}); auto inferred_status_error5 = ShapeInference::InferConcatOpShape( {&vector_32_, &vector_s32}, /*dimension=*/0); ASSERT_FALSE(inferred_status_error5.ok()); - ASSERT_MATCH(inferred_status_error5.status().error_message(), - testing::ContainsRegex( - "cannot concatenate arrays with different element types")); + ASSERT_THAT( + inferred_status_error5.status().error_message(), + HasSubstr("cannot concatenate arrays with different element types")); auto inferred_status_error6 = ShapeInference::InferConcatOpShape( {&matrix_32_48_, &matrix_32_64_}, /*dimension=*/0); ASSERT_FALSE(inferred_status_error6.ok()); - ASSERT_MATCH( - inferred_status_error6.status().error_message(), - testing::ContainsRegex("cannot concatenate arrays that differ in " - "dimensions other than the one being " - "concatenated")); + ASSERT_THAT(inferred_status_error6.status().error_message(), + HasSubstr("cannot concatenate arrays that differ in " + "dimensions other than the one being " + "concatenated")); } TEST_F(ShapeInferenceTest, Pad) { @@ -1156,27 +1174,27 @@ TEST_F(ShapeInferenceTest, ReverseInvalidDimension) { auto inferred_status_error0 = ShapeInference::InferReverseShape(input_shape, {0, 2}); ASSERT_FALSE(inferred_status_error0.ok()); - ASSERT_MATCH(inferred_status_error0.status().error_message(), - testing::ContainsRegex("out-of-bounds")); + ASSERT_THAT(inferred_status_error0.status().error_message(), + HasSubstr("out-of-bounds")); auto inferred_status_error1 = ShapeInference::InferReverseShape(input_shape, {0, -1}); ASSERT_FALSE(inferred_status_error1.ok()); - ASSERT_MATCH(inferred_status_error1.status().error_message(), - testing::ContainsRegex("out-of-bounds")); + ASSERT_THAT(inferred_status_error1.status().error_message(), + HasSubstr("out-of-bounds")); auto inferred_status_error2 = ShapeInference::InferReverseShape(input_shape, {0, 0}); ASSERT_FALSE(inferred_status_error2.ok()); - ASSERT_MATCH(inferred_status_error2.status().error_message(), - testing::ContainsRegex("duplicated")); + ASSERT_THAT(inferred_status_error2.status().error_message(), + HasSubstr("duplicated")); Shape tuple_shape = ShapeUtil::MakeTupleShape({input_shape, input_shape}); auto inferred_status_error3 = ShapeInference::InferReverseShape(tuple_shape, {0}); ASSERT_FALSE(inferred_status_error3.ok()); - ASSERT_MATCH(inferred_status_error3.status().error_message(), - testing::ContainsRegex("Expected non-tuple argument")); + ASSERT_THAT(inferred_status_error3.status().error_message(), + HasSubstr("Expected non-tuple argument")); } TEST_F(ShapeInferenceTest, Call) { @@ -1196,20 +1214,20 @@ TEST_F(ShapeInferenceTest, Call) { auto inferred_status_error0 = ShapeInference::InferCallShape( {}, ShapeUtil::MakeProgramShape({f32_}, f32_)); EXPECT_FALSE(inferred_status_error0.ok()); - EXPECT_MATCH(inferred_status_error0.status().error_message(), - testing::ContainsRegex("arity must match")); + EXPECT_THAT(inferred_status_error0.status().error_message(), + HasSubstr("arity must match")); auto inferred_status_error1 = ShapeInference::InferCallShape( {&f32_}, ShapeUtil::MakeProgramShape({}, f32_)); EXPECT_FALSE(inferred_status_error1.ok()); - EXPECT_MATCH(inferred_status_error1.status().error_message(), - testing::ContainsRegex("arity must match")); + EXPECT_THAT(inferred_status_error1.status().error_message(), + HasSubstr("arity must match")); auto inferred_status_error2 = ShapeInference::InferCallShape( {&f32_}, ShapeUtil::MakeProgramShape({s32_}, f32_)); EXPECT_FALSE(inferred_status_error2.ok()); - EXPECT_MATCH(inferred_status_error2.status().error_message(), - testing::ContainsRegex("parameter must match argument")); + EXPECT_THAT(inferred_status_error2.status().error_message(), + HasSubstr("parameter must match argument")); } TEST_F(ShapeInferenceTest, Transpose) { diff --git a/tensorflow/compiler/xla/service/shaped_buffer.cc b/tensorflow/compiler/xla/service/shaped_buffer.cc index cf49fd72b7d1f31fd09b39a9286b62300d8bd50e..865be1b84f2fd599b68f09fdad0323076e637906 100644 --- a/tensorflow/compiler/xla/service/shaped_buffer.cc +++ b/tensorflow/compiler/xla/service/shaped_buffer.cc @@ -73,16 +73,13 @@ ShapedBuffer::MakeUnnestedTupleShapedBuffer( } TF_ASSIGN_OR_RETURN(std::unique_ptr shaped_buffer, MakeShapedBuffer(shape, platform, device_ordinal)); - TF_CHECK_OK(shaped_buffer->mutable_shape_index_to_buffer_entry() - ->ForEachMutableElement( - [](const ShapeIndex& index, bool is_leaf, - size_t* buffer_element) -> tensorflow::Status { - if (is_leaf) { - CHECK_EQ(index.size(), 1); - *buffer_element = index[0]; - } - return tensorflow::Status::OK(); - })); + shaped_buffer->mutable_shape_index_to_buffer_entry()->ForEachMutableElement( + [&shaped_buffer](const ShapeIndex& index, size_t* buffer_element) { + if (ShapeUtil::IsLeafIndex(shaped_buffer->shape(), index)) { + CHECK_EQ(index.size(), 1); + *buffer_element = index[0]; + } + }); shaped_buffer->mutable_buffers()->reserve(buffers.size()); for (const perftools::gputools::DeviceMemoryBase& memory_base : buffers) { shaped_buffer->mutable_buffers()->push_back(memory_base); @@ -126,10 +123,12 @@ ScopedShapedBuffer::MakeScopedShapedBuffer(const Shape& shape, // Allocate an appropriate sized buffer for each array element in the shape. TF_RETURN_IF_ERROR( - shaped_buffer->shape_index_to_buffer_entry_.ForEachMutableElement( - [&shaped_buffer](const ShapeIndex& index, bool is_leaf, - size_t* buffer_entry) -> tensorflow::Status { - if (is_leaf) { + shaped_buffer->shape_index_to_buffer_entry_ + .ForEachMutableElementWithStatus([&shaped_buffer]( + const ShapeIndex& index, + size_t* buffer_entry) + -> tensorflow::Status { + if (ShapeUtil::IsLeafIndex(shaped_buffer->shape(), index)) { TF_ASSIGN_OR_RETURN( perftools::gputools::DeviceMemoryBase memory_base, shaped_buffer->allocator_->Allocate( diff --git a/tensorflow/compiler/xla/service/transfer_manager.cc b/tensorflow/compiler/xla/service/transfer_manager.cc index c7f6a13023d32f48d357430bb62241cd537422ca..4da0a0d36841a6dfaed5c7eebdfb9e6980ad1090 100644 --- a/tensorflow/compiler/xla/service/transfer_manager.cc +++ b/tensorflow/compiler/xla/service/transfer_manager.cc @@ -72,7 +72,7 @@ TransferManager::GetPlatformTransferManagers() { it->second.manager = (*it->second.creation_function)(); } - return it->second.manager; + return it->second.manager.get(); } Status TransferManager::TransferBufferFromDevice( diff --git a/tensorflow/compiler/xla/service/transfer_manager.h b/tensorflow/compiler/xla/service/transfer_manager.h index b052bb814693c2e9364c94154ca223fe98526622..c79ffa9cd73950b1653f72b1c6286346f76c10fb 100644 --- a/tensorflow/compiler/xla/service/transfer_manager.h +++ b/tensorflow/compiler/xla/service/transfer_manager.h @@ -20,6 +20,7 @@ limitations under the License. #include #include +#include "tensorflow/compiler/xla/literal_util.h" #include "tensorflow/compiler/xla/statusor.h" #include "tensorflow/compiler/xla/types.h" #include "tensorflow/compiler/xla/xla_data.pb.h" @@ -64,6 +65,17 @@ class TransferManager { perftools::gputools::StreamExecutor* executor, const Literal& literal) = 0; + // Transfer a memory block of the given size from 'source' buffer to the + // Infeed interface of the device using the given executor. + // + // size is the size to transfer from source in bytes. + // + // source is the source data that must be in the target-dependent layout that + // the Infeed HLO used in the computation expects. + virtual Status TransferBufferToInfeed( + perftools::gputools::StreamExecutor* executor, int64 size, + const void* source) = 0; + // Transfers the given literal from the Outfeed interface of the device, // using the given executor. virtual Status TransferLiteralFromOutfeed( @@ -99,13 +111,6 @@ class TransferManager { // region for a host-to-device transfer. virtual int64 GetByteSizeRequirement(const Shape& shape) = 0; - // Returns whether tuple elements are distinct buffers (in which case each of - // the elements of a tuple should be deallocated, in addition to the tuple's - // buffer itself). - // - // TODO(b/36256956) Ideally tuple elements could always be distinct buffers. - virtual bool TupleElementsAreDistinctBuffers() const { return true; } - // Transfer a memory block of the given size from the device source into the // 'destination' buffer. // @@ -123,7 +128,7 @@ class TransferManager { perftools::gputools::StreamExecutor* executor, int64 size, const void* source, perftools::gputools::DeviceMemoryBase* destination); - typedef TransferManager* (*TransferManagerCreationFunction)(); + typedef std::unique_ptr (*TransferManagerCreationFunction)(); ///// // The TransferManager class also serves as a point to register objects for @@ -153,7 +158,7 @@ class TransferManager { // set up creation_function, and then we use that to lazily create // "manager" the first time GetForPlatform is invoked for a particular id. struct State { - TransferManager* manager = nullptr; + std::unique_ptr manager; TransferManagerCreationFunction creation_function = nullptr; }; diff --git a/tensorflow/compiler/xla/service/transfer_manager_test.cc b/tensorflow/compiler/xla/service/transfer_manager_test.cc index 564111c4f2b22df6652fadb7cd3cd237ba957905..29ecef9510cfe6b8764c2e5fe1216255ca1dc983 100644 --- a/tensorflow/compiler/xla/service/transfer_manager_test.cc +++ b/tensorflow/compiler/xla/service/transfer_manager_test.cc @@ -55,7 +55,7 @@ class CpuTransferManagerTest : public ::testing::Test { TEST_F(CpuTransferManagerTest, TransferR0U32ToDevice) { std::vector storage(sizeof(uint32), '\x00'); se::DeviceMemoryBase memptr(storage.data(), storage.size()); - std::unique_ptr literal = LiteralUtil::CreateR0(42); + std::unique_ptr literal = Literal::CreateR0(42); TF_CHECK_OK(transfer_manager_.TransferLiteralToDevice(stream_exec_, *literal, &memptr)); @@ -66,7 +66,7 @@ TEST_F(CpuTransferManagerTest, TransferR1F32ToDevice) { std::vector storage(4 * sizeof(float), '\x00'); se::DeviceMemoryBase memptr(storage.data(), storage.size()); std::unique_ptr literal = - LiteralUtil::CreateR1({1.25f, 2.5f, -17.0f, -20.125f}); + Literal::CreateR1({1.25f, 2.5f, -17.0f, -20.125f}); TF_CHECK_OK(transfer_manager_.TransferLiteralToDevice(stream_exec_, *literal, &memptr)); @@ -80,7 +80,7 @@ TEST_F(CpuTransferManagerTest, TransferR1U8ToDevice) { std::vector storage(16, '\x00'); se::DeviceMemoryBase memptr(storage.data(), storage.size()); const char* str = "0123456789abcdef"; - std::unique_ptr literal = LiteralUtil::CreateR1U8(str); + std::unique_ptr literal = Literal::CreateR1U8(str); TF_CHECK_OK(transfer_manager_.TransferLiteralToDevice(stream_exec_, *literal, &memptr)); @@ -121,7 +121,7 @@ TEST_F(CpuTransferManagerTest, TransferR1U8FromDevice) { const Shape shape = ShapeUtil::MakeShape(U8, {4}); TF_CHECK_OK(transfer_manager_.TransferLiteralFromDevice( stream_exec_, memptr, shape, shape, &literal)); - CHECK_EQ("klmn", literal.u8s()); + CHECK_EQ("klmn", literal.u8s_string()); } TEST_F(CpuTransferManagerTest, TransferBufferFromDevice) { diff --git a/tensorflow/compiler/xla/service/transpose_folding.cc b/tensorflow/compiler/xla/service/transpose_folding.cc index 07e0ce89f6ad2ba194832096de2399ab618422a4..585833573606058514d20fa396b433497ec65bd6 100644 --- a/tensorflow/compiler/xla/service/transpose_folding.cc +++ b/tensorflow/compiler/xla/service/transpose_folding.cc @@ -21,7 +21,9 @@ limitations under the License. #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/shape_util.h" #include "tensorflow/compiler/xla/status_macros.h" +#include "tensorflow/compiler/xla/util.h" #include "tensorflow/core/lib/core/errors.h" #include "tensorflow/core/lib/core/status.h" #include "tensorflow/core/platform/logging.h" @@ -30,43 +32,55 @@ namespace xla { namespace { -bool IsOperandFoldableToDot(const HloInstruction& hlo) { - return hlo.IsRank2Transpose() && - hlo.user_count() == 1; // The dot is its only user. -} - -bool CanFoldOperandsIntoDot( +TransposeFolding::OperandIndices CanFoldOperandsIntoDot( const HloInstruction& dot, - const TransposeFolding::IsTransposableGemmFn& is_transposable_gemm) { + const TransposeFolding::TransposableGemmOperandsFn& + transposable_gemm_operands) { if (HloOpcode::kDot != dot.opcode()) { - return false; + return {}; } - if (!is_transposable_gemm(dot)) { - return false; + TransposeFolding::OperandIndices operand_set; + for (int64 i = 0; i < dot.operand_count(); ++i) { + auto& operand = *dot.operand(i); + if (operand.IsRank2Transpose() && operand.user_count() == 1) { + operand_set.push_back(i); + } } - const HloInstruction* lhs = dot.operand(0); - const HloInstruction* rhs = dot.operand(1); - bool lhs_foldable = IsOperandFoldableToDot(*lhs); - bool rhs_foldable = IsOperandFoldableToDot(*rhs); - if (!lhs_foldable && !rhs_foldable) { - return false; + return transposable_gemm_operands(dot, operand_set); +} + +TransposeFolding::OperandIndices CanFoldOperandsIntoConvolution( + const HloInstruction& convolution, + const TransposeFolding::TransposableConvOperandsFn& + transposable_conv_operands) { + if (HloOpcode::kConvolution != convolution.opcode()) { + return {}; } - return true; + + // We only support folding the RHS. + const int64 kRhsOperandIndex = 1; + auto& operand = *convolution.operand(kRhsOperandIndex); + if (operand.opcode() == HloOpcode::kTranspose && operand.user_count() == 1) { + return transposable_conv_operands(convolution, {kRhsOperandIndex}); + } + + return {}; } +using InstructionOperandsPair = + std::pair; + // Folds the operands of `dot` that are foldable transposes. `computation` is -// the parent HLO computation of `dot`. `module` is the parent HloModule of -// `computation`. +// the parent HLO computation of `dot`. // // Returns whether the module is changed. -bool FoldTransposeIntoDot(HloInstruction* dot, HloComputation* computation) { +bool FoldTransposeIntoDot(InstructionOperandsPair pair) { + auto* dot = pair.first; std::vector instructions_to_fuse(1, dot); - for (HloInstruction* operand : dot->operands()) { - if (IsOperandFoldableToDot(*operand)) { - instructions_to_fuse.push_back(operand); - } + for (const int64 operand_index : pair.second) { + instructions_to_fuse.push_back(dot->mutable_operand(operand_index)); } // Early-exit if no operands are foldable. @@ -74,33 +88,107 @@ bool FoldTransposeIntoDot(HloInstruction* dot, HloComputation* computation) { return false; } - computation->CreateFusionInstruction( + dot->parent()->CreateFusionInstruction( instructions_to_fuse, HloInstruction::FusionKind::kTransposeDot); return true; } +// Folds the operands of `convolution` that are foldable transposes. +// `computation` is the parent HLO computation of `convolution`. +// +// Returns whether the module is changed. +bool FoldTransposeIntoConvolution(InstructionOperandsPair pair) { + auto& convolution = *pair.first; + + // We only support fusing the RHS transpose into convolution. + // + // ConvolutionDimensionNumbers doesn't make enough of a distinction between + // the output and the activations. + // + // TODO(b/37125184): Support transposing the LHS too. + if (pair.second.size() != 1 || pair.second.front() != 1) { + return false; + } + + const ConvolutionDimensionNumbers& dnums = + convolution.convolution_dimension_numbers(); + HloInstruction& transpose = *convolution.mutable_operand(1); + CHECK_EQ(transpose.opcode(), HloOpcode::kTranspose); + const auto& transpose_dimensions = transpose.dimensions(); + HloInstruction& transpose_operand = *transpose.mutable_operand(0); + + // Everything remains the same except for the kernel dimension numbers. We + // need to apply the transpose permutation to the original shape to figure out + // what the new logical dimensions are. + ConvolutionDimensionNumbers new_dnums = dnums; + new_dnums.set_kernel_input_feature_dimension( + transpose_dimensions[dnums.kernel_input_feature_dimension()]); + new_dnums.set_kernel_output_feature_dimension( + transpose_dimensions[dnums.kernel_output_feature_dimension()]); + for (auto& kernel_spatial_dimension : + *new_dnums.mutable_kernel_spatial_dimensions()) { + kernel_spatial_dimension = transpose_dimensions[kernel_spatial_dimension]; + } + + auto new_conv = HloInstruction::CreateConvolve( + convolution.shape(), convolution.mutable_operand(0), &transpose_operand, + convolution.window(), new_dnums); + TF_CHECK_OK(convolution.parent()->ReplaceWithNewInstruction( + &convolution, std::move(new_conv))); + + return true; +} + } // namespace -TransposeFolding::TransposeFolding(IsTransposableGemmFn is_transposable_gemm) - : is_transposable_gemm_(std::move(is_transposable_gemm)) {} +TransposeFolding::TransposeFolding( + TransposableGemmOperandsFn transposable_gemm_operands, + TransposableConvOperandsFn transposable_conv_operands) + : transposable_gemm_operands_(std::move(transposable_gemm_operands)), + transposable_conv_operands_(std::move(transposable_conv_operands)) {} StatusOr TransposeFolding::Run(HloModule* module) { // Modifying the graph while traversing is dangerous, so we find all folding // opportunities before actually folding them. - HloComputation* entry_computation = module->entry_computation(); - - std::vector foldable_dots; - auto visit_fn = [this, &foldable_dots](HloInstruction* instruction) { - if (CanFoldOperandsIntoDot(*instruction, is_transposable_gemm_)) { - foldable_dots.emplace_back(instruction); + std::vector> foldable_dots; + std::vector> foldable_convolutions; + auto visit_fn = [this, &foldable_dots, + &foldable_convolutions](HloInstruction* instruction) { + { + OperandIndices operand_indices = + CanFoldOperandsIntoDot(*instruction, transposable_gemm_operands_); + if (!operand_indices.empty()) { + foldable_dots.emplace_back(instruction, operand_indices); + } + } + { + OperandIndices operand_indices = CanFoldOperandsIntoConvolution( + *instruction, transposable_conv_operands_); + if (!operand_indices.empty()) { + foldable_convolutions.emplace_back( + std::make_pair(instruction, operand_indices)); + } } return tensorflow::Status::OK(); }; - TF_RETURN_IF_ERROR(entry_computation->root_instruction()->Accept(visit_fn)); + + std::vector computations; + for (auto& computation : module->computations()) { + if (computation->IsFusionComputation()) { + continue; + } + computations.push_back(computation.get()); + } + for (auto& comp : computations) { + TF_RETURN_IF_ERROR(comp->Accept(visit_fn)); + } bool changed = false; - for (HloInstruction* dot : foldable_dots) { - changed |= FoldTransposeIntoDot(dot, entry_computation); + for (InstructionOperandsPair& pair : foldable_dots) { + changed |= FoldTransposeIntoDot(pair); + } + for (InstructionOperandsPair& pair : foldable_convolutions) { + changed |= FoldTransposeIntoConvolution(pair); } return changed; } diff --git a/tensorflow/compiler/xla/service/transpose_folding.h b/tensorflow/compiler/xla/service/transpose_folding.h index d857c04ed8d0c0d9d6c005db0f29ab0c5abd3bb2..71e8446452f072c22bb730cbda65a1743a95cd4c 100644 --- a/tensorflow/compiler/xla/service/transpose_folding.h +++ b/tensorflow/compiler/xla/service/transpose_folding.h @@ -25,16 +25,37 @@ namespace xla { // operator is implemented by a GEMM kernel that can transpose its inputs. class TransposeFolding : public HloPassInterface { public: - // IsTransposableGemmFn should return true iff the instruction argument is - // implemented as a GEMM kernel that supports transposing its arguments. - typedef std::function IsTransposableGemmFn; - explicit TransposeFolding(IsTransposableGemmFn is_transposable_gemm); + using OperandIndices = std::vector; + + // Returns the set of foldable operands for a given HLO and some candidate + // operands. + using FoldableOperands = std::function; + using TransposableGemmOperandsFn = FoldableOperands; + using TransposableConvOperandsFn = FoldableOperands; + + // Helper function to explicitly not fold transposes. + static OperandIndices NeverFoldTranspose(const HloInstruction&, + const OperandIndices&) { + return {}; + } + // transposable_gemm_operands returns the set of operands it wants to fold if + // the instruction argument is implemented as a GEMM kernel that supports + // transposing its arguments. + // + // transposable_conv_operands returns the set of operands it wants to fold if + // the instruction argument is implemented as a convolution that supports + // transposing its arguments. + explicit TransposeFolding( + TransposableGemmOperandsFn transposable_gemm_operands, + TransposableConvOperandsFn transposable_conv_operands); tensorflow::StringPiece name() const override { return "transpose-folding"; } StatusOr Run(HloModule* module) override; private: - IsTransposableGemmFn is_transposable_gemm_; + TransposableGemmOperandsFn transposable_gemm_operands_; + TransposableConvOperandsFn transposable_conv_operands_; }; } // namespace xla diff --git a/tensorflow/compiler/xla/service/transpose_folding_test.cc b/tensorflow/compiler/xla/service/transpose_folding_test.cc index 09f932e29e61a24b178e7ced0d2643aa484bea02..9520c42d280968e3f21a110089583c94277ef1a6 100644 --- a/tensorflow/compiler/xla/service/transpose_folding_test.cc +++ b/tensorflow/compiler/xla/service/transpose_folding_test.cc @@ -16,16 +16,19 @@ limitations under the License. #include "tensorflow/compiler/xla/service/transpose_folding.h" #include -#include +#include #include +#include "tensorflow/compiler/xla/client/computation_builder.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_opcode.h" +#include "tensorflow/compiler/xla/service/shape_inference.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/xla_data.pb.h" #include "tensorflow/core/platform/logging.h" @@ -35,12 +38,20 @@ namespace xla { class TransposeFoldingTest : public ::testing::Test { protected: void FoldTranspose(HloModule* module) { - TransposeFolding transpose_folding(gpu::ImplementedAsGemm); + TransposeFolding transpose_folding( + [](const HloInstruction& dot, + const TransposeFolding::OperandIndices& candidate_operands) { + return candidate_operands; + }, + [](const HloInstruction& convolution, + const TransposeFolding::OperandIndices& candidate_operands) { + return candidate_operands; + }); EXPECT_IS_OK(transpose_folding.Run(module).status()); } }; -TEST_F(TransposeFoldingTest, FoldTranspose) { +TEST_F(TransposeFoldingTest, FoldDotTranspose) { auto builder = HloComputation::Builder("entry_computation"); HloInstruction* x = builder.AddInstruction(HloInstruction::CreateParameter( /*parameter_number=*/0, ShapeUtil::MakeShape(F32, {2, 3}), @@ -61,7 +72,7 @@ TEST_F(TransposeFoldingTest, FoldTranspose) { FoldTranspose(&module); // Instructions after folding: x, y, and the fusion. - std::set instruction_set; + std::unordered_set instruction_set; for (auto& instruction : entry_computation->instructions()) { instruction_set.insert(instruction.get()); } @@ -77,15 +88,15 @@ TEST_F(TransposeFoldingTest, FoldTranspose) { EXPECT_EQ(4, fusion->fused_instructions().size()); } -TEST_F(TransposeFoldingTest, FoldTransposeConstant) { +TEST_F(TransposeFoldingTest, FoldDotTransposeConstant) { auto builder = HloComputation::Builder("entry_computation"); // 2x1 HloInstruction* const0 = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR2({{1}, {2}}))); + HloInstruction::CreateConstant(Literal::CreateR2({{1}, {2}}))); // 3x2 HloInstruction* const1 = builder.AddInstruction(HloInstruction::CreateConstant( - LiteralUtil::CreateR2({{1, 2}, {3, 4}, {5, 6}}))); + Literal::CreateR2({{1, 2}, {3, 4}, {5, 6}}))); HloInstruction* transpose0 = builder.AddInstruction(HloInstruction::CreateTranspose( ShapeUtil::MakeShape(F32, {1, 2}), const0, {1, 0})); @@ -115,15 +126,15 @@ TEST_F(TransposeFoldingTest, FoldTransposeConstant) { entry_computation->root_instruction()->fused_instructions().size()); } -TEST_F(TransposeFoldingTest, FuseWithConstantOperands) { +TEST_F(TransposeFoldingTest, FuseDotWithConstantOperands) { auto builder = HloComputation::Builder("entry"); // (1.0 + 2.0) * (2.0 - 3.0) HloInstruction* const1 = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(1.0))); + HloInstruction::CreateConstant(Literal::CreateR0(1.0))); HloInstruction* const2 = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(2.0))); + HloInstruction::CreateConstant(Literal::CreateR0(2.0))); HloInstruction* const3 = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(3.0))); + HloInstruction::CreateConstant(Literal::CreateR0(3.0))); HloInstruction* add = builder.AddInstruction(HloInstruction::CreateBinary( const1->shape(), HloOpcode::kAdd, const1, const2)); HloInstruction* sub = builder.AddInstruction(HloInstruction::CreateBinary( @@ -139,11 +150,219 @@ TEST_F(TransposeFoldingTest, FuseWithConstantOperands) { EXPECT_EQ(call, entry_computation->root_instruction()); HloComputation* callee_computation = call->to_apply(); // The arguments to the call should be const1, const2, and const3. - EXPECT_MATCH(call->operands(), testing::UnorderedMatcher( - const1, const2, const3)); + EXPECT_THAT(call->operands(), + ::testing::UnorderedElementsAre(const1, const2, const3)); // The callee should contain 3 parameters and 3 binary operators. EXPECT_EQ(6, callee_computation->instructions().size()); } +TEST_F(TransposeFoldingTest, FoldDotTransposeInWhile) { + auto builder = HloComputation::Builder("entry_computation"); + HloInstruction* x = builder.AddInstruction(HloInstruction::CreateParameter( + /*parameter_number=*/0, ShapeUtil::MakeShape(F32, {2, 3}), + /*name=*/"x")); + HloInstruction* y = builder.AddInstruction(HloInstruction::CreateParameter( + /*parameter_number=*/1, ShapeUtil::MakeShape(F32, {2, 3}), + /*name=*/"y")); + HloInstruction* transpose_y = + builder.AddInstruction(HloInstruction::CreateTranspose( + ShapeUtil::MakeShape(F32, {3, 2}), y, {1, 0})); + HloInstruction* dot = builder.AddInstruction(HloInstruction::CreateBinary( + ShapeUtil::MakeShape(F32, {2, 2}), /*opcode=*/HloOpcode::kDot, + /*lhs=*/x, /*rhs=*/transpose_y)); + + HloModule module("test_module"); + HloComputation* entry_computation = + module.AddEntryComputation(builder.Build(dot)); + + HloInstruction* call = module.OutlineExpressionFromComputation( + {transpose_y, dot}, "outlined", entry_computation); + + FoldTranspose(&module); + + // Instructions after folding: x, y, and the fusion. + std::unordered_set instruction_set; + for (auto& instruction : entry_computation->instructions()) { + instruction_set.insert(instruction.get()); + } + CHECK_EQ(1, instruction_set.erase(x)) << "x is not in entry_computation."; + CHECK_EQ(1, instruction_set.erase(y)) << "y is not in entry_computation."; + CHECK_EQ(1, instruction_set.erase(call)) + << "call is not in entry_computation."; + CHECK(instruction_set.empty()) + << "entry_computation should contain exactly 3 instructions."; + HloInstruction* fusion = + call->called_computations().front()->root_instruction(); + EXPECT_EQ(HloOpcode::kFusion, fusion->opcode()); + + // The fusion instruction should contain two parameters, one transpose and + // one dot. + EXPECT_EQ(4, fusion->fused_instructions().size()); +} + +// Test that a two dimension swap of the kernel gets folded into convolution. +TEST_F(TransposeFoldingTest, FoldConvDimSwapTransposeRhs) { + auto builder = HloComputation::Builder("entry_computation"); + HloInstruction* x = builder.AddInstruction(HloInstruction::CreateParameter( + /*parameter_number=*/0, ShapeUtil::MakeShape(F32, {2, 3, 1, 1}), + /*name=*/"x")); + HloInstruction* y = builder.AddInstruction(HloInstruction::CreateParameter( + /*parameter_number=*/1, ShapeUtil::MakeShape(F32, {3, 2, 1, 1}), + /*name=*/"y")); + HloInstruction* transpose_y = + builder.AddInstruction(HloInstruction::CreateTranspose( + ShapeUtil::MakeShape(F32, {2, 3, 1, 1}), y, {1, 0, 2, 3})); + auto dnums = ComputationBuilder::CreateDefaultConvDimensionNumbers(); + Window window; + for (int i = 0; i < 2; ++i) { + WindowDimension* dim = window.add_dimensions(); + dim->set_padding_low(0); + dim->set_padding_high(0); + dim->set_base_dilation(1); + dim->set_window_dilation(1); + dim->set_stride(1); + dim->set_size( + transpose_y->shape().dimensions(dnums.kernel_spatial_dimensions(i))); + } + StatusOr conv_shape = ShapeInference::InferConvolveShape( + x->shape(), transpose_y->shape(), window, dnums); + EXPECT_IS_OK(conv_shape); + HloInstruction* conv = builder.AddInstruction(HloInstruction::CreateConvolve( + conv_shape.ValueOrDie(), x, transpose_y, window, dnums)); + + HloModule module("test_module"); + HloComputation* entry_computation = + module.AddEntryComputation(builder.Build(conv)); + FoldTranspose(&module); + + // Instructions after folding: x, y, and the convolution. + std::unordered_set instruction_set; + for (auto& instruction : entry_computation->instructions()) { + instruction_set.insert(instruction.get()); + } + CHECK_EQ(1, instruction_set.erase(x)) << "x is not in entry_computation."; + CHECK_EQ(1, instruction_set.erase(y)) << "y is not in entry_computation."; + CHECK_EQ(1, instruction_set.size()) + << "entry_computation should contain exactly 3 instructions."; + HloInstruction* new_conv = *instruction_set.begin(); + EXPECT_EQ(HloOpcode::kConvolution, new_conv->opcode()); + EXPECT_EQ(dnums.kernel_input_feature_dimension(), + new_conv->convolution_dimension_numbers() + .kernel_output_feature_dimension()); + EXPECT_EQ(dnums.kernel_output_feature_dimension(), + new_conv->convolution_dimension_numbers() + .kernel_input_feature_dimension()); +} + +// Test that a complex transpose of the kernel gets folded into convolution. +TEST_F(TransposeFoldingTest, FoldConvComplexTransposeRhs) { + auto builder = HloComputation::Builder("entry_computation"); + HloInstruction* x = builder.AddInstruction(HloInstruction::CreateParameter( + /*parameter_number=*/0, ShapeUtil::MakeShape(F32, {2, 3, 1, 1}), + /*name=*/"x")); + HloInstruction* y = builder.AddInstruction(HloInstruction::CreateParameter( + /*parameter_number=*/1, ShapeUtil::MakeShape(F32, {1, 2, 1, 3}), + /*name=*/"y")); + HloInstruction* transpose_y = + builder.AddInstruction(HloInstruction::CreateTranspose( + ShapeUtil::MakeShape(F32, {2, 3, 1, 1}), y, {1, 3, 0, 2})); + auto dnums = ComputationBuilder::CreateDefaultConvDimensionNumbers(); + Window window; + for (int i = 0; i < 2; ++i) { + WindowDimension* dim = window.add_dimensions(); + dim->set_padding_low(0); + dim->set_padding_high(0); + dim->set_base_dilation(1); + dim->set_window_dilation(1); + dim->set_stride(1); + dim->set_size( + transpose_y->shape().dimensions(dnums.kernel_spatial_dimensions(i))); + } + StatusOr conv_shape = ShapeInference::InferConvolveShape( + x->shape(), transpose_y->shape(), window, dnums); + EXPECT_IS_OK(conv_shape); + HloInstruction* conv = builder.AddInstruction(HloInstruction::CreateConvolve( + conv_shape.ValueOrDie(), x, transpose_y, window, dnums)); + + HloModule module("test_module"); + HloComputation* entry_computation = + module.AddEntryComputation(builder.Build(conv)); + FoldTranspose(&module); + + // Instructions after folding: x, y, and the convolution. + std::unordered_set instruction_set; + for (auto& instruction : entry_computation->instructions()) { + instruction_set.insert(instruction.get()); + } + CHECK_EQ(1, instruction_set.erase(x)) << "x is not in entry_computation."; + CHECK_EQ(1, instruction_set.erase(y)) << "y is not in entry_computation."; + CHECK_EQ(1, instruction_set.size()) + << "entry_computation should contain exactly 3 instructions."; + HloInstruction* new_conv = *instruction_set.begin(); + EXPECT_EQ(HloOpcode::kConvolution, new_conv->opcode()); + EXPECT_EQ(dnums.kernel_input_feature_dimension(), + new_conv->convolution_dimension_numbers() + .kernel_output_feature_dimension()); + EXPECT_EQ(dnums.kernel_spatial_dimensions(1), + new_conv->convolution_dimension_numbers() + .kernel_input_feature_dimension()); + EXPECT_EQ( + dnums.kernel_output_feature_dimension(), + new_conv->convolution_dimension_numbers().kernel_spatial_dimensions(0)); + EXPECT_EQ( + dnums.kernel_spatial_dimensions(0), + new_conv->convolution_dimension_numbers().kernel_spatial_dimensions(1)); +} + +// Test that a transpose of the activations does not get folded into +// convolution. +TEST_F(TransposeFoldingTest, FoldConvTransposeLhs) { + auto builder = HloComputation::Builder("entry_computation"); + HloInstruction* x = builder.AddInstruction(HloInstruction::CreateParameter( + /*parameter_number=*/0, ShapeUtil::MakeShape(F32, {3, 2, 1, 1}), + /*name=*/"x")); + HloInstruction* y = builder.AddInstruction(HloInstruction::CreateParameter( + /*parameter_number=*/1, ShapeUtil::MakeShape(F32, {2, 3, 1, 1}), + /*name=*/"y")); + HloInstruction* transpose_x = + builder.AddInstruction(HloInstruction::CreateTranspose( + ShapeUtil::MakeShape(F32, {2, 3, 1, 1}), x, {1, 0, 2, 3})); + auto dnums = ComputationBuilder::CreateDefaultConvDimensionNumbers(); + Window window; + for (int i = 0; i < 2; ++i) { + WindowDimension* dim = window.add_dimensions(); + dim->set_padding_low(0); + dim->set_padding_high(0); + dim->set_base_dilation(1); + dim->set_window_dilation(1); + dim->set_stride(1); + dim->set_size(y->shape().dimensions(dnums.kernel_spatial_dimensions(i))); + } + StatusOr conv_shape = ShapeInference::InferConvolveShape( + transpose_x->shape(), y->shape(), window, dnums); + EXPECT_IS_OK(conv_shape); + HloInstruction* conv = builder.AddInstruction(HloInstruction::CreateConvolve( + conv_shape.ValueOrDie(), transpose_x, y, window, dnums)); + + HloModule module("test_module"); + HloComputation* entry_computation = + module.AddEntryComputation(builder.Build(conv)); + FoldTranspose(&module); + + // Instructions after folding: transpose_x, y, and the convolution. + std::unordered_set instruction_set; + for (auto& instruction : entry_computation->instructions()) { + instruction_set.insert(instruction.get()); + } + CHECK_EQ(1, instruction_set.erase(x)) << "x is not in entry_computation."; + CHECK_EQ(1, instruction_set.erase(y)) << "y is not in entry_computation."; + CHECK_EQ(1, instruction_set.erase(transpose_x)) + << "transpose_x is not in entry_computation."; + CHECK_EQ(1, instruction_set.erase(conv)) + << "transpose_x is not in entry_computation."; + CHECK_EQ(0, instruction_set.size()) + << "entry_computation should contain exactly 4 instructions."; +} + } // namespace xla diff --git a/tensorflow/compiler/xla/service/tuple_points_to_analysis.cc b/tensorflow/compiler/xla/service/tuple_points_to_analysis.cc index 98c51b48f9022c5f2d1e23b59a6ce775f3a48e0b..9fc288d3017137c8a0741a9a69c7a20396ce4af1 100644 --- a/tensorflow/compiler/xla/service/tuple_points_to_analysis.cc +++ b/tensorflow/compiler/xla/service/tuple_points_to_analysis.cc @@ -33,10 +33,9 @@ limitations under the License. namespace xla { string BufferAlias::ToString() const { - return tensorflow::strings::StrCat("BufferAlias(", - instruction_->FullyQualifiedName(), "[", + return tensorflow::strings::StrCat("BufferAlias(", instruction_->name(), "[", tensorflow::str_util::Join(index_, ","), - "] => ", buffer_->ToString(), ")"); + "])"); } std::ostream& operator<<(std::ostream& out, const BufferAlias& buffer_alias) { @@ -46,29 +45,25 @@ std::ostream& operator<<(std::ostream& out, const BufferAlias& buffer_alias) { bool PointsToSet::IsAmbiguous() const { bool ambiguous = false; - TF_CHECK_OK(ForEachElement( - [&ambiguous](const ShapeIndex& /*index*/, bool /*is_leaf*/, - const std::vector& points_to) { + ForEachElement( + [&ambiguous](const ShapeIndex& /*index*/, const BufferList& points_to) { ambiguous |= points_to.size() > 1; - return Status::OK(); - })); + }); return ambiguous; } bool PointsToSet::IsDistinct() const { bool distinct = true; std::set all_points_to; - TF_CHECK_OK(ForEachElement([&distinct, &all_points_to]( - const ShapeIndex& /*index*/, bool /*is_leaf*/, - const std::vector& points_to) { + ForEachElement([&distinct, &all_points_to](const ShapeIndex& /*index*/, + const BufferList& points_to) { for (auto& buffer : points_to) { if (all_points_to.count(buffer) != 0) { distinct = false; } all_points_to.insert(buffer); } - return Status::OK(); - })); + }); return distinct; } @@ -78,36 +73,31 @@ size_t PointsToSet::size() const { return CreateFlattenedSet().size(); } -tensorflow::gtl::FlatSet PointsToSet::CreateFlattenedSet() - const { - tensorflow::gtl::FlatSet flat_set; - TF_CHECK_OK(ForEachElement( - [&flat_set](const ShapeIndex& /*index*/, bool /*is_leaf*/, - const std::vector& buffers) { +PointsToSet::BufferSet PointsToSet::CreateFlattenedSet() const { + BufferSet flat_set; + ForEachElement( + [&flat_set](const ShapeIndex& /*index*/, const BufferList& buffers) { flat_set.insert(buffers.begin(), buffers.end()); - return Status::OK(); - })); + }); return flat_set; } bool PointsToSet::ContainsBuffer(const LogicalBuffer& buffer) const { bool found = false; - TF_CHECK_OK(ForEachElement([&found, &buffer]( - const ShapeIndex& /*index*/, bool /*is_leaf*/, - const std::vector& pointed_to_buffers) { + 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()) { found = true; } - return Status::OK(); - })); + }); return found; } bool PointsToSet::ContainsBufferAtIndex(const LogicalBuffer& buffer, const ShapeIndex& index) const { - const std::vector& pointed_to_buffers = element(index); + const auto& pointed_to_buffers = element(index); return std::find(pointed_to_buffers.begin(), pointed_to_buffers.end(), &buffer) != pointed_to_buffers.end(); } @@ -120,42 +110,48 @@ void PointsToSet::AddPointedToBuffer(const LogicalBuffer& buffer, mutable_element(index)->push_back(&buffer); } -const std::set& PointsToSet::tuple_sources( +const PointsToSet::SourceSet& PointsToSet::tuple_sources( const ShapeIndex& index) const { - return tuple_sources_.element(index); + return tree_.element(index).tuple_sources; } void PointsToSet::add_tuple_source(const ShapeIndex& index, HloInstruction* tuple) { - tuple_sources_.mutable_element(index)->insert(tuple); + tree_.mutable_element(index)->tuple_sources.insert(tuple); } /* static */ StatusOr> -TuplePointsToAnalysis::Run(const HloModule* module, - const bool include_loop_fusion_instructions) { - std::unique_ptr analysis( - new TuplePointsToAnalysis(module, include_loop_fusion_instructions)); +TuplePointsToAnalysis::Run(const HloModule* module) { + auto logical_buffer_analysis = LogicalBufferAnalysis::Run(module); + std::unique_ptr analysis(new TuplePointsToAnalysis( + module, logical_buffer_analysis.ConsumeValueOrDie())); TF_RETURN_IF_ERROR(analysis->Analyze()); return std::move(analysis); } Status TuplePointsToAnalysis::Analyze() { - points_to_.clear(); + per_instruction_.clear(); + per_instruction_.resize(module_->NumUniqueInstructionIds()); + + logical_buffer_aliases_.clear(); + logical_buffer_aliases_.resize( + logical_buffer_analysis_->num_logical_buffers()); + for (auto& computation : module_->computations()) { + if (computation->IsFusionComputation()) { + continue; + } TF_RETURN_IF_ERROR(computation->Accept(this)); TF_RETURN_IF_ERROR( PopulateDefinedBuffersAndAliases(computation->instructions())); - if (include_loop_fusion_instructions_) { - // Run points-to analysis on loop fusion instructions in 'computation'. - for (auto& instruction : computation->instructions()) { - if (instruction->opcode() != HloOpcode::kFusion || - instruction->fusion_kind() != HloInstruction::FusionKind::kLoop) { - continue; - } - TF_RETURN_IF_ERROR(instruction->fused_expression_root()->Accept(this)); - TF_RETURN_IF_ERROR(PopulateDefinedBuffersAndAliases( - instruction->fused_instructions())); + // Run points-to analysis on fusion instructions in 'computation'. + for (auto& instruction : computation->instructions()) { + if (instruction->opcode() != HloOpcode::kFusion) { + continue; } + TF_RETURN_IF_ERROR(instruction->fused_expression_root()->Accept(this)); + TF_RETURN_IF_ERROR( + PopulateDefinedBuffersAndAliases(instruction->fused_instructions())); } } @@ -167,46 +163,35 @@ Status TuplePointsToAnalysis::Analyze() { Status TuplePointsToAnalysis::PopulateDefinedBuffersAndAliases( const std::list>& instructions) { for (auto& instruction : instructions) { + PerInstruction* pi = PerInst(instruction.get()); TF_RETURN_IF_ERROR(GatherBuffersDefinedByInstruction( - instruction.get(), &instruction_defined_buffers_[instruction.get()])); + instruction.get(), &pi->instruction_defined_buffers)); const PointsToSet& points_to_set = GetPointsToSet(instruction.get()); - TF_RETURN_IF_ERROR(points_to_set.ForEachElement([this, &instruction]( - const ShapeIndex& index, bool /*is_leaf*/, - const std::vector& pointed_to_buffers) { - for (const LogicalBuffer* buffer : pointed_to_buffers) { - if (buffer_aliases_.count(buffer) == 0) { - buffer_aliases_.insert({buffer, std::vector()}); - } - buffer_aliases_[buffer].emplace_back(*buffer, instruction.get(), index); - } - return Status::OK(); - })); + points_to_set.ForEachElement( + [this, &instruction]( + const ShapeIndex& index, + const PointsToSet::BufferList& pointed_to_buffers) { + for (const LogicalBuffer* buffer : pointed_to_buffers) { + logical_buffer_aliases_[buffer->id()].emplace_back( + instruction.get(), index); + } + }); } return Status::OK(); } -const LogicalBuffer& TuplePointsToAnalysis::NewLogicalBuffer( - HloInstruction* instruction, const ShapeIndex& index) { - CHECK_EQ(logical_buffers_.size(), next_buffer_id_); - logical_buffers_.push_back( - MakeUnique(instruction, index, next_buffer_id_)); - ++next_buffer_id_; - return *logical_buffers_.back(); -} - Status TuplePointsToAnalysis::DefaultAction(HloInstruction* hlo_instruction) { // Create trivial points-to set for instruction. Each points-to set at index i // contains a single element LogicalBuffer(hlo_instruction, i). This indicates // that this instruction is the source of all buffers in its own output. PointsToSet& points_to_set = CreateEmptyPointsToSet(hlo_instruction); - TF_RETURN_IF_ERROR(points_to_set.ForEachMutableElement( - [this, hlo_instruction](const ShapeIndex& index, bool /*is_leaf*/, - std::vector* buffers) { - const LogicalBuffer& buffer = NewLogicalBuffer(hlo_instruction, index); - buffers->push_back(&buffer); - return Status::OK(); - })); + points_to_set.ForEachMutableElement( + [this, hlo_instruction](const ShapeIndex& index, + PointsToSet::BufferList* buffers) { + buffers->push_back( + &logical_buffer_analysis_->GetBuffer(hlo_instruction, index)); + }); if (ShapeUtil::IsTuple(hlo_instruction->shape())) { // If the hlo instruction is a tuple-shaped, then trivially the instruction @@ -224,41 +209,40 @@ Status TuplePointsToAnalysis::HandleGetTupleElement( int64 element_index = get_tuple_element->tuple_index(); PointsToSet& points_to_set = CreateEmptyPointsToSet(get_tuple_element); - const PointsToSet& operand_points_to_set = *FindOrDie(points_to_, operand); + const PointsToSet& operand_points_to_set = *PerInst(operand)->points_to_set; // Copy the points-to set (and tuple sources) at index {element_index} of the // operand to the points-to set for this GetTupleElement instruction. - TF_RETURN_IF_ERROR(points_to_set.ForEachMutableElement([&, this]( - const ShapeIndex& target_index, bool /*is_leaf*/, - std::vector* points_to) { - // Construct an index into the operand by prepending element_index to the - // index for the GetTupleElement instruction's points-to set. - ShapeIndex src_index; - src_index.push_back(element_index); - for (auto element : target_index) { - src_index.push_back(element); - } + points_to_set.ForEachMutableElement( + [&, this](const ShapeIndex& target_index, + PointsToSet::BufferList* points_to) { + // Construct an index into the operand by prepending element_index to + // the index for the GetTupleElement instruction's points-to set. + ShapeIndex src_index; + src_index.push_back(element_index); + for (auto element : target_index) { + src_index.push_back(element); + } - *points_to = operand_points_to_set.element(src_index); - for (HloInstruction* tuple : - operand_points_to_set.tuple_sources(src_index)) { - points_to_set.add_tuple_source(target_index, tuple); - } - return Status::OK(); - })); + *points_to = operand_points_to_set.element(src_index); + for (HloInstruction* tuple : + operand_points_to_set.tuple_sources(src_index)) { + points_to_set.add_tuple_source(target_index, tuple); + } + }); return Status::OK(); } -Status TuplePointsToAnalysis::HandleCopy(HloInstruction* copy, - HloInstruction* operand) { +Status TuplePointsToAnalysis::HandleCopy(HloInstruction* copy) { // A kCopy instruction performs a shallow copy of the operand. The top-level // buffer (index={}) is newly created, but all other buffers (in the case of a // tuple shape) come from the operand - PointsToSet& points_to_set = CreateCopiedPointsToSet(copy, operand); + PointsToSet& points_to_set = CreateCopiedPointsToSet(copy, copy->operand(0)); points_to_set.mutable_element(/*index=*/{})->clear(); - points_to_set.AddPointedToBuffer(NewLogicalBuffer(copy, /*index=*/{}), - /*index=*/{}); + points_to_set.AddPointedToBuffer( + logical_buffer_analysis_->GetBuffer(copy, /*index=*/{}), + /*index=*/{}); return Status::OK(); } @@ -275,21 +259,22 @@ Status TuplePointsToAnalysis::HandleTuple( HloInstruction* tuple, tensorflow::gtl::ArraySlice operands) { PointsToSet& points_to_set = CreateEmptyPointsToSet(tuple); - points_to_set.AddPointedToBuffer(NewLogicalBuffer(tuple, /*index=*/{}), - /*index=*/{}); + points_to_set.AddPointedToBuffer( + logical_buffer_analysis_->GetBuffer(tuple, /*index=*/{}), + /*index=*/{}); // A tuple contains references to all input operands and transitively any // references in those operands. for (int64 i = 0; i < operands.size(); ++i) { const PointsToSet& operand_points_to_set = - *FindOrDie(points_to_, operands[i]); + *PerInst(operands[i])->points_to_set; // Copy the points-to set (and tuple sources) of the operand into the // respective subtree of the tuple instructions points-to set. - TF_RETURN_IF_ERROR(operand_points_to_set.ForEachElement( + operand_points_to_set.ForEachElement( [&points_to_set, &operand_points_to_set, i]( - const ShapeIndex& src_index, bool /*is_leaf*/, - const std::vector& points_to) { + const ShapeIndex& src_index, + const PointsToSet::BufferList& points_to) { ShapeIndex target_index; target_index.push_back(i); for (auto element : src_index) { @@ -302,8 +287,7 @@ Status TuplePointsToAnalysis::HandleTuple( operand_points_to_set.tuple_sources(src_index)) { points_to_set.add_tuple_source(target_index, tuple); } - return Status::OK(); - })); + }); } points_to_set.add_tuple_source({}, tuple); @@ -323,10 +307,9 @@ Status TuplePointsToAnalysis::HandleSelect(HloInstruction* select, // First create a copy of the on_true points-to set (and tuple sources), then // add in elements of the on_false points-to set (tuple sources). PointsToSet& points_to_set = CreateCopiedPointsToSet(select, on_true); - const PointsToSet& false_points_to_set = *FindOrDie(points_to_, on_false); - TF_RETURN_IF_ERROR(points_to_set.ForEachMutableElement( - [&](const ShapeIndex& index, bool /*is_leaf*/, - std::vector* buffers) { + const PointsToSet& false_points_to_set = *PerInst(on_false)->points_to_set; + points_to_set.ForEachMutableElement( + [&](const ShapeIndex& index, PointsToSet::BufferList* buffers) { for (const LogicalBuffer* false_buffer : false_points_to_set.element(index)) { points_to_set.AddPointedToBuffer(*false_buffer, index); @@ -335,39 +318,36 @@ Status TuplePointsToAnalysis::HandleSelect(HloInstruction* select, for (HloInstruction* tuple : false_points_to_set.tuple_sources(index)) { points_to_set.add_tuple_source(index, tuple); } - return Status::OK(); - })); + }); // Select creates a new (top-level) buffer to store its result, so its // respective element in the points-to set should contain only itself. points_to_set.mutable_element({})->clear(); - points_to_set.AddPointedToBuffer(NewLogicalBuffer(select, /*index=*/{}), - /*index=*/{}); + points_to_set.AddPointedToBuffer( + logical_buffer_analysis_->GetBuffer(select, /*index=*/{}), + /*index=*/{}); return Status::OK(); } -Status TuplePointsToAnalysis::HandleFusion(HloInstruction* fusion) { - return ShapeUtil::IsTuple(fusion->shape()) - ? Unimplemented("HandleFusion with tuple output") - : DefaultAction(fusion); -} - const PointsToSet& TuplePointsToAnalysis::GetPointsToSet( const HloInstruction* hlo_instruction) const { - return *FindOrDie(points_to_, hlo_instruction); + return *PerInst(hlo_instruction)->points_to_set; } PointsToSet& TuplePointsToAnalysis::CreateEmptyPointsToSet( const HloInstruction* instruction) { - CHECK_EQ(0, points_to_.count(instruction)); - points_to_[instruction] = MakeUnique(instruction->shape()); - return *FindOrDie(points_to_, instruction); + PerInstruction* pi = PerInst(instruction); + CHECK(pi->points_to_set == nullptr) + << "instruction should not have been present in the map."; + auto set = MakeUnique(&instruction->shape()); + pi->points_to_set = std::move(set); + // Return *set using the iterator returned by emplace. + return *pi->points_to_set; } bool TuplePointsToAnalysis::InstructionDefinesBufferAtIndex( const HloInstruction* instruction, const ShapeIndex& index) const { - const std::vector& buffers = - GetPointsToSet(instruction).element(index); + const auto& buffers = GetPointsToSet(instruction).element(index); return (buffers.size() == 1 && buffers[0]->instruction() == instruction); } @@ -379,7 +359,8 @@ Status TuplePointsToAnalysis::VerifyBuffer(const LogicalBuffer& buffer) const { buffer.ToString().c_str(), buffer.instruction()->name().c_str()); } - if (buffer.id() < 0 || buffer.id() >= next_buffer_id_) { + if (buffer.id() < 0 || + buffer.id() >= logical_buffer_analysis_->num_logical_buffers()) { return FailedPrecondition( "LogicalBuffer %s is ill-defined: invalid id %lld", buffer.ToString().c_str(), buffer.id()); @@ -397,14 +378,13 @@ Status TuplePointsToAnalysis::VerifyBuffer(const LogicalBuffer& buffer) const { const LogicalBuffer& TuplePointsToAnalysis::GetBuffer( LogicalBuffer::Id id) const { CHECK_GE(id, 0); - CHECK_LT(id, logical_buffers_.size()); - return *logical_buffers_[id]; + CHECK_LT(id, logical_buffer_analysis_->num_logical_buffers()); + return logical_buffer_analysis_->GetBuffer(id); } StatusOr TuplePointsToAnalysis::GetBufferDefinedAt( const HloInstruction* instruction, const ShapeIndex& index) const { - const std::vector& buffers = - GetPointsToSet(instruction).element(index); + const auto& buffers = GetPointsToSet(instruction).element(index); if (buffers.size() != 1 || buffers[0]->instruction() != instruction) { return FailedPrecondition( "instruction %s does not define buffer at index {%s}", @@ -414,24 +394,24 @@ StatusOr TuplePointsToAnalysis::GetBufferDefinedAt( return buffers[0]; } -const std::vector& TuplePointsToAnalysis::GetBufferAliases( - const LogicalBuffer& buffer) const { - return buffer_aliases_.at(&buffer); +const TuplePointsToAnalysis::BufferAliasVector& +TuplePointsToAnalysis::GetBufferAliases(const LogicalBuffer& buffer) const { + return logical_buffer_aliases_.at(buffer.id()); } -const std::vector& +const TuplePointsToAnalysis::BufferDefinitionVector& TuplePointsToAnalysis::GetBuffersDefinedByInstruction( const HloInstruction* instruction) const { - return instruction_defined_buffers_.at(instruction); + return PerInst(instruction)->instruction_defined_buffers; } Status TuplePointsToAnalysis::GatherBuffersDefinedByInstruction( const HloInstruction* instruction, - std::vector* buffers) { - return GetPointsToSet(instruction) + TuplePointsToAnalysis::BufferDefinitionVector* buffers) { + GetPointsToSet(instruction) .ForEachElement([this, buffers, instruction]( - const ShapeIndex& index, bool /*is_leaf*/, - const std::vector& source_buffers) { + const ShapeIndex& index, + const PointsToSet::BufferList& source_buffers) { // Add buffers which 'instruction' is the source of. CHECK(!source_buffers.empty()); if (source_buffers.size() == 1 && @@ -448,8 +428,8 @@ Status TuplePointsToAnalysis::GatherBuffersDefinedByInstruction( DCHECK(source_buffer->instruction() != instruction); } } - return Status::OK(); }); + return Status::OK(); } PointsToSet& TuplePointsToAnalysis::CreateCopiedPointsToSet( @@ -458,23 +438,24 @@ PointsToSet& TuplePointsToAnalysis::CreateCopiedPointsToSet( // from src PointsToSet. PointsToSet& dst_points_to_set = CreateEmptyPointsToSet(instruction); const PointsToSet& src_points_to_set = GetPointsToSet(src); - TF_CHECK_OK(dst_points_to_set.ForEachMutableElement( + dst_points_to_set.ForEachMutableElement( [this, &dst_points_to_set, &src_points_to_set]( - const ShapeIndex& index, bool /*is_leaf*/, - std::vector* buffers) { + const ShapeIndex& index, PointsToSet::BufferList* buffers) { *buffers = src_points_to_set.element(index); for (auto& tuple_source : src_points_to_set.tuple_sources(index)) { dst_points_to_set.add_tuple_source(index, tuple_source); } - return Status::OK(); - })); - return *FindOrDie(points_to_, instruction); + }); + return *PerInst(instruction)->points_to_set; } string TuplePointsToAnalysis::ToString() const { string output = tensorflow::strings::Printf( "TuplePointsToSet for module %s:\n", module_->name().c_str()); for (const auto& computation : module_->computations()) { + if (computation->IsFusionComputation()) { + continue; + } const char* entry = computation.get() == module_->entry_computation() ? "entry " : ""; tensorflow::strings::StrAppend(&output, entry, "computation ", @@ -482,9 +463,7 @@ string TuplePointsToAnalysis::ToString() const { for (const HloInstruction* instruction : computation->MakeInstructionPostOrder()) { InstructionToString(instruction, &output); - if (include_loop_fusion_instructions_ && - instruction->opcode() == HloOpcode::kFusion && - instruction->fusion_kind() == HloInstruction::FusionKind::kLoop) { + if (instruction->opcode() == HloOpcode::kFusion) { for (auto& fused : instruction->fused_instructions()) { InstructionToString(fused.get(), &output); } @@ -493,12 +472,11 @@ string TuplePointsToAnalysis::ToString() const { } tensorflow::strings::StrAppend(&output, "LogicalBuffers:\n"); - for (auto& buffer : logical_buffers_) { - tensorflow::strings::StrAppend(&output, " buffer ", buffer->ToString(), - ":\n"); - for (const BufferAlias& buffer_alias : buffer_aliases_.at(buffer.get())) { - tensorflow::strings::StrAppend(&output, " alias ", - buffer_alias.ToString(), "\n"); + for (const auto& b : logical_buffer_analysis_->logical_buffers()) { + tensorflow::strings::StrAppend(&output, " buffer ", b->ToString(), ":\n"); + for (const BufferAlias& alias : logical_buffer_aliases_.at(b->id())) { + tensorflow::strings::StrAppend(&output, " alias ", alias.ToString(), + "\n"); } } return output; @@ -510,9 +488,9 @@ void TuplePointsToAnalysis::InstructionToString( tensorflow::strings::StrAppend(output, prefix, " instruction ", instruction->ToShortString(), ":\n"); const PointsToSet& points_to_set = GetPointsToSet(instruction); - TF_CHECK_OK(points_to_set.ForEachElement([&prefix, &output]( - const ShapeIndex& index, bool /*is_leaf*/, - const std::vector& points_to) { + points_to_set.ForEachElement([&prefix, &output]( + const ShapeIndex& index, + const PointsToSet::BufferList& points_to) { tensorflow::strings::StrAppend( output, prefix, " {", tensorflow::str_util::Join(index, ","), "}: ", tensorflow::str_util::Join( @@ -521,8 +499,7 @@ void TuplePointsToAnalysis::InstructionToString( out->append(source->ToString()); }), "\n"); - return Status::OK(); - })); + }); } } // namespace xla diff --git a/tensorflow/compiler/xla/service/tuple_points_to_analysis.h b/tensorflow/compiler/xla/service/tuple_points_to_analysis.h index a384529171a7371c848ca8949d22cb6717d83a78..3b3a046e498f0e5fdd6c0a18caadab856f5db676 100644 --- a/tensorflow/compiler/xla/service/tuple_points_to_analysis.h +++ b/tensorflow/compiler/xla/service/tuple_points_to_analysis.h @@ -27,12 +27,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/logical_buffer.h" +#include "tensorflow/compiler/xla/service/logical_buffer_analysis.h" #include "tensorflow/compiler/xla/shape_tree.h" #include "tensorflow/compiler/xla/statusor.h" #include "tensorflow/compiler/xla/types.h" #include "tensorflow/compiler/xla/xla_data.pb.h" #include "tensorflow/core/lib/core/status.h" #include "tensorflow/core/lib/gtl/array_slice.h" +#include "tensorflow/core/lib/gtl/compactptrset.h" #include "tensorflow/core/lib/gtl/flatmap.h" #include "tensorflow/core/lib/gtl/flatset.h" #include "tensorflow/core/platform/macros.h" @@ -46,11 +48,12 @@ namespace xla { // nested tuple). Each node in this tree corresponds to a single buffer in the // instruction's output and contains the set of Buffers which might define // the corresponding buffer. -class PointsToSet : public ShapeTree> { +class PointsToSet { public: - explicit PointsToSet(const Shape& shape) - : ShapeTree>(shape), - tuple_sources_(shape) {} + // Construct our ShapeTree with a pointer rather than a reference to a Shape + // because this is very hot code, and copying (and then destroying) all these + // Shapes is slow. + explicit PointsToSet(const Shape* shape) : tree_(shape) {} // Returns true if any points-to sets for any subshape element is not a // singleton. @@ -66,7 +69,8 @@ class PointsToSet : public ShapeTree> { // Creates a set containing the union of all LogicalBuffers contained in the // PointsToSet. - tensorflow::gtl::FlatSet CreateFlattenedSet() const; + using BufferSet = tensorflow::gtl::CompactPointerSet; + BufferSet CreateFlattenedSet() const; // Returns true if the given buffer is in the points-to set at the given // index. @@ -99,13 +103,49 @@ class PointsToSet : public ShapeTree> { // tuple_sources() at the index of an array shape (not a tuple) returns the // empty set. The instructions in the set returned by tuple_sources // necessarily are either Tuple instructions, constants, or parameters. - const std::set& tuple_sources(const ShapeIndex& index) const; + using SourceSet = tensorflow::gtl::CompactPointerSet; + const SourceSet& tuple_sources(const ShapeIndex& index) const; // Add a tuple source instruction for the given index. void add_tuple_source(const ShapeIndex& index, HloInstruction* tuple); + using BufferList = tensorflow::gtl::InlinedVector; + + // Return the list of logical buffers for the subshape at index. + const BufferList& element(const ShapeIndex& index) const { + return tree_.element(index).buffers; + } + BufferList* mutable_element(const ShapeIndex& index) { + return &tree_.mutable_element(index)->buffers; + } + + // Call fn(index, buflist) for every subshape index. + template + void ForEachElement(const Fn& fn) const { + tree_.ForEachElement([&fn](const ShapeIndex& index, const Elem& elem) { + fn(index, elem.buffers); + }); + } + template + void ForEachMutableElement(const Fn& fn) { + tree_.ForEachMutableElement([&fn](const ShapeIndex& index, Elem* elem) { + fn(index, &elem->buffers); + }); + } + template + Status ForEachElementWithStatus(const Fn& fn) const { + return tree_.ForEachElementWithStatus( + [&fn](const ShapeIndex& index, const Elem& elem) { + return fn(index, elem.buffers); + }); + } + private: - ShapeTree> tuple_sources_; + struct Elem { + BufferList buffers; + SourceSet tuple_sources; + }; + ShapeTree tree_; // PointsToSet contains references (const LogicalBuffer*) to elements within // TuplePointsToAnalysis so disable copying. @@ -117,29 +157,23 @@ class PointsToSet : public ShapeTree> { // value. class BufferAlias { public: - BufferAlias(const LogicalBuffer& buffer, HloInstruction* instruction, - const ShapeIndex& index) - : buffer_(&buffer), instruction_(instruction), index_(index) {} - - // Return the logical buffer aliased at the instruction and index. - const LogicalBuffer& buffer() const { return *buffer_; } + BufferAlias(HloInstruction* instruction, const ShapeIndex& index) + : instruction_(instruction), index_(index) {} // Return the instruction/index of the subshape. HloInstruction* instruction() const { return instruction_; } const ShapeIndex& index() const { return index_; } bool operator==(const BufferAlias& other) const { - return buffer_ == other.buffer_ && instruction_ == other.instruction_ && - index_ == other.index_; + return instruction_ == other.instruction_ && index_ == other.index_; } bool operator!=(const BufferAlias& other) const { return !(*this == other); } string ToString() const; private: - const LogicalBuffer* buffer_; HloInstruction* instruction_; - const ShapeIndex index_; + ShapeIndex index_; }; std::ostream& operator<<(std::ostream& out, const BufferAlias& buffer_alias); @@ -148,12 +182,9 @@ std::ostream& operator<<(std::ostream& out, const BufferAlias& buffer_alias); // the potential sources of each buffer in each instruction's output. class TuplePointsToAnalysis : public DfsHloVisitorWithDefault { public: - // Runs points-to analysis on 'module'. If 'include_loop_fusion_instructions' - // is true, includes fused instructions from each loop fusion instruction - // in 'module' in the points-to analysis. + // Runs points-to analysis on 'module'. static StatusOr> Run( - const HloModule* module, - const bool include_loop_fusion_instructions = false); + const HloModule* module); // Return the points-to set of an instruction. This describes the potential // sources of each buffer in the instruction's output. @@ -174,19 +205,32 @@ class TuplePointsToAnalysis : public DfsHloVisitorWithDefault { // buffer alias set is the inverse of the points-to set. That is, // LogicalBuffer B is in the points-to set of instruction I at index N iff // instruction I, index N is a BufferAlias of B. - const std::vector& GetBufferAliases( - const LogicalBuffer& buffer) const; + using BufferAliasVector = tensorflow::gtl::InlinedVector; + const BufferAliasVector& GetBufferAliases(const LogicalBuffer& buffer) const; - // Return a vector containing all logical buffers in the module. - const std::vector>& logical_buffers() const { - return logical_buffers_; + // Returns the number of logical buffers in the module + LogicalBuffer::Id num_logical_buffers() const { + return logical_buffer_analysis_->num_logical_buffers(); + } + + // Return a the logical buffer with id "id" in the module. Iteration + // over all logical buffers is usually done with something like: + // + // for (LogicalBuffer:Id id = 0; id < points_to.num_logical_buffers(); id++){ + // const auto& buffer = points_to.logical_buffer(id); + // ... do something with buffer ... + // } + LogicalBuffer& logical_buffer(LogicalBuffer::Id id) const { + return logical_buffer_analysis_->GetBuffer(id); } // Returns a vector of buffers that the instruction produces. Most // instructions produce a single buffer (the top-level buffer), some produce // no buffers (eg bitcast), and some produce more than one buffer (eg, // tuple-shaped parameters). - const std::vector& GetBuffersDefinedByInstruction( + using BufferDefinitionVector = + tensorflow::gtl::InlinedVector; + const BufferDefinitionVector& GetBuffersDefinedByInstruction( const HloInstruction* instruction) const; // Returns true if the given instruction defines a buffer at the given index. @@ -209,8 +253,7 @@ class TuplePointsToAnalysis : public DfsHloVisitorWithDefault { Status HandleGetTupleElement(HloInstruction* get_tuple_element, HloInstruction* operand) override; Status HandleBitcast(HloInstruction* bitcast) override; - Status HandleCopy(HloInstruction* copy, HloInstruction* operand) override; - Status HandleFusion(HloInstruction* fusion) override; + Status HandleCopy(HloInstruction* copy) override; Status HandleSelect(HloInstruction* select, HloInstruction* pred, HloInstruction* on_true, HloInstruction* on_false) override; @@ -218,10 +261,11 @@ class TuplePointsToAnalysis : public DfsHloVisitorWithDefault { string ToString() const; private: - explicit TuplePointsToAnalysis(const HloModule* module, - const bool include_loop_fusion_instructions) + explicit TuplePointsToAnalysis( + const HloModule* module, + std::unique_ptr logical_buffer_analysis) : module_(module), - include_loop_fusion_instructions_(include_loop_fusion_instructions) {} + logical_buffer_analysis_(std::move(logical_buffer_analysis)) {} // Perform the analysis. Should be called immediately after constructing the // object and before calling GetPointsToSet. @@ -234,12 +278,6 @@ class TuplePointsToAnalysis : public DfsHloVisitorWithDefault { Status PopulateDefinedBuffersAndAliases( const std::list>& instructions); - // Create a new logical buffer and return a reference to it. The newly created - // buffer is stored in an internal vector of LogicalBuffers and can be - // accessed with GetBuffer. - const LogicalBuffer& NewLogicalBuffer(HloInstruction* instruction, - const ShapeIndex& index); - // Creates an empty PointsToSet in the points_to_ map for the given // instruction. PointsToSet& CreateEmptyPointsToSet(const HloInstruction* instruction); @@ -250,38 +288,46 @@ class TuplePointsToAnalysis : public DfsHloVisitorWithDefault { const HloInstruction* src); // Adds the buffers defined by the given instruction to the given vector. - Status GatherBuffersDefinedByInstruction( - const HloInstruction* instruction, - std::vector* buffers); + Status GatherBuffersDefinedByInstruction(const HloInstruction* instruction, + BufferDefinitionVector* buffers); // Print points-to set for 'instruction' to 'output'. void InstructionToString(const HloInstruction* instruction, string* output) const; + // Information kept per instruction + struct PerInstruction { + std::unique_ptr points_to_set; + // Empircally, ~92% of instructions have 1 + // instruction_defined_buffer, and 99% have 0 or 1 + BufferDefinitionVector instruction_defined_buffers; + }; + + const PerInstruction* PerInst(const HloInstruction* inst) const { + int id = inst->unique_id(); + DCHECK_GE(id, 0); + DCHECK_LT(id, per_instruction_.size()); + return &per_instruction_[id]; + } + PerInstruction* PerInst(const HloInstruction* inst) { + int id = inst->unique_id(); + DCHECK_GE(id, 0); + DCHECK_LT(id, per_instruction_.size()); + return &per_instruction_[id]; + } + // The module this analysis is performed on. const HloModule* module_; - // Whether to run points-to analysis on loop fusion instructions in 'module_'. - const bool include_loop_fusion_instructions_; - - // A map containing a PointsToSet for every HLO instruction. - tensorflow::gtl::FlatMap> - points_to_; - - // A map containing the LogicalBuffers defined by each HLO instruction. - tensorflow::gtl::FlatMap> - instruction_defined_buffers_; - - tensorflow::gtl::FlatMap> - buffer_aliases_; + // The logical buffers for this module. + const std::unique_ptr logical_buffer_analysis_; - // All logical buffers in the module, indexed by LogicalBuffer::Id. Keep as - // vector of std::unique_ptr to keep the underlying pointer values stable. - std::vector> logical_buffers_; + // A map from instruction->unique_id() to + std::vector per_instruction_; - // The ID of the next logical buffer created. - LogicalBuffer::Id next_buffer_id_ = 0; + // A map from LogicalBuffer->id() to alias information about that logical + // buffer + std::vector logical_buffer_aliases_; TF_DISALLOW_COPY_AND_ASSIGN(TuplePointsToAnalysis); }; 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 4a4a6e64ffae265bc143cfd7adb9f7d53b2b0359..5a23553d4ec7ea5d34a3d0957fbf0ad04ad801fb 100644 --- a/tensorflow/compiler/xla/service/tuple_points_to_analysis_test.cc +++ b/tensorflow/compiler/xla/service/tuple_points_to_analysis_test.cc @@ -19,18 +19,25 @@ limitations under the License. #include #include "tensorflow/compiler/xla/literal_util.h" +#include "tensorflow/compiler/xla/service/hlo_matchers.h" #include "tensorflow/compiler/xla/service/hlo_opcode.h" #include "tensorflow/compiler/xla/service/instruction_fusion.h" #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/platform/logging.h" #include "tensorflow/core/platform/test.h" +namespace op = xla::testing::opcode_matchers; + namespace xla { namespace { +using ::testing::UnorderedElementsAre; +using ::testing::UnorderedElementsAreArray; + class TuplePointsToAnalysisTest : public HloTestBase { protected: // Builds a module with the given entry computation and runs points to @@ -41,22 +48,21 @@ class TuplePointsToAnalysisTest : public HloTestBase { } void BuildModule(std::unique_ptr computation) { - module_.reset(new HloModule(TestName())); + module_ = CreateNewModule(); module_->AddEntryComputation(std::move(computation)); } - void RunAnalysis(const bool include_loop_fusion_instructions = false) { + void RunAnalysis() { CHECK_NOTNULL(module_.get()); - points_to_analysis_ = TuplePointsToAnalysis::Run( - module_.get(), include_loop_fusion_instructions) - .ConsumeValueOrDie(); + points_to_analysis_ = + TuplePointsToAnalysis::Run(module_.get()).ConsumeValueOrDie(); } // Returns the LogicalBuffer defined at the given instruction and // index. CHECKs if no buffer is defined at that point. const LogicalBuffer* const GetBuffer(const HloInstruction* instruction, const ShapeIndex& index) { - const std::vector& pointed_to = + const auto& pointed_to = points_to_analysis_->GetPointsToSet(instruction).element(index); CHECK_EQ(1, pointed_to.size()); CHECK_EQ(instruction, pointed_to[0]->instruction()); @@ -67,31 +73,31 @@ class TuplePointsToAnalysisTest : public HloTestBase { // Checks that the given points-to set contains exactly (unordered) the given // LogicalBuffers. void ExpectHasBuffers( - const std::vector& points_to_set, + const PointsToSet::BufferList& points_to_set, tensorflow::gtl::ArraySlice buffers) { std::vector vec(buffers.begin(), buffers.end()); - EXPECT_MATCH(points_to_set, testing::UnorderedElementsAre(vec)); + EXPECT_THAT(points_to_set, UnorderedElementsAreArray(vec)); } // Checks that the given points-to set contains exactly (unordered) the // top-level buffers of the given instructions. void ExpectHasTopLevelBuffers( - const std::vector& points_to_set, + const PointsToSet::BufferList& points_to_set, tensorflow::gtl::ArraySlice instructions) { - std::vector buffers; + PointsToSet::BufferList buffers; for (auto instruction : instructions) { buffers.push_back(GetBuffer(instruction, /*index=*/{})); } ExpectHasBuffers(points_to_set, buffers); } - // Overload which takes a std::set instead of a std::vector. + // Overload which takes a set instead of a vector. void ExpectHasTopLevelBuffers( - const tensorflow::gtl::FlatSet& points_to_set, + const PointsToSet::BufferSet& points_to_set, tensorflow::gtl::ArraySlice instructions) { - ExpectHasTopLevelBuffers(std::vector( - points_to_set.begin(), points_to_set.end()), - instructions); + ExpectHasTopLevelBuffers( + PointsToSet::BufferList(points_to_set.begin(), points_to_set.end()), + instructions); } // Checks that the buffer defined at the given instruction and index has @@ -105,28 +111,22 @@ class TuplePointsToAnalysisTest : public HloTestBase { .ValueOrDie(); std::vector expected_aliases; for (auto& pair : expected) { - expected_aliases.push_back(BufferAlias(*buffer, pair.first, pair.second)); + expected_aliases.push_back(BufferAlias(pair.first, pair.second)); } - EXPECT_MATCH(points_to_analysis_->GetBufferAliases(*buffer), - testing::UnorderedElementsAre(expected_aliases)); + EXPECT_THAT(points_to_analysis_->GetBufferAliases(*buffer), + UnorderedElementsAreArray(expected_aliases)); } std::unique_ptr module_; std::unique_ptr points_to_analysis_; }; -// Expect the given std::set as A contains exactly the given -// HloInstruction*s as __VA_ARGS__. -#define EXPECT_ISET(A, ...) \ - EXPECT_MATCH(testing::SetToVec(A), \ - testing::UnorderedMatcher(__VA_ARGS__)) - TEST_F(TuplePointsToAnalysisTest, SimpleTuple) { auto builder = HloComputation::Builder(TestName()); auto constant1 = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(1.0))); + HloInstruction::CreateConstant(Literal::CreateR0(1.0))); auto constant2 = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(2.0))); + HloInstruction::CreateConstant(Literal::CreateR0(2.0))); auto tuple = builder.AddInstruction( HloInstruction::CreateTuple({constant1, constant2})); @@ -146,8 +146,8 @@ TEST_F(TuplePointsToAnalysisTest, SimpleTuple) { EXPECT_EQ(3, points_to_analysis_->GetPointsToSet(tuple).size()); EXPECT_FALSE(points_to_analysis_->GetPointsToSet(tuple).IsAmbiguous()); - EXPECT_ISET(points_to_analysis_->GetPointsToSet(tuple).tuple_sources({}), - tuple); + EXPECT_THAT(points_to_analysis_->GetPointsToSet(tuple).tuple_sources({}), + UnorderedElementsAre(tuple)); ExpectHasTopLevelBuffers( points_to_analysis_->GetPointsToSet(tuple).CreateFlattenedSet(), @@ -177,14 +177,14 @@ TEST_F(TuplePointsToAnalysisTest, NestedTuple) { // tuple. auto builder = HloComputation::Builder(TestName()); auto constant1 = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(1.0))); + HloInstruction::CreateConstant(Literal::CreateR0(1.0))); auto constant2 = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(2.0))); + HloInstruction::CreateConstant(Literal::CreateR0(2.0))); auto inner_tuple = builder.AddInstruction( HloInstruction::CreateTuple({constant1, constant2})); auto constant3 = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(3.0))); + HloInstruction::CreateConstant(Literal::CreateR0(3.0))); auto tuple = builder.AddInstruction( HloInstruction::CreateTuple({inner_tuple, constant3})); @@ -205,9 +205,9 @@ TEST_F(TuplePointsToAnalysisTest, NestedTuple) { ExpectHasTopLevelBuffers( points_to_analysis_->GetPointsToSet(inner_tuple).element({}), {inner_tuple}); - EXPECT_ISET( + EXPECT_THAT( points_to_analysis_->GetPointsToSet(inner_tuple).tuple_sources({}), - inner_tuple); + UnorderedElementsAre(inner_tuple)); EXPECT_EQ(5, points_to_analysis_->GetPointsToSet(tuple).size()); EXPECT_FALSE(points_to_analysis_->GetPointsToSet(tuple).IsAmbiguous()); @@ -215,10 +215,10 @@ TEST_F(TuplePointsToAnalysisTest, NestedTuple) { points_to_analysis_->GetPointsToSet(tuple).CreateFlattenedSet(), {constant1, constant2, constant3, inner_tuple, tuple}); - EXPECT_ISET(points_to_analysis_->GetPointsToSet(tuple).tuple_sources({}), - tuple); - EXPECT_ISET(points_to_analysis_->GetPointsToSet(tuple).tuple_sources({0}), - inner_tuple); + EXPECT_THAT(points_to_analysis_->GetPointsToSet(tuple).tuple_sources({}), + UnorderedElementsAre(tuple)); + EXPECT_THAT(points_to_analysis_->GetPointsToSet(tuple).tuple_sources({0}), + UnorderedElementsAre(inner_tuple)); EXPECT_TRUE( points_to_analysis_->GetPointsToSet(tuple).tuple_sources({1}).empty()); @@ -238,14 +238,14 @@ TEST_F(TuplePointsToAnalysisTest, GetTupleElement) { // tuple. auto builder = HloComputation::Builder(TestName()); auto constant1 = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(1.0))); + HloInstruction::CreateConstant(Literal::CreateR0(1.0))); auto constant2 = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(2.0))); + HloInstruction::CreateConstant(Literal::CreateR0(2.0))); auto inner_tuple = builder.AddInstruction( HloInstruction::CreateTuple({constant1, constant2})); auto constant3 = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(3.0))); + HloInstruction::CreateConstant(Literal::CreateR0(3.0))); auto tuple = builder.AddInstruction( HloInstruction::CreateTuple({inner_tuple, constant3})); @@ -262,14 +262,15 @@ TEST_F(TuplePointsToAnalysisTest, GetTupleElement) { {constant1, constant2, inner_tuple}); ExpectHasTopLevelBuffers(points_to_set.element({}), {inner_tuple}); - EXPECT_ISET(points_to_set.tuple_sources({}), inner_tuple); + EXPECT_THAT(points_to_set.tuple_sources({}), + UnorderedElementsAre(inner_tuple)); } TEST_F(TuplePointsToAnalysisTest, DuplicatedElement) { // Create a tuple which contains duplicate elements. auto builder = HloComputation::Builder(TestName()); auto constant = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(1.0))); + HloInstruction::CreateConstant(Literal::CreateR0(1.0))); auto tuple = builder.AddInstruction( HloInstruction::CreateTuple({constant, constant, constant})); @@ -290,9 +291,9 @@ TEST_F(TuplePointsToAnalysisTest, TupleCopy) { // the same. auto builder = HloComputation::Builder(TestName()); auto constant1 = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(1.0))); + HloInstruction::CreateConstant(Literal::CreateR0(1.0))); auto constant2 = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(2.0))); + HloInstruction::CreateConstant(Literal::CreateR0(2.0))); auto tuple = builder.AddInstruction( HloInstruction::CreateTuple({constant1, constant2})); auto copy = builder.AddInstruction( @@ -317,16 +318,16 @@ TEST_F(TuplePointsToAnalysisTest, TupleSelect) { // set containing the union of both sides. auto builder = HloComputation::Builder(TestName()); auto constant1 = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(1.0))); + HloInstruction::CreateConstant(Literal::CreateR0(1.0))); auto constant2 = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(2.0))); + HloInstruction::CreateConstant(Literal::CreateR0(2.0))); auto tuple1 = builder.AddInstruction( HloInstruction::CreateTuple({constant1, constant2})); auto tuple2 = builder.AddInstruction( HloInstruction::CreateTuple({constant2, constant2})); auto pred = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(false))); + HloInstruction::CreateConstant(Literal::CreateR0(false))); auto select = builder.AddInstruction(HloInstruction::CreateTernary( tuple1->shape(), HloOpcode::kSelect, pred, tuple1, tuple2)); @@ -355,7 +356,7 @@ TEST_F(TuplePointsToAnalysisTest, SelectTupleParameters) { auto param1 = builder.AddInstruction( HloInstruction::CreateParameter(1, tuple_shape, "param1")); auto pred = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(false))); + HloInstruction::CreateConstant(Literal::CreateR0(false))); auto select = builder.AddInstruction(HloInstruction::CreateTernary( tuple_shape, HloOpcode::kSelect, pred, param0, param1)); auto copy = builder.AddInstruction( @@ -395,16 +396,16 @@ TEST_F(TuplePointsToAnalysisTest, UnambiguousTupleSelect) { // Select from two identical tuples. The result should not be ambiguous. auto builder = HloComputation::Builder(TestName()); auto constant1 = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(1.0))); + HloInstruction::CreateConstant(Literal::CreateR0(1.0))); auto constant2 = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(2.0))); + HloInstruction::CreateConstant(Literal::CreateR0(2.0))); auto tuple1 = builder.AddInstruction( HloInstruction::CreateTuple({constant1, constant2})); auto tuple2 = builder.AddInstruction( HloInstruction::CreateTuple({constant1, constant2})); auto pred = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(false))); + HloInstruction::CreateConstant(Literal::CreateR0(false))); auto select = builder.AddInstruction(HloInstruction::CreateTernary( tuple1->shape(), HloOpcode::kSelect, pred, tuple1, tuple2)); @@ -426,9 +427,9 @@ TEST_F(TuplePointsToAnalysisTest, NestedTupleSelect) { // the right values. auto builder = HloComputation::Builder(TestName()); auto constant1 = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(1.0))); + HloInstruction::CreateConstant(Literal::CreateR0(1.0))); auto constant2 = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(2.0))); + HloInstruction::CreateConstant(Literal::CreateR0(2.0))); auto inner_tuple1 = builder.AddInstruction( HloInstruction::CreateTuple({constant1, constant2})); auto inner_tuple2 = builder.AddInstruction( @@ -440,7 +441,7 @@ TEST_F(TuplePointsToAnalysisTest, NestedTupleSelect) { builder.AddInstruction(HloInstruction::CreateTuple({inner_tuple2})); auto pred = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(false))); + HloInstruction::CreateConstant(Literal::CreateR0(false))); auto select = builder.AddInstruction(HloInstruction::CreateTernary( tuple1->shape(), HloOpcode::kSelect, pred, tuple1, tuple2)); @@ -460,8 +461,10 @@ TEST_F(TuplePointsToAnalysisTest, NestedTupleSelect) { ExpectHasTopLevelBuffers(points_to_set.element({0, 1}), {constant2}); // Verify tuple sources. - EXPECT_ISET(points_to_set.tuple_sources({}), tuple1, tuple2); - EXPECT_ISET(points_to_set.tuple_sources({0}), inner_tuple1, inner_tuple2); + EXPECT_THAT(points_to_set.tuple_sources({}), + UnorderedElementsAre(tuple1, tuple2)); + EXPECT_THAT(points_to_set.tuple_sources({0}), + UnorderedElementsAre(inner_tuple1, inner_tuple2)); EXPECT_EQ(0, points_to_set.tuple_sources({0, 0}).size()); EXPECT_EQ(0, points_to_set.tuple_sources({0, 1}).size()); } @@ -471,9 +474,9 @@ TEST_F(TuplePointsToAnalysisTest, TupleWithBitcast) { // have the operand of the bitcast in its points-to set. auto builder = HloComputation::Builder(TestName()); auto constant1 = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(1.0))); + HloInstruction::CreateConstant(Literal::CreateR0(1.0))); auto constant2 = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(2.0))); + HloInstruction::CreateConstant(Literal::CreateR0(2.0))); auto bitcast = builder.AddInstruction(HloInstruction::CreateUnary( constant2->shape(), HloOpcode::kBitcast, constant2)); auto tuple = @@ -489,8 +492,8 @@ TEST_F(TuplePointsToAnalysisTest, TupleWithBitcast) { EXPECT_EQ(3, points_to_analysis_->GetPointsToSet(tuple).size()); EXPECT_FALSE(points_to_analysis_->GetPointsToSet(tuple).IsAmbiguous()); - EXPECT_ISET(points_to_analysis_->GetPointsToSet(tuple).tuple_sources({}), - tuple); + EXPECT_THAT(points_to_analysis_->GetPointsToSet(tuple).tuple_sources({}), + UnorderedElementsAre(tuple)); ExpectHasTopLevelBuffers( points_to_analysis_->GetPointsToSet(tuple).CreateFlattenedSet(), @@ -507,10 +510,9 @@ TEST_F(TuplePointsToAnalysisTest, PointsToTupleConstantElements) { // Construct a tuple constant and kCopy it. Verify the points-to set of the // copy correctly correctly points into the nested elements of the constant. auto builder = HloComputation::Builder(TestName()); - auto tuple_constant = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::MakeTuple( - {LiteralUtil::CreateR2({{1.0}, {2.0}}).get(), - LiteralUtil::CreateR1({2.0, 42}).get()}))); + auto tuple_constant = builder.AddInstruction(HloInstruction::CreateConstant( + Literal::MakeTuple({Literal::CreateR2({{1.0}, {2.0}}).get(), + Literal::CreateR1({2.0, 42}).get()}))); auto copy = builder.AddInstruction(HloInstruction::CreateUnary( tuple_constant->shape(), HloOpcode::kCopy, tuple_constant)); @@ -530,9 +532,9 @@ TEST_F(TuplePointsToAnalysisTest, BufferAliases) { // times. Verify buffer alias sets. auto builder = HloComputation::Builder(TestName()); auto constant1 = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(1.0))); + HloInstruction::CreateConstant(Literal::CreateR0(1.0))); auto constant2 = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(2.0))); + HloInstruction::CreateConstant(Literal::CreateR0(2.0))); auto inner_tuple = builder.AddInstruction( HloInstruction::CreateTuple({constant1, constant2})); auto tuple = builder.AddInstruction( @@ -571,7 +573,7 @@ class FusionPointsToAnalysisTest : public TuplePointsToAnalysisTest { auto tuple_element1 = builder.AddInstruction( HloInstruction::CreateGetTupleElement(update_shape, tuple_param0, 1)); auto ones = builder.AddInstruction(HloInstruction::CreateConstant( - LiteralUtil::CreateR1({1.f, 1.f, 1.f, 1.f}))); + Literal::CreateR1({1.f, 1.f, 1.f, 1.f}))); // Create 'update' = Add(GetTupleElement(tuple_param0, 1), ones) auto update = builder.AddInstruction(HloInstruction::CreateBinary( update_shape, HloOpcode::kAdd, tuple_element1, ones)); @@ -582,7 +584,7 @@ class FusionPointsToAnalysisTest : public TuplePointsToAnalysisTest { if (add_additional_gte0_user) { // Create 'slice' as an additional user of 'input'. auto slice = builder.AddInstruction( - HloInstruction::CreateSlice(update_shape, input, {0}, {3})); + HloInstruction::CreateSlice(update_shape, input, {0}, {3}, {1})); // Modify 'update' to take 'slice' output. update = builder.AddInstruction(HloInstruction::CreateBinary( update_shape, HloOpcode::kAdd, update, slice)); @@ -603,9 +605,9 @@ class FusionPointsToAnalysisTest : public TuplePointsToAnalysisTest { .ValueOrDie()); // Get computation root instruction (should be a kFusion). auto* fusion = module_->entry_computation()->root_instruction(); - EXPECT_EQ(HloOpcode::kFusion, fusion->opcode()); + EXPECT_THAT(fusion, op::Fusion(tuple_param0)); // Run points-to analysis (should include fused instructions from 'fusion'). - RunAnalysis(/*include_loop_fusion_instructions=*/true); + RunAnalysis(); // Check points-to set of fusion parameter associated with 'tuple_param0'. auto* fusion_param = GetFusionParameterForOperand(fusion, tuple_param0); @@ -761,3 +763,7 @@ TEST_F(FusionPointsToAnalysisTest, FusionParam0TwoUsers) { } // namespace } // namespace xla + +int main(int argc, char** argv) { + return xla::ParseDebugOptionsFlagsAndRunTests(argc, argv); +} diff --git a/tensorflow/compiler/xla/service/user_computation.cc b/tensorflow/compiler/xla/service/user_computation.cc index a77788e0b63b984328c0ea52ebbb94cb8583e6e3..297bfd93d122d6190efddcc159f6f40894964ebf 100644 --- a/tensorflow/compiler/xla/service/user_computation.cc +++ b/tensorflow/compiler/xla/service/user_computation.cc @@ -17,6 +17,8 @@ limitations under the License. #include #include +#include +#include #include #include "tensorflow/compiler/xla/layout_util.h" @@ -46,6 +48,8 @@ HloOpcode UnaryOperationToHloOpcode(UnaryOperation unop) { return HloOpcode::kAbs; case UNOP_CEIL: return HloOpcode::kCeil; + case UNOP_COS: + return HloOpcode::kCos; case UNOP_EXP: return HloOpcode::kExp; case UNOP_FLOOR: @@ -60,6 +64,8 @@ HloOpcode UnaryOperationToHloOpcode(UnaryOperation unop) { return HloOpcode::kNegate; case UNOP_SIGN: return HloOpcode::kSign; + case UNOP_SIN: + return HloOpcode::kSin; case UNOP_SORT: return HloOpcode::kSort; case UNOP_TANH: @@ -462,6 +468,137 @@ StatusOr UserComputation::AddReduceInstruction( return handle; } +StatusOr +UserComputation::AddBatchNormTrainingInstruction( + const BatchNormTrainingRequest& batch_norm_training_request) { + tensorflow::mutex_lock lock(mutex_); + + TF_ASSIGN_OR_RETURN(const OperationRequest* operand, + LookUpRequest(batch_norm_training_request.operand())); + + TF_ASSIGN_OR_RETURN(const OperationRequest* scale, + LookUpRequest(batch_norm_training_request.scale())); + + TF_ASSIGN_OR_RETURN(const OperationRequest* offset, + LookUpRequest(batch_norm_training_request.offset())); + + ComputationDataHandle handle = CreateComputationDataHandle(); + + OperationRequest& request = + (*session_computation_.mutable_requests())[handle.handle()]; + + TF_ASSIGN_OR_RETURN( + Shape inferred_shape, + ShapeInference::InferBatchNormTrainingShape( + operand->output_shape(), scale->output_shape(), + offset->output_shape(), batch_norm_training_request.feature_index())); + + *request.mutable_output_shape() = inferred_shape; + + *request.mutable_output_handle() = handle; + + *request.mutable_request()->mutable_batch_norm_training_request() = + batch_norm_training_request; + + VLOG(1) << "AddBatchNormTrainingInstruction (" << GetVersionedHandleInternal() + << "), data handle " << handle.handle() << ": " + << batch_norm_training_request.ShortDebugString(); + + return handle; +} + +StatusOr +UserComputation::AddBatchNormInferenceInstruction( + const BatchNormInferenceRequest& batch_norm_inference_request) { + tensorflow::mutex_lock lock(mutex_); + + TF_ASSIGN_OR_RETURN(const OperationRequest* operand, + LookUpRequest(batch_norm_inference_request.operand())); + + TF_ASSIGN_OR_RETURN(const OperationRequest* scale, + LookUpRequest(batch_norm_inference_request.scale())); + + TF_ASSIGN_OR_RETURN(const OperationRequest* offset, + LookUpRequest(batch_norm_inference_request.offset())); + + TF_ASSIGN_OR_RETURN(const OperationRequest* mean, + LookUpRequest(batch_norm_inference_request.mean())); + + TF_ASSIGN_OR_RETURN(const OperationRequest* variance, + LookUpRequest(batch_norm_inference_request.variance())); + + ComputationDataHandle handle = CreateComputationDataHandle(); + + OperationRequest& request = + (*session_computation_.mutable_requests())[handle.handle()]; + + TF_ASSIGN_OR_RETURN(Shape inferred_shape, + ShapeInference::InferBatchNormInferenceShape( + operand->output_shape(), scale->output_shape(), + offset->output_shape(), mean->output_shape(), + variance->output_shape(), + batch_norm_inference_request.feature_index())); + + *request.mutable_output_shape() = inferred_shape; + + *request.mutable_output_handle() = handle; + + *request.mutable_request()->mutable_batch_norm_inference_request() = + batch_norm_inference_request; + + VLOG(1) << "AddBatchNormInferenceInstruction (" + << GetVersionedHandleInternal() << "), data handle " + << handle.handle() << ": " + << batch_norm_inference_request.ShortDebugString(); + + return handle; +} + +StatusOr UserComputation::AddBatchNormGradInstruction( + const BatchNormGradRequest& batch_norm_grad_request) { + tensorflow::mutex_lock lock(mutex_); + + TF_ASSIGN_OR_RETURN(const OperationRequest* operand, + LookUpRequest(batch_norm_grad_request.operand())); + + TF_ASSIGN_OR_RETURN(const OperationRequest* scale, + LookUpRequest(batch_norm_grad_request.scale())); + + TF_ASSIGN_OR_RETURN(const OperationRequest* mean, + LookUpRequest(batch_norm_grad_request.mean())); + + TF_ASSIGN_OR_RETURN(const OperationRequest* variance, + LookUpRequest(batch_norm_grad_request.variance())); + + TF_ASSIGN_OR_RETURN(const OperationRequest* grad_output, + LookUpRequest(batch_norm_grad_request.grad_output())); + + ComputationDataHandle handle = CreateComputationDataHandle(); + + OperationRequest& request = + (*session_computation_.mutable_requests())[handle.handle()]; + + TF_ASSIGN_OR_RETURN( + Shape inferred_shape, + ShapeInference::InferBatchNormGradShape( + operand->output_shape(), scale->output_shape(), mean->output_shape(), + variance->output_shape(), grad_output->output_shape(), + batch_norm_grad_request.feature_index())); + + *request.mutable_output_shape() = inferred_shape; + + *request.mutable_output_handle() = handle; + + *request.mutable_request()->mutable_batch_norm_grad_request() = + batch_norm_grad_request; + + VLOG(1) << "AddBatchNormGradInstruction (" << GetVersionedHandleInternal() + << "), data handle " << handle.handle() << ": " + << batch_norm_grad_request.ShortDebugString(); + + return handle; +} + StatusOr UserComputation::AddReduceWindowInstruction( const ReduceWindowRequest& reduce_window_request, const UserComputation& to_apply_computation) { @@ -700,7 +837,8 @@ StatusOr UserComputation::AddSliceInstruction( Shape new_shape, ShapeInference::InferSliceShape( operand->output_shape(), AsInt64Slice(slice_request.start_indices()), - AsInt64Slice(slice_request.limit_indices()))); + AsInt64Slice(slice_request.limit_indices()), + AsInt64Slice(slice_request.strides()))); ComputationDataHandle handle = CreateComputationDataHandle(); @@ -837,6 +975,34 @@ StatusOr UserComputation::AddConvertInstruction( return handle; } +StatusOr UserComputation::AddReducePrecisionInstruction( + const ReducePrecisionRequest& reduce_precision_request) { + tensorflow::mutex_lock lock(mutex_); + + TF_ASSIGN_OR_RETURN(const OperationRequest* operand, + LookUpRequest(reduce_precision_request.operand())); + + TF_ASSIGN_OR_RETURN( + Shape new_shape, + ShapeInference::InferReducePrecisionShape( + operand->output_shape(), reduce_precision_request.exponent_bits(), + reduce_precision_request.mantissa_bits())); + + ComputationDataHandle handle = CreateComputationDataHandle(); + + OperationRequest& request = + (*session_computation_.mutable_requests())[handle.handle()]; + *request.mutable_output_handle() = handle; + *request.mutable_output_shape() = new_shape; + *request.mutable_request()->mutable_reduce_precision_request() = + reduce_precision_request; + + VLOG(1) << "AddReducePrecisionInstruction (" << GetVersionedHandleInternal() + << "), data handle " << handle.handle() << ": " + << reduce_precision_request.ShortDebugString(); + return handle; +} + StatusOr UserComputation::AddConvolveInstruction( const ConvolveRequest& convolve_request) { tensorflow::mutex_lock lock(mutex_); @@ -893,9 +1059,6 @@ StatusOr UserComputation::AddInfeedInstruction( tensorflow::mutex_lock lock(mutex_); const Shape& shape = infeed_request.shape(); - if (ShapeUtil::IsNestedTuple(shape)) { - return InvalidArgument("Infeed does not support nested tuple shapes"); - } if (!LayoutUtil::HasLayout(shape)) { return InvalidArgument("Given shape to Infeed must have a layout"); } @@ -919,9 +1082,6 @@ Status UserComputation::AddOutfeedInstruction( tensorflow::mutex_lock lock(mutex_); const Shape& shape = outfeed_request.shape(); - if (ShapeUtil::IsNestedTuple(shape)) { - return InvalidArgument("Outfeed does not support nested tuple shapes"); - } if (!LayoutUtil::HasLayout(shape)) { return InvalidArgument("Given shape to Outfeed must have a layout"); } @@ -1510,6 +1670,7 @@ void ConstantVisitor(const SessionComputation& session_computation, is_constant); // TODO(b/32495713): We aren't checking the condition and body // computations themselves. + *is_constant = false; break; } @@ -1551,6 +1712,55 @@ void ConstantVisitor(const SessionComputation& session_computation, break; } + case OpRequest::kBatchNormTrainingRequest: { + const BatchNormTrainingRequest& batch_norm_training_request = + request.request().batch_norm_training_request(); + ConstantVisitor(session_computation, + batch_norm_training_request.operand(), visited, + is_constant); + ConstantVisitor(session_computation, batch_norm_training_request.scale(), + visited, is_constant); + ConstantVisitor(session_computation, batch_norm_training_request.offset(), + visited, is_constant); + break; + } + + case OpRequest::kBatchNormInferenceRequest: { + const BatchNormInferenceRequest& batch_norm_inference_request = + request.request().batch_norm_inference_request(); + ConstantVisitor(session_computation, + batch_norm_inference_request.operand(), visited, + is_constant); + ConstantVisitor(session_computation, batch_norm_inference_request.scale(), + visited, is_constant); + ConstantVisitor(session_computation, + batch_norm_inference_request.offset(), visited, + is_constant); + ConstantVisitor(session_computation, batch_norm_inference_request.mean(), + visited, is_constant); + ConstantVisitor(session_computation, + batch_norm_inference_request.variance(), visited, + is_constant); + break; + } + + case OpRequest::kBatchNormGradRequest: { + const BatchNormGradRequest& batch_norm_grad_request = + request.request().batch_norm_grad_request(); + ConstantVisitor(session_computation, batch_norm_grad_request.operand(), + visited, is_constant); + ConstantVisitor(session_computation, batch_norm_grad_request.scale(), + visited, is_constant); + ConstantVisitor(session_computation, batch_norm_grad_request.mean(), + visited, is_constant); + ConstantVisitor(session_computation, batch_norm_grad_request.variance(), + visited, is_constant); + ConstantVisitor(session_computation, + batch_norm_grad_request.grad_output(), visited, + is_constant); + break; + } + case OpRequest::kBinaryOpRequest: { const BinaryOpRequest& binary_op_request = request.request().binary_op_request(); @@ -1819,7 +2029,6 @@ Status UserComputation::CheckParametersAreContiguous( } } - auto program_shape = MakeUnique(); for (int64 i = 0; i < parameter_requests.size(); ++i) { auto it = parameter_requests.find(i); if (it == parameter_requests.end()) { @@ -1845,32 +2054,43 @@ class ComputationLowerer { const SessionComputation& session_computation, VersionedComputationHandle::Version version, UserComputation::HloComputationResolver hlo_resolver, + const DebugOptions& debug_options, bool include_unreachable_instructions) { ComputationLowerer lowerer(computation_name, session_computation, version, - std::move(hlo_resolver)); - return lowerer.Lower(include_unreachable_instructions); + std::move(hlo_resolver), debug_options, + include_unreachable_instructions); + return lowerer.Lower(); } private: ComputationLowerer(const string& computation_name, const SessionComputation& session_computation, VersionedComputationHandle::Version version, - UserComputation::HloComputationResolver hlo_resolver) + UserComputation::HloComputationResolver hlo_resolver, + const DebugOptions& debug_options, + bool include_unreachable_instructions) : hlo_builder_(computation_name), session_computation_(session_computation), version_(version), - hlo_resolver_(std::move(hlo_resolver)) {} + hlo_resolver_(std::move(hlo_resolver)), + debug_options_(debug_options), + include_unreachable_instructions_(include_unreachable_instructions) {} // Build an HLO computation from the SessionComputation at the given // version. - StatusOr> Lower( - bool include_unreachable_instructions); + StatusOr> Lower(); private: + // Traverses the computation 'root' using a DFS, calling 'visit' in postorder. + void TraversePostorder( + const ComputationDataHandle& root, + std::unordered_map* visited, + const std::function& visit); + // DFS visitor of the UserComputation operations which lowers the operations // to HLO instructions. - HloInstruction* Visit(const ComputationDataHandle& handle, - std::map* visited); + void Visit(const ComputationDataHandle& handle, + std::unordered_map* instructions); // Resolves a ComputationHandle and Version to a previously lowered // HloComputation using the hlo_resolver_ function. @@ -1878,32 +2098,358 @@ class ComputationLowerer { const ComputationHandle& handle, VersionedComputationHandle::Version version); + // This function takes an input value which is being implicitly broadcast into + // an output shape and figures out the right kBroadcast instruction(s) + // necessary to replicate the implicit broadcast semantics explicitly. + HloInstruction* ImplicitBroadcastToExplicitBroadcast( + HloInstruction* operand, const Shape& output_shape); + HloComputation::Builder hlo_builder_; const SessionComputation& session_computation_; const VersionedComputationHandle::Version version_; const UserComputation::HloComputationResolver hlo_resolver_; + const DebugOptions& debug_options_; + const bool include_unreachable_instructions_; }; -StatusOr> ComputationLowerer::Lower( - bool include_unreachable_instructions) { +// Calls 'apply' on each operand of 'request'. +static void ForEachOperand( + const OperationRequest& request, + const std::function& apply) { + switch (request.request().op_case()) { + case OpRequest::kRngRequest: { + const RngRequest& rng_request = request.request().rng_request(); + for (const ComputationDataHandle& param : rng_request.parameter()) { + apply(param); + } + break; + } + + case OpRequest::kConstantRequest: + break; + case OpRequest::kGetTupleElementRequest: { + const GetTupleElementRequest& get_tuple_element_request = + request.request().get_tuple_element_request(); + apply(get_tuple_element_request.operand()); + break; + } + + case OpRequest::kSliceRequest: { + const SliceRequest& slice_request = request.request().slice_request(); + apply(slice_request.operand()); + break; + } + + case OpRequest::kDynamicSliceRequest: { + const DynamicSliceRequest& dynamic_slice_request = + request.request().dynamic_slice_request(); + apply(dynamic_slice_request.operand()); + apply(dynamic_slice_request.start_indices()); + break; + } + + case OpRequest::kDynamicUpdateSliceRequest: { + const DynamicUpdateSliceRequest& dynamic_update_slice_request = + request.request().dynamic_update_slice_request(); + apply(dynamic_update_slice_request.operand()); + apply(dynamic_update_slice_request.update()); + apply(dynamic_update_slice_request.start_indices()); + break; + } + + case OpRequest::kConcatenateRequest: { + const ConcatenateRequest& concatenate_request = + request.request().concatenate_request(); + for (const ComputationDataHandle& handle : + concatenate_request.operands()) { + apply(handle); + } + break; + } + + case OpRequest::kConvolveRequest: { + const ConvolveRequest& convolve_request = + request.request().convolve_request(); + apply(convolve_request.lhs()); + apply(convolve_request.rhs()); + break; + } + + case OpRequest::kBatchNormTrainingRequest: { + const BatchNormTrainingRequest& batch_norm_training_request = + request.request().batch_norm_training_request(); + + apply(batch_norm_training_request.operand()); + apply(batch_norm_training_request.scale()); + apply(batch_norm_training_request.offset()); + break; + } + + case OpRequest::kBatchNormInferenceRequest: { + const BatchNormInferenceRequest& batch_norm_inference_request = + request.request().batch_norm_inference_request(); + + apply(batch_norm_inference_request.operand()); + apply(batch_norm_inference_request.scale()); + apply(batch_norm_inference_request.offset()); + apply(batch_norm_inference_request.mean()); + apply(batch_norm_inference_request.variance()); + break; + } + + case OpRequest::kBatchNormGradRequest: { + const BatchNormGradRequest& batch_norm_grad_request = + request.request().batch_norm_grad_request(); + + apply(batch_norm_grad_request.operand()); + apply(batch_norm_grad_request.scale()); + apply(batch_norm_grad_request.mean()); + apply(batch_norm_grad_request.variance()); + apply(batch_norm_grad_request.grad_output()); + break; + } + + case OpRequest::kCrossReplicaSumRequest: { + const CrossReplicaSumRequest& cross_replica_sum_request = + request.request().cross_replica_sum_request(); + apply(cross_replica_sum_request.operand()); + break; + } + + case OpRequest::kInfeedRequest: + break; + + case OpRequest::kOutfeedRequest: { + const OutfeedRequest& outfeed_request = + request.request().outfeed_request(); + apply(outfeed_request.operand()); + break; + } + + case OpRequest::kMapRequest: { + const MapRequest& map_request = request.request().map_request(); + for (const ComputationDataHandle& handle : map_request.operands()) { + apply(handle); + } + break; + } + + case OpRequest::kReduceRequest: { + const ReduceRequest& reduce_request = request.request().reduce_request(); + apply(reduce_request.operand()); + apply(reduce_request.init_value()); + break; + } + + case OpRequest::kReduceWindowRequest: { + const ReduceWindowRequest& reduce_window_request = + request.request().reduce_window_request(); + apply(reduce_window_request.operand()); + apply(reduce_window_request.init_value()); + break; + } + + case OpRequest::kSelectAndScatterRequest: { + const SelectAndScatterRequest& select_and_scatter_request = + request.request().select_and_scatter_request(); + apply(select_and_scatter_request.operand()); + apply(select_and_scatter_request.source()); + apply(select_and_scatter_request.init_value()); + + break; + } + + case OpRequest::kBroadcastRequest: { + const BroadcastRequest& broadcast_request = + request.request().broadcast_request(); + apply(broadcast_request.operand()); + break; + } + + case OpRequest::kReshapeRequest: { + const ReshapeRequest& reshape_request = + request.request().reshape_request(); + apply(reshape_request.operand()); + break; + } + + case OpRequest::kTransposeRequest: { + const TransposeRequest& transpose_request = + request.request().transpose_request(); + apply(transpose_request.operand()); + break; + } + + case OpRequest::kReverseRequest: { + const ReverseRequest& reverse_request = + request.request().reverse_request(); + apply(reverse_request.operand()); + break; + } + + case OpRequest::kPadRequest: { + const PadRequest& pad_request = request.request().pad_request(); + apply(pad_request.operand()); + apply(pad_request.padding_value()); + break; + } + + case OpRequest::kRecvRequest: + case OpRequest::kParameterRequest: + break; + + case OpRequest::kConvertRequest: { + const ConvertRequest& convert_request = + request.request().convert_request(); + apply(convert_request.operand()); + break; + } + + case OpRequest::kWhileRequest: { + const WhileRequest& while_request = request.request().while_request(); + apply(while_request.init()); + break; + } + + case OpRequest::kTernaryOpRequest: { + const TernaryOpRequest& ternary_op_request = + request.request().ternary_op_request(); + apply(ternary_op_request.lhs()); + apply(ternary_op_request.rhs()); + apply(ternary_op_request.ehs()); + break; + } + + case OpRequest::kVariadicOpRequest: { + const VariadicOpRequest& variadic_op_request = + request.request().variadic_op_request(); + for (const ComputationDataHandle& handle : + variadic_op_request.operands()) { + apply(handle); + } + break; + } + + case OpRequest::kCallRequest: { + const CallRequest& call_request = request.request().call_request(); + for (const ComputationDataHandle& handle : call_request.operands()) { + apply(handle); + } + break; + } + + case OpRequest::kCustomCallRequest: { + const CustomCallRequest& cc_request = + request.request().custom_call_request(); + for (const ComputationDataHandle& operand : cc_request.operands()) { + apply(operand); + } + break; + } + + case OpRequest::kUnaryOpRequest: { + const UnaryOpRequest& unary_op_request = + request.request().unary_op_request(); + apply(unary_op_request.operand()); + break; + } + + case OpRequest::kBinaryOpRequest: { + const BinaryOpRequest& binary_op_request = + request.request().binary_op_request(); + apply(binary_op_request.rhs()); + apply(binary_op_request.lhs()); + break; + } + + case OpRequest::kReducePrecisionRequest: { + const ReducePrecisionRequest& reduce_precision_request = + request.request().reduce_precision_request(); + apply(reduce_precision_request.operand()); + break; + } + + case OpRequest::kTraceRequest: { + const TraceRequest& trace_request = request.request().trace_request(); + apply(trace_request.operand()); + break; + } + + case OpRequest::kSendRequest: { + const SendRequest& send_request = request.request().send_request(); + apply(send_request.operand()); + break; + } + + case OpRequest::OP_NOT_SET: + LOG(FATAL) << "OperationRequest doesn't contain a request"; + + default: + LOG(FATAL) << "Unexpected request type: " << request.request().op_case(); + } +} + +void ComputationLowerer::TraversePostorder( + const ComputationDataHandle& root, + std::unordered_map* visited, + const std::function& visit) { + // Stack containing {handle, enter} pairs. The 'enter' value describes whether + // we are entering or leaving 'handle'. + std::stack> work; + work.push({root, true}); + while (!work.empty()) { + ComputationDataHandle handle; + bool enter; + std::tie(handle, enter) = work.top(); + work.pop(); + + if (enter) { + // We are entering 'handle'. The first time we enter 'handle', we add it + // to 'visited' with a nullptr value. If 'handle' is already in 'visited', + // we do not visit it again. This algorithm only uses the presence of + // a handle in 'visited', but we use a map so we can use the same data + // structure to store the HloInstruction outputs. + if (visited->emplace(handle.handle(), nullptr).second) { + const OperationRequest& request = + session_computation_.requests().at(handle.handle()); + // Push the corresponding 'leave' action onto the stack, followed by + // the operands. + work.push({handle, false}); + ForEachOperand(request, [&work](const ComputationDataHandle& child) { + work.push({child, true}); + }); + } + } else { + // We are leaving 'handle'. We have visited the operands of 'handle', and + // now can visit the 'handle' itself. + visit(handle); + } + } +} + +StatusOr> ComputationLowerer::Lower() { // Map from ComputationDataHandle to HLO instruction. Serves as a record of // which operations have been visited as well as a cache for looking up // ComputationDataHandles as HloInstructions. - std::map visited; + std::unordered_map instructions; TF_ASSIGN_OR_RETURN(const OperationRequest* root_request, GetRoot(version_, session_computation_)); - HloInstruction* hlo_root = Visit(root_request->output_handle(), &visited); - if (include_unreachable_instructions) { + auto visit = [&](const ComputationDataHandle& handle) { + Visit(handle, &instructions); + }; + TraversePostorder(root_request->output_handle(), &instructions, visit); + HloInstruction* hlo_root = + instructions.at(root_request->output_handle().handle()); + + if (include_unreachable_instructions_) { // Iterate through all computation data handles, and visit any unvisited // operations. for (int64 request_num = 1; request_num <= version_; ++request_num) { TF_ASSIGN_OR_RETURN(const OperationRequest* request, LookUpRequest(request_num, session_computation_)); - if (visited.count(request->output_handle().handle()) == 0) { - Visit(request->output_handle(), &visited); - } + TraversePostorder(request->output_handle(), &instructions, visit); } } @@ -1917,25 +2463,62 @@ HloComputation* ComputationLowerer::ResolveComputation( return hlo_resolver_(checked_handle); } -HloInstruction* ComputationLowerer::Visit( - const ComputationDataHandle& handle, - std::map* visited) { - CHECK_LE(handle.handle(), version_); - if (visited->count(handle.handle()) != 0) { - return (*visited)[handle.handle()]; +HloInstruction* ComputationLowerer::ImplicitBroadcastToExplicitBroadcast( + HloInstruction* operand, const Shape& output_shape) { + CHECK(ShapeUtil::IsScalar(operand->shape()) || + ShapeUtil::Rank(operand->shape()) == ShapeUtil::Rank(output_shape)); + Shape broadcast_shape = ShapeUtil::MakeShape( + operand->shape().element_type(), AsInt64Slice(output_shape.dimensions())); + // Do explicit broadcast for scalar. + if (ShapeUtil::IsScalar(operand->shape())) { + return hlo_builder_.AddInstruction(HloInstruction::CreateBroadcast( + broadcast_shape, operand, AsInt64Slice(broadcast_shape.dimensions()))); } + // Do explicit broadcast for degenerate broadcast. + std::vector broadcast_dimensions; + std::vector reshaped_dimensions; + for (int i = 0; i < ShapeUtil::Rank(operand->shape()); i++) { + if (operand->shape().dimensions(i) == output_shape.dimensions(i)) { + broadcast_dimensions.push_back(i); + reshaped_dimensions.push_back(operand->shape().dimensions(i)); + } + } + // Eliminate the size one dimensions. + HloInstruction* reshaped_operand = + hlo_builder_.AddInstruction(HloInstruction::CreateReshape( + ShapeUtil::MakeShape(operand->shape().element_type(), + reshaped_dimensions), + operand)); + // Broadcast 'reshape' up to the larger size. + return hlo_builder_.AddInstruction(HloInstruction::CreateBroadcast( + broadcast_shape, reshaped_operand, broadcast_dimensions)); +} +void ComputationLowerer::Visit( + const ComputationDataHandle& handle, + std::unordered_map* instructions) { + CHECK_LE(handle.handle(), version_); + CHECK(instructions->at(handle.handle()) == nullptr); const OperationRequest& request = session_computation_.requests().at(handle.handle()); + auto add_instruction = [&](std::unique_ptr instruction) { + HloInstruction* hlo_instruction = + hlo_builder_.AddInstruction(std::move(instruction)); + hlo_instruction->set_metadata(request.request().metadata()); + return hlo_instruction; + }; + auto lookup_instruction = [&](const ComputationDataHandle& handle) { + return instructions->at(handle.handle()); + }; HloInstruction* hlo_instruction; switch (request.request().op_case()) { case OpRequest::kRngRequest: { const RngRequest& rng_request = request.request().rng_request(); std::vector parameters; for (const ComputationDataHandle& param : rng_request.parameter()) { - parameters.push_back(Visit(param, visited)); + parameters.push_back(lookup_instruction(param)); } - hlo_instruction = hlo_builder_.AddInstruction(HloInstruction::CreateRng( + hlo_instruction = add_instruction(HloInstruction::CreateRng( request.output_shape(), rng_request.distribution(), parameters)); break; } @@ -1943,9 +2526,8 @@ HloInstruction* ComputationLowerer::Visit( case OpRequest::kConstantRequest: { const ConstantRequest& constant_request = request.request().constant_request(); - hlo_instruction = - hlo_builder_.AddInstruction(HloInstruction::CreateConstant( - LiteralUtil::CloneToUnique(constant_request.literal()))); + hlo_instruction = add_instruction(HloInstruction::CreateConstant( + Literal(constant_request.literal()).CloneToUnique())); break; } @@ -1953,35 +2535,34 @@ HloInstruction* ComputationLowerer::Visit( const GetTupleElementRequest& get_tuple_element_request = request.request().get_tuple_element_request(); HloInstruction* operand = - Visit(get_tuple_element_request.operand(), visited); - hlo_instruction = - hlo_builder_.AddInstruction(HloInstruction::CreateGetTupleElement( - request.output_shape(), operand, - get_tuple_element_request.index())); + lookup_instruction(get_tuple_element_request.operand()); + hlo_instruction = add_instruction(HloInstruction::CreateGetTupleElement( + request.output_shape(), operand, get_tuple_element_request.index())); break; } case OpRequest::kSliceRequest: { const SliceRequest& slice_request = request.request().slice_request(); - HloInstruction* operand = Visit(slice_request.operand(), visited); - hlo_instruction = hlo_builder_.AddInstruction(HloInstruction::CreateSlice( + HloInstruction* operand = lookup_instruction(slice_request.operand()); + hlo_instruction = add_instruction(HloInstruction::CreateSlice( request.output_shape(), operand, AsInt64Slice(slice_request.start_indices()), - AsInt64Slice(slice_request.limit_indices()))); + AsInt64Slice(slice_request.limit_indices()), + AsInt64Slice(slice_request.strides()))); break; } case OpRequest::kDynamicSliceRequest: { const DynamicSliceRequest& dynamic_slice_request = request.request().dynamic_slice_request(); - HloInstruction* operand = Visit(dynamic_slice_request.operand(), visited); + HloInstruction* operand = + lookup_instruction(dynamic_slice_request.operand()); HloInstruction* start_indices = - Visit(dynamic_slice_request.start_indices(), visited); + lookup_instruction(dynamic_slice_request.start_indices()); - hlo_instruction = - hlo_builder_.AddInstruction(HloInstruction::CreateDynamicSlice( - request.output_shape(), operand, start_indices, - AsInt64Slice(dynamic_slice_request.slice_sizes()))); + hlo_instruction = add_instruction(HloInstruction::CreateDynamicSlice( + request.output_shape(), operand, start_indices, + AsInt64Slice(dynamic_slice_request.slice_sizes()))); break; } @@ -1989,13 +2570,13 @@ HloInstruction* ComputationLowerer::Visit( const DynamicUpdateSliceRequest& dynamic_update_slice_request = request.request().dynamic_update_slice_request(); HloInstruction* operand = - Visit(dynamic_update_slice_request.operand(), visited); + lookup_instruction(dynamic_update_slice_request.operand()); HloInstruction* update = - Visit(dynamic_update_slice_request.update(), visited); + lookup_instruction(dynamic_update_slice_request.update()); HloInstruction* start_indices = - Visit(dynamic_update_slice_request.start_indices(), visited); + lookup_instruction(dynamic_update_slice_request.start_indices()); hlo_instruction = - hlo_builder_.AddInstruction(HloInstruction::CreateDynamicUpdateSlice( + add_instruction(HloInstruction::CreateDynamicUpdateSlice( request.output_shape(), operand, update, start_indices)); break; } @@ -2006,24 +2587,22 @@ HloInstruction* ComputationLowerer::Visit( std::vector operands; for (const ComputationDataHandle& handle : concatenate_request.operands()) { - HloInstruction* operand = Visit(handle, visited); + HloInstruction* operand = lookup_instruction(handle); operands.push_back(operand); } - hlo_instruction = hlo_builder_.AddInstruction( - HloInstruction::CreateConcatenate(request.output_shape(), operands, - concatenate_request.dimension())); + hlo_instruction = add_instruction(HloInstruction::CreateConcatenate( + request.output_shape(), operands, concatenate_request.dimension())); break; } case OpRequest::kConvolveRequest: { const ConvolveRequest& convolve_request = request.request().convolve_request(); - HloInstruction* lhs = Visit(convolve_request.lhs(), visited); - HloInstruction* rhs = Visit(convolve_request.rhs(), visited); - hlo_instruction = - hlo_builder_.AddInstruction(HloInstruction::CreateConvolve( - request.output_shape(), lhs, rhs, convolve_request.window(), - convolve_request.dimension_numbers())); + HloInstruction* lhs = lookup_instruction(convolve_request.lhs()); + HloInstruction* rhs = lookup_instruction(convolve_request.rhs()); + hlo_instruction = add_instruction(HloInstruction::CreateConvolve( + request.output_shape(), lhs, rhs, convolve_request.window(), + convolve_request.dimension_numbers())); break; } @@ -2031,28 +2610,25 @@ HloInstruction* ComputationLowerer::Visit( const CrossReplicaSumRequest& cross_replica_sum_request = request.request().cross_replica_sum_request(); HloInstruction* operand = - Visit(cross_replica_sum_request.operand(), visited); - hlo_instruction = - hlo_builder_.AddInstruction(HloInstruction::CreateCrossReplicaSum( - request.output_shape(), operand)); + lookup_instruction(cross_replica_sum_request.operand()); + hlo_instruction = add_instruction(HloInstruction::CreateCrossReplicaSum( + request.output_shape(), operand)); break; } case OpRequest::kInfeedRequest: { const InfeedRequest& infeed_request = request.request().infeed_request(); - hlo_instruction = - hlo_builder_.AddInstruction(HloInstruction::CreateInfeed( - request.output_shape(), infeed_request.config())); + hlo_instruction = add_instruction(HloInstruction::CreateInfeed( + request.output_shape(), infeed_request.config())); break; } case OpRequest::kOutfeedRequest: { const OutfeedRequest& outfeed_request = request.request().outfeed_request(); - HloInstruction* operand = Visit(outfeed_request.operand(), visited); - hlo_instruction = hlo_builder_.AddInstruction( - HloInstruction::CreateOutfeed(outfeed_request.shape(), operand, - outfeed_request.outfeed_config())); + HloInstruction* operand = lookup_instruction(outfeed_request.operand()); + hlo_instruction = add_instruction(HloInstruction::CreateOutfeed( + outfeed_request.shape(), operand, outfeed_request.outfeed_config())); break; } @@ -2060,7 +2636,7 @@ HloInstruction* ComputationLowerer::Visit( const MapRequest& map_request = request.request().map_request(); std::vector operands; for (const ComputationDataHandle& handle : map_request.operands()) { - HloInstruction* operand = Visit(handle, visited); + HloInstruction* operand = lookup_instruction(handle); operands.push_back(operand); } CHECK_EQ(1, request.embedded_computation_versions_size()); @@ -2068,42 +2644,42 @@ HloInstruction* ComputationLowerer::Visit( request.embedded_computation_versions(0); HloComputation* map_computation = ResolveComputation(map_request.to_apply(), map_version); - hlo_instruction = hlo_builder_.AddInstruction(HloInstruction::CreateMap( + hlo_instruction = add_instruction(HloInstruction::CreateMap( request.output_shape(), operands, map_computation)); break; } case OpRequest::kReduceRequest: { const ReduceRequest& reduce_request = request.request().reduce_request(); - HloInstruction* operand = Visit(reduce_request.operand(), visited); - HloInstruction* init_value = Visit(reduce_request.init_value(), visited); + HloInstruction* operand = lookup_instruction(reduce_request.operand()); + HloInstruction* init_value = + lookup_instruction(reduce_request.init_value()); CHECK_EQ(1, request.embedded_computation_versions_size()); VersionedComputationHandle::Version reduce_version = request.embedded_computation_versions(0); HloComputation* reduce_computation = ResolveComputation(reduce_request.to_apply(), reduce_version); - hlo_instruction = - hlo_builder_.AddInstruction(HloInstruction::CreateReduce( - request.output_shape(), operand, init_value, - AsInt64Slice(reduce_request.dimensions()), reduce_computation)); + hlo_instruction = add_instruction(HloInstruction::CreateReduce( + request.output_shape(), operand, init_value, + AsInt64Slice(reduce_request.dimensions()), reduce_computation)); break; } case OpRequest::kReduceWindowRequest: { const ReduceWindowRequest& reduce_window_request = request.request().reduce_window_request(); - HloInstruction* operand = Visit(reduce_window_request.operand(), visited); + HloInstruction* operand = + lookup_instruction(reduce_window_request.operand()); HloInstruction* init_value = - Visit(reduce_window_request.init_value(), visited); + lookup_instruction(reduce_window_request.init_value()); CHECK_EQ(1, request.embedded_computation_versions_size()); VersionedComputationHandle::Version reduce_window_version = request.embedded_computation_versions(0); HloComputation* reduce_window_computation = ResolveComputation( reduce_window_request.to_apply(), reduce_window_version); - hlo_instruction = - hlo_builder_.AddInstruction(HloInstruction::CreateReduceWindow( - request.output_shape(), operand, init_value, - reduce_window_request.window(), reduce_window_computation)); + hlo_instruction = add_instruction(HloInstruction::CreateReduceWindow( + request.output_shape(), operand, init_value, + reduce_window_request.window(), reduce_window_computation)); break; } @@ -2111,11 +2687,11 @@ HloInstruction* ComputationLowerer::Visit( const SelectAndScatterRequest& select_and_scatter_request = request.request().select_and_scatter_request(); HloInstruction* operand = - Visit(select_and_scatter_request.operand(), visited); + lookup_instruction(select_and_scatter_request.operand()); HloInstruction* source = - Visit(select_and_scatter_request.source(), visited); + lookup_instruction(select_and_scatter_request.source()); HloInstruction* init_value = - Visit(select_and_scatter_request.init_value(), visited); + lookup_instruction(select_and_scatter_request.init_value()); CHECK_EQ(2, request.embedded_computation_versions_size()); VersionedComputationHandle::Version select_version = request.embedded_computation_versions(0); @@ -2125,18 +2701,77 @@ HloInstruction* ComputationLowerer::Visit( select_and_scatter_request.select(), select_version); HloComputation* scatter_computation = ResolveComputation( select_and_scatter_request.scatter(), scatter_version); + hlo_instruction = add_instruction(HloInstruction::CreateSelectAndScatter( + request.output_shape(), operand, select_computation, + select_and_scatter_request.window(), source, init_value, + scatter_computation)); + break; + } + + case OpRequest::kBatchNormTrainingRequest: { + const BatchNormTrainingRequest& batch_norm_training_request = + request.request().batch_norm_training_request(); + HloInstruction* operand = + lookup_instruction(batch_norm_training_request.operand()); + HloInstruction* scale = + lookup_instruction(batch_norm_training_request.scale()); + HloInstruction* offset = + lookup_instruction(batch_norm_training_request.offset()); + + hlo_instruction = add_instruction(HloInstruction::CreateBatchNormTraining( + request.output_shape(), operand, scale, offset, + batch_norm_training_request.epsilon(), + batch_norm_training_request.feature_index())); + break; + } + + case OpRequest::kBatchNormInferenceRequest: { + const BatchNormInferenceRequest& batch_norm_inference_request = + request.request().batch_norm_inference_request(); + HloInstruction* operand = + lookup_instruction(batch_norm_inference_request.operand()); + HloInstruction* scale = + lookup_instruction(batch_norm_inference_request.scale()); + HloInstruction* offset = + lookup_instruction(batch_norm_inference_request.offset()); + HloInstruction* mean = + lookup_instruction(batch_norm_inference_request.mean()); + HloInstruction* variance = + lookup_instruction(batch_norm_inference_request.variance()); + hlo_instruction = - hlo_builder_.AddInstruction(HloInstruction::CreateSelectAndScatter( - request.output_shape(), operand, select_computation, - select_and_scatter_request.window(), source, init_value, - scatter_computation)); + add_instruction(HloInstruction::CreateBatchNormInference( + request.output_shape(), operand, scale, offset, mean, variance, + batch_norm_inference_request.epsilon(), + batch_norm_inference_request.feature_index())); + break; + } + + case OpRequest::kBatchNormGradRequest: { + const BatchNormGradRequest& batch_norm_grad_request = + request.request().batch_norm_grad_request(); + + HloInstruction* operand = + lookup_instruction(batch_norm_grad_request.operand()); + HloInstruction* scale = + lookup_instruction(batch_norm_grad_request.scale()); + HloInstruction* mean = lookup_instruction(batch_norm_grad_request.mean()); + HloInstruction* variance = + lookup_instruction(batch_norm_grad_request.variance()); + HloInstruction* grad_output = + lookup_instruction(batch_norm_grad_request.grad_output()); + + hlo_instruction = add_instruction(HloInstruction::CreateBatchNormGrad( + request.output_shape(), operand, scale, mean, variance, grad_output, + batch_norm_grad_request.epsilon(), + batch_norm_grad_request.feature_index())); break; } case OpRequest::kBroadcastRequest: { const BroadcastRequest& broadcast_request = request.request().broadcast_request(); - HloInstruction* operand = Visit(broadcast_request.operand(), visited); + HloInstruction* operand = lookup_instruction(broadcast_request.operand()); std::vector broadcast_dimensions; // The client-level broadcast instruction just appends dimensions on the // left (adds lowest numbered dimensions). The HLO broadcast op is more @@ -2145,33 +2780,32 @@ HloInstruction* ComputationLowerer::Visit( // to append dimensions on the left the broadcast_dimensions should just // be the n highest dimension numbers of the output shape where n is // the number of input dimensions. + broadcast_dimensions.reserve(ShapeUtil::Rank(operand->shape())); for (int i = 0; i < ShapeUtil::Rank(operand->shape()); ++i) { broadcast_dimensions.push_back(i + ShapeUtil::Rank(request.output_shape()) - ShapeUtil::Rank(operand->shape())); } - hlo_instruction = - hlo_builder_.AddInstruction(HloInstruction::CreateBroadcast( - request.output_shape(), operand, broadcast_dimensions)); + hlo_instruction = add_instruction(HloInstruction::CreateBroadcast( + request.output_shape(), operand, broadcast_dimensions)); break; } case OpRequest::kReshapeRequest: { const ReshapeRequest& reshape_request = request.request().reshape_request(); - HloInstruction* operand = Visit(reshape_request.operand(), visited); + HloInstruction* operand = lookup_instruction(reshape_request.operand()); HloInstruction* transposed; if (IsIdentityPermutation(AsInt64Slice(reshape_request.dimensions()))) { transposed = operand; } else { - transposed = - hlo_builder_.AddInstruction(HloInstruction::CreateTranspose( - ShapeUtil::PermuteDimensions(InversePermutation(AsInt64Slice( - reshape_request.dimensions())), - operand->shape()), - operand, AsInt64Slice(reshape_request.dimensions()))); + transposed = add_instruction(HloInstruction::CreateTranspose( + ShapeUtil::PermuteDimensions( + InversePermutation(AsInt64Slice(reshape_request.dimensions())), + operand->shape()), + operand, AsInt64Slice(reshape_request.dimensions()))); } - hlo_instruction = hlo_builder_.AddInstruction( + hlo_instruction = add_instruction( HloInstruction::CreateReshape(request.output_shape(), transposed)); break; } @@ -2179,33 +2813,31 @@ HloInstruction* ComputationLowerer::Visit( case OpRequest::kTransposeRequest: { const TransposeRequest& transpose_request = request.request().transpose_request(); - HloInstruction* operand = Visit(transpose_request.operand(), visited); - hlo_instruction = - hlo_builder_.AddInstruction(HloInstruction::CreateTranspose( - ShapeUtil::PermuteDimensions(InversePermutation(AsInt64Slice( - transpose_request.dimensions())), - operand->shape()), - operand, AsInt64Slice(transpose_request.dimensions()))); + HloInstruction* operand = lookup_instruction(transpose_request.operand()); + hlo_instruction = add_instruction(HloInstruction::CreateTranspose( + ShapeUtil::PermuteDimensions( + InversePermutation(AsInt64Slice(transpose_request.dimensions())), + operand->shape()), + operand, AsInt64Slice(transpose_request.dimensions()))); break; } case OpRequest::kReverseRequest: { const ReverseRequest& reverse_request = request.request().reverse_request(); - HloInstruction* operand = Visit(reverse_request.operand(), visited); - hlo_instruction = - hlo_builder_.AddInstruction(HloInstruction::CreateReverse( - request.output_shape(), operand, - AsInt64Slice(reverse_request.dimensions()))); + HloInstruction* operand = lookup_instruction(reverse_request.operand()); + hlo_instruction = add_instruction(HloInstruction::CreateReverse( + request.output_shape(), operand, + AsInt64Slice(reverse_request.dimensions()))); break; } case OpRequest::kPadRequest: { const PadRequest& pad_request = request.request().pad_request(); - HloInstruction* operand = Visit(pad_request.operand(), visited); + HloInstruction* operand = lookup_instruction(pad_request.operand()); HloInstruction* padding_value = - Visit(pad_request.padding_value(), visited); - hlo_instruction = hlo_builder_.AddInstruction(HloInstruction::CreatePad( + lookup_instruction(pad_request.padding_value()); + hlo_instruction = add_instruction(HloInstruction::CreatePad( request.output_shape(), operand, padding_value, pad_request.padding_config())); break; @@ -2213,7 +2845,7 @@ HloInstruction* ComputationLowerer::Visit( case OpRequest::kRecvRequest: { const RecvRequest& recv_request = request.request().recv_request(); - hlo_instruction = hlo_builder_.AddInstruction(HloInstruction::CreateRecv( + hlo_instruction = add_instruction(HloInstruction::CreateRecv( request.output_shape(), recv_request.channel_handle().handle())); break; } @@ -2221,18 +2853,17 @@ HloInstruction* ComputationLowerer::Visit( case OpRequest::kParameterRequest: { const ParameterRequest& parameter_request = request.request().parameter_request(); - hlo_instruction = - hlo_builder_.AddInstruction(HloInstruction::CreateParameter( - parameter_request.parameter(), request.output_shape(), - parameter_request.name())); + hlo_instruction = add_instruction(HloInstruction::CreateParameter( + parameter_request.parameter(), request.output_shape(), + parameter_request.name())); break; } case OpRequest::kConvertRequest: { const ConvertRequest& convert_request = request.request().convert_request(); - HloInstruction* operand = Visit(convert_request.operand(), visited); - hlo_instruction = hlo_builder_.AddInstruction( + HloInstruction* operand = lookup_instruction(convert_request.operand()); + hlo_instruction = add_instruction( HloInstruction::CreateConvert(request.output_shape(), operand)); break; } @@ -2248,8 +2879,8 @@ HloInstruction* ComputationLowerer::Visit( request.embedded_computation_versions(1); HloComputation* body = ResolveComputation(while_request.body(), body_version); - HloInstruction* init = Visit(while_request.init(), visited); - hlo_instruction = hlo_builder_.AddInstruction(HloInstruction::CreateWhile( + HloInstruction* init = lookup_instruction(while_request.init()); + hlo_instruction = add_instruction(HloInstruction::CreateWhile( request.output_shape(), condition, body, init)); break; } @@ -2257,13 +2888,12 @@ HloInstruction* ComputationLowerer::Visit( case OpRequest::kTernaryOpRequest: { const TernaryOpRequest& ternary_op_request = request.request().ternary_op_request(); - HloInstruction* lhs = Visit(ternary_op_request.lhs(), visited); - HloInstruction* rhs = Visit(ternary_op_request.rhs(), visited); - HloInstruction* ehs = Visit(ternary_op_request.ehs(), visited); + HloInstruction* lhs = lookup_instruction(ternary_op_request.lhs()); + HloInstruction* rhs = lookup_instruction(ternary_op_request.rhs()); + HloInstruction* ehs = lookup_instruction(ternary_op_request.ehs()); auto hlo_opcode = TernaryOperationToHloOpcode(ternary_op_request.triop()); - hlo_instruction = - hlo_builder_.AddInstruction(HloInstruction::CreateTernary( - request.output_shape(), hlo_opcode, lhs, rhs, ehs)); + hlo_instruction = add_instruction(HloInstruction::CreateTernary( + request.output_shape(), hlo_opcode, lhs, rhs, ehs)); break; } @@ -2273,14 +2903,13 @@ HloInstruction* ComputationLowerer::Visit( std::vector operands; for (const ComputationDataHandle& handle : variadic_op_request.operands()) { - HloInstruction* operand = Visit(handle, visited); + HloInstruction* operand = lookup_instruction(handle); operands.push_back(operand); } auto hlo_opcode = VariadicOperationToHloOpcode(variadic_op_request.varop()); - hlo_instruction = - hlo_builder_.AddInstruction(HloInstruction::CreateVariadic( - request.output_shape(), hlo_opcode, operands)); + hlo_instruction = add_instruction(HloInstruction::CreateVariadic( + request.output_shape(), hlo_opcode, operands)); break; } @@ -2288,14 +2917,14 @@ HloInstruction* ComputationLowerer::Visit( const CallRequest& call_request = request.request().call_request(); std::vector operands; for (const ComputationDataHandle& handle : call_request.operands()) { - operands.push_back(Visit(handle, visited)); + operands.push_back(lookup_instruction(handle)); } CHECK_EQ(1, request.embedded_computation_versions_size()); VersionedComputationHandle::Version call_version = request.embedded_computation_versions(0); HloComputation* call_computation = ResolveComputation(call_request.to_apply(), call_version); - hlo_instruction = hlo_builder_.AddInstruction(HloInstruction::CreateCall( + hlo_instruction = add_instruction(HloInstruction::CreateCall( request.output_shape(), operands, call_computation)); break; } @@ -2305,20 +2934,19 @@ HloInstruction* ComputationLowerer::Visit( request.request().custom_call_request(); std::vector operands; for (const ComputationDataHandle& operand : cc_request.operands()) { - operands.push_back(Visit(operand, visited)); + operands.push_back(lookup_instruction(operand)); } - hlo_instruction = - hlo_builder_.AddInstruction(HloInstruction::CreateCustomCall( - cc_request.shape(), operands, cc_request.call_target_name())); + hlo_instruction = add_instruction(HloInstruction::CreateCustomCall( + cc_request.shape(), operands, cc_request.call_target_name())); break; } case OpRequest::kUnaryOpRequest: { const UnaryOpRequest& unary_op_request = request.request().unary_op_request(); - HloInstruction* operand = Visit(unary_op_request.operand(), visited); + HloInstruction* operand = lookup_instruction(unary_op_request.operand()); auto hlo_opcode = UnaryOperationToHloOpcode(unary_op_request.unop()); - hlo_instruction = hlo_builder_.AddInstruction(HloInstruction::CreateUnary( + hlo_instruction = add_instruction(HloInstruction::CreateUnary( request.output_shape(), hlo_opcode, operand)); break; } @@ -2326,8 +2954,8 @@ HloInstruction* ComputationLowerer::Visit( case OpRequest::kBinaryOpRequest: { const BinaryOpRequest& binary_op_request = request.request().binary_op_request(); - HloInstruction* lhs = Visit(binary_op_request.lhs(), visited); - HloInstruction* rhs = Visit(binary_op_request.rhs(), visited); + HloInstruction* lhs = lookup_instruction(binary_op_request.lhs()); + HloInstruction* rhs = lookup_instruction(binary_op_request.rhs()); auto hlo_opcode = BinaryOperationToHloOpcode(binary_op_request.binop()); if (binary_op_request.broadcast_dimensions_size() > 0) { // Emit a broadcast instruction to perform the "broadcast in dimension" @@ -2346,23 +2974,47 @@ HloInstruction* ComputationLowerer::Visit( // identical to the HLO broadcast semantics so the broadcast_dimensions // field can just be passed to the instruction builder. HloInstruction* broadcasted_operand = - hlo_builder_.AddInstruction(HloInstruction::CreateBroadcast( + add_instruction(HloInstruction::CreateBroadcast( broadcast_shape, operand_to_broadcast, AsInt64Slice(binary_op_request.broadcast_dimensions()))); lhs = (lhs == operand_to_broadcast) ? broadcasted_operand : lhs; rhs = (rhs == operand_to_broadcast) ? broadcasted_operand : rhs; } - hlo_instruction = - hlo_builder_.AddInstruction(HloInstruction::CreateBinary( - request.output_shape(), hlo_opcode, lhs, rhs)); + if (debug_options_.xla_eliminate_hlo_implicit_broadcast() && + binary_op_request.binop() != BINOP_DOT) { + if (!ShapeUtil::SameDimensions(request.output_shape(), lhs->shape())) { + // lhs side is being implicitly broadcast. Change to explicit. + lhs = + ImplicitBroadcastToExplicitBroadcast(lhs, request.output_shape()); + } + + if (!ShapeUtil::SameDimensions(request.output_shape(), rhs->shape())) { + rhs = + ImplicitBroadcastToExplicitBroadcast(rhs, request.output_shape()); + } + } + hlo_instruction = add_instruction(HloInstruction::CreateBinary( + request.output_shape(), hlo_opcode, lhs, rhs)); + break; + } + + case OpRequest::kReducePrecisionRequest: { + const ReducePrecisionRequest& reduce_precision_request = + request.request().reduce_precision_request(); + HloInstruction* operand = + lookup_instruction(reduce_precision_request.operand()); + auto exponent_bits = reduce_precision_request.exponent_bits(); + auto mantissa_bits = reduce_precision_request.mantissa_bits(); + hlo_instruction = add_instruction(HloInstruction::CreateReducePrecision( + request.output_shape(), operand, exponent_bits, mantissa_bits)); break; } case OpRequest::kTraceRequest: { const TraceRequest& trace_request = request.request().trace_request(); - HloInstruction* operand = Visit(trace_request.operand(), visited); - hlo_instruction = hlo_builder_.AddInstruction( + HloInstruction* operand = lookup_instruction(trace_request.operand()); + hlo_instruction = add_instruction( HloInstruction::CreateTrace(trace_request.tag(), operand)); operand->set_tracing(hlo_instruction); break; @@ -2370,8 +3022,8 @@ HloInstruction* ComputationLowerer::Visit( case OpRequest::kSendRequest: { const SendRequest& send_request = request.request().send_request(); - HloInstruction* operand = Visit(send_request.operand(), visited); - hlo_instruction = hlo_builder_.AddInstruction(HloInstruction::CreateSend( + HloInstruction* operand = lookup_instruction(send_request.operand()); + hlo_instruction = add_instruction(HloInstruction::CreateSend( operand, send_request.channel_handle().handle())); break; } @@ -2382,16 +3034,14 @@ HloInstruction* ComputationLowerer::Visit( default: LOG(FATAL) << "Unexpected request type: " << request.request().op_case(); } - hlo_instruction->set_metadata(request.request().metadata()); - (*visited)[handle.handle()] = hlo_instruction; - return hlo_instruction; + (*instructions)[handle.handle()] = hlo_instruction; } } // namespace StatusOr> UserComputation::BuildHloComputation( VersionedComputationHandle::Version version, - HloComputationResolver hlo_resolver, + HloComputationResolver hlo_resolver, const DebugOptions& debug_options, bool include_unreachable_instructions) const { tensorflow::mutex_lock lock(mutex_); @@ -2403,10 +3053,9 @@ StatusOr> UserComputation::BuildHloComputation( std::unique_ptr hlo_computation, ComputationLowerer::Lower( tensorflow::strings::StrCat(name(), ".v", version), - session_computation_, version, std::move(hlo_resolver), + session_computation_, version, std::move(hlo_resolver), debug_options, include_unreachable_instructions)); - XLA_VLOG_LINES(2, hlo_computation->ToString()); return std::move(hlo_computation); } diff --git a/tensorflow/compiler/xla/service/user_computation.h b/tensorflow/compiler/xla/service/user_computation.h index fb5425ae61ab1edcd00aac493c9e2ac3c430cb72..b779b1f76c82a43c4d5aee4b1f01d0941d104c5c 100644 --- a/tensorflow/compiler/xla/service/user_computation.h +++ b/tensorflow/compiler/xla/service/user_computation.h @@ -27,6 +27,7 @@ limitations under the License. #include "tensorflow/compiler/xla/service/versioned_computation_handle.h" #include "tensorflow/compiler/xla/statusor.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/platform/macros.h" #include "tensorflow/core/platform/mutex.h" @@ -84,6 +85,18 @@ class UserComputation { StatusOr AddUnaryInstruction( const UnaryOpRequest& unary_request); + // Enqueues a batch norm training instruction onto this user computation. + StatusOr AddBatchNormTrainingInstruction( + const BatchNormTrainingRequest& batch_norm_training_request); + + // Enqueues a batch norm inference instruction onto this user computation. + StatusOr AddBatchNormInferenceInstruction( + const BatchNormInferenceRequest& batch_norm_inference_request); + + // Enqueues a batch norm grad instruction onto this user computation. + StatusOr AddBatchNormGradInstruction( + const BatchNormGradRequest& batch_norm_grad_request); + // Enqueues a binary instruction onto this user computation. // Returns an error status if the operand indices are out of bounds. StatusOr AddBinaryInstruction( @@ -112,6 +125,10 @@ class UserComputation { const MapRequest& map_request, const UserComputation& to_apply_computation); + // Enqueues a reduce-precision instruction onto this user computation. + StatusOr AddReducePrecisionInstruction( + const ReducePrecisionRequest& reduce_precision_request); + // Enqueues a convolution instruction onto this user computation. StatusOr AddConvolveInstruction( const ConvolveRequest& convolve_request); @@ -256,7 +273,7 @@ class UserComputation { std::function; StatusOr> BuildHloComputation( VersionedComputationHandle::Version version, - HloComputationResolver hlo_resolver, + HloComputationResolver hlo_resolver, const DebugOptions& debug_options, bool include_unreachable_instructions = true) const; // Return a vector containing the embedded computations used by this diff --git a/tensorflow/compiler/xla/service/user_computation_test.cc b/tensorflow/compiler/xla/service/user_computation_test.cc index e67254328ad6973ee63a83d45cd3c2618e39ff56..6b0d6b9e11cd638b8f8a2d6f6be7e5a96b351382 100644 --- a/tensorflow/compiler/xla/service/user_computation_test.cc +++ b/tensorflow/compiler/xla/service/user_computation_test.cc @@ -17,12 +17,16 @@ limitations under the License. #include "tensorflow/compiler/xla/literal_util.h" #include "tensorflow/compiler/xla/service/hlo_computation.h" +#include "tensorflow/compiler/xla/service/hlo_matchers.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/xla_data.pb.h" #include "tensorflow/core/lib/core/status_test_util.h" +namespace op = xla::testing::opcode_matchers; + namespace xla { namespace { @@ -45,16 +49,19 @@ TEST_F(UserComputationTest, SimpleComputation) { ConstantRequest constant_request; *constant_request.mutable_literal() = - *LiteralUtil::CreateR1({123.0f, 42.0f}); - TF_ASSIGN_OR_ASSERT_OK(ComputationDataHandle constant_handle, - computation.AddConstantInstruction(constant_request)); + Literal::CreateR1({123.0f, 42.0f})->ToProto(); + TF_ASSERT_OK_AND_ASSIGN(ComputationDataHandle constant_handle, + computation.AddConstantInstruction(constant_request)); ParameterRequest param_request; *param_request.mutable_shape() = kScalarShape; param_request.set_parameter(0); param_request.set_name("param0"); - TF_ASSIGN_OR_ASSERT_OK(ComputationDataHandle param_handle, - computation.AddParameterInstruction(param_request)); + TF_ASSERT_OK_AND_ASSIGN(ComputationDataHandle param_handle, + computation.AddParameterInstruction(param_request)); + OpMetadata metadata; + metadata.set_op_name("meta"); + TF_ASSERT_OK(computation.SetOpMetadata(param_handle, metadata)); OutfeedRequest outfeed_request; *outfeed_request.mutable_operand() = constant_handle; @@ -73,7 +80,7 @@ TEST_F(UserComputationTest, SimpleComputation) { // Program shape should have a single scalar parameter and scalar // result. The outfeed instruction should not affect the program shape. - TF_ASSIGN_OR_ASSERT_OK( + TF_ASSERT_OK_AND_ASSIGN( std::shared_ptr program_shape, computation.ComputeProgramShape(latest_version.version)); ASSERT_EQ(1, program_shape->parameters_size()); @@ -82,15 +89,15 @@ TEST_F(UserComputationTest, SimpleComputation) { EXPECT_TRUE(ShapeUtil::Compatible(kScalarShape, program_shape->result())); // Build the HLO computation. - TF_ASSIGN_OR_ASSERT_OK( + TF_ASSERT_OK_AND_ASSIGN( std::unique_ptr hlo_computation, - computation.BuildHloComputation(latest_version.version, hlo_resolver)); + computation.BuildHloComputation(latest_version.version, hlo_resolver, + DebugOptions())); // There should be one HloInstruction per UserComputation operation. EXPECT_EQ(3, hlo_computation->instruction_count()); // The root of the instruction should be the parameter instruction (not the // outfeed). - EXPECT_EQ(HloOpcode::kParameter, - hlo_computation->root_instruction()->opcode()); + EXPECT_THAT(hlo_computation->root_instruction(), op::Parameter()); } { @@ -100,7 +107,7 @@ TEST_F(UserComputationTest, SimpleComputation) { computation.GetVersionedHandleAtOperation(param_handle); // Program shape should have a single scalar parameter, and scalar result. - TF_ASSIGN_OR_ASSERT_OK( + TF_ASSERT_OK_AND_ASSIGN( std::shared_ptr program_shape, computation.ComputeProgramShape(version_at_param.version)); ASSERT_EQ(1, program_shape->parameters_size()); @@ -110,12 +117,12 @@ TEST_F(UserComputationTest, SimpleComputation) { // There should be two instructions, one for the constant and one for the // parameter. The outfeed instruction should not be included. - TF_ASSIGN_OR_ASSERT_OK(std::unique_ptr hlo_computation, - computation.BuildHloComputation( - version_at_param.version, hlo_resolver)); + TF_ASSERT_OK_AND_ASSIGN( + std::unique_ptr hlo_computation, + computation.BuildHloComputation(version_at_param.version, hlo_resolver, + DebugOptions())); EXPECT_EQ(2, hlo_computation->instruction_count()); - EXPECT_EQ(HloOpcode::kParameter, - hlo_computation->root_instruction()->opcode()); + EXPECT_THAT(hlo_computation->root_instruction(), op::Parameter()); } { // Test the computation at the latest version, but lowered with @@ -124,18 +131,240 @@ TEST_F(UserComputationTest, SimpleComputation) { computation.GetVersionedHandle(); // Build the HLO computation. - TF_ASSIGN_OR_ASSERT_OK(std::unique_ptr hlo_computation, - computation.BuildHloComputation( - latest_version.version, hlo_resolver, - /*include_unreachable_instructions=*/false)); + TF_ASSERT_OK_AND_ASSIGN( + std::unique_ptr hlo_computation, + computation.BuildHloComputation( + latest_version.version, hlo_resolver, DebugOptions(), + /*include_unreachable_instructions=*/false)); // There is only one reachable instruction, the parameter. EXPECT_EQ(1, hlo_computation->instruction_count()); // The root of the instruction should be the parameter instruction (not the // outfeed). - EXPECT_EQ(HloOpcode::kParameter, - hlo_computation->root_instruction()->opcode()); + EXPECT_THAT(hlo_computation->root_instruction(), op::Parameter()); + EXPECT_EQ(hlo_computation->root_instruction()->metadata().op_name(), + "meta"); } } +TEST_F(UserComputationTest, EliminateScalarBroadcast) { + auto debug_options = DebugOptions(); + debug_options.set_xla_eliminate_hlo_implicit_broadcast(true); + + // Build a binary computation with scalar broadcast. + // + // %a = Constant({123, 42}) + // %b = Constant(1) + // %add = Add(%a, %b) + ComputationHandle handle; + handle.set_handle(123); + UserComputation computation("TheComputation", handle); + + ConstantRequest a_request; + *a_request.mutable_literal() = + Literal::CreateR1({123.0f, 42.0f})->ToProto(); + TF_ASSERT_OK_AND_ASSIGN(ComputationDataHandle a_handle, + computation.AddConstantInstruction(a_request)); + + ConstantRequest b_request; + *b_request.mutable_literal() = Literal::CreateR0(1.0f)->ToProto(); + TF_ASSERT_OK_AND_ASSIGN(ComputationDataHandle b_handle, + computation.AddConstantInstruction(b_request)); + + BinaryOpRequest add; + add.set_binop(BINOP_ADD); + *add.mutable_lhs() = a_handle; + *add.mutable_rhs() = b_handle; + TF_ASSERT_OK(computation.AddBinaryInstruction(add).status()); + + auto hlo_resolver = [](const VersionedComputationHandle& handle) { + return nullptr; + }; + VersionedComputationHandle latest_version = computation.GetVersionedHandle(); + + // Build the HLO computation. + TF_ASSERT_OK_AND_ASSIGN( + std::unique_ptr hlo_computation, + computation.BuildHloComputation(latest_version.version, hlo_resolver, + debug_options)); + // The binary operation has implicit scalar broadcast, should be converted + // to an explicit broadcast intruction and a binary instruction. + EXPECT_EQ(4, hlo_computation->instruction_count()); + EXPECT_THAT(hlo_computation->root_instruction(), op::Add()); + LOG(INFO) << hlo_computation->root_instruction()->ToString(); + const auto& operands = hlo_computation->root_instruction()->operands(); + ASSERT_EQ(2, operands.size()); + EXPECT_TRUE(operands[0]->opcode() == HloOpcode::kBroadcast || + operands[1]->opcode() == HloOpcode::kBroadcast); +} + +TEST_F(UserComputationTest, CheckImplicitBroadcastToExplicitBroadcast) { + auto debug_options = DebugOptions(); + debug_options.set_xla_eliminate_hlo_implicit_broadcast(true); + + // Build a binary computation with degenerate broadcast. + // + // %a = Param({1, 2, 3}); + // %b = Param({1, 2, 1}); + // %add = Add(%a, %b, {}); + ComputationHandle handle; + handle.set_handle(123); + UserComputation computation("TheComputation", handle); + + ParameterRequest a_request; + *a_request.mutable_shape() = ShapeUtil::MakeShape(F32, {1, 2, 3}); + a_request.set_name("a"); + a_request.set_parameter(0); + TF_ASSERT_OK_AND_ASSIGN(ComputationDataHandle a_handle, + computation.AddParameterInstruction(a_request)); + + ParameterRequest b_request; + *b_request.mutable_shape() = ShapeUtil::MakeShape(F32, {1, 2, 1}); + b_request.set_name("b"); + b_request.set_parameter(1); + TF_ASSERT_OK_AND_ASSIGN(ComputationDataHandle b_handle, + computation.AddParameterInstruction(b_request)); + + BinaryOpRequest add; + add.set_binop(BINOP_ADD); + *add.mutable_lhs() = a_handle; + *add.mutable_rhs() = b_handle; + TF_ASSERT_OK(computation.AddBinaryInstruction(add).status()); + + auto hlo_resolver = [](const VersionedComputationHandle& handle) { + return nullptr; + }; + VersionedComputationHandle latest_version = computation.GetVersionedHandle(); + + // Build the HLO computation. + TF_ASSERT_OK_AND_ASSIGN( + std::unique_ptr hlo_computation, + computation.BuildHloComputation(latest_version.version, hlo_resolver, + debug_options)); + + // b a + // | | + // reshape | + // | | + // broadcast | + // \ / + // add + EXPECT_EQ(5, hlo_computation->instruction_count()); + EXPECT_THAT(hlo_computation->root_instruction(), op::Add()); + const auto& operands = hlo_computation->root_instruction()->operands(); + ASSERT_EQ(2, operands.size()); + EXPECT_TRUE(operands[0]->opcode() == HloOpcode::kParameter && + operands[1]->opcode() == HloOpcode::kBroadcast); +} + +TEST_F(UserComputationTest, EliminateDegenerateBroadcastAfterIndimBroadcast) { + auto debug_options = DebugOptions(); + debug_options.set_xla_eliminate_hlo_implicit_broadcast(true); + + // Build a binary computation with in-dim broadcast and degenerate broadcast. + // + // %a = Param({2, 3}); + // %b = Param({2, 1, 4}); + // %add = Add(%a, %b, {0, 1}); + ComputationHandle handle; + handle.set_handle(123); + UserComputation computation("TheComputation", handle); + + ParameterRequest a_request; + *a_request.mutable_shape() = ShapeUtil::MakeShape(F32, {2, 3}); + a_request.set_name("a"); + a_request.set_parameter(0); + TF_ASSERT_OK_AND_ASSIGN(ComputationDataHandle a_handle, + computation.AddParameterInstruction(a_request)); + + ParameterRequest b_request; + *b_request.mutable_shape() = ShapeUtil::MakeShape(F32, {2, 1, 4}); + b_request.set_name("b"); + b_request.set_parameter(1); + TF_ASSERT_OK_AND_ASSIGN(ComputationDataHandle b_handle, + computation.AddParameterInstruction(b_request)); + + BinaryOpRequest add; + add.set_binop(BINOP_ADD); + *add.mutable_lhs() = a_handle; + *add.mutable_rhs() = b_handle; + add.add_broadcast_dimensions(0); + add.add_broadcast_dimensions(1); + TF_ASSERT_OK(computation.AddBinaryInstruction(add).status()); + + auto hlo_resolver = [](const VersionedComputationHandle& handle) { + return nullptr; + }; + VersionedComputationHandle latest_version = computation.GetVersionedHandle(); + + // Build the HLO computation. + TF_ASSERT_OK_AND_ASSIGN( + std::unique_ptr hlo_computation, + computation.BuildHloComputation(latest_version.version, hlo_resolver, + debug_options)); + + // The binary operation has in-dim broadcast and degenerate broadcast, should + // first do the in-dim broadcast then convert the degnerate broadcast into a + // reshape and a broadcast. + // + // b a + // | | + // broadcast reshape + // | | + // | broadcast + // \ / + // add + EXPECT_EQ(6, hlo_computation->instruction_count()); + EXPECT_THAT(hlo_computation->root_instruction(), op::Add()); + const auto& operands = hlo_computation->root_instruction()->operands(); + ASSERT_EQ(2, operands.size()); + EXPECT_TRUE(operands[0]->opcode() == HloOpcode::kBroadcast && + operands[1]->opcode() == HloOpcode::kBroadcast); +} + +TEST_F(UserComputationTest, SkipDotInEliminatingImplicitBroadcast) { + auto debug_options = DebugOptions(); + debug_options.set_xla_eliminate_hlo_implicit_broadcast(true); + + // %a = Param({1, 3}); + // %b = Param({3, 1}); + // %dot = Dot(%a, %b); + ComputationHandle handle; + handle.set_handle(123); + UserComputation computation("TheComputation", handle); + + ParameterRequest a_request; + *a_request.mutable_shape() = ShapeUtil::MakeShape(F32, {1, 3}); + a_request.set_name("a"); + a_request.set_parameter(0); + TF_ASSERT_OK_AND_ASSIGN(ComputationDataHandle a_handle, + computation.AddParameterInstruction(a_request)); + + ParameterRequest b_request; + *b_request.mutable_shape() = ShapeUtil::MakeShape(F32, {3, 1}); + b_request.set_name("b"); + b_request.set_parameter(1); + TF_ASSERT_OK_AND_ASSIGN(ComputationDataHandle b_handle, + computation.AddParameterInstruction(b_request)); + + BinaryOpRequest dot; + dot.set_binop(BINOP_DOT); + *dot.mutable_lhs() = a_handle; + *dot.mutable_rhs() = b_handle; + TF_ASSERT_OK(computation.AddBinaryInstruction(dot).status()); + + auto hlo_resolver = [](const VersionedComputationHandle& handle) { + return nullptr; + }; + VersionedComputationHandle latest_version = computation.GetVersionedHandle(); + + // Build the HLO computation. + TF_ASSERT_OK_AND_ASSIGN( + std::unique_ptr hlo_computation, + computation.BuildHloComputation(latest_version.version, hlo_resolver, + debug_options)); + + EXPECT_EQ(3, hlo_computation->instruction_count()); +} + } // namespace } // namespace xla diff --git a/tensorflow/compiler/xla/service_interface.h b/tensorflow/compiler/xla/service_interface.h index 2159386152b34e4f9b59ca14faa756e37551d724..809941d8fe1f63d66bf104e66eea66167a0f509d 100644 --- a/tensorflow/compiler/xla/service_interface.h +++ b/tensorflow/compiler/xla/service_interface.h @@ -21,7 +21,10 @@ limitations under the License. namespace xla { -// Defines the interface for an XLA service. +// Defines the interface for an XLA service on the client side. This service +// helps abstract around the actual implementation of a service - the service +// can be local (running in the same process), or remote - in which case an RPC +// stub is used as the implementation. class ServiceInterface { public: ServiceInterface() {} @@ -31,10 +34,6 @@ class ServiceInterface { virtual tensorflow::Status TransferToClient( const TransferToClientRequest* arg, TransferToClientResponse* result) = 0; - virtual tensorflow::Status TransferToClientInProcess( - const TransferToClientInProcessRequest* arg, - TransferToClientInProcessResponse* result) = 0; - virtual tensorflow::Status TransferToServer( const TransferToServerRequest* arg, TransferToServerResponse* result) = 0; @@ -48,10 +47,6 @@ class ServiceInterface { virtual tensorflow::Status ResetDevice(const ResetDeviceRequest* arg, ResetDeviceResponse* result) = 0; - virtual tensorflow::Status TransferToServerInProcess( - const TransferToServerInProcessRequest* arg, - TransferToServerInProcessResponse* result) = 0; - virtual tensorflow::Status LoadComputationSnapshot( const LoadComputationSnapshotRequest* request, LoadComputationSnapshotResponse* result) = 0; diff --git a/tensorflow/compiler/xla/shape_tree.h b/tensorflow/compiler/xla/shape_tree.h index 6963a68d10d527acebde65f30f9caf87608950cb..64a36471b9f1b35517c29c01554e02c5d1035086 100644 --- a/tensorflow/compiler/xla/shape_tree.h +++ b/tensorflow/compiler/xla/shape_tree.h @@ -17,6 +17,7 @@ limitations under the License. #define TENSORFLOW_COMPILER_XLA_SHAPE_TREE_H_ #include +#include #include #include @@ -28,27 +29,108 @@ limitations under the License. #include "tensorflow/core/lib/core/errors.h" #include "tensorflow/core/lib/core/status.h" #include "tensorflow/core/lib/gtl/array_slice.h" +#include "tensorflow/core/lib/gtl/iterator_range.h" +#include "tensorflow/core/lib/gtl/optional.h" #include "tensorflow/core/platform/logging.h" #include "tensorflow/core/platform/types.h" namespace xla { +namespace internal { + +// Internal representation of each node in a ShapeTree. +template +struct ShapeTreeNode { + // Data corresponding to this node. + T data; + + // Children of this node. + std::vector> children; + + ShapeTreeNode() = default; + explicit ShapeTreeNode(const T& data) : data(data) {} + + ShapeTreeNode(const ShapeTreeNode& other) + : data(other.data), children(other.children.size()) { + for (size_t i = 0; i < children.size(); ++i) { + children[i] = MakeUnique(*other.children[i]); + } + } + + ShapeTreeNode& operator=(const ShapeTreeNode& other) { + if (this != &other) { + data = other.data; + children.resize(other.children.size()); + for (size_t i = 0; i < children.size(); ++i) { + children[i] = MakeUnique(*other.children[i]); + } + } + return *this; + } +}; + +} // namespace internal + +template +class ShapeTreeIterator; + // A ShapeTree is a recursive data structure which mirrors the structure of a -// XLA shape and holds a value of type T for each array in the shape. For -// array shapes, a ShapeTree trivially holds a single value of type T. For tuple -// shapes which can be an arbitrary tree with arrays at the leaves, a ShapeTree -// is an identically structured tree with data elements of type T at the leaves. +// XLA shape and holds a value of type T for each subshape (i.e. tuple or array) +// in the shape. For array shapes, a ShapeTree trivially holds a single value of +// type T. +// +// For tuple shapes which can be an arbitrary tree with arrays at the leaves, a +// ShapeTree is an identically structured tree with data elements of type T at +// every node. I.e. the root is a tuple by definition, all interior nodes are +// also tuples, and all leaves are arrays. // // Like the Shape data structure, this is a tree and tuple elements cannot be -// duplicated. That is, every distinct element position in the Shape has a -// unique T object. +// duplicated. That is, every distinct ShapeIndex in the Shape has a unique T +// object. +// +// Normally a ShapeTree owns its Shape, but for efficiency reasons, sometimes +// it's helpful not to copy a Shape just to make a ShapeTree. In these cases, +// you can pass a Shape* instead of a Shape& to the ShapeTree constructor. It's +// then up to you to ensure that the pointed-to Shape doesn't die or mutate +// before its ShapeTree goes away. template class ShapeTree { + friend class ShapeTreeIterator; + friend class ShapeTreeIterator; + public: - explicit ShapeTree(const Shape& shape); - ShapeTree(const Shape& shape, const T& init_value); - ShapeTree(const ShapeTree& other); - ShapeTree& operator=(const ShapeTree& other); + // Default constructor creates a tree with a nil shape (i.e. an empty tuple). + ShapeTree() : ShapeTree(ShapeUtil::MakeNil()) {} + + // Create ShapeTree with the given shape, and default-constructed T values for + // all nodes. + // + // The version that takes a pointer may be cheaper because it doesn't require + // any Shape copies, but then it's up to you to ensure that the pointer stays + // alive longer than this ShapeTree. + explicit ShapeTree(Shape shape); + explicit ShapeTree(const Shape* shape); + + // Create ShapeTree with the given shape, and init_value for all nodes. + ShapeTree(Shape shape, const T& init_value); + ShapeTree(const Shape* shape, const T& init_value); + + ShapeTree(const ShapeTree& other) { *this = other; } + + ShapeTree& operator=(const ShapeTree& other) { + root_ = other.root_; + + // Fix up internal pointer if necessary. + if (other.shape_storage_) { + CHECK_EQ(other.shape_, other.shape_storage_.get()); + shape_storage_.reset(new Shape(*other.shape_)); + shape_ = shape_storage_.get(); + } else { + shape_ = other.shape_; + } + + return *this; + } // Returns the data element associated with the array in the shape at the // given index (see ShapeUtil::GetSubshape for how indexes are defined). @@ -61,198 +143,417 @@ class ShapeTree { // Returns true if the node at the given index is a leaf node (an array // shape). bool IsLeaf(const ShapeIndex& index) const { - return Lookup(index).elements_.empty(); + return Lookup(index)->children.empty(); + } + + // iterator implements a forward_iterator with value_type = + // std::pair + using iterator = ShapeTreeIterator; + using const_iterator = ShapeTreeIterator; + + // begin/end for iterating over all nodes. + iterator begin() { return iterator(&root_, /*iterate_leaves_only=*/false); } + iterator end() { return iterator(nullptr, /*iterate_leaves_only=*/false); } + const_iterator begin() const { + return const_iterator(&root_, /*iterate_leaves_only=*/false); + } + const_iterator end() const { + return const_iterator(nullptr, /*iterate_leaves_only=*/false); + } + + // leaf_begin()/leaf_end() iterates over all leaf nodes (nodes with no + // children). + iterator leaf_begin() { + return iterator(&root_, /*iterate_leaves_only=*/true); + } + iterator leaf_end() { + return iterator(nullptr, /*iterate_leaves_only=*/true); + } + const_iterator leaf_begin() const { + return const_iterator(&root_, /*iterate_leaves_only=*/true); + } + const_iterator leaf_end() const { + return const_iterator(nullptr, /*iterate_leaves_only=*/true); + } + // range-based iterator for leaf_begin()/leaf_end(). + tensorflow::gtl::iterator_range leaves() { + return tensorflow::gtl::make_range(leaf_begin(), leaf_end()); + } + tensorflow::gtl::iterator_range leaves() const { + return tensorflow::gtl::make_range(leaf_begin(), leaf_end()); } // Recursively traverses the shape and calls the given function at each // element. The function has the following arguments: // + // Fn : A callable of type void(const ShapeIndex& index, const T& data) + // (or compatible). // index : the index of the element in the shape. See ShapeUtil::GetSubshape // for definition of index. - // is_leaf : Whether this element is a leaf element in the shape. That is, - // whether this index corresponds to an array and not a (nested) - // tuple element. // data : The data value at this elemnt. + template + void ForEachElement(const Fn& func) const; + + // Like ForEachElement, but the callable has type // - // If any call to the given function returns a non-OK status, then traversal - // is aborted and the status value is returned. - using VisitorFunction = std::function; - tensorflow::Status ForEachElement(VisitorFunction func) const; + // void (const ShapeIndex& index, T* data). + // + template + void ForEachMutableElement(const Fn& func); + + // Like ForEach(Mutable)Element, but the callable returns a Status instead of + // void. The first non-OK return value is returned by the ForEach* function. + template + Status ForEachElementWithStatus(const Fn& func) const; + template + Status ForEachMutableElementWithStatus(const Fn& func); + + // Copy the subtree of values from 'other' rooted at ShapeIndex + // 'source_base_index' into the subtree of value in this ShapeTree rooted at + // 'target_base_index'. + // + // Precondition: The subshape of other.shape() at index source_base_index must + // be compatible with the subshape of shape() at index target_base_index. + void CopySubtreeFrom(const ShapeTree& other, + const ShapeIndex& source_base_index, + const ShapeIndex& target_base_index); - using MutableVisitorFunction = std::function; - tensorflow::Status ForEachMutableElement(MutableVisitorFunction func); + bool operator==(const ShapeTree& other) const; + bool operator!=(const ShapeTree& other) const { return !(*this == other); } private: - // Private default constructor for non-root nodes of the tree. - ShapeTree() = default; + using Node = internal::ShapeTreeNode; + + // Initialize node->children based on 'shape'. All children are assigned the + // the given 'init_value'. + void InitChildren(const Shape& shape, const T& init_value, Node* node); + + // Initialize node->children based on 'shape'. All children have + // default-constructed data values. + void InitChildren(const Shape& shape, Node* node); // Helpers for traversing the shape via ForEachElement. The helpers // recursively traverse the subtree rooted at "index" (defined as in // ShapeUtil::GetSubshape). - static tensorflow::Status ForEachHelperMutable(ShapeIndex* index, - ShapeTree* shape_tree, - MutableVisitorFunction func); - static tensorflow::Status ForEachHelper(ShapeIndex* index, - const ShapeTree& shape_tree, - VisitorFunction func); + template + static Status ForEachHelper(const Fn& func, const Node& node, + ShapeIndex* index); + template + static Status ForEachMutableHelper(const Fn& func, Node* node, + ShapeIndex* index); + + // Return the tree node at the given index. + Node* Lookup(const ShapeIndex& index); + const Node* Lookup(const ShapeIndex& index) const; - // Copy all the data elements (of type T) from "other" into "this". "this" - // must have the same tree structure as "other" prior to calling this method. - void CopyDataElements(const ShapeTree& other); + // The root node, which contains all other nodes. + Node root_; - // Recursive helper for constructing a subtree beneath "this" node. - void BuildTree(const Shape& shape); + // If we own our Shape, this field contains it, and shape_ is a pointer into + // here. Otherwise if we don't own our shape, this is nullptr. + std::unique_ptr shape_storage_; - // Return the tree node at the given index. - ShapeTree& Lookup(const ShapeIndex& index); - const ShapeTree& Lookup(const ShapeIndex& index) const; + // The XLA shape mirrored in this ShapeTree. This is either + // shape_storage_.get() or the Shape pointer passed to our constructor. + const Shape* shape_; +}; - // The data corresponding to the array at this node. - T data_; +// Internal iterator that performs a pre-order walk. This is copyable, but +// contains a vector so isn't cheap to copy. This also means post-increment is +// expensive. The iterator value_type is equivalent to a std::pair, similar to std::map. The non-const iterator's T& type can be mutated +// in-place. +template +class ShapeTreeIterator : public std::iterator> { + public: + using value_type = + typename std::conditional, + std::pair>::type; + using NodeType = + typename std::conditional::Node, + typename ShapeTree::Node>::type; + + // Construct an iterator pointing at node. Node must either be the tree root + // or nullptr (which is equivalent to end() and should not be dereferenced or + // incremented). If iterate_leaves_only is true, the iterator will not include + // interior tree nodes, only leaves. + ShapeTreeIterator(NodeType* node, bool iterate_leaves_only) + : node_(node), iterate_leaves_only_(iterate_leaves_only) { + if (node_ && !node_->children.empty() && iterate_leaves_only) { + ++*this; + } + } + ShapeTreeIterator(const ShapeTreeIterator& other) + : node_(other.node_), + stack_(other.stack_), + iterate_leaves_only_(other.iterate_leaves_only_) {} + + ShapeTreeIterator& operator++() { + CHECK_NE(nullptr, node_) << "walking off the end() of an iterator!"; + // We're doing a pre-order walk, so if our current node has children take + // the first child. + if (!node_->children.empty()) { + stack_.push_back({node_, /*child-index=*/0}); + node_ = node_->children[0].get(); + if (node_->children.empty() || !iterate_leaves_only_) { + return *this; + } else { + // This is a non-leaf; tail-recurse. + return ++(*this); + } + } + // Otherwise we are currently at a leaf. Walk back up until a node contains + // a child we haven't visited yet. + while (!stack_.empty()) { + node_ = stack_.back().first; + int64 next_child_index = stack_.back().second + 1; + stack_.pop_back(); + if (node_->children.size() > next_child_index) { + stack_.push_back({node_, next_child_index}); + node_ = node_->children[next_child_index].get(); + + if (node_->children.empty() || !iterate_leaves_only_) { + return *this; + } else { + // This is a non-leaf; tail-recurse. + return ++(*this); + } + } + } + // We've walked off the end of the tree. Set node_ to nullptr to signify + // end(). + node_ = nullptr; + current_.reset(); + return *this; + } + ShapeTreeIterator operator++(int) { + auto i = *this; + ++(*this); + return i; + } + bool operator==(const ShapeTreeIterator& other) const { + return node_ == other.node_; + } + bool operator!=(const ShapeTreeIterator& other) const { + return node_ != other.node_; + } + value_type& operator*() { return UpdateCurrent(); } + value_type* operator->() { return &UpdateCurrent(); } - // The XLA shape mirrored in this ShapeTree. Only the root of the - // ShapeTree has this member set. - std::unique_ptr shape_; + private: + // Updates the current_ member to reflect the current state. + value_type& UpdateCurrent() { + ShapeIndex index; + for (auto& node_and_index : stack_) { + index.push_back(node_and_index.second); + } + current_ = MakeUnique(index, node_->data); + return *current_; + } - // The children of this node in the tree. - std::vector> elements_; + // The node to which this iterator is pointing. This is the source of truth in + // the iterator - the stack only exists to facilitate walking back from + // children to parents. + NodeType* node_; + // Stack of {node, child-index} pairs of the path taken from the root to get + // to node_. This allows us to backtrack and know where to go next. + std::vector> stack_; + // True if we should not include interior nodes in our walk. + bool iterate_leaves_only_; + // Placeholder for the current value. Ideally this wouldn't exist and would + // just be an rvalue, but operator -> needs to return a pointer to something. + // We cannot just use a plain old value_type as it contains a reference so + // cannot be default-constructed. + std::unique_ptr current_; }; template -void ShapeTree::BuildTree(const Shape& shape) { +void ShapeTree::InitChildren(const Shape& shape, const T& init_value, + Node* node) { if (ShapeUtil::IsTuple(shape)) { for (int i = 0; i < ShapeUtil::TupleElementCount(shape); ++i) { - elements_.emplace_back(new ShapeTree()); - elements_.back()->BuildTree(shape.tuple_shapes(i)); + node->children.emplace_back(new Node(init_value)); + InitChildren(shape.tuple_shapes(i), init_value, + node->children.back().get()); } } } template -ShapeTree::ShapeTree(const Shape& shape) : shape_(MakeUnique(shape)) { - // The shape_ field is just used to hold the structure of the shape. It should - // not be relied upon to store layout information. - LayoutUtil::ClearLayout(shape_.get()); - BuildTree(*shape_); +void ShapeTree::InitChildren(const Shape& shape, Node* node) { + if (ShapeUtil::IsTuple(shape)) { + for (int i = 0; i < ShapeUtil::TupleElementCount(shape); ++i) { + node->children.emplace_back(new Node()); + InitChildren(shape.tuple_shapes(i), node->children.back().get()); + } + } } template -ShapeTree::ShapeTree(const Shape& shape, const T& init_value) - : shape_(MakeUnique(shape)) { - LayoutUtil::ClearLayout(shape_.get()); - BuildTree(*shape_); - TF_CHECK_OK(ForEachMutableElement( - [&init_value](const ShapeIndex& /*index*/, bool /*is_leaf*/, bool* data) { - *data = init_value; - return tensorflow::Status::OK(); - })); +ShapeTree::ShapeTree(Shape shape) + : root_(), + shape_storage_(MakeUnique(std::move(shape))), + shape_(shape_storage_.get()) { + // The shape_ field is just used to hold the structure of the shape. + // It should not be relied upon to store layout information. + LayoutUtil::ClearLayout(shape_storage_.get()); + InitChildren(*shape_, &root_); } template -ShapeTree::ShapeTree(const ShapeTree& other) - : shape_(MakeUnique(other.shape())) { - LayoutUtil::ClearLayout(shape_.get()); - BuildTree(*shape_); - CopyDataElements(other); +ShapeTree::ShapeTree(const Shape* shape) : root_(), shape_(shape) { + InitChildren(*shape_, &root_); } template -ShapeTree& ShapeTree::operator=(const ShapeTree& other) { - if (this == &other) { - return *this; - } - elements_.clear(); - shape_ = MakeUnique(other.shape()); - LayoutUtil::ClearLayout(shape_.get()); - - BuildTree(*shape_); - CopyDataElements(other); - return *this; +ShapeTree::ShapeTree(Shape shape, const T& init_value) + : root_(init_value), + shape_storage_(MakeUnique(std::move(shape))), + shape_(shape_storage_.get()) { + // The shape_ field is just used to hold the structure of the shape. + // It should not be relied upon to store layout information. + LayoutUtil::ClearLayout(shape_storage_.get()); + InitChildren(*shape_, init_value, &root_); } template -void ShapeTree::CopyDataElements(const ShapeTree& other) { - CHECK(ShapeUtil::Compatible(shape(), other.shape())); - TF_CHECK_OK(ForEachMutableElement( - [&other](const ShapeIndex& index, bool /*is_leaf*/, T* data) { - *data = other.element(index); - return tensorflow::Status::OK(); - })); +ShapeTree::ShapeTree(const Shape* shape, const T& init_value) + : root_(init_value), shape_(shape) { + InitChildren(*shape_, init_value, &root_); } template const T& ShapeTree::element(const ShapeIndex& index) const { - return Lookup(index).data_; + return Lookup(index)->data; } template T* ShapeTree::mutable_element(const ShapeIndex& index) { - return &Lookup(index).data_; + return &Lookup(index)->data; } template -ShapeTree& ShapeTree::Lookup(const ShapeIndex& index) { - ShapeTree* node = this; - for (auto& i : index) { +internal::ShapeTreeNode* ShapeTree::Lookup(const ShapeIndex& index) { + Node* node = &root_; + for (const int64 i : index) { CHECK_GE(i, 0); - CHECK_LT(i, node->elements_.size()); - node = node->elements_[i].get(); + CHECK_LT(i, node->children.size()); + node = node->children[i].get(); } - return *node; + return node; } template -const ShapeTree& ShapeTree::Lookup(const ShapeIndex& index) const { - return const_cast*>(this)->Lookup(index); +const internal::ShapeTreeNode* ShapeTree::Lookup( + const ShapeIndex& index) const { + return const_cast(this)->Lookup(index); } /* static */ template -tensorflow::Status ShapeTree::ForEachHelperMutable( - ShapeIndex* index, ShapeTree* shape_tree, - ShapeTree::MutableVisitorFunction func) { - TF_RETURN_IF_ERROR( - func(*index, shape_tree->elements_.empty(), &shape_tree->data_)); - for (int i = 0; i < shape_tree->elements_.size(); ++i) { +template +Status ShapeTree::ForEachHelper(const Fn& func, const Node& node, + ShapeIndex* index) { + TF_RETURN_IF_ERROR(func(*index, node.data)); + for (int64 i = 0; i < node.children.size(); ++i) { index->push_back(i); - TF_RETURN_IF_ERROR( - ForEachHelperMutable(index, shape_tree->elements_[i].get(), func)); + TF_RETURN_IF_ERROR(ForEachHelper(func, *node.children[i], index)); index->pop_back(); } - - return tensorflow::Status::OK(); + return Status::OK(); } /* static */ template -tensorflow::Status ShapeTree::ForEachHelper( - ShapeIndex* index, const ShapeTree& shape_tree, - ShapeTree::VisitorFunction func) { - TF_RETURN_IF_ERROR( - func(*index, shape_tree.elements_.empty(), shape_tree.data_)); - for (int i = 0; i < shape_tree.elements_.size(); ++i) { +template +Status ShapeTree::ForEachMutableHelper(const Fn& func, Node* node, + ShapeIndex* index) { + TF_RETURN_IF_ERROR(func(*index, &node->data)); + for (int64 i = 0; i < node->children.size(); ++i) { index->push_back(i); - TF_RETURN_IF_ERROR(ForEachHelper(index, *shape_tree.elements_[i], func)); + TF_RETURN_IF_ERROR( + ForEachMutableHelper(func, node->children[i].get(), index)); index->pop_back(); } + return Status::OK(); +} + +template +template +Status ShapeTree::ForEachElementWithStatus(const Fn& func) const { + ShapeIndex index; + return ForEachHelper(func, root_, &index); +} - return tensorflow::Status::OK(); +template +template +Status ShapeTree::ForEachMutableElementWithStatus(const Fn& func) { + ShapeIndex index; + return ForEachMutableHelper(func, &root_, &index); } template -tensorflow::Status ShapeTree::ForEachElement( - ShapeTree::VisitorFunction func) const { +template +void ShapeTree::ForEachElement(const Fn& func) const { ShapeIndex index; - return ForEachHelper(&index, *this, func); + return ForEachHelper( + [&func](const ShapeIndex& index, const T& data) { + func(index, data); + return Status::OK(); + }, + root_, &index) + .IgnoreError(); } template -tensorflow::Status ShapeTree::ForEachMutableElement( - ShapeTree::MutableVisitorFunction func) { +template +void ShapeTree::ForEachMutableElement(const Fn& func) { ShapeIndex index; - return ForEachHelperMutable(&index, this, func); + return ForEachMutableHelper( + [&func](const ShapeIndex& index, T* data) { + func(index, data); + return Status::OK(); + }, + &root_, &index) + .IgnoreError(); +} + +template +void ShapeTree::CopySubtreeFrom(const ShapeTree& other, + const ShapeIndex& source_base_index, + const ShapeIndex& target_base_index) { + CHECK(ShapeUtil::Compatible( + ShapeUtil::GetSubshape(shape(), target_base_index), + ShapeUtil::GetSubshape(other.shape(), source_base_index))); + ForEachMutableElement([this, &other, &source_base_index, &target_base_index]( + const ShapeIndex& index, T* data) { + // Copy the data element only if index is in the + // subtree rooted at target_base_index. + for (int i = 0; i < target_base_index.size(); ++i) { + if (i >= index.size() || index[i] != target_base_index[i]) { + return; + } + } + // Construct source element index to copy from. + ShapeIndex source_index = source_base_index; + for (int i = target_base_index.size(); i < index.size(); ++i) { + source_index.push_back(index[i]); + } + *data = other.element(source_index); + }); +} + +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; + } + }); + return equal; } } // namespace xla diff --git a/tensorflow/compiler/xla/shape_tree_test.cc b/tensorflow/compiler/xla/shape_tree_test.cc index d37f536b755d1feca57360edf950329197ba2dd4..7b4b5cb0fb5e1564ca12ac6e3b901e94ea4c8db6 100644 --- a/tensorflow/compiler/xla/shape_tree_test.cc +++ b/tensorflow/compiler/xla/shape_tree_test.cc @@ -16,8 +16,8 @@ limitations under the License. #include "tensorflow/compiler/xla/shape_tree.h" #include "tensorflow/compiler/xla/shape_util.h" +#include "tensorflow/compiler/xla/test.h" #include "tensorflow/compiler/xla/xla_data.pb.h" -#include "tensorflow/core/platform/test.h" namespace xla { namespace { @@ -35,6 +35,9 @@ class ShapeTreeTest : public ::testing::Test { array_shape_})}); } + void TestShapeConstructor(const Shape& shape, int expected_num_nodes); + void TestInitValueConstructor(const Shape& shape, int expected_num_nodes); + // An array shape (non-tuple). Shape array_shape_; @@ -45,6 +48,73 @@ class ShapeTreeTest : public ::testing::Test { Shape nested_tuple_shape_; }; +TEST_F(ShapeTreeTest, DefaultConstructor) { + ShapeTree int_tree; + EXPECT_TRUE(ShapeUtil::IsNil(int_tree.shape())); + + ShapeTree bool_tree; + EXPECT_TRUE(ShapeUtil::IsNil(bool_tree.shape())); +} + +void ShapeTreeTest::TestShapeConstructor(const Shape& shape, + int expected_num_nodes) { + ShapeTree int_tree(shape); + int num_nodes = 0; + int_tree.ForEachElement([&num_nodes](const ShapeIndex& /*index*/, int data) { + EXPECT_EQ(0, data); + ++num_nodes; + }); + EXPECT_EQ(expected_num_nodes, num_nodes); + + ShapeTree bool_tree(shape); + num_nodes = 0; + bool_tree.ForEachElement( + [&num_nodes](const ShapeIndex& /*index*/, bool data) { + EXPECT_EQ(false, data); + ++num_nodes; + }); + EXPECT_EQ(expected_num_nodes, num_nodes); +} + +TEST_F(ShapeTreeTest, ShapeConstructor) { + TestShapeConstructor(array_shape_, 1); + TestShapeConstructor(tuple_shape_, 4); + TestShapeConstructor(nested_tuple_shape_, 10); +} + +void ShapeTreeTest::TestInitValueConstructor(const Shape& shape, + int expected_num_nodes) { + ShapeTree tree(shape, 42); + int num_nodes = 0; + tree.ForEachElement([&num_nodes](const ShapeIndex& /*index*/, int data) { + EXPECT_EQ(42, data); + ++num_nodes; + }); + EXPECT_EQ(expected_num_nodes, num_nodes); + + num_nodes = 0; + tree.ForEachMutableElement( + [&num_nodes](const ShapeIndex& /*index*/, int* data) { + EXPECT_EQ(42, *data); + *data = num_nodes; + ++num_nodes; + }); + EXPECT_EQ(expected_num_nodes, num_nodes); + + num_nodes = 0; + tree.ForEachElement([&num_nodes](const ShapeIndex& /*index*/, int data) { + EXPECT_EQ(num_nodes, data); + ++num_nodes; + }); + EXPECT_EQ(expected_num_nodes, num_nodes); +} + +TEST_F(ShapeTreeTest, InitValueConstructor) { + TestInitValueConstructor(array_shape_, 1); + TestInitValueConstructor(tuple_shape_, 4); + TestInitValueConstructor(nested_tuple_shape_, 10); +} + TEST_F(ShapeTreeTest, ArrayShape) { ShapeTree shape_tree{array_shape_}; *shape_tree.mutable_element({}) = 42; @@ -57,6 +127,15 @@ TEST_F(ShapeTreeTest, ArrayShape) { // Test the copy constructor. ShapeTree copy{shape_tree}; EXPECT_EQ(123, copy.element({})); + + // Mutate the copy, and ensure the original doesn't change. + *copy.mutable_element({}) = 99; + EXPECT_EQ(99, copy.element({})); + EXPECT_EQ(123, shape_tree.element({})); + + // Test the assignment operator. + copy = shape_tree; + EXPECT_EQ(123, copy.element({})); } TEST_F(ShapeTreeTest, TupleShape) { @@ -74,11 +153,8 @@ TEST_F(ShapeTreeTest, TupleShape) { // Sum all elements in the shape. int sum = 0; - TF_CHECK_OK(shape_tree.ForEachElement( - [&sum](const ShapeIndex& /*index*/, bool /*is_leaf*/, int data) { - sum += data; - return tensorflow::Status::OK(); - })); + shape_tree.ForEachElement( + [&sum](const ShapeIndex& /*index*/, int data) { sum += data; }); EXPECT_EQ(66, sum); // Test the copy constructor. @@ -89,15 +165,23 @@ TEST_F(ShapeTreeTest, TupleShape) { EXPECT_EQ(-100, copy.element({2})); // Write zero to all data elements. - TF_CHECK_OK(shape_tree.ForEachMutableElement( - [&sum](const ShapeIndex& /*index*/, bool /*is_leaf*/, int* data) { - *data = 0; - return tensorflow::Status::OK(); - })); + shape_tree.ForEachMutableElement( + [&sum](const ShapeIndex& /*index*/, int* data) { *data = 0; }); EXPECT_EQ(0, shape_tree.element({})); EXPECT_EQ(0, shape_tree.element({0})); EXPECT_EQ(0, shape_tree.element({1})); EXPECT_EQ(0, shape_tree.element({2})); + EXPECT_EQ(1, copy.element({})); + EXPECT_EQ(42, copy.element({0})); + EXPECT_EQ(123, copy.element({1})); + EXPECT_EQ(-100, copy.element({2})); + + // Test the assignment operator. + copy = shape_tree; + EXPECT_EQ(0, copy.element({})); + EXPECT_EQ(0, copy.element({0})); + EXPECT_EQ(0, copy.element({1})); + EXPECT_EQ(0, copy.element({2})); } TEST_F(ShapeTreeTest, NestedTupleShape) { @@ -116,6 +200,23 @@ TEST_F(ShapeTreeTest, NestedTupleShape) { EXPECT_EQ(42, copy.element({0})); EXPECT_EQ(123, copy.element({1, 1})); EXPECT_EQ(-100, copy.element({2, 0, 1})); + + // Mutate the copy, and ensure the original doesn't change. + *copy.mutable_element({0}) = 1; + *copy.mutable_element({1, 1}) = 2; + *copy.mutable_element({2, 0, 1}) = 3; + EXPECT_EQ(1, copy.element({0})); + EXPECT_EQ(2, copy.element({1, 1})); + EXPECT_EQ(3, copy.element({2, 0, 1})); + EXPECT_EQ(42, shape_tree.element({0})); + EXPECT_EQ(123, shape_tree.element({1, 1})); + EXPECT_EQ(-100, shape_tree.element({2, 0, 1})); + + // Test the assignment operator. + copy = shape_tree; + EXPECT_EQ(42, copy.element({0})); + EXPECT_EQ(123, copy.element({1, 1})); + EXPECT_EQ(-100, copy.element({2, 0, 1})); } TEST_F(ShapeTreeTest, InvalidIndexingTuple) { @@ -130,5 +231,240 @@ TEST_F(ShapeTreeTest, InvalidIndexingNestedTuple) { EXPECT_DEATH(shape_tree.element({0, 0}), ""); } +TEST_F(ShapeTreeTest, ShapeTreeOfNonCopyableType) { + ShapeTree> shape_tree{tuple_shape_}; + EXPECT_EQ(shape_tree.element({2}).get(), nullptr); + *shape_tree.mutable_element({2}) = MakeUnique(42); + EXPECT_EQ(*shape_tree.element({2}), 42); +} + +TEST_F(ShapeTreeTest, CopySubtreeFromArrayShape) { + // Test CopySubtreeFrom method for a single value copied between array-shaped + // ShapeTrees. + ShapeTree source(array_shape_); + *source.mutable_element(/*index=*/{}) = 42; + ShapeTree destination(array_shape_, 123); + + EXPECT_EQ(destination.element(/*index=*/{}), 123); + destination.CopySubtreeFrom(source, /*source_base_index=*/{}, + /*target_base_index=*/{}); + EXPECT_EQ(destination.element(/*index=*/{}), 42); +} + +TEST_F(ShapeTreeTest, FullCopySubtreeFromTupleShape) { + // Test CopySubtreeFrom method for a copy of all elements from one + // tuple-shaped ShapeTree to another. + ShapeTree source(tuple_shape_); + *source.mutable_element(/*index=*/{}) = 10; + *source.mutable_element(/*index=*/{0}) = 11; + *source.mutable_element(/*index=*/{1}) = 12; + *source.mutable_element(/*index=*/{2}) = 13; + + ShapeTree destination(tuple_shape_, 0); + + destination.CopySubtreeFrom(source, /*source_base_index=*/{}, + /*target_base_index=*/{}); + EXPECT_EQ(destination.element(/*index=*/{}), 10); + EXPECT_EQ(destination.element(/*index=*/{0}), 11); + EXPECT_EQ(destination.element(/*index=*/{1}), 12); + EXPECT_EQ(destination.element(/*index=*/{2}), 13); +} + +TEST_F(ShapeTreeTest, SingleElementCopySubtreeFromTupleShape) { + // Test CopySubtreeFrom method for a copy of a single element from one + // tuple-shaped ShapeTree to another. + ShapeTree source(tuple_shape_); + *source.mutable_element(/*index=*/{}) = 10; + *source.mutable_element(/*index=*/{0}) = 11; + *source.mutable_element(/*index=*/{1}) = 12; + *source.mutable_element(/*index=*/{2}) = 13; + + ShapeTree destination(tuple_shape_, 0); + + destination.CopySubtreeFrom(source, /*source_base_index=*/{0}, + /*target_base_index=*/{1}); + EXPECT_EQ(destination.element(/*index=*/{}), 0); + EXPECT_EQ(destination.element(/*index=*/{0}), 0); + EXPECT_EQ(destination.element(/*index=*/{1}), 11); + EXPECT_EQ(destination.element(/*index=*/{2}), 0); +} + +TEST_F(ShapeTreeTest, CopySubtreeIntoNestedShape) { + // Test CopySubtreeFrom method for a copy of a tuple-shaped ShapeTree into a + // nested-tuple-shaped ShapeTree. + ShapeTree source( + ShapeUtil::MakeTupleShape({array_shape_, array_shape_})); + *source.mutable_element(/*index=*/{}) = 10; + *source.mutable_element(/*index=*/{0}) = 11; + *source.mutable_element(/*index=*/{1}) = 12; + + ShapeTree destination(nested_tuple_shape_, 0); + + destination.CopySubtreeFrom(source, /*source_base_index=*/{}, + /*target_base_index=*/{2, 0}); + + EXPECT_EQ(destination.element(/*index=*/{}), 0); + EXPECT_EQ(destination.element(/*index=*/{0}), 0); + EXPECT_EQ(destination.element(/*index=*/{1}), 0); + EXPECT_EQ(destination.element(/*index=*/{1, 0}), 0); + EXPECT_EQ(destination.element(/*index=*/{1, 1}), 0); + EXPECT_EQ(destination.element(/*index=*/{2}), 0); + EXPECT_EQ(destination.element(/*index=*/{2, 0}), 10); + EXPECT_EQ(destination.element(/*index=*/{2, 0, 0}), 11); + EXPECT_EQ(destination.element(/*index=*/{2, 0, 1}), 12); + EXPECT_EQ(destination.element(/*index=*/{2, 1}), 0); +} + +TEST_F(ShapeTreeTest, CopySubtreeFromNestedShape) { + // Test CopySubtreeFrom method for a copy from a nested-tuple-shape. + ShapeTree source(nested_tuple_shape_, 42); + *source.mutable_element(/*index=*/{1}) = 10; + *source.mutable_element(/*index=*/{1, 0}) = 11; + *source.mutable_element(/*index=*/{1, 1}) = 12; + + ShapeTree destination( + ShapeUtil::MakeTupleShape({array_shape_, array_shape_}), 0); + + destination.CopySubtreeFrom(source, /*source_base_index=*/{1}, + /*target_base_index=*/{}); + + EXPECT_EQ(destination.element(/*index=*/{}), 10); + EXPECT_EQ(destination.element(/*index=*/{0}), 11); + EXPECT_EQ(destination.element(/*index=*/{1}), 12); +} + +TEST_F(ShapeTreeTest, OperatorEquals) { + { + ShapeTree a(array_shape_, 123); + ShapeTree b(array_shape_, 42); + ShapeTree c(array_shape_, 42); + EXPECT_FALSE(a == b); + EXPECT_TRUE(a != b); + EXPECT_TRUE(b == c); + } + { + ShapeTree a(tuple_shape_); + *a.mutable_element(/*index=*/{}) = 10; + *a.mutable_element(/*index=*/{0}) = 11; + *a.mutable_element(/*index=*/{1}) = 12; + + ShapeTree b(tuple_shape_); + *b.mutable_element(/*index=*/{}) = 10; + *b.mutable_element(/*index=*/{0}) = 42; + *b.mutable_element(/*index=*/{1}) = 11; + + ShapeTree c(tuple_shape_); + *c.mutable_element(/*index=*/{}) = 10; + *c.mutable_element(/*index=*/{0}) = 42; + *c.mutable_element(/*index=*/{1}) = 11; + + EXPECT_FALSE(a == b); + EXPECT_TRUE(a != b); + EXPECT_TRUE(b == c); + EXPECT_FALSE(b != c); + } +} + +TEST_F(ShapeTreeTest, ConstructWithPointerToShape) { + // Construct a ShapeTree using a pointer to a shape, rather than a reference + // to a shape. This constructor is an optimization to let us avoid + // constructing and destroying temporary shapes when we have many ShapeTrees. + ShapeTree t(&nested_tuple_shape_, 42); + int num_nodes = 0; + t.ForEachElement([&num_nodes](const ShapeIndex& /*index*/, int data) { + EXPECT_EQ(42, data); + ++num_nodes; + }); + EXPECT_EQ(10, num_nodes); +} + +TEST_F(ShapeTreeTest, CopyWithPointerToShape) { + ShapeTree source(&nested_tuple_shape_, 0); + ShapeTree dest(source); + EXPECT_EQ(&dest.shape(), &nested_tuple_shape_); +} + +TEST_F(ShapeTreeTest, CopyAssignWithPointerToShape) { + ShapeTree source(&nested_tuple_shape_, 0); + ShapeTree dest; + dest = source; + EXPECT_EQ(&dest.shape(), &nested_tuple_shape_); +} + +TEST_F(ShapeTreeTest, IterateSimple) { + ShapeTree t(nested_tuple_shape_, 42); + int num_nodes = 0; + for (auto index_to_data : t) { + EXPECT_EQ(42, index_to_data.second); + ++num_nodes; + } + EXPECT_EQ(10, num_nodes); +} + +TEST_F(ShapeTreeTest, ConstIterate) { + const ShapeTree t(nested_tuple_shape_, 42); + int num_nodes = 0; + for (const auto& index_to_data : t) { + EXPECT_EQ(42, index_to_data.second); + ++num_nodes; + } + EXPECT_EQ(10, num_nodes); +} + +TEST_F(ShapeTreeTest, IterateAndMutate) { + ShapeTree t(nested_tuple_shape_, 42); + int i = 0; + for (auto& index_to_data : t) { + EXPECT_EQ(42, index_to_data.second); + if (i == 1) { + index_to_data.second = 98; + } + ++i; + } + t.begin()->second = 78; + EXPECT_EQ(78, t.begin()->second); + i = 0; + for (auto& index_to_data : t) { + if (i == 0) { + EXPECT_EQ(78, index_to_data.second); + } else if (i == 1) { + EXPECT_EQ(98, index_to_data.second); + } else { + EXPECT_EQ(42, index_to_data.second); + } + ++i; + } + EXPECT_EQ(78, t.begin()->second); + EXPECT_EQ(98, std::next(t.begin())->second); +} + +TEST_F(ShapeTreeTest, IterateOrder) { + ShapeTree t(nested_tuple_shape_, 42); + std::vector v; + for (auto& index_to_data : t) { + v.push_back(index_to_data.first); + } + EXPECT_EQ(v, (std::vector{{}, + {0}, + {1}, + {1, 0}, + {1, 1}, + {2}, + {2, 0}, + {2, 0, 0}, + {2, 0, 1}, + {2, 1}})); +} + +TEST_F(ShapeTreeTest, IterateOrderLeaves) { + ShapeTree t(nested_tuple_shape_, 42); + std::vector v; + for (auto& index_to_data : t.leaves()) { + v.push_back(index_to_data.first); + } + EXPECT_EQ(v, (std::vector{ + {0}, {1, 0}, {1, 1}, {2, 0, 0}, {2, 0, 1}, {2, 1}})); +} + } // namespace } // namespace xla diff --git a/tensorflow/compiler/xla/shape_util.cc b/tensorflow/compiler/xla/shape_util.cc index 57d91e4bfc1145faa25c9b5c57422c7653d4a163..b71b3a9e1316522be0490e4c15ae696c1cbc8391 100644 --- a/tensorflow/compiler/xla/shape_util.cc +++ b/tensorflow/compiler/xla/shape_util.cc @@ -18,6 +18,7 @@ limitations under the License. #include #include #include +#include #include #include "tensorflow/compiler/xla/index_util.h" @@ -28,6 +29,7 @@ limitations under the License. #include "tensorflow/compiler/xla/util.h" #include "tensorflow/core/lib/core/errors.h" #include "tensorflow/core/lib/core/stringpiece.h" +#include "tensorflow/core/lib/gtl/optional.h" #include "tensorflow/core/lib/strings/numbers.h" #include "tensorflow/core/lib/strings/str_util.h" #include "tensorflow/core/lib/strings/strcat.h" @@ -103,6 +105,11 @@ bool CompareShapes(const Shape& lhs, const Shape& rhs, bool compare_layouts) { return equal; } +/* static */ int64 ShapeUtil::Rank(const Shape& shape) { + CHECK(!ShapeUtil::IsTuple(shape)) << "Tuples do not have a rank"; + return shape.dimensions_size(); +} + /* static */ int64 ShapeUtil::TrueRank(const Shape& shape) { int64 accum = 0; for (int64 dimension : shape.dimensions()) { @@ -120,7 +127,7 @@ bool CompareShapes(const Shape& lhs, const Shape& rhs, bool compare_layouts) { for (const auto& shape : parameters) { *program_shape.add_parameters() = shape; } - *program_shape.mutable_result() = result; + *program_shape.mutable_result() = std::move(result); return program_shape; } @@ -163,6 +170,7 @@ bool CompareShapes(const Shape& lhs, const Shape& rhs, bool compare_layouts) { } return MakeShapeWithMonotonicDim0MajorLayout(shape.element_type(), dims); } + /* static */ void ShapeUtil::PopulateShape( PrimitiveType element_type, tensorflow::gtl::ArraySlice dimensions, Shape* shape) { @@ -200,7 +208,7 @@ bool CompareShapes(const Shape& lhs, const Shape& rhs, bool compare_layouts) { } /* static */ void ShapeUtil::AppendMajorDimension(int bound, Shape* shape) { - shape->mutable_layout()->add_minor_to_major(ShapeUtil::Rank(*shape)); + shape->mutable_layout()->add_minor_to_major(Rank(*shape)); shape->add_dimensions(bound); TF_DCHECK_OK(ValidateShape(*shape)); } @@ -250,14 +258,6 @@ bool CompareShapes(const Shape& lhs, const Shape& rhs, bool compare_layouts) { return primitive_util::IsFloatingPointType(shape.element_type()); } -/* static */ bool ShapeUtil::IsTuple(const Shape& shape) { - return shape.element_type() == TUPLE; -} - -/* static */ bool ShapeUtil::IsArray(const Shape& shape) { - return !IsTuple(shape) && !IsOpaque(shape); -} - /* static */ bool ShapeUtil::IsNestedTuple(const Shape& shape) { return IsTuple(shape) && std::any_of(shape.tuple_shapes().begin(), shape.tuple_shapes().end(), IsTuple); @@ -268,7 +268,7 @@ bool CompareShapes(const Shape& lhs, const Shape& rhs, bool compare_layouts) { } /* static */ bool ShapeUtil::IsNil(const Shape& shape) { - return IsEmptyTuple(shape) || HasZeroElements(shape); + return IsTuple(shape) ? IsEmptyTuple(shape) : HasZeroElements(shape); } /* static */ int64 ShapeUtil::TupleElementCount(const Shape& shape) { @@ -293,11 +293,7 @@ bool CompareShapes(const Shape& lhs, const Shape& rhs, bool compare_layouts) { std::vector new_elements(tuple.tuple_shapes().begin() + start, tuple.tuple_shapes().begin() + limit); - return ShapeUtil::MakeTupleShape(new_elements); -} - -/* static */ bool ShapeUtil::IsOpaque(const Shape& shape) { - return shape.element_type() == OPAQUE; + return MakeTupleShape(new_elements); } /* static */ bool ShapeUtil::ShapeIs(const Shape& shape, @@ -307,7 +303,7 @@ bool CompareShapes(const Shape& lhs, const Shape& rhs, bool compare_layouts) { if (shape.element_type() != element_type) { return false; } - if (shape.dimensions_size() != ShapeUtil::Rank(shape)) { + if (shape.dimensions_size() != Rank(shape)) { return false; } int64 i = 0; @@ -321,7 +317,8 @@ bool CompareShapes(const Shape& lhs, const Shape& rhs, bool compare_layouts) { } /* static */ int64 ShapeUtil::ElementsIn(const Shape& shape) { - CHECK_EQ(shape.dimensions_size(), ShapeUtil::Rank(shape)); + CHECK(!IsTuple(shape)); + CHECK_EQ(shape.dimensions_size(), Rank(shape)); return std::accumulate( shape.dimensions().begin(), shape.dimensions().end(), 1LL, std::multiplies()); @@ -332,7 +329,7 @@ bool CompareShapes(const Shape& lhs, const Shape& rhs, bool compare_layouts) { } /* static */ bool ShapeUtil::IsScalarF32(const Shape& shape) { - return shape.element_type() == F32 && ShapeUtil::Rank(shape) == 0; + return shape.element_type() == F32 && Rank(shape) == 0; } /* static */ string ShapeUtil::HumanString(const Shape& shape) { @@ -353,6 +350,35 @@ bool CompareShapes(const Shape& lhs, const Shape& rhs, bool compare_layouts) { } } +namespace { + +// Class to memoize the computation of +// tensorflow::str_util::Lowercase(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] = tensorflow::str_util::Lowercase( + 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); +} +} // namespace + /* static */ string ShapeUtil::HumanStringWithLayout(const Shape& shape) { if (shape.element_type() == TUPLE) { string text = "("; @@ -365,18 +391,22 @@ bool CompareShapes(const Shape& lhs, const Shape& rhs, bool compare_layouts) { text += ")"; return text; } else { - string layout; + string result = tensorflow::strings::StrCat( + LowercasePrimitiveTypeName(shape.element_type()), "["); + for (int i = 0; i < shape.dimensions().size(); i++) { + tensorflow::strings::StrAppend(&result, (i > 0) ? "," : "", + shape.dimensions(i)); + } + result += "]"; if (!IsScalar(shape) && !IsOpaque(shape)) { if (LayoutUtil::HasLayout(shape)) { - layout = LayoutUtil::HumanString(shape.layout()); + tensorflow::strings::StrAppend(&result, + LayoutUtil::HumanString(shape.layout())); } else { - layout = "{no layout}"; + tensorflow::strings::StrAppend(&result, "{no layout}"); } } - return tensorflow::strings::StrCat( - tensorflow::str_util::Lowercase( - PrimitiveType_Name(shape.element_type())), - "[", tensorflow::str_util::Join(shape.dimensions(), ","), "]", layout); + return result; } } @@ -395,13 +425,42 @@ bool CompareShapes(const Shape& lhs, const Shape& rhs, bool compare_layouts) { HumanString(program_shape.result())); } -/* static */ StatusOr ShapeUtil::ParseShapeString(const string& s) { +namespace { +// Parses shapes with simple recursive descent structure -- consumes from the +// front of s and passes that view recursively as required. +StatusOr ParseShapeStringInternal(tensorflow::StringPiece* s) { + tensorflow::str_util::RemoveLeadingWhitespace(s); + + if (s->Consume("(")) { // Tuple. + std::vector shapes; + bool must_end = false; + while (true) { + if (s->Consume(")")) { + break; + } else if (must_end) { + return InvalidArgument("Expected end of tuple; got: \"%s\"", + s->ToString().c_str()); + } + shapes.emplace_back(); + TF_ASSIGN_OR_RETURN(shapes.back(), ParseShapeStringInternal(s)); + tensorflow::str_util::RemoveLeadingWhitespace(s); + must_end = !s->Consume(","); + } + return ShapeUtil::MakeTupleShape(shapes); + } + string element_type_string; string dimensions_string; string layout_string; - if (RE2::FullMatch(s, "([fsu]32)\\[([\\d,]*)\\](?: {([\\d,]*)})?", - &element_type_string, &dimensions_string, - &layout_string)) { + // tensorflow::StringPiece is not compatible with internal RE2 StringPiece, so + // we convert in to the RE2-consumable type and then consume the corresponding + // amount from our StringPiece type. + tensorflow::RegexpStringPiece s_consumable(s->data(), s->size()); + if (RE2::Consume(&s_consumable, + "^(\\w*\\d*)\\[([\\d,]*)\\](?:\\s*{([\\d,]*)})?", + &element_type_string, &dimensions_string, &layout_string)) { + size_t consumed = s->size() - s_consumable.size(); + s->remove_prefix(consumed); auto comma_list_to_int64s = [&s](const string& input) -> StatusOr> { std::vector results; @@ -409,40 +468,58 @@ bool CompareShapes(const Shape& lhs, const Shape& rhs, bool compare_layouts) { int64 element; if (!tensorflow::strings::safe_strto64(piece.c_str(), &element)) { return InvalidArgument( - "invalid value in parsed shape string: \"%s\" in \"%s\"", - piece.c_str(), s.c_str()); + "Invalid s64 value in parsed shape string: \"%s\" in \"%s\"", + piece.c_str(), s->ToString().c_str()); } results.push_back(element); } return results; }; + + // Extract the dimensions. TF_ASSIGN_OR_RETURN(std::vector dimensions, comma_list_to_int64s(dimensions_string)); - PrimitiveType primitive_type; - if (element_type_string == "f32") { - primitive_type = F32; - } else if (element_type_string == "s32") { - primitive_type = S32; - } else if (element_type_string == "u32") { - primitive_type = U32; - } else { - LOG(FATAL) << "unhandled element type string: " << element_type_string; + + // Extract the primitive element type. + PrimitiveType primitive_type = PRIMITIVE_TYPE_INVALID; + for (PrimitiveType i = + static_cast(PRIMITIVE_TYPE_INVALID + 1); + i < TUPLE; i = static_cast(i + 1)) { + if (tensorflow::str_util::Lowercase(PrimitiveType_Name(i)) == + element_type_string) { + primitive_type = i; + break; + } + } + if (primitive_type == PRIMITIVE_TYPE_INVALID) { + return InvalidArgument("Invalid element type string: \"%s\".", + element_type_string.c_str()); } + Shape result; if (layout_string.empty()) { + // Create a shape without a layout set. result = ShapeUtil::MakeShape(primitive_type, dimensions); } else { + // Extract the layout minor-to-major and set it. TF_ASSIGN_OR_RETURN(std::vector min2maj, comma_list_to_int64s(layout_string)); TF_RET_CHECK(dimensions.size() == min2maj.size()); result = ShapeUtil::MakeShapeWithLayout(primitive_type, dimensions, min2maj); } - TF_DCHECK_OK(ValidateShape(result)); - return result; + TF_DCHECK_OK(ShapeUtil::ValidateShape(result)); + return std::move(result); } - return InvalidArgument("invalid shape string to parse: \"%s\"", s.c_str()); + return InvalidArgument("Invalid shape string to parse: \"%s\"", + s->ToString().c_str()); +} +} // namespace + +/* static */ StatusOr ShapeUtil::ParseShapeString( + tensorflow::StringPiece s) { + return ParseShapeStringInternal(&s); } /* static */ bool ShapeUtil::SameDimensions(const Shape& lhs, @@ -466,7 +543,7 @@ bool CompareShapes(const Shape& lhs, const Shape& rhs, bool compare_layouts) { /* static */ int64 ShapeUtil::GetDimensionNumber(const Shape& shape, int64 dimension_number) { if (dimension_number < 0) { - dimension_number += ShapeUtil::Rank(shape); + dimension_number += Rank(shape); } CHECK_GE(dimension_number, 0); return dimension_number; @@ -518,7 +595,7 @@ bool CompareShapes(const Shape& lhs, const Shape& rhs, bool compare_layouts) { } int64 allocated_element_count; if (shape.layout().padded_dimensions_size() > 0) { - CHECK_EQ(ShapeUtil::Rank(shape), shape.layout().padded_dimensions_size()); + CHECK_EQ(Rank(shape), shape.layout().padded_dimensions_size()); allocated_element_count = 1; for (int64 dimension_size : shape.layout().padded_dimensions()) { allocated_element_count *= dimension_size; @@ -533,11 +610,6 @@ bool CompareShapes(const Shape& lhs, const Shape& rhs, bool compare_layouts) { /* static */ Status ShapeUtil::ValidateShapeWithOptionalLayoutInternal( const Shape& shape) { if (shape.element_type() == TUPLE) { - // Tuple shape. - if (ShapeUtil::Rank(shape) != 0) { - return InvalidArgument("tuples must be rank-0; got rank %lld", - ShapeUtil::Rank(shape)); - } if (shape.dimensions_size() != 0) { return InvalidArgument("tuples must not have dimensions specified"); } @@ -556,13 +628,13 @@ bool CompareShapes(const Shape& lhs, const Shape& rhs, bool compare_layouts) { return InvalidArgument("shape has invalid element type: %s", shape.ShortDebugString().c_str()); } - if (ShapeUtil::Rank(shape) != shape.dimensions_size()) { + if (Rank(shape) != shape.dimensions_size()) { return InvalidArgument( "shape's rank is mismatched with dimension count; rank=%lld " "dimensions_size=%d", - ShapeUtil::Rank(shape), shape.dimensions_size()); + Rank(shape), shape.dimensions_size()); } - for (int64 i = 0; i < ShapeUtil::Rank(shape); ++i) { + for (int64 i = 0; i < Rank(shape); ++i) { int64 dimension = shape.dimensions(i); if (dimension < 0) { return InvalidArgument( @@ -617,6 +689,11 @@ bool CompareShapes(const Shape& lhs, const Shape& rhs, bool compare_layouts) { return return_shape; } +/* static */ +bool ShapeUtil::IsLeafIndex(const Shape& shape, const ShapeIndex& index) { + return !IsTuple(GetSubshape(shape, index)); +} + /* static */ Shape ShapeUtil::StripDegenerateDimensions(const Shape& shape) { std::vector dimension_sizes; std::vector degenerate_dimensions; @@ -675,7 +752,7 @@ namespace { // Helper for ForEachSubshape which visits the subshapes of the given shape in // DFS pre-order starting with the index. Status ForEachSubshapeHelper(const Shape& shape, - const ShapeUtil::VisitorFunction func, + const ShapeUtil::StatusVisitorFunction& func, ShapeIndex* index) { TF_RETURN_IF_ERROR(func(shape, *index)); if (ShapeUtil::IsTuple(shape)) { @@ -692,7 +769,7 @@ Status ForEachSubshapeHelper(const Shape& shape, // Helper for ForEachMutableSubshape which visits the subshapes of the given // shape in DFS pre-order starting with the index. Status ForEachMutableSubshapeHelper( - Shape* shape, const ShapeUtil::MutatingVisitorFunction func, + Shape* shape, const ShapeUtil::MutatingStatusVisitorFunction& func, ShapeIndex* index) { TF_RETURN_IF_ERROR(func(shape, *index)); if (ShapeUtil::IsTuple(*shape)) { @@ -708,14 +785,40 @@ Status ForEachMutableSubshapeHelper( } // namespace -/* static */ Status ShapeUtil::ForEachSubshape(const Shape& shape, - VisitorFunction func) { +/* static */ void ShapeUtil::ForEachSubshape(const Shape& shape, + const VisitorFunction& func) { + ShapeIndex index; + ForEachSubshapeHelper( + shape, + [&func](const Shape& subshape, const ShapeIndex& index) { + func(subshape, index); + return Status::OK(); + }, + &index) + .IgnoreError(); +} + +/* static */ void ShapeUtil::ForEachMutableSubshape( + Shape* shape, const MutatingVisitorFunction& func) { + ShapeIndex index; + ForEachMutableSubshapeHelper( + shape, + [&func](Shape* subshape, const ShapeIndex& index) { + func(subshape, index); + return Status::OK(); + }, + &index) + .IgnoreError(); +} + +/* static */ Status ShapeUtil::ForEachSubshapeWithStatus( + const Shape& shape, const StatusVisitorFunction& func) { ShapeIndex index; return ForEachSubshapeHelper(shape, func, &index); } -/* static */ Status ShapeUtil::ForEachMutableSubshape( - Shape* shape, MutatingVisitorFunction func) { +/* static */ Status ShapeUtil::ForEachMutableSubshapeWithStatus( + Shape* shape, const MutatingStatusVisitorFunction& func) { ShapeIndex index; return ForEachMutableSubshapeHelper(shape, func, &index); } @@ -728,9 +831,17 @@ Status ForEachMutableSubshapeHelper( new_shape.add_dimensions(dim); } if (shape.has_layout()) { - new_shape.mutable_layout()->clear_minor_to_major(); + Layout* new_layout = new_shape.mutable_layout(); + new_layout->clear_minor_to_major(); for (auto index : Permute(permutation, shape.layout().minor_to_major())) { - new_shape.mutable_layout()->add_minor_to_major(index); + 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); + } } } return new_shape; @@ -747,27 +858,28 @@ ShapeUtil::InsertedOrDeleted1SizedDimensions(const Shape& shape_pre, // and unmodified_dim_pair have size >1. Otherwise, returns true and appends // the degerenate input/output dimensions in the gap to // deleted_indices/inserted_indices respectively. - auto check_modified_dims = [&shape_pre, &shape_post, &deleted_indices, - &inserted_indices]( - std::pair prior_unmodified_dim_pair, - std::pair unmodified_dim_pair) { - for (int64 modified_input_dim = prior_unmodified_dim_pair.first + 1; - modified_input_dim < unmodified_dim_pair.first; ++modified_input_dim) { - if (shape_pre.dimensions(modified_input_dim) > 1) { - return false; - } - deleted_indices.push_back(modified_input_dim); - } - for (int64 modified_output_dim = prior_unmodified_dim_pair.second + 1; - modified_output_dim < unmodified_dim_pair.second; - ++modified_output_dim) { - if (shape_post.dimensions(modified_output_dim) > 1) { - return false; - } - inserted_indices.push_back(modified_output_dim); - } - return true; - }; + auto check_modified_dims = + [&shape_pre, &shape_post, &deleted_indices, &inserted_indices]( + std::pair prior_unmodified_dim_pair, + std::pair unmodified_dim_pair) { + for (int64 modified_input_dim = prior_unmodified_dim_pair.first + 1; + modified_input_dim < unmodified_dim_pair.first; + ++modified_input_dim) { + if (shape_pre.dimensions(modified_input_dim) > 1) { + return false; + } + deleted_indices.push_back(modified_input_dim); + } + for (int64 modified_output_dim = prior_unmodified_dim_pair.second + 1; + modified_output_dim < unmodified_dim_pair.second; + ++modified_output_dim) { + if (shape_post.dimensions(modified_output_dim) > 1) { + return false; + } + inserted_indices.push_back(modified_output_dim); + } + return true; + }; std::vector> unmodified_dims = DimensionsUnmodifiedByReshape(shape_pre, shape_post); @@ -783,8 +895,7 @@ ShapeUtil::InsertedOrDeleted1SizedDimensions(const Shape& shape_pre, auto unmodified_dim_pair = i < unmodified_dims.size() ? unmodified_dims[i] - : std::make_pair(ShapeUtil::Rank(shape_pre), - ShapeUtil::Rank(shape_post)); + : std::make_pair(Rank(shape_pre), Rank(shape_post)); if (!check_modified_dims(prior_unmodified_dim_pair, unmodified_dim_pair)) { return nil; } @@ -859,9 +970,8 @@ ShapeUtil::DimensionsUnmodifiedByReshape(const Shape& input_shape, return false; } - CHECK_EQ(ShapeUtil::ElementsIn(input_shape), - ShapeUtil::ElementsIn(output_shape)); - if (ShapeUtil::ElementsIn(input_shape) == 0) { + CHECK_EQ(ElementsIn(input_shape), ElementsIn(output_shape)); + if (ElementsIn(input_shape) == 0) { return true; } @@ -975,21 +1085,17 @@ ShapeUtil::DimensionsUnmodifiedByReshape(const Shape& input_shape, // as input_shape/output_shape and the dimension-0-major layout. These two // shapes are used for conversion between logical linear indices and // multi-dimensional indices. - Shape input_shape_dim0_major = - ShapeUtil::MakeShapeWithMonotonicDim0MajorLayout( - input_shape.element_type(), AsInt64Slice(input_shape.dimensions())); - Shape output_shape_dim0_major = - ShapeUtil::MakeShapeWithMonotonicDim0MajorLayout( - output_shape.element_type(), - AsInt64Slice(output_shape.dimensions())); - - for (int64 input_dim = 0; input_dim < ShapeUtil::Rank(input_shape); - ++input_dim) { + Shape input_shape_dim0_major = MakeShapeWithMonotonicDim0MajorLayout( + input_shape.element_type(), AsInt64Slice(input_shape.dimensions())); + Shape output_shape_dim0_major = MakeShapeWithMonotonicDim0MajorLayout( + output_shape.element_type(), AsInt64Slice(output_shape.dimensions())); + + for (int64 input_dim = 0; input_dim < Rank(input_shape); ++input_dim) { if (input_shape.dimensions(input_dim) <= 1) { continue; } - std::vector input_unit_index(ShapeUtil::Rank(input_shape), 0); + std::vector input_unit_index(Rank(input_shape), 0); input_unit_index[input_dim] = 1; int64 logical_linear_index = IndexUtil::MultidimensionalIndexToLinearIndex(input_shape_dim0_major, @@ -1013,6 +1119,140 @@ ShapeUtil::DimensionsUnmodifiedByReshape(const Shape& input_shape, check_input_unit_indices(output_shape, input_shape); } +/* static */ tensorflow::gtl::optional ShapeUtil::AlignLayouts( + const Shape& input_shape, const Shape& output_shape) { + int64 input_rank = Rank(input_shape); + int64 output_rank = Rank(output_shape); + + // First, calculate an alignment of the dimensions. A consecutive sequence of + // input dimensions and output dimensions belong to the same alignment part if + // the products of their dimension bounds are the same. In the easiest case, + // an alignment part consists of one input dimension and one output dimension + // which both have the same dimension bound. An alignment part specifies which + // dimensions need to be kept together in a physical layout if we want a + // reshape to be a bitcast. The order of the alignment parts is defined by the + // physical layout of the input shape, so when we construct the layout for the + // output shape we just process the alignment parts in this order, and then + // layout the dimensions belonging to each part in descending (major to minor) + // order. + + // Stores the input and output dimension numbers where each alignment part + // starts. + std::vector> alignment; + alignment.push_back({0, 0}); + + // Stores a mapping from the input dimension to the alignment part it belongs + // to. + std::vector dimension_to_alignment_index(input_rank); + int64 input_dimension_product = 1, output_dimension_product = 1; + for (int64 i = 0, j = 0; i < input_rank || j < output_rank;) { + // Check if we have reached the end of an alignment part. + if (input_dimension_product == output_dimension_product && + input_dimension_product > 1) { + alignment.push_back({i, j}); + input_dimension_product = output_dimension_product = 1; + } + if (input_dimension_product < output_dimension_product || + j == output_rank) { + if (i == input_rank) { + return tensorflow::gtl::nullopt; + } + dimension_to_alignment_index[i] = alignment.size() - 1; + input_dimension_product *= input_shape.dimensions(i); + ++i; + } else { + output_dimension_product *= output_shape.dimensions(j); + ++j; + } + } + if (input_dimension_product != output_dimension_product) { + return tensorflow::gtl::nullopt; + } + // We also need to store an end element so that we know where the last + // alignment part ends. + alignment.push_back({input_rank, output_rank}); + + // Now check if the physical layout can potentially be aligned to the output + // shape by changing the physical layout of the output shape. We need to check + // that all dimension numbers that belong to the same alignment part appear + // consecutively, and are in descending order. However we can ignore any + // trivial dimension bounds of 1, because they can be placed anywhere. + auto input_dimension_numbers = input_shape.layout().minor_to_major(); + std::vector output_layout; + output_layout.reserve(output_rank); + for (int64 i = 0; i < input_rank;) { + int64 current_dimension_number = input_dimension_numbers[i]; + + // Skip trivial dimensions with a bound of 1. + if (input_shape.dimensions(current_dimension_number) == 1) { + ++i; + continue; + } + + // Calculate the number of non-trivial dimension bounds in the input shape + // belonging to the current alignment part. + const int64 current_alignment_index = + dimension_to_alignment_index[current_dimension_number]; + // Because of the special end element that we added, we can be sure that + // 'current_alignment_index' is < alignment.size() - 1. + CHECK_LT(current_alignment_index, alignment.size() - 1); + int64 num_non_trivial_dimensions_in_alignment_part = 0; + for (int64 j = alignment[current_alignment_index].first; + j < alignment[current_alignment_index + 1].first; ++j) { + if (input_shape.dimensions(j) != 1) { + ++num_non_trivial_dimensions_in_alignment_part; + } + } + + // Check that the following 'num_non_trivial_dimensions_in_alignment_part' + // dimension numbers (ignoring dimension numbers with dimension bound 1) are + // in descending order and belong to the current alignment part. + for (int64 j = 0; j < num_non_trivial_dimensions_in_alignment_part; + ++i, ++j) { + if (i == input_rank) { + return tensorflow::gtl::nullopt; + } + // Skip trivial dimensions with a bound of 1. + if (input_shape.dimensions(input_dimension_numbers[i]) == 1) { + --j; + continue; + } + // If the current dimension number belongs to a different alignment part, + // or the dimension numbers are not in descending order, we can return + // early. + if (dimension_to_alignment_index[input_dimension_numbers[i]] != + current_alignment_index || + input_dimension_numbers[i] > current_dimension_number) { + return tensorflow::gtl::nullopt; + } + current_dimension_number = input_dimension_numbers[i]; + } + + // The output dimension numbers that belong to the current alignment part + // need to appear in the same descending order as in the input. Again, we + // can skip dimensions with a bound of 1. + for (int64 j = alignment[current_alignment_index + 1].second - 1; + j >= alignment[current_alignment_index].second; --j) { + if (output_shape.dimensions(j) != 1) { + output_layout.push_back(j); + } + } + } + // Now add all the dimensions with dimension bound 1 at the end of + // 'output_layout'. + for (int64 i = 0; i < output_rank; ++i) { + if (output_shape.dimensions(i) == 1) { + output_layout.push_back(i); + } + } + CHECK_EQ(output_layout.size(), output_rank); + Shape output_shape_with_layout = MakeShapeWithLayout( + output_shape.element_type(), AsInt64Slice(output_shape.dimensions()), + output_layout); + CHECK(ReshapeIsBitcast(input_shape, output_shape_with_layout)); + return output_shape_with_layout; +} + /* static */ Shape ShapeUtil::DeleteDimension(int64 dim_to_delete, Shape shape) { shape.mutable_dimensions()->erase(shape.dimensions().begin() + dim_to_delete); diff --git a/tensorflow/compiler/xla/shape_util.h b/tensorflow/compiler/xla/shape_util.h index 68e138e6aca9d2cf157466eca1ea6960e3c448e8..e347313837622c4fceb850a34028ab530233524d 100644 --- a/tensorflow/compiler/xla/shape_util.h +++ b/tensorflow/compiler/xla/shape_util.h @@ -26,6 +26,7 @@ limitations under the License. #include "tensorflow/compiler/xla/types.h" #include "tensorflow/compiler/xla/xla_data.pb.h" #include "tensorflow/core/lib/gtl/array_slice.h" +#include "tensorflow/core/lib/gtl/optional.h" #include "tensorflow/core/platform/macros.h" #include "tensorflow/core/platform/types.h" @@ -92,6 +93,7 @@ class ShapeUtil { public: // Returns the number of elements are contained within the provided shape; // e.g. for rank 0 (scalars) the result is always 1. + // Precondition: !IsTuple(shape) static int64 ElementsIn(const Shape& shape); // Returns true if 'shape' has zero elements. @@ -123,7 +125,7 @@ class ShapeUtil { // Parses a ShapeUtil::HumanString-format shape string back into a shape // object. - static StatusOr ParseShapeString(const string& s); + static StatusOr ParseShapeString(tensorflow::StringPiece s); // Returns whether the LHS and RHS shapes have the same dimensions; note: does // not check element type. @@ -143,7 +145,8 @@ class ShapeUtil { static bool Equal(const Shape& lhs, const Shape& rhs); // Returns the rank (number of dimensions) of the given shape. - static int64 Rank(const Shape& shape) { return shape.dimensions_size(); } + // 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 @@ -254,13 +257,20 @@ class ShapeUtil { static bool ElementIsSigned(const Shape& shape); // Returns whether the shape is a tuple. - static bool IsTuple(const Shape& shape); + static bool IsTuple(const Shape& shape) { + return shape.element_type() == TUPLE; + } - // Returns whether the shape is an array. - static bool IsArray(const Shape& shape); + // 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 opaque. - static bool IsOpaque(const Shape& shape); + // Returns whether the shape is an array. + static bool IsArray(const Shape& shape) { + return !IsTuple(shape) && !IsOpaque(shape); + } // Returns whether the shape is a tuple with at least one element which is // also a tuple. @@ -294,18 +304,31 @@ class ShapeUtil { static const Shape& GetSubshape(const Shape& shape, const ShapeIndex& index); static Shape* GetMutableSubshape(Shape* shape, const ShapeIndex& index); - // Calls the given visitor function for each subshape of the given shape. - // Returns early if an error status is returned. Subshapes are visited in DFS - // pre-order starting with the entire shape (index {}). - using VisitorFunction = std::function; - static Status ForEachSubshape(const Shape& shape, VisitorFunction func); + // Returns whether the given index in the given shape is a leaf element of the + // shape. + static bool IsLeafIndex(const Shape& shape, const ShapeIndex& index); - // Mutating variant of ForEachSubshape. + // Calls the given visitor function for each subshape of the given shape. + // Subshapes are visited in DFS pre-order starting with the entire shape + // (index {}). + using VisitorFunction = std::function; + static void ForEachSubshape(const Shape& shape, const VisitorFunction& func); using MutatingVisitorFunction = + std::function; + static void ForEachMutableSubshape(Shape* shape, + const MutatingVisitorFunction& func); + + // Variants of ForEach(Mutable)Subshape which propagate Status from the + // visitor function. + using StatusVisitorFunction = std::function; + static Status ForEachSubshapeWithStatus(const Shape& shape, + const StatusVisitorFunction& func); + using MutatingStatusVisitorFunction = std::function; - static Status ForEachMutableSubshape(Shape* shape, - MutatingVisitorFunction func); + static Status ForEachMutableSubshapeWithStatus( + Shape* shape, const MutatingStatusVisitorFunction& func); // Removes all degenerate dimensions (size one) from the given shape. The // stripped minor_to_major preserves the relative ordering of non-degenerate @@ -377,6 +400,15 @@ class ShapeUtil { static bool ReshapeIsBitcast(const Shape& input_shape, const Shape& output_shape); + // Find a physical layout for 'output_shape' such that + // ShapeUtil::ReshapeIsBitcast(input_shape, output_shape_with_layout) returns + // true (where 'output_shape_with_layout' is 'output_shape' with the found + // layout). The layout of 'input_shape' is kept fixed. Returns + // 'output_shape_with_layout' if such a layout can be found, and an error + // otherwise. + static tensorflow::gtl::optional AlignLayouts( + const Shape& input_shape, const Shape& output_shape); + // Returns a shape with the given dimension deleted. // For example: // • `DeleteDimension(1, T[m, n, k]) = T[m, k]` @@ -390,6 +422,46 @@ class ShapeUtil { static Shape FilterDimensions(const std::function& p, Shape shape); + // Iterates through all the shape indexes, in minor to major order, starting + // from the base indexes, incrementing by the incr steps, up to count + // (index[i] < base[i] + count[i]), and calls the visitor_function with the + // current index. + // The visitor_function visitor function should return true if it wants to + // continue, or false otherwise. + // + // visitor_function must be a callable of type bool(const std::vector&) + // or compatible. + template + static void ForEachIndex(const Shape& shape, + tensorflow::gtl::ArraySlice base, + tensorflow::gtl::ArraySlice count, + tensorflow::gtl::ArraySlice incr, + const FnType& visitor_function) { + if (ShapeUtil::HasZeroElements(shape)) { + return; + } + CHECK_EQ(Rank(shape), base.size()); + CHECK_EQ(incr.size(), base.size()); + CHECK_EQ(count.size(), base.size()); + const Layout& layout = shape.layout(); + const int64 rank = layout.minor_to_major_size(); + // Allows handling R0 arrays, such that the visitor function will be called + // once with the proper empty indexes. + int64 n = -1; + std::vector indexes(base.begin(), base.end()); + while (n < rank && visitor_function(indexes)) { + // Increments dimensions in minor to major order. + for (n = 0; n < rank; ++n) { + int64 dim = layout.minor_to_major(n); + indexes[dim] += incr[dim]; + if (indexes[dim] < base[dim] + count[dim]) { + break; + } + indexes[dim] = base[dim]; + } + } + } + private: // Validates all of the non-layout properties of the shape -- this is a helper // used by both the layout-optional and layout-required public method. diff --git a/tensorflow/compiler/xla/shape_util_test.cc b/tensorflow/compiler/xla/shape_util_test.cc index 9e6b243611b57d38339a8f6460c655255f60899d..9635e5ad2eb704a08423e31cd53ef5b79d2558fd 100644 --- a/tensorflow/compiler/xla/shape_util_test.cc +++ b/tensorflow/compiler/xla/shape_util_test.cc @@ -16,14 +16,17 @@ limitations under the License. #include "tensorflow/compiler/xla/shape_util.h" #include "tensorflow/compiler/xla/layout_util.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/core/platform/test.h" +#include "tensorflow/compiler/xla/xla_data.pb.h" namespace xla { namespace { +using ::testing::ElementsAre; + TEST(ShapeUtilTest, GetDimensionHelperCanNegativeIndex) { Shape matrix = ShapeUtil::MakeShape(F32, {2, 3}); EXPECT_EQ(3, ShapeUtil::GetDimension(matrix, -1)); @@ -75,6 +78,30 @@ TEST(ShapeUtilTest, ParseShapeStringR2F32) { << "actual: " << ShapeUtil::HumanString(actual); } +TEST(ShapeUtilTest, ParseShapeStringTupleOfArrays) { + string shape_string = "(f32[1572864],s8[5120,1024])"; + Shape actual = ShapeUtil::ParseShapeString(shape_string).ValueOrDie(); + 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]), f32[3])"; + Shape actual = ShapeUtil::ParseShapeString(shape_string).ValueOrDie(); + Shape expected = ShapeUtil::MakeTupleShape({ + ShapeUtil::MakeShape(F32, {1}), + ShapeUtil::MakeTupleShape({ShapeUtil::MakeShape(F32, {2})}), + ShapeUtil::MakeShape(F32, {3}), + }); + ASSERT_TRUE(ShapeUtil::Equal(expected, actual)) + << "expected: " << ShapeUtil::HumanString(expected) + << "actual: " << ShapeUtil::HumanString(actual); +} + TEST(ShapeUtilTest, CompatibleIdenticalShapes) { Shape shape1 = ShapeUtil::MakeShape(F32, {3, 2}); Shape shape2 = ShapeUtil::MakeShape(F32, {3, 2}); @@ -319,6 +346,30 @@ TEST(ShapeUtilTest, GetSubshape) { ShapeUtil::GetSubshape(nested_tuple_shape, {2, 0}))); } +TEST(ShapeUtilTest, IsLeafIndex) { + // Test array shape. + Shape array_shape = ShapeUtil::MakeShape(F32, {42, 42, 123}); + EXPECT_TRUE(ShapeUtil::IsLeafIndex(array_shape, {})); + + // Test tuple shape. + Shape tuple_shape = ShapeUtil::MakeTupleShape({array_shape, array_shape}); + EXPECT_FALSE(ShapeUtil::IsLeafIndex(tuple_shape, {})); + EXPECT_TRUE(ShapeUtil::IsLeafIndex(tuple_shape, {0})); + EXPECT_TRUE(ShapeUtil::IsLeafIndex(tuple_shape, {1})); + + // Test nested tuple shape. + Shape nested_tuple_shape = ShapeUtil::MakeTupleShape( + {array_shape, ShapeUtil::MakeTupleShape({array_shape, array_shape}), + ShapeUtil::MakeTupleShape( + {ShapeUtil::MakeTupleShape({array_shape, array_shape}), + array_shape})}); + EXPECT_FALSE(ShapeUtil::IsLeafIndex(nested_tuple_shape, {})); + EXPECT_TRUE(ShapeUtil::IsLeafIndex(nested_tuple_shape, {0})); + EXPECT_FALSE(ShapeUtil::IsLeafIndex(nested_tuple_shape, {1})); + EXPECT_TRUE(ShapeUtil::IsLeafIndex(nested_tuple_shape, {1, 0})); + EXPECT_TRUE(ShapeUtil::IsLeafIndex(nested_tuple_shape, {1, 1})); +} + TEST(ShapeUtilTest, HumanString) { Shape opaque = ShapeUtil::MakeOpaqueShape(); Shape scalar = ShapeUtil::MakeShape(F32, {}); @@ -377,13 +428,12 @@ TEST(ShapeUtilTest, HumanString) { TEST(ShapeUtilTest, ForEachSubshapeArray) { const Shape shape = ShapeUtil::MakeShape(F32, {2, 3}); int calls = 0; - EXPECT_IS_OK(ShapeUtil::ForEachSubshape( + ShapeUtil::ForEachSubshape( shape, [&calls, &shape](const Shape& subshape, const ShapeIndex& index) { EXPECT_EQ(&shape, &subshape); EXPECT_TRUE(index.empty()); ++calls; - return tensorflow::Status::OK(); - })); + }); EXPECT_EQ(1, calls); } @@ -393,7 +443,7 @@ TEST(ShapeUtilTest, ForEachSubshapeNestedTuple) { ShapeUtil::MakeTupleShape({ShapeUtil::MakeShape(F32, {101}), ShapeUtil::MakeShape(PRED, {33})})}); int calls = 0; - EXPECT_IS_OK(ShapeUtil::ForEachSubshape( + ShapeUtil::ForEachSubshape( shape, [&calls, &shape](const Shape& subshape, const ShapeIndex& index) { EXPECT_TRUE( ShapeUtil::Equal(subshape, ShapeUtil::GetSubshape(shape, index))); @@ -405,8 +455,7 @@ TEST(ShapeUtilTest, ForEachSubshapeNestedTuple) { EXPECT_EQ(33, ShapeUtil::ElementsIn(subshape)); } ++calls; - return tensorflow::Status::OK(); - })); + }); EXPECT_EQ(5, calls); } @@ -416,7 +465,7 @@ TEST(ShapeUtilTest, ForEachMutableSubshapeNestedTuple) { ShapeUtil::MakeTupleShape({ShapeUtil::MakeShape(F32, {101}), ShapeUtil::MakeShape(PRED, {33})})}); int calls = 0; - EXPECT_IS_OK(ShapeUtil::ForEachMutableSubshape( + ShapeUtil::ForEachMutableSubshape( &shape, [&calls, &shape](const Shape* subshape, const ShapeIndex& index) { // Pointer values should be equal EXPECT_EQ(subshape, ShapeUtil::GetMutableSubshape(&shape, index)); @@ -428,8 +477,7 @@ TEST(ShapeUtilTest, ForEachMutableSubshapeNestedTuple) { EXPECT_EQ(33, ShapeUtil::ElementsIn(*subshape)); } ++calls; - return tensorflow::Status::OK(); - })); + }); EXPECT_EQ(5, calls); } @@ -443,24 +491,52 @@ TEST(ShapeUtilTest, InsertedOrDeleted1SizedDimensions) { ShapeUtil::InsertedOrDeleted1SizedDimensions(shape0, shape2))); } +TEST(ShapeUtilTest, ForEachIndex) { + struct ShapeDimensionAndNumberInvocations { + std::vector dimensions; + int invocations; + } test_data[] = { + {{}, 1}, {{0}, 0}, {{16}, 16}, {{3, 0}, 0}, + {{0, 2}, 0}, {{4, 16}, 64}, {{6, 11, 17}, 1122}, {{6, 11, 5, 17}, 5610}, + }; + + for (const auto& data : test_data) { + Shape shape = ShapeUtil::MakeShape(F32, data.dimensions); + // Increments at every invocation. + int invocations = 0; + auto increment_func = [&invocations](const std::vector& indexes) { + invocations++; + return true; + }; + + std::vector zero_base(data.dimensions.size(), 0); + std::vector step(data.dimensions.size(), 1); + + ShapeUtil::ForEachIndex(shape, zero_base, data.dimensions, step, + increment_func); + + EXPECT_EQ(invocations, data.invocations); + } +} + TEST(ShapeUtilTest, DimensionsUnmodifiedByReshape_1x1x1x1_to_1x1x1) { // All output dimensions should be unmodified. One of the input dimensions is // modified because the input rank is larger by one. - EXPECT_EQ(3, - ShapeUtil::DimensionsUnmodifiedByReshape( - ShapeUtil::MakeShape(S32, {1, 1, 1, 1}), - ShapeUtil::MakeShape(S32, {1, 1, 1})) - .size()); + EXPECT_THAT(ShapeUtil::DimensionsUnmodifiedByReshape( + ShapeUtil::MakeShape(S32, {1, 1, 1, 1}), + ShapeUtil::MakeShape(S32, {1, 1, 1})), + ElementsAre(std::make_pair(0, 0), std::make_pair(1, 1), + std::make_pair(2, 2))); } TEST(ShapeUtilTest, DimensionsUnmodifiedByReshape_1x1x1_to_1x1x1x1) { // All input dimensions should be unmodified. One of the output dimensions is // modified because the output rank is larger by one. - EXPECT_EQ(3, - ShapeUtil::DimensionsUnmodifiedByReshape( - ShapeUtil::MakeShape(S32, {1, 1, 1}), - ShapeUtil::MakeShape(S32, {1, 1, 1, 1})) - .size()); + EXPECT_THAT(ShapeUtil::DimensionsUnmodifiedByReshape( + ShapeUtil::MakeShape(S32, {1, 1, 1}), + ShapeUtil::MakeShape(S32, {1, 1, 1, 1})), + ElementsAre(std::make_pair(0, 0), std::make_pair(1, 1), + std::make_pair(2, 2))); } TEST(ShapeUtilTest, DimensionsUnmodifiedByReshape_4x1x3x5x6x7_to_2x6x1x5x1x42) { @@ -468,11 +544,10 @@ TEST(ShapeUtilTest, DimensionsUnmodifiedByReshape_4x1x3x5x6x7_to_2x6x1x5x1x42) { // 4, 1, 3, 5, 6, 7 // | // 2, 6, 1, 5, 1, 42 - EXPECT_TRUE( - ContainersEqual(ShapeUtil::DimensionsUnmodifiedByReshape( - ShapeUtil::MakeShape(S32, {4, 1, 3, 5, 6, 7}), - ShapeUtil::MakeShape(S32, {2, 6, 1, 5, 1, 42})), - std::vector>({{3, 3}}))); + EXPECT_THAT(ShapeUtil::DimensionsUnmodifiedByReshape( + ShapeUtil::MakeShape(S32, {4, 1, 3, 5, 6, 7}), + ShapeUtil::MakeShape(S32, {2, 6, 1, 5, 1, 42})), + ElementsAre(std::make_pair(3, 3))); } TEST(ShapeUtilTest, ReshapeIsBitcast_3x4_6x2) { @@ -521,5 +596,58 @@ TEST(AlgebraicSimplifierTest, ReshapeIsBitcast_3x2x2_6x2_Dim0IsMostMinor) { ShapeUtil::MakeShapeWithLayout(F32, {6, 2}, {0, 1}))); } +TEST(AlignmentTest, AlignLayoutsWithoutTrivialDimensions) { + Shape input = ShapeUtil::MakeShapeWithLayout(xla::F32, {3, 8, 5, 7, 11}, + {3, 2, 1, 0, 4}); + auto aligned_shape = ShapeUtil::AlignLayouts( + input, ShapeUtil::MakeShape(xla::F32, {4, 3, 2, 7, 5, 11})); + EXPECT_TRUE(aligned_shape); + EXPECT_THAT(aligned_shape.value().layout().minor_to_major(), + ElementsAre(4, 3, 2, 1, 0, 5)); + EXPECT_TRUE(ShapeUtil::ReshapeIsBitcast(input, aligned_shape.value())); + + aligned_shape = ShapeUtil::AlignLayouts( + input, ShapeUtil::MakeShape(xla::F32, {3, 2, 4, 35, 11})); + EXPECT_TRUE(aligned_shape); + EXPECT_THAT(aligned_shape.value().layout().minor_to_major(), + ElementsAre(3, 2, 1, 0, 4)); + EXPECT_TRUE(ShapeUtil::ReshapeIsBitcast(input, aligned_shape.value())); +} + +TEST(AlignmentTest, AlignLayoutsWithTrivialDimensions) { + Shape input = + ShapeUtil::MakeShapeWithLayout(xla::F32, {1, 3, 8, 1, 5, 7, 1, 11, 1, 1}, + {5, 0, 4, 2, 1, 3, 6, 7, 9, 8}); + auto aligned_shape = ShapeUtil::AlignLayouts( + input, ShapeUtil::MakeShape(xla::F32, {1, 4, 1, 3, 2, 7, 5, 11, 1})); + EXPECT_TRUE(aligned_shape); + EXPECT_THAT(aligned_shape.value().layout().minor_to_major(), + ElementsAre(6, 5, 4, 3, 1, 7, 0, 2, 8)); + EXPECT_TRUE(ShapeUtil::ReshapeIsBitcast(input, aligned_shape.value())); +} + +// A test case where the consecutive elements of the input shape belonging to +// the same layout part are not in descending order. +TEST(AlignmentTest, AlignLayoutsWithoutTrivialDimensionsWrongInputLayout) { + // Same physical layout as in AlignLayoutsWithoutTrivialDimensions, except + // that the first two dimension numbers are exchanged. + Shape input = ShapeUtil::MakeShapeWithLayout(xla::F32, {3, 8, 5, 7, 11}, + {2, 3, 1, 0, 4}); + auto aligned_shape = ShapeUtil::AlignLayouts( + input, ShapeUtil::MakeShape(xla::F32, {4, 3, 2, 7, 5, 11})); + EXPECT_FALSE(aligned_shape); +} + +// A test case where the physical layout of the input shape does not place all +// dimensions that belong to the same alignment part consecutively. +TEST(AlignmentTest, + AlignLayoutsWithoutTrivialDimensionsNonConsecutiveAlignmentPart) { + Shape input = ShapeUtil::MakeShapeWithLayout(xla::F32, {3, 8, 5, 7, 11}, + {3, 2, 1, 0, 4}); + auto aligned_shape = ShapeUtil::AlignLayouts( + input, ShapeUtil::MakeShape(xla::F32, {4, 3, 2, 5, 77})); + EXPECT_FALSE(aligned_shape); +} + } // namespace } // namespace xla diff --git a/tensorflow/compiler/xla/status_macros.h b/tensorflow/compiler/xla/status_macros.h index aa12cda666c4abfbf7ec38f0aa640df3b51ea106..5e5550563d02de99ddefbeb8ee8e1bf98afdcdbf 100644 --- a/tensorflow/compiler/xla/status_macros.h +++ b/tensorflow/compiler/xla/status_macros.h @@ -183,15 +183,15 @@ class StatusAdaptorForMacros { .with_log_stack_trace() \ .add_ret_check_failure(#condition) -#define TF_ASSIGN_OR_ASSERT_OK(lhs, rexpr) \ - TF_ASSIGN_OR_ASSERT_OK_IMPL( \ +#define TF_ASSERT_OK_AND_ASSIGN(lhs, rexpr) \ + TF_ASSERT_OK_AND_ASSIGN_IMPL( \ TF_STATUS_MACROS_CONCAT_NAME(_status_or_value, __COUNTER__), lhs, \ rexpr); -#define TF_ASSIGN_OR_ASSERT_OK_IMPL(statusor, lhs, rexpr) \ +#define TF_ASSERT_OK_AND_ASSIGN_IMPL(statusor, lhs, rexpr) \ auto statusor = (rexpr); \ ASSERT_TRUE(statusor.status().ok()) << statusor.status(); \ - lhs = statusor.ConsumeValueOrDie() + lhs = std::move(statusor.ValueOrDie()) #define TF_STATUS_MACROS_CONCAT_NAME(x, y) TF_STATUS_MACROS_CONCAT_IMPL(x, y) #define TF_STATUS_MACROS_CONCAT_IMPL(x, y) x##y diff --git a/tensorflow/compiler/xla/status_macros_test.cc b/tensorflow/compiler/xla/status_macros_test.cc index 4e7b9161db5c7e01a4b80da49bdded025eaf298a..4b0740dad72f5d96e5ae153abf9232553ff834c2 100644 --- a/tensorflow/compiler/xla/status_macros_test.cc +++ b/tensorflow/compiler/xla/status_macros_test.cc @@ -16,9 +16,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/test_helpers.h" #include "tensorflow/core/lib/core/errors.h" -#include "tensorflow/core/platform/test.h" namespace xla { @@ -40,15 +40,15 @@ Status RetCheckSuccess() { TEST(StatusMacros, RetCheckFailing) { Status status = RetCheckFail(); EXPECT_EQ(status.code(), tensorflow::error::INTERNAL); - EXPECT_MATCH(status.error_message(), - xla::testing::ContainsRegex("RET_CHECK failure.*2 > 3")); + EXPECT_THAT(status.error_message(), + ::testing::ContainsRegex("RET_CHECK failure.*2 > 3")); } TEST(StatusMacros, RetCheckFailingWithExtraMessage) { Status status = RetCheckFailWithExtraMessage(); EXPECT_EQ(status.code(), tensorflow::error::INTERNAL); - EXPECT_MATCH(status.error_message(), - xla::testing::ContainsRegex("RET_CHECK.*2 > 3 extra message")); + EXPECT_THAT(status.error_message(), + ::testing::ContainsRegex("RET_CHECK.*2 > 3 extra message")); } TEST(StatusMacros, RetCheckSucceeding) { @@ -63,7 +63,7 @@ StatusOr CreateIntUnsuccessfully() { } TEST(StatusMacros, AssignOrAssertOnOK) { - TF_ASSIGN_OR_ASSERT_OK(int result, CreateIntSuccessfully()); + TF_ASSERT_OK_AND_ASSIGN(int result, CreateIntSuccessfully()); EXPECT_EQ(42, result); } @@ -73,7 +73,7 @@ Status ReturnStatusError() { return (tensorflow::errors::Internal("foobar")); } using StatusReturningFunction = std::function; -StatusOr CallStatusReturningFunction(StatusReturningFunction func) { +StatusOr CallStatusReturningFunction(const StatusReturningFunction& func) { TF_RETURN_IF_ERROR(func()); return 42; } diff --git a/tensorflow/compiler/xla/statusor.cc b/tensorflow/compiler/xla/statusor.cc index 36f08fc99f45a7c82f086d04fa60014343d574da..72ab67ff810e0ec384a22da092363cc7446435bb 100644 --- a/tensorflow/compiler/xla/statusor.cc +++ b/tensorflow/compiler/xla/statusor.cc @@ -19,28 +19,20 @@ limitations under the License. #include "tensorflow/core/platform/logging.h" namespace xla { -namespace internal { +namespace internal_statusor { -Status StatusOrHelper::HandleInvalidStatusCtorArg() { +void Helper::HandleInvalidStatusCtorArg(Status* status) { const char* kMessage = - "Status::OK is not a valid constructor argument to StatusOr"; + "An OK status is not a valid constructor argument to StatusOr"; LOG(ERROR) << kMessage; - // In optimized builds, we will fall back to tensorflow::error::INTERNAL. - return Status(tensorflow::error::INTERNAL, kMessage); + // Fall back to tensorflow::error::INTERNAL. + *status = ::tensorflow::errors::Internal(kMessage); } -Status StatusOrHelper::HandleNullObjectCtorArg() { - const char* kMessage = - "NULL is not a valid constructor argument to StatusOr"; - LOG(ERROR) << kMessage; - // In optimized builds, we will fall back to tensorflow::error::INTERNAL. - return Status(tensorflow::error::INTERNAL, kMessage); -} - -void StatusOrHelper::Crash(const Status& status) { +void Helper::Crash(const Status& status) { LOG(FATAL) << "Attempting to fetch value instead of handling error " << status; } -} // namespace internal +} // namespace internal_statusor } // namespace xla diff --git a/tensorflow/compiler/xla/statusor.h b/tensorflow/compiler/xla/statusor.h index d8cd736238c19cc00d0302daa54fc7417740001a..641b5e9a6accc0a2e7737f79bcd485d317e4e521 100644 --- a/tensorflow/compiler/xla/statusor.h +++ b/tensorflow/compiler/xla/statusor.h @@ -72,216 +72,233 @@ limitations under the License. #define TENSORFLOW_COMPILER_XLA_STATUSOR_H_ #include "tensorflow/compiler/xla/status.h" +#include "tensorflow/compiler/xla/statusor_internals.h" #include "tensorflow/core/platform/macros.h" namespace xla { #if defined(__clang__) // Only clang supports warn_unused_result as a type annotation. -template +template class TF_MUST_USE_RESULT StatusOr; #endif -template ::value> -class StatusOr { - template +template +class StatusOr : private internal_statusor::StatusOrData, + private internal_statusor::TraitsBase< + std::is_copy_constructible::value, + std::is_move_constructible::value> { + template friend class StatusOr; + typedef internal_statusor::StatusOrData Base; + public: typedef T element_type; - // Construct a new StatusOr with Status::UNKNOWN status - StatusOr(); + // Constructs a new StatusOr with Status::UNKNOWN status. This is marked + // 'explicit' to try to catch cases like 'return {};', where people think + // StatusOr> will be initialized with an empty vector, + // instead of a Status::UNKNOWN status. + explicit StatusOr(); + + // StatusOr will be copy constructible/assignable if T is copy + // constructible. + StatusOr(const StatusOr&) = default; + StatusOr& operator=(const StatusOr&) = default; + + // StatusOr will be move constructible/assignable if T is move + // constructible. + StatusOr(StatusOr&&) = default; + StatusOr& operator=(StatusOr&&) = default; + + // Conversion copy/move constructor, T must be convertible from U. + // TODO(b/62186717): These should not participate in overload resolution if U + // is not convertible to T. + template + StatusOr(const StatusOr& other); + template + StatusOr(StatusOr&& other); - // Construct a new StatusOr with the given non-ok status. After calling - // this constructor, calls to ValueOrDie() will CHECK-fail. - // - // NOTE: Not explicit - we want to use StatusOr as a return - // value, so it is convenient and sensible to be able to do 'return - // Status()' when the return type is StatusOr. - // - // REQUIRES: status != Status::OK. This requirement is DCHECKed. - // In optimized builds, passing Status::OK here will have the effect - // of passing tensorflow::error::INTERNAL as a fallback. - StatusOr(Status status); // NOLINT + // Conversion copy/move assignment operator, T must be convertible from U. + template + StatusOr& operator=(const StatusOr& other); + template + StatusOr& operator=(StatusOr&& other); - // Construct a new StatusOr with the given value. If T is a plain pointer, - // value must not be NULL. After calling this constructor, calls to - // ValueOrDie() will succeed, and calls to status() will return OK. + // Constructs a new StatusOr with the given value. After calling this + // constructor, calls to ValueOrDie() will succeed, and calls to status() will + // return OK. // // NOTE: Not explicit - we want to use StatusOr as a return type // so it is convenient and sensible to be able to do 'return T()' // when the return type is StatusOr. // - // REQUIRES: if T is a plain pointer, value != NULL. This requirement is - // DCHECKed. In optimized builds, passing a NULL pointer here will have - // the effect of passing tensorflow::error::INTERNAL as a fallback. - StatusOr(const T& value); // NOLINT - - // Copy constructor. - StatusOr(const StatusOr& other) = default; - - // Conversion copy constructor, T must be copy constructible from U - template - StatusOr(const StatusOr& other); - - // Assignment operator. - StatusOr& operator=(const StatusOr& other) = default; + // REQUIRES: T is copy constructible. + StatusOr(const T& value); - // Conversion assignment operator, T must be assignable from U - template - StatusOr& operator=(const StatusOr& other); + // Constructs a new StatusOr with the given non-ok status. After calling + // this constructor, calls to ValueOrDie() will CHECK-fail. + // + // NOTE: Not explicit - we want to use StatusOr as a return + // value, so it is convenient and sensible to be able to do 'return + // Status()' when the return type is StatusOr. + // + // REQUIRES: !status.ok(). This requirement is DCHECKed. + // In optimized builds, passing Status::OK() here will have the effect + // of passing tensorflow::error::INTERNAL as a fallback. + StatusOr(const Status& status); + StatusOr& operator=(const Status& status); - // Move constructor and move-assignment operator. - StatusOr(StatusOr&& other) = default; - StatusOr& operator=(StatusOr&& other) = default; + // TODO(b/62186997): Add operator=(T) overloads. - // Rvalue-reference overloads of the other constructors and assignment - // operators, to support move-only types and avoid unnecessary copying. + // Similar to the `const T&` overload. // - // Implementation note: we could avoid all these rvalue-reference overloads - // if the existing lvalue-reference overloads took their arguments by value - // instead. I think this would also let us omit the conversion assignment - // operator altogether, since we'd get the same functionality for free - // from the implicit conversion constructor and ordinary assignment. - // However, this could result in extra copy operations unless we use - // std::move to avoid them, and we can't use std::move because this code - // needs to be portable to C++03. - StatusOr(T&& value); // NOLINT - template - StatusOr(StatusOr&& other); + // REQUIRES: T is move constructible. + StatusOr(T&& value); - // Returns a reference to our status. If this contains a T, then - // returns Status::OK. - const Status& status() const { return status_; } + // RValue versions of the operations declared above. + StatusOr(Status&& status); + StatusOr& operator=(Status&& status); // Returns this->status().ok() - bool ok() const { return status_.ok(); } + bool ok() const { return this->status_.ok(); } + + // Returns a reference to our status. If this contains a T, then + // returns Status::OK(). + const Status& status() const &; + Status status() &&; // Returns a reference to our current value, or CHECK-fails if !this->ok(). - const T& ValueOrDie() const; - T& ValueOrDie(); + // + // Note: for value types that are cheap to copy, prefer simple code: + // + // T value = statusor.ValueOrDie(); + // + // Otherwise, if the value type is expensive to copy, but can be left + // in the StatusOr, simply assign to a reference: + // + // T& value = statusor.ValueOrDie(); // or `const T&` + // + // Otherwise, if the value type supports an efficient move, it can be + // used as follows: + // + // T value = std::move(statusor).ValueOrDie(); + // + // The std::move on statusor instead of on the whole expression enables + // warnings about possible uses of the statusor object after the move. + // C++ style guide waiver for ref-qualified overloads granted in cl/143176389 + // See go/ref-qualifiers for more details on such overloads. + const T& ValueOrDie() const &; + T& ValueOrDie() &; + const T&& ValueOrDie() const &&; + T&& ValueOrDie() &&; - // Moves our current value out of this object and returns it, or CHECK-fails - // if !this->ok(). - // Use of this method is discouraged; prefer std::move(statusor.ValueOrDie()) - // instead. T ConsumeValueOrDie() { return std::move(ValueOrDie()); } - private: - Status status_; - T value_; -}; - -// Partial specialization for when T is not copy-constructible. This uses all -// methods from the core implementation, but removes copy assignment and copy -// construction. -template -class StatusOr : public StatusOr { - public: - // Remove copies. - StatusOr(const StatusOr& other) = delete; - StatusOr& operator=(const StatusOr& other) = delete; - template - StatusOr(const StatusOr& other) = delete; - StatusOr(const T& value) = delete; - - // Use the superclass version for other constructors and operators. - StatusOr() = default; - StatusOr(StatusOr&& other) = default; - StatusOr& operator=(StatusOr&& other) = default; - StatusOr(T&& value) // NOLINT - : StatusOr::StatusOr(std::move(value)) {} - StatusOr(Status status) // NOLINT - : StatusOr::StatusOr(std::move(status)) {} - template - StatusOr(StatusOr&& other) // NOLINT - : StatusOr::StatusOr(std::move(other)) {} + // Ignores any errors. This method does nothing except potentially suppress + // complaints from any tools that are checking that errors are not dropped on + // the floor. + void IgnoreError() const; }; //////////////////////////////////////////////////////////////////////////////// // Implementation details for StatusOr -namespace internal { +template +StatusOr::StatusOr() : Base(Status(tensorflow::error::UNKNOWN, "")) {} -class StatusOrHelper { - public: - // Move type-agnostic error handling to the .cc. - static Status HandleInvalidStatusCtorArg(); - static Status HandleNullObjectCtorArg(); - static void Crash(const Status& status); - - // Customized behavior for StatusOr vs. StatusOr - template - struct Specialize; -}; +template +StatusOr::StatusOr(const T& value) : Base(value) {} template -struct StatusOrHelper::Specialize { - // For non-pointer T, a reference can never be NULL. - static inline bool IsValueNull(const T& t) { return false; } -}; +StatusOr::StatusOr(const Status& status) : Base(status) {} template -struct StatusOrHelper::Specialize { - static inline bool IsValueNull(const T* t) { return t == NULL; } -}; +StatusOr& StatusOr::operator=(const Status& status) { + this->Assign(status); + return *this; +} -} // namespace internal +template +StatusOr::StatusOr(T&& value) : Base(std::move(value)) {} -template -inline StatusOr::StatusOr() - : status_(tensorflow::error::UNKNOWN, "") {} +template +StatusOr::StatusOr(Status&& status) : Base(std::move(status)) {} -template -inline StatusOr::StatusOr(Status status) - : status_(std::move(status)) { - if (status_.ok()) { - status_ = internal::StatusOrHelper::HandleInvalidStatusCtorArg(); - } +template +StatusOr& StatusOr::operator=(Status&& status) { + this->Assign(std::move(status)); + return *this; } -template -inline StatusOr::StatusOr(const T& value) - : value_(value) { - if (internal::StatusOrHelper::Specialize::IsValueNull(value)) { - status_ = internal::StatusOrHelper::HandleNullObjectCtorArg(); - } -} +template +template +inline StatusOr::StatusOr(const StatusOr& other) + : Base(static_cast::Base&>(other)) {} -template +template template -inline StatusOr::StatusOr(const StatusOr& other) - : status_(other.status_), value_(other.value_) {} - -template -inline StatusOr::StatusOr(T&& value) - : value_(std::move(value)) { - if (internal::StatusOrHelper::Specialize::IsValueNull(value_)) { - status_ = internal::StatusOrHelper::HandleNullObjectCtorArg(); - } +inline StatusOr& StatusOr::operator=(const StatusOr& other) { + if (other.ok()) + this->Assign(other.ValueOrDie()); + else + this->Assign(other.status()); + return *this; } -template +template template -inline StatusOr::StatusOr(StatusOr&& other) - : status_(std::move(other.status_)), value_(std::move(other.value_)) {} +inline StatusOr::StatusOr(StatusOr&& other) + : Base(static_cast::Base&&>(other)) {} -template -inline const T& StatusOr::ValueOrDie() const { - if (!ok()) { - internal::StatusOrHelper::Crash(status()); +template +template +inline StatusOr& StatusOr::operator=(StatusOr&& other) { + if (other.ok()) { + this->Assign(std::move(other).ValueOrDie()); + } else { + this->Assign(std::move(other).status()); } - return value_; + return *this; } -template -inline T& StatusOr::ValueOrDie() { - if (!status_.ok()) { - internal::StatusOrHelper::Crash(status()); - } - return value_; +template +const Status& StatusOr::status() const & { + return this->status_; +} +template +Status StatusOr::status() && { + return ok() ? Status::OK() : std::move(this->status_); +} + +template +const T& StatusOr::ValueOrDie() const & { + this->EnsureOk(); + return this->data_; +} + +template +T& StatusOr::ValueOrDie() & { + this->EnsureOk(); + return this->data_; +} + +template +const T&& StatusOr::ValueOrDie() const && { + this->EnsureOk(); + return std::move(this->data_); +} + +template +T&& StatusOr::ValueOrDie() && { + this->EnsureOk(); + return std::move(this->data_); +} + +template +void StatusOr::IgnoreError() const { + // no-op } } // namespace xla diff --git a/tensorflow/compiler/xla/statusor_internals.h b/tensorflow/compiler/xla/statusor_internals.h new file mode 100644 index 0000000000000000000000000000000000000000..a2fda5bb3c6f11c20fc45c57885b1ce7523db81d --- /dev/null +++ b/tensorflow/compiler/xla/statusor_internals.h @@ -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. +==============================================================================*/ + +#ifndef THIRD_PARTY_TENSORFLOW_COMPILER_XLA_STATUSOR_INTERNALS_H_ +#define THIRD_PARTY_TENSORFLOW_COMPILER_XLA_STATUSOR_INTERNALS_H_ + +#include "tensorflow/compiler/xla/status.h" +#include "tensorflow/core/platform/macros.h" + +namespace xla { +namespace internal_statusor { + +class Helper { + public: + // Move type-agnostic error handling to the .cc. + static void HandleInvalidStatusCtorArg(Status*); + TF_ATTRIBUTE_NORETURN static void Crash(const Status& status); +}; + +// Construct an instance of T in `p` through placement new, passing Args... to +// the constructor. +// This abstraction is here mostly for the gcc performance fix. +template +void PlacementNew(void* p, Args&&... args) { +#if defined(__GNUC__) && !defined(__clang__) + // Teach gcc that 'p' cannot be null, fixing code size issues. + if (p == nullptr) __builtin_unreachable(); +#endif + new (p) T(std::forward(args)...); +} + +// Helper base class to hold the data and all operations. +// We move all this to a base class to allow mixing with the appropriate +// TraitsBase specialization. +template +class StatusOrData { + template + friend class StatusOrData; + + public: + StatusOrData() = delete; + + StatusOrData(const StatusOrData& other) { + if (other.ok()) { + MakeValue(other.data_); + MakeStatus(); + } else { + MakeStatus(other.status_); + } + } + + StatusOrData(StatusOrData&& other) noexcept { + if (other.ok()) { + MakeValue(std::move(other.data_)); + MakeStatus(); + } else { + MakeStatus(std::move(other.status_)); + } + } + + template + StatusOrData(const StatusOrData& other) { + if (other.ok()) { + MakeValue(other.data_); + MakeStatus(); + } else { + MakeStatus(other.status_); + } + } + + template + StatusOrData(StatusOrData&& other) { + if (other.ok()) { + MakeValue(std::move(other.data_)); + MakeStatus(); + } else { + MakeStatus(std::move(other.status_)); + } + } + + explicit StatusOrData(const T& value) : data_(value) { MakeStatus(); } + explicit StatusOrData(T&& value) : data_(std::move(value)) { MakeStatus(); } + + explicit StatusOrData(const Status& status) : status_(status) { + EnsureNotOk(); + } + explicit StatusOrData(Status&& status) : status_(std::move(status)) { + EnsureNotOk(); + } + + StatusOrData& operator=(const StatusOrData& other) { + if (this == &other) return *this; + if (other.ok()) + Assign(other.data_); + else + Assign(other.status_); + return *this; + } + + StatusOrData& operator=(StatusOrData&& other) { + if (this == &other) return *this; + if (other.ok()) + Assign(std::move(other.data_)); + else + Assign(std::move(other.status_)); + return *this; + } + + ~StatusOrData() { + if (ok()) { + status_.~Status(); + data_.~T(); + } else { + status_.~Status(); + } + } + + void Assign(const T& value) { + if (ok()) { + data_.~T(); + MakeValue(value); + } else { + MakeValue(value); + status_ = Status::OK(); + } + } + + void Assign(T&& value) { + if (ok()) { + data_.~T(); + MakeValue(std::move(value)); + } else { + MakeValue(std::move(value)); + status_ = Status::OK(); + } + } + + void Assign(const Status& status) { + Clear(); + status_ = status; + EnsureNotOk(); + } + + void Assign(Status&& status) { + Clear(); + status_ = std::move(status); + EnsureNotOk(); + } + + bool ok() const { return status_.ok(); } + + protected: + // status_ will always be active after the constructor. + // We make it a union to be able to initialize exactly how we need without + // waste. + // Eg. in the copy constructor we use the default constructor of Status in + // the ok() path to avoid an extra Ref call. + union { + Status status_; + }; + + // data_ is active iff status_.ok()==true + struct Dummy {}; + union { + // When T is const, we need some non-const object we can cast to void* for + // the placement new. dummy_ is that object. + Dummy dummy_; + T data_; + }; + + void Clear() { + if (ok()) data_.~T(); + } + + void EnsureOk() const { + if (!ok()) Helper::Crash(status_); + } + + void EnsureNotOk() { + if (ok()) Helper::HandleInvalidStatusCtorArg(&status_); + } + + // Construct the value (ie. data_) through placement new with the passed + // argument. + template + void MakeValue(Arg&& arg) { + internal_statusor::PlacementNew(&dummy_, std::forward(arg)); + } + + // Construct the status (ie. status_) through placement new with the passed + // argument. + template + void MakeStatus(Args&&... args) { + internal_statusor::PlacementNew(&status_, + std::forward(args)...); + } +}; + +// Helper base class to allow implicitly deleted constructors and assignment +// operations in StatusOr. +// TraitsBase will explicitly delete what it can't support and StatusOr will +// inherit that behavior implicitly. +template +struct TraitsBase { + TraitsBase() = default; + TraitsBase(const TraitsBase&) = default; + TraitsBase(TraitsBase&&) = default; + TraitsBase& operator=(const TraitsBase&) = default; + TraitsBase& operator=(TraitsBase&&) = default; +}; + +template <> +struct TraitsBase { + TraitsBase() = default; + TraitsBase(const TraitsBase&) = delete; + TraitsBase(TraitsBase&&) = default; + TraitsBase& operator=(const TraitsBase&) = delete; + TraitsBase& operator=(TraitsBase&&) = default; +}; + +template <> +struct TraitsBase { + TraitsBase() = default; + TraitsBase(const TraitsBase&) = delete; + TraitsBase(TraitsBase&&) = delete; + TraitsBase& operator=(const TraitsBase&) = delete; + TraitsBase& operator=(TraitsBase&&) = delete; +}; + +} // namespace internal_statusor +} // namespace xla + +#endif // THIRD_PARTY_TENSORFLOW_COMPILER_XLA_STATUSOR_INTERNALS_H_ diff --git a/tensorflow/compiler/xla/statusor_test.cc b/tensorflow/compiler/xla/statusor_test.cc index d98eb2793363ac855b43f88eb4201f34a3b7693b..5fa2211ac66177514ac8ecabfa8791e7c8c014a2 100644 --- a/tensorflow/compiler/xla/statusor_test.cc +++ b/tensorflow/compiler/xla/statusor_test.cc @@ -20,17 +20,15 @@ limitations under the License. #include #include +#include "tensorflow/compiler/xla/test.h" #include "tensorflow/compiler/xla/types.h" #include "tensorflow/core/lib/core/errors.h" #include "tensorflow/core/platform/macros.h" -#include "tensorflow/core/platform/test.h" #include "tensorflow/core/platform/test_benchmark.h" namespace xla { namespace { -using tensorflow::Status; - class Base1 { public: virtual ~Base1() {} @@ -45,7 +43,7 @@ class Base2 { class Derived : public Base1, public Base2 { public: - virtual ~Derived() {} + ~Derived() override {} int evenmorepad; }; @@ -59,6 +57,14 @@ class CopyNoAssign { const CopyNoAssign& operator=(const CopyNoAssign&); }; +class NoDefaultConstructor { + public: + explicit NoDefaultConstructor(int foo); +}; + +static_assert(!std::is_default_constructible(), + "Should not be default-constructible."); + StatusOr> ReturnUniquePtr() { // Uses implicit constructor from T&& return std::unique_ptr(new int(0)); @@ -69,6 +75,18 @@ TEST(StatusOr, ElementType) { static_assert(std::is_same::element_type, char>(), ""); } +TEST(StatusOr, TestNoDefaultConstructorInitialization) { + // Explicitly initialize it with an error code. + StatusOr statusor(tensorflow::errors::Cancelled("")); + EXPECT_FALSE(statusor.ok()); + EXPECT_EQ(statusor.status().code(), tensorflow::error::CANCELLED); + + // Default construction of StatusOr initializes it with an UNKNOWN error code. + StatusOr statusor2; + EXPECT_FALSE(statusor2.ok()); + EXPECT_EQ(statusor2.status().code(), tensorflow::error::UNKNOWN); +} + TEST(StatusOr, TestMoveOnlyInitialization) { StatusOr> thing(ReturnUniquePtr()); ASSERT_TRUE(thing.ok()); @@ -436,17 +454,17 @@ class BenchmarkFactory { } Status ArgumentFactoryFail(T** result) TF_ATTRIBUTE_NOINLINE { - *result = NULL; + *result = nullptr; return Status(tensorflow::error::CANCELLED, ""); } Status ArgumentFactoryFailShortMsg(T** result) TF_ATTRIBUTE_NOINLINE { - *result = NULL; + *result = nullptr; return Status(::tensorflow::error::INTERNAL, ""); } Status ArgumentFactoryFailLongMsg(T** result) TF_ATTRIBUTE_NOINLINE { - *result = NULL; + *result = nullptr; return Status(::tensorflow::error::INTERNAL, "a big string of message junk that will never be read"); } @@ -489,26 +507,30 @@ class BenchmarkType { // Calibrate the amount of time spent just calling DoWork, since each of our // tests will do this, we can subtract this out of benchmark results. -static void BM_CalibrateWorkLoop(int iters) { +void BM_CalibrateWorkLoop(int iters) { tensorflow::testing::StopTiming(); BenchmarkFactory factory; BenchmarkType* result = factory.TrivialFactory(); tensorflow::testing::StartTiming(); for (int i = 0; i != iters; ++i) { - if (result != NULL) result->DoWork(); + if (result != nullptr) { + result->DoWork(); + } } } BENCHMARK(BM_CalibrateWorkLoop); // Measure the time taken to call into the factory, return the value, // determine that it is OK, and invoke a trivial function. -static void BM_TrivialFactory(int iters) { +void BM_TrivialFactory(int iters) { tensorflow::testing::StopTiming(); BenchmarkFactory factory; tensorflow::testing::StartTiming(); for (int i = 0; i != iters; ++i) { BenchmarkType* result = factory.TrivialFactory(); - if (result != NULL) result->DoWork(); + if (result != nullptr) { + result->DoWork(); + } } } BENCHMARK(BM_TrivialFactory); @@ -516,14 +538,14 @@ BENCHMARK(BM_TrivialFactory); // Measure the time taken to call into the factory, providing an // out-param for the result, evaluating the status result and the // result pointer, and invoking the trivial function. -static void BM_ArgumentFactory(int iters) { +void BM_ArgumentFactory(int iters) { tensorflow::testing::StopTiming(); BenchmarkFactory factory; tensorflow::testing::StartTiming(); for (int i = 0; i != iters; ++i) { - BenchmarkType* result = NULL; + BenchmarkType* result = nullptr; Status status = factory.ArgumentFactory(&result); - if (status.ok() && result != NULL) { + if (status.ok() && result != nullptr) { result->DoWork(); } } @@ -532,7 +554,7 @@ BENCHMARK(BM_ArgumentFactory); // Measure the time to use the StatusOr factory, evaluate the result, // and invoke the trivial function. -static void BM_StatusOrFactory(int iters) { +void BM_StatusOrFactory(int iters) { tensorflow::testing::StopTiming(); BenchmarkFactory factory; tensorflow::testing::StartTiming(); @@ -548,14 +570,14 @@ BENCHMARK(BM_StatusOrFactory); // Measure the time taken to call into the factory, providing an // out-param for the result, evaluating the status result and the // result pointer, and invoking the trivial function. -static void BM_ArgumentFactoryFail(int iters) { +void BM_ArgumentFactoryFail(int iters) { tensorflow::testing::StopTiming(); BenchmarkFactory factory; tensorflow::testing::StartTiming(); for (int i = 0; i != iters; ++i) { - BenchmarkType* result = NULL; + BenchmarkType* result = nullptr; Status status = factory.ArgumentFactoryFail(&result); - if (status.ok() && result != NULL) { + if (status.ok() && result != nullptr) { result->DoWork(); } } @@ -564,7 +586,7 @@ BENCHMARK(BM_ArgumentFactoryFail); // Measure the time to use the StatusOr factory, evaluate the result, // and invoke the trivial function. -static void BM_StatusOrFactoryFail(int iters) { +void BM_StatusOrFactoryFail(int iters) { tensorflow::testing::StopTiming(); BenchmarkFactory factory; tensorflow::testing::StartTiming(); @@ -580,14 +602,14 @@ BENCHMARK(BM_StatusOrFactoryFail); // Measure the time taken to call into the factory, providing an // out-param for the result, evaluating the status result and the // result pointer, and invoking the trivial function. -static void BM_ArgumentFactoryFailShortMsg(int iters) { +void BM_ArgumentFactoryFailShortMsg(int iters) { tensorflow::testing::StopTiming(); BenchmarkFactory factory; tensorflow::testing::StartTiming(); for (int i = 0; i != iters; ++i) { - BenchmarkType* result = NULL; + BenchmarkType* result = nullptr; Status status = factory.ArgumentFactoryFailShortMsg(&result); - if (status.ok() && result != NULL) { + if (status.ok() && result != nullptr) { result->DoWork(); } } @@ -596,7 +618,7 @@ BENCHMARK(BM_ArgumentFactoryFailShortMsg); // Measure the time to use the StatusOr factory, evaluate the result, // and invoke the trivial function. -static void BM_StatusOrFactoryFailShortMsg(int iters) { +void BM_StatusOrFactoryFailShortMsg(int iters) { tensorflow::testing::StopTiming(); BenchmarkFactory factory; tensorflow::testing::StartTiming(); @@ -612,14 +634,14 @@ BENCHMARK(BM_StatusOrFactoryFailShortMsg); // Measure the time taken to call into the factory, providing an // out-param for the result, evaluating the status result and the // result pointer, and invoking the trivial function. -static void BM_ArgumentFactoryFailLongMsg(int iters) { +void BM_ArgumentFactoryFailLongMsg(int iters) { tensorflow::testing::StopTiming(); BenchmarkFactory factory; tensorflow::testing::StartTiming(); for (int i = 0; i != iters; ++i) { - BenchmarkType* result = NULL; + BenchmarkType* result = nullptr; Status status = factory.ArgumentFactoryFailLongMsg(&result); - if (status.ok() && result != NULL) { + if (status.ok() && result != nullptr) { result->DoWork(); } } @@ -628,7 +650,7 @@ BENCHMARK(BM_ArgumentFactoryFailLongMsg); // Measure the time to use the StatusOr factory, evaluate the result, // and invoke the trivial function. -static void BM_StatusOrFactoryFailLongMsg(int iters) { +void BM_StatusOrFactoryFailLongMsg(int iters) { tensorflow::testing::StopTiming(); BenchmarkFactory factory; tensorflow::testing::StartTiming(); diff --git a/tensorflow/compiler/xla/test.h b/tensorflow/compiler/xla/test.h new file mode 100644 index 0000000000000000000000000000000000000000..87a8c5f3a528289d47c1729ae6719aae47037c36 --- /dev/null +++ b/tensorflow/compiler/xla/test.h @@ -0,0 +1,48 @@ +/* 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_COMPLIER_XLA_TEST_H_ +#define TENSORFLOW_COMPLIER_XLA_TEST_H_ + +// This header includes gmock.h and enables the use of gmock matchers in tests +// in third_party/tensorflow/compiler/xla. +// +// Test including this header can use the macros EXPECT_THAT(...) and +// ASSERT_THAT(...) in combination with gmock matchers. +// Example: +// std::vector vec = Foo(); +// EXPECT_THAT(vec, ::testing::ElementsAre(1,2,3)); +// +// For more details on gmock matchers see: +// https://github.com/google/googletest/blob/master/googlemock/docs/CheatSheet.md#matchers +// +// The advantages of using gmock matchers instead of self defined matchers are +// better error messages, more maintainable tests and more test coverage. +// +// Note that while the use of gmock matchers is allowed in the xla project, the +// use of mocks is disallowed in the whole tensorflow project! + +#include "tensorflow/core/platform/platform.h" + +#if defined(PLATFORM_GOOGLE) || defined(PLATFORM_GOOGLE_ANDROID) +#include "testing/base/public/gmock.h" +#else +#include +#include +#endif + +#include "tensorflow/core/platform/test.h" + +#endif // TENSORFLOW_COMPLIER_XLA_TEST_H_ diff --git a/tensorflow/compiler/xla/test_helpers.cc b/tensorflow/compiler/xla/test_helpers.cc deleted file mode 100644 index 02abfdeab80ee34c79e8d54b825937d6fc4b4053..0000000000000000000000000000000000000000 --- a/tensorflow/compiler/xla/test_helpers.cc +++ /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. -==============================================================================*/ - -#include "tensorflow/compiler/xla/test_helpers.h" -#include "tensorflow/compiler/xla/types.h" -#include "tensorflow/core/platform/regexp.h" - -namespace xla { -namespace testing { - -AssertionResult::AssertionResult(const AssertionResult& other) - : success_(other.success_), - message_(other.message_ != nullptr ? new std::string(*other.message_) - : static_cast(nullptr)) { -} - -// Returns the assertion's negation. Used with EXPECT/ASSERT_FALSE. -AssertionResult AssertionResult::operator!() const { - AssertionResult negation(!success_); - if (message_ != nullptr) negation << *message_; - return negation; -} - -AssertionResult& AssertionResult::operator=(const AssertionResult& ar) { - success_ = ar.success_; - message_.reset(ar.message_ != nullptr ? new std::string(*ar.message_) - : nullptr); - return *this; -} - -AssertionResult AssertionFailure() { return AssertionResult(false); } - -AssertionResult AssertionSuccess() { return AssertionResult(true); } - -std::function ContainsRegex( - const tensorflow::StringPiece regex) { - return [regex](const tensorflow::StringPiece to_test) { - if (RE2::PartialMatch( - tensorflow::RegexpStringPiece(to_test.data(), to_test.size()), - tensorflow::RegexpStringPiece(regex.data(), regex.size()))) { - return true; - } else { - LOG(ERROR) << "Expected to find " << regex << " in " << to_test; - return false; - } - }; -} - -std::function HasSubstr( - const tensorflow::StringPiece part) { - return [part](const tensorflow::StringPiece whole) { - return whole.contains(part); - }; -} - -} // namespace testing -} // namespace xla diff --git a/tensorflow/compiler/xla/test_helpers.h b/tensorflow/compiler/xla/test_helpers.h index f923d9f36c878c1ae4e37f052a84e9c2a279b4ed..634cdb5aa29651b08090ff99f0a6cafb9facb645 100644 --- a/tensorflow/compiler/xla/test_helpers.h +++ b/tensorflow/compiler/xla/test_helpers.h @@ -39,286 +39,6 @@ class Literal; namespace testing { -class AssertionResult { - public: - explicit AssertionResult(bool success) : success_(success) {} - - // Returns true iff the assertion succeeded. - operator bool() const { return success_; } // NOLINT - - // Returns the assertion's negation. Used with EXPECT/ASSERT_FALSE. - AssertionResult operator!() const; - - // Returns the text streamed into this AssertionResult. Test assertions - // use it when they fail (i.e., the predicate's outcome doesn't match the - // assertion's expectation). When nothing has been streamed into the - // object, returns an empty string. - const char* message() const { - return message_ != nullptr ? message_->c_str() : ""; - } - - // Streams a custom failure message into this object. - template - AssertionResult& operator<<(const T& value) { - AppendMessage(::testing::Message() << value); - return *this; - } - - // Allows streaming basic output manipulators such as endl or flush into - // this object. - AssertionResult& operator<<( - std::ostream& (*basic_manipulator)(std::ostream& stream)) { - AppendMessage(::testing::Message() << basic_manipulator); - return *this; - } - - // Copy operator. - AssertionResult(const AssertionResult& ar); - - // Assignment operator. - AssertionResult& operator=(const AssertionResult&); - - private: - // Appends the contents of message to message_. - void AppendMessage(const ::testing::Message& a_message) { - if (message_ == nullptr) message_.reset(new std::string); - message_->append(a_message.GetString().c_str()); - } - - bool success_ = false; - - // Stores the message describing the condition in case the - // expectation construct is not satisfied with the predicate's - // outcome. Referenced via a pointer to avoid taking too much stack - // frame space with test assertions. - std::unique_ptr message_; -}; - -AssertionResult AssertionFailure(); - -AssertionResult AssertionSuccess(); - -std::function ContainsRegex( - const tensorflow::StringPiece regex); - -std::function HasSubstr( - const tensorflow::StringPiece part); - -// Matcher for a vector of same-type values for which operator= is -// defined. -template -std::function& actual)> VectorMatcher( - const std::vector& expected) { - return [expected](const std::vector& actual) -> AssertionResult { - int len = expected.size(); - if (actual.size() != len) { - return AssertionFailure() << "Actual values len of " << actual.size() - << " != expected.size " << len; - } - for (int i = 0; i < len; ++i) { - if (actual[i] != expected[i]) { - return AssertionFailure() << "Element " << i << " actual " << actual[i] - << " != " << expected[i]; - } - } - return AssertionSuccess(); - }; -} - -// Approximate matcher for a vector of floats or similar. -template -std::function& actual)> -ApproxVectorMatcher(const std::vector& expected, float abs_diff, - float rel_diff) { - return [abs_diff, rel_diff, - expected](const std::vector& actual) -> AssertionResult { - int len = expected.size(); - if (actual.size() != len) { - AssertionResult ar = AssertionFailure() << "Actual values len of " - << actual.size() - << " != expected.size " << len; - LOG(ERROR) << ar.message(); - return ar; - } - for (int i = 0; i < len; ++i) { - T diff = actual[i] - expected[i]; - if (diff < 0) { - diff *= -1; - } - if (diff > abs_diff) { - T rdiff = (expected[i] != 0 ? diff / expected[i] : 0.0 * expected[i]); - if (rdiff > rel_diff) { - AssertionResult ar = AssertionFailure() - << "Element " << i << " actual " << actual[i] - << " != " << expected[i] - << "( abs_diff = " << diff - << ", rel_diff = " << rdiff << ")"; - LOG(ERROR) << ar.message(); - return ar; - } - } - } - return AssertionSuccess(); - }; -} - -// Matches a vector of same-type values against another, succeeding so -// long as they have the same length and every value in 'actual' -// matches one in 'expected.' Does not verify an exhaustive -// one-to-one mapping between the two. -template -std::function& actual)> -UnorderedElementsAre(const std::vector& expected) { - return [expected](const std::vector& actual) -> AssertionResult { - if (actual.size() != expected.size()) { - return AssertionFailure() << "sizes don't match"; - } - for (auto a : actual) { - bool found = false; - for (auto e : expected) { - if (a == e) { - found = true; - break; - } - } - if (!found) { - return AssertionFailure() << "actual element " << a - << " not in expected"; - } - } - return AssertionSuccess(); - }; -} - -// Overloaded cover functions for UnorderedElementsAre, for the numbers -// of values used in practice. -template -std::function& actual)> UnorderedMatcher( - T a) { - std::vector expected; - expected.push_back(a); - return testing::UnorderedElementsAre(expected); -} - -template -std::function& actual)> UnorderedMatcher( - T a, T b) { - std::vector expected; - expected.push_back(a); - expected.push_back(b); - return testing::UnorderedElementsAre(expected); -} - -template -std::function& actual)> UnorderedMatcher( - T a, T b, T c) { - std::vector expected; - expected.push_back(a); - expected.push_back(b); - expected.push_back(c); - return testing::UnorderedElementsAre(expected); -} - -template -std::function& actual)> UnorderedMatcher( - T a, T b, T c, T d) { - std::vector expected; - expected.push_back(a); - expected.push_back(b); - expected.push_back(c); - expected.push_back(d); - return testing::UnorderedElementsAre(expected); -} - -template -std::function& actual)> UnorderedMatcher( - T a, T b, T c, T d, T e) { - std::vector expected; - expected.push_back(a); - expected.push_back(b); - expected.push_back(c); - expected.push_back(d); - expected.push_back(e); - return testing::UnorderedElementsAre(expected); -} - -template -std::function& actual)> UnorderedMatcher( - T a, T b, T c, T d, T e, T f) { - std::vector expected; - expected.push_back(a); - expected.push_back(b); - expected.push_back(c); - expected.push_back(d); - expected.push_back(e); - expected.push_back(f); - return testing::UnorderedElementsAre(expected); -} - -// Overloaded cover functions for VectorMatcher for the numbers of -// elements used in practice. -template -std::function& actual)> OrderedMatcher( - T a) { - std::vector expected; - expected.push_back(a); - return testing::VectorMatcher(expected); -} - -template -std::function& actual)> OrderedMatcher( - T a, T b) { - std::vector expected; - expected.push_back(a); - expected.push_back(b); - return testing::VectorMatcher(expected); -} - -template -std::function& actual)> OrderedMatcher( - T a, T b, T c) { - std::vector expected; - expected.push_back(a); - expected.push_back(b); - expected.push_back(c); - return testing::VectorMatcher(expected); -} - -template -std::function& actual)> OrderedMatcher( - T a, T b, T c, T d) { - std::vector expected; - expected.push_back(a); - expected.push_back(b); - expected.push_back(c); - expected.push_back(d); - return testing::VectorMatcher(expected); -} - -// Convert a RepeatedField to a flat vector. -template -std::vector PBToVec(const tensorflow::protobuf::RepeatedField rf) { - return std::vector(rf.begin(), rf.end()); -} - -// Convert a List to a flat vector. -template -std::vector ListToVec(const std::list& l) { - return std::vector(l.begin(), l.end()); -} - -// Convert a Set to a flat vector. -template -std::vector SetToVec(const std::set& c) { - return std::vector(c.begin(), c.end()); -} - -// Convert an Array to a flat vector. -template -std::vector Array2DToVec(const Array2D& a) { - return std::vector(a.data(), a.data() + a.num_elements()); -} - namespace internal_status { inline const ::tensorflow::Status& GetStatus( const ::tensorflow::Status& status) { @@ -347,9 +67,4 @@ inline const ::tensorflow::Status& GetStatus(const StatusOr& status) { ASSERT_EQ(tensorflow::Status::OK(), \ xla::testing::internal_status::GetStatus(expression)) -// Macros that apply a Matcher to a Value, returning an -// AssertionResult which gets digested by a standard gunit macro. -#define EXPECT_MATCH(V, M) EXPECT_TRUE((M)((V))) -#define ASSERT_MATCH(V, M) ASSERT_TRUE(M(V)) - #endif // TENSORFLOW_COMPILER_XLA_TEST_HELPERS_H_ diff --git a/tensorflow/compiler/xla/tests/BUILD b/tensorflow/compiler/xla/tests/BUILD index c7cbbdf4999970b0a09660ddadc31a068c752a55..52b2027aece62aba4d21dca0a61046dd59ee7db0 100644 --- a/tensorflow/compiler/xla/tests/BUILD +++ b/tensorflow/compiler/xla/tests/BUILD @@ -5,7 +5,6 @@ licenses(["notice"]) # Apache 2.0 package( default_visibility = [":friends"], - features = ["no_layering_check"], ) package_group( @@ -26,12 +25,26 @@ filegroup( load("//tensorflow/compiler/xla:xla.bzl", "export_dynamic_linkopts") 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") # Generate test_suites for all backends, named "${backend}_tests". generate_backend_suites() +# Target to add main for tests. Do not link this target and +# //third_party/tensorflow/core:test_main into the same target. +cc_library( + name = "xla_internal_test_main", + testonly = True, + srcs = ["xla_internal_test_main.cc"], + deps = [ + "//tensorflow/compiler/xla/legacy_flags:debug_options_flags", + "//tensorflow/core:lib", + "//tensorflow/core:test", + ], +) + cc_library( name = "test_macros_header", testonly = True, @@ -69,6 +82,7 @@ cc_library( "//tensorflow/compiler/xla:array4d", "//tensorflow/compiler/xla:literal_util", "//tensorflow/compiler/xla:shape_util", + "//tensorflow/compiler/xla:test", "//tensorflow/compiler/xla:test_helpers", "//tensorflow/compiler/xla:types", "//tensorflow/compiler/xla:util", @@ -92,17 +106,18 @@ cc_library( "//tensorflow/compiler/xla:types", "//tensorflow/compiler/xla:util", "//tensorflow/compiler/xla:xla_data_proto", - "//tensorflow/compiler/xla/legacy_flags:hlo_test_base_flags", + "//tensorflow/compiler/xla/legacy_flags:debug_options_flags", "//tensorflow/compiler/xla/service", "//tensorflow/compiler/xla/service:backend", "//tensorflow/compiler/xla/service:compiler", "//tensorflow/compiler/xla/service:computation_layout", + "//tensorflow/compiler/xla/service:computation_placer", "//tensorflow/compiler/xla/service:executable", "//tensorflow/compiler/xla/service:hlo", "//tensorflow/compiler/xla/service:hlo_execution_profile", "//tensorflow/compiler/xla/service:hlo_graph_dumper", - "//tensorflow/compiler/xla/service:hlo_module_config", "//tensorflow/compiler/xla/service:transfer_manager", + "//tensorflow/core:core_cpu_internal", "//tensorflow/core:lib", "//tensorflow/core:stream_executor_no_cuda", "//tensorflow/core:test", @@ -114,9 +129,14 @@ cc_binary( name = "local_client_aot_test_helper", srcs = ["local_client_aot_test_helper.cc"], deps = [ + "//tensorflow/compiler/xla:types", + "//tensorflow/compiler/xla:util", "//tensorflow/compiler/xla/client:client_library", "//tensorflow/compiler/xla/client:computation_builder", + "//tensorflow/compiler/xla/service/cpu:cpu_compiler", + "//tensorflow/compiler/xla/service/llvm_ir:llvm_util", "//tensorflow/core:lib", + "@llvm//:support", ], ) @@ -137,6 +157,7 @@ cc_library( "//tensorflow/compiler/xla:array2d", "//tensorflow/compiler/xla:array3d", "//tensorflow/compiler/xla:array4d", + "//tensorflow/compiler/xla:execution_options_util", "//tensorflow/compiler/xla:literal_util", "//tensorflow/compiler/xla:shape_util", "//tensorflow/compiler/xla:status_macros", @@ -149,7 +170,6 @@ cc_library( "//tensorflow/compiler/xla/client:computation_builder", "//tensorflow/compiler/xla/client:global_data", "//tensorflow/compiler/xla/client:local_client", - "//tensorflow/compiler/xla/legacy_flags:hlo_pass_pipeline_flags", "//tensorflow/compiler/xla/tests:literal_test_util", "//tensorflow/compiler/xla/tests:test_utils", "//tensorflow/core:lib", @@ -158,23 +178,44 @@ cc_library( ], ) +cc_library( + name = "llvm_irgen_test_base", + testonly = True, + srcs = ["llvm_irgen_test_base.cc"], + hdrs = ["llvm_irgen_test_base.h"], + deps = [ + ":codegen_test_base", + ":filecheck", + "//tensorflow/compiler/xla/service:llvm_compiler", + "//tensorflow/compiler/xla/service/llvm_ir:llvm_util", + "//tensorflow/core:test", + ], +) + cc_library( name = "codegen_test_base", testonly = True, srcs = ["codegen_test_base.cc"], hdrs = ["codegen_test_base.h"], + deps = [ + ":hlo_test_base", + "//tensorflow/compiler/xla/service:compiler", + "//tensorflow/compiler/xla/service:executable", + "//tensorflow/compiler/xla/service:hlo", + ], +) + +cc_library( + name = "filecheck", + testonly = True, + srcs = ["filecheck.cc"], + hdrs = ["filecheck.h"], data = [ "@llvm//:FileCheck", ], deps = [ "//tensorflow/compiler/xla:statusor", "//tensorflow/compiler/xla:util", - "//tensorflow/compiler/xla/service:backend", - "//tensorflow/compiler/xla/service:compiler", - "//tensorflow/compiler/xla/service:executable", - "//tensorflow/compiler/xla/service:hlo", - "//tensorflow/compiler/xla/service:hlo_module_config", - "//tensorflow/compiler/xla/tests:hlo_test_base", "//tensorflow/core:lib", "//tensorflow/core:test", ], @@ -195,14 +236,17 @@ cc_library( "//tensorflow/compiler/xla/client:client_library", "//tensorflow/compiler/xla/client:computation", "//tensorflow/compiler/xla/client:local_client", + "//tensorflow/compiler/xla/service:computation_placer", "//tensorflow/compiler/xla/service:device_memory_allocator", "//tensorflow/compiler/xla/service:local_service", "//tensorflow/compiler/xla/service:platform_util", "//tensorflow/compiler/xla/service:shaped_buffer", "//tensorflow/compiler/xla/service:transfer_manager", "//tensorflow/compiler/xla/tests:client_library_test_base", + "//tensorflow/core:core_cpu_internal", "//tensorflow/core:lib", "//tensorflow/core:stream_executor_no_cuda", + "//third_party/eigen3", ], ) @@ -211,13 +255,15 @@ xla_test( srcs = ["bad_rng_shape_validation_test.cc"], deps = [ "//tensorflow/compiler/xla:statusor", + "//tensorflow/compiler/xla:test", "//tensorflow/compiler/xla:test_helpers", + "//tensorflow/compiler/xla:types", "//tensorflow/compiler/xla:xla_data_proto", "//tensorflow/compiler/xla/client:computation", "//tensorflow/compiler/xla/client:computation_builder", "//tensorflow/compiler/xla/client:local_client", - "//tensorflow/compiler/xla/legacy_flags:cpu_compiler_flags", "//tensorflow/compiler/xla/tests:client_library_test_base", + "//tensorflow/compiler/xla/tests:xla_internal_test_main", "//tensorflow/core:lib", "//tensorflow/core:test", ], @@ -230,13 +276,14 @@ xla_test( "//tensorflow/compiler/xla:literal_util", "//tensorflow/compiler/xla:shape_util", "//tensorflow/compiler/xla:statusor", + "//tensorflow/compiler/xla:test", "//tensorflow/compiler/xla:test_helpers", "//tensorflow/compiler/xla:xla_data_proto", "//tensorflow/compiler/xla/client:computation_builder", "//tensorflow/compiler/xla/client:global_data", "//tensorflow/compiler/xla/client:local_client", - "//tensorflow/compiler/xla/legacy_flags:cpu_compiler_flags", "//tensorflow/compiler/xla/tests:client_library_test_base", + "//tensorflow/compiler/xla/tests:xla_internal_test_main", "//tensorflow/core:test", ], ) @@ -251,8 +298,8 @@ xla_test( "//tensorflow/compiler/xla:xla_data_proto", "//tensorflow/compiler/xla/client:computation_builder", "//tensorflow/compiler/xla/client:local_client", - "//tensorflow/compiler/xla/legacy_flags:cpu_compiler_flags", "//tensorflow/compiler/xla/tests:client_library_test_base", + "//tensorflow/compiler/xla/tests:xla_internal_test_main", "//tensorflow/core:test", ], ) @@ -263,6 +310,7 @@ xla_test( deps = [ "//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:client_library", @@ -270,10 +318,10 @@ xla_test( "//tensorflow/compiler/xla/client:computation_builder", "//tensorflow/compiler/xla/client:local_client", "//tensorflow/compiler/xla/client/lib:arithmetic", - "//tensorflow/compiler/xla/legacy_flags:cpu_compiler_flags", "//tensorflow/compiler/xla/service:platform_util", "//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", ], @@ -285,9 +333,9 @@ xla_test( deps = [ "//tensorflow/compiler/xla/client:computation_builder", "//tensorflow/compiler/xla/client:local_client", - "//tensorflow/compiler/xla/legacy_flags:cpu_compiler_flags", "//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", ], ) @@ -300,6 +348,7 @@ xla_test( "//tensorflow/compiler/xla:literal_util", "//tensorflow/compiler/xla:shape_util", "//tensorflow/compiler/xla:statusor", + "//tensorflow/compiler/xla:test", "//tensorflow/compiler/xla:test_helpers", "//tensorflow/compiler/xla:xla_data_proto", "//tensorflow/compiler/xla:xla_proto", @@ -308,10 +357,10 @@ xla_test( "//tensorflow/compiler/xla/client:global_data", "//tensorflow/compiler/xla/client:local_client", "//tensorflow/compiler/xla/client/lib:arithmetic", - "//tensorflow/compiler/xla/legacy_flags:cpu_compiler_flags", "//tensorflow/compiler/xla/tests:client_library_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:stream_executor_no_cuda", "//tensorflow/core:test", @@ -331,9 +380,9 @@ xla_test( "//tensorflow/compiler/xla/client:computation_builder", "//tensorflow/compiler/xla/client:global_data", "//tensorflow/compiler/xla/client:local_client", - "//tensorflow/compiler/xla/legacy_flags:cpu_compiler_flags", "//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", ], @@ -347,8 +396,8 @@ xla_test( "//tensorflow/compiler/xla/client:computation_builder", "//tensorflow/compiler/xla/client:local_client", "//tensorflow/compiler/xla/client/lib:arithmetic", - "//tensorflow/compiler/xla/legacy_flags:cpu_compiler_flags", "//tensorflow/compiler/xla/tests:client_library_test_base", + "//tensorflow/compiler/xla/tests:xla_internal_test_main", "//tensorflow/core:test", ], ) @@ -361,9 +410,9 @@ xla_test( "//tensorflow/compiler/xla/client:computation_builder", "//tensorflow/compiler/xla/client:global_data", "//tensorflow/compiler/xla/client:local_client", - "//tensorflow/compiler/xla/legacy_flags:cpu_compiler_flags", "//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", ], @@ -378,9 +427,9 @@ xla_test( "//tensorflow/compiler/xla/client:global_data", "//tensorflow/compiler/xla/client:local_client", "//tensorflow/compiler/xla/client/lib:arithmetic", - "//tensorflow/compiler/xla/legacy_flags:cpu_compiler_flags", "//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", ], @@ -399,9 +448,9 @@ xla_test( "//tensorflow/compiler/xla/client:computation_builder", "//tensorflow/compiler/xla/client:global_data", "//tensorflow/compiler/xla/client:local_client", - "//tensorflow/compiler/xla/legacy_flags:cpu_compiler_flags", "//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", ], @@ -412,13 +461,14 @@ xla_test( srcs = ["deallocation_test.cc"], deps = [ "//tensorflow/compiler/xla:statusor", + "//tensorflow/compiler/xla:test", "//tensorflow/compiler/xla:test_helpers", "//tensorflow/compiler/xla/client:computation", "//tensorflow/compiler/xla/client:computation_builder", "//tensorflow/compiler/xla/client:global_data", "//tensorflow/compiler/xla/client:local_client", - "//tensorflow/compiler/xla/legacy_flags:cpu_compiler_flags", "//tensorflow/compiler/xla/tests:client_library_test_base", + "//tensorflow/compiler/xla/tests:xla_internal_test_main", "//tensorflow/core:lib", "//tensorflow/core:test", ], @@ -431,14 +481,15 @@ xla_test( "//tensorflow/compiler/xla:literal_util", "//tensorflow/compiler/xla:shape_util", "//tensorflow/compiler/xla:statusor", + "//tensorflow/compiler/xla:test", "//tensorflow/compiler/xla:test_helpers", "//tensorflow/compiler/xla:xla_data_proto", "//tensorflow/compiler/xla/client:computation", "//tensorflow/compiler/xla/client:computation_builder", "//tensorflow/compiler/xla/client:global_data", "//tensorflow/compiler/xla/client:local_client", - "//tensorflow/compiler/xla/legacy_flags:cpu_compiler_flags", "//tensorflow/compiler/xla/tests:client_library_test_base", + "//tensorflow/compiler/xla/tests:xla_internal_test_main", "//tensorflow/core:lib", "//tensorflow/core:test", ], @@ -455,50 +506,44 @@ xla_test( "//tensorflow/compiler/xla:literal_util", "//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:computation_builder", "//tensorflow/compiler/xla/client:global_data", "//tensorflow/compiler/xla/client:local_client", - "//tensorflow/compiler/xla/legacy_flags:cpu_compiler_flags", "//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", ], ) xla_test( - name = "dot_operation_test", - srcs = ["dot_operation_test.cc"], + name = "reduce_precision_test", + srcs = ["reduce_precision_test.cc"], deps = [ "//tensorflow/compiler/xla:array2d", - "//tensorflow/compiler/xla:array3d", - "//tensorflow/compiler/xla:reference_util", + "//tensorflow/compiler/xla:literal_util", "//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:computation_builder", + "//tensorflow/compiler/xla/client:global_data", "//tensorflow/compiler/xla/client:local_client", - "//tensorflow/compiler/xla/legacy_flags:cpu_compiler_flags", - "//tensorflow/compiler/xla/legacy_flags:cpu_runtime_flags", - "//tensorflow/compiler/xla/legacy_flags:layout_util_flags", + "//tensorflow/compiler/xla/service:reduce_precision_insertion", "//tensorflow/compiler/xla/tests:client_library_test_base", "//tensorflow/compiler/xla/tests:literal_test_util", - "//tensorflow/compiler/xla/tests:test_utils", - "//tensorflow/core:framework_internal", + "//tensorflow/compiler/xla/tests:xla_internal_test_main", "//tensorflow/core:lib", - "//tensorflow/core:test", ], ) -# Tests the dot operation in some cases that can be performed via a -# runtime call on some backends - e.g. a runtime call to to Eigen. xla_test( - name = "dot_operation_runtime_test", + name = "dot_operation_test", srcs = ["dot_operation_test.cc"], - backend_args = { - "cpu": ["--xla_cpu_use_eigen"], - "cpu_parallel": ["--xla_cpu_use_eigen"], - }, deps = [ "//tensorflow/compiler/xla:array2d", "//tensorflow/compiler/xla:array3d", @@ -506,32 +551,21 @@ xla_test( "//tensorflow/compiler/xla:shape_util", "//tensorflow/compiler/xla/client:computation_builder", "//tensorflow/compiler/xla/client:local_client", - "//tensorflow/compiler/xla/legacy_flags:cpu_compiler_flags", - "//tensorflow/compiler/xla/legacy_flags:cpu_runtime_flags", - "//tensorflow/compiler/xla/legacy_flags:layout_util_flags", "//tensorflow/compiler/xla/tests:client_library_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:framework_internal", "//tensorflow/core:lib", "//tensorflow/core:test", ], ) -# Repeat dot_operation_runtime_test with single-threded eigen. +# Tests the dot operation in some cases that can be performed via a +# runtime call on some backends - e.g. a runtime call to Eigen. xla_test( - name = "dot_operation_single_threaded_runtime_test", + name = "dot_operation_runtime_test", srcs = ["dot_operation_test.cc"], - backend_args = { - "cpu": [ - "--xla_cpu_use_eigen", - "--xla_cpu_multi_thread_eigen=false", - ], - "cpu_parallel": [ - "--xla_cpu_use_eigen", - "--xla_cpu_multi_thread_eigen=false", - ], - }, deps = [ "//tensorflow/compiler/xla:array2d", "//tensorflow/compiler/xla:array3d", @@ -539,29 +573,26 @@ xla_test( "//tensorflow/compiler/xla:shape_util", "//tensorflow/compiler/xla/client:computation_builder", "//tensorflow/compiler/xla/client:local_client", - "//tensorflow/compiler/xla/legacy_flags:cpu_compiler_flags", - "//tensorflow/compiler/xla/legacy_flags:cpu_runtime_flags", - "//tensorflow/compiler/xla/legacy_flags:layout_util_flags", "//tensorflow/compiler/xla/tests:client_library_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:framework_internal", "//tensorflow/core:lib", "//tensorflow/core:test", ], ) +# Repeat dot_operation_runtime_test with single-threded eigen. xla_test( - name = "dot_operation_rowmajor_runtime_test", + name = "dot_operation_single_threaded_runtime_test", srcs = ["dot_operation_test.cc"], backend_args = { "cpu": [ - "--xla_cpu_use_eigen", - "--xla_default_layout=major2minor", + "--xla_cpu_multi_thread_eigen=false", ], "cpu_parallel": [ - "--xla_cpu_use_eigen", - "--xla_default_layout=major2minor", + "--xla_cpu_multi_thread_eigen=false", ], }, deps = [ @@ -571,12 +602,10 @@ xla_test( "//tensorflow/compiler/xla:shape_util", "//tensorflow/compiler/xla/client:computation_builder", "//tensorflow/compiler/xla/client:local_client", - "//tensorflow/compiler/xla/legacy_flags:cpu_compiler_flags", - "//tensorflow/compiler/xla/legacy_flags:cpu_runtime_flags", - "//tensorflow/compiler/xla/legacy_flags:layout_util_flags", "//tensorflow/compiler/xla/tests:client_library_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:framework_internal", "//tensorflow/core:lib", "//tensorflow/core:test", @@ -593,10 +622,10 @@ xla_test( "//tensorflow/compiler/xla:xla_data_proto", "//tensorflow/compiler/xla/client:computation_builder", "//tensorflow/compiler/xla/client:local_client", - "//tensorflow/compiler/xla/legacy_flags:cpu_compiler_flags", "//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:test", ], ) @@ -612,9 +641,9 @@ xla_test( "//tensorflow/compiler/xla:xla_data_proto", "//tensorflow/compiler/xla/client:computation_builder", "//tensorflow/compiler/xla/client:local_client", - "//tensorflow/compiler/xla/legacy_flags:cpu_compiler_flags", "//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", ], @@ -638,9 +667,9 @@ xla_test( "//tensorflow/compiler/xla/client:global_data", "//tensorflow/compiler/xla/client:local_client", "//tensorflow/compiler/xla/client:padding", - "//tensorflow/compiler/xla/legacy_flags:cpu_compiler_flags", "//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", ], @@ -665,9 +694,9 @@ xla_test( "//tensorflow/compiler/xla/client:computation_builder", "//tensorflow/compiler/xla/client:local_client", "//tensorflow/compiler/xla/client:padding", - "//tensorflow/compiler/xla/legacy_flags:cpu_compiler_flags", "//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", ], @@ -682,14 +711,15 @@ xla_test( "//tensorflow/compiler/xla:array4d", "//tensorflow/compiler/xla:reference_util", "//tensorflow/compiler/xla:statusor", + "//tensorflow/compiler/xla:test", "//tensorflow/compiler/xla:test_helpers", "//tensorflow/compiler/xla:util", "//tensorflow/compiler/xla/client:computation_builder", "//tensorflow/compiler/xla/client:local_client", "//tensorflow/compiler/xla/client:padding", - "//tensorflow/compiler/xla/legacy_flags:cpu_compiler_flags", "//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", ], @@ -698,20 +728,30 @@ xla_test( xla_test( name = "batch_normalization_test", srcs = ["batch_normalization_test.cc"], + shard_count = 40, deps = [ + ":test_utils", "//tensorflow/compiler/xla:array2d", "//tensorflow/compiler/xla:array4d", "//tensorflow/compiler/xla:literal_util", + "//tensorflow/compiler/xla:reference_util", "//tensorflow/compiler/xla:shape_util", + "//tensorflow/compiler/xla:status_macros", "//tensorflow/compiler/xla:statusor", + "//tensorflow/compiler/xla:test", + "//tensorflow/compiler/xla:test_helpers", + "//tensorflow/compiler/xla:util", "//tensorflow/compiler/xla:xla_data_proto", "//tensorflow/compiler/xla/client:computation", "//tensorflow/compiler/xla/client:computation_builder", + "//tensorflow/compiler/xla/client:global_data", "//tensorflow/compiler/xla/client:local_client", "//tensorflow/compiler/xla/client/lib:arithmetic", - "//tensorflow/compiler/xla/legacy_flags:cpu_compiler_flags", + "//tensorflow/compiler/xla/service:hlo", "//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", "//tensorflow/core:test", ], @@ -726,9 +766,9 @@ xla_test( "//tensorflow/compiler/xla:reference_util", "//tensorflow/compiler/xla/client:computation_builder", "//tensorflow/compiler/xla/client:local_client", - "//tensorflow/compiler/xla/legacy_flags:cpu_compiler_flags", "//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", ], @@ -742,9 +782,9 @@ xla_test( "//tensorflow/compiler/xla:array3d", "//tensorflow/compiler/xla/client:computation_builder", "//tensorflow/compiler/xla/client:local_client", - "//tensorflow/compiler/xla/legacy_flags:cpu_compiler_flags", "//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", ], ) @@ -756,11 +796,12 @@ xla_test( deps = [ "//tensorflow/compiler/xla:array2d", "//tensorflow/compiler/xla:reference_util", + "//tensorflow/compiler/xla:test_helpers", "//tensorflow/compiler/xla/client:client_library", "//tensorflow/compiler/xla/client:computation", "//tensorflow/compiler/xla/client:computation_builder", "//tensorflow/compiler/xla/client:local_client", - "//tensorflow/compiler/xla/legacy_flags:cpu_compiler_flags", + "//tensorflow/compiler/xla/service:computation_placer", "//tensorflow/compiler/xla/service:device_memory_allocator", "//tensorflow/compiler/xla/service:local_service", "//tensorflow/compiler/xla/service:platform_util", @@ -768,6 +809,7 @@ xla_test( "//tensorflow/compiler/xla/service:transfer_manager", "//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:stream_executor_no_cuda", "//tensorflow/core:test", @@ -787,9 +829,9 @@ xla_test( "//tensorflow/compiler/xla/client:computation", "//tensorflow/compiler/xla/client:computation_builder", "//tensorflow/compiler/xla/client:local_client", - "//tensorflow/compiler/xla/legacy_flags:cpu_compiler_flags", "//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", ], ) @@ -804,9 +846,9 @@ xla_test( "//tensorflow/compiler/xla/client:computation_builder", "//tensorflow/compiler/xla/client:local_client", "//tensorflow/compiler/xla/client/lib:arithmetic", - "//tensorflow/compiler/xla/legacy_flags:cpu_compiler_flags", "//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", ], ) @@ -830,19 +872,22 @@ xla_test( "//tensorflow/compiler/xla/client:global_data", "//tensorflow/compiler/xla/client:local_client", "//tensorflow/compiler/xla/client/lib:arithmetic", - "//tensorflow/compiler/xla/legacy_flags:cpu_compiler_flags", "//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", ], ) -xla_test( - name = "reduce_window_test", - timeout = "long", +# External xla_test targets can add "reduce_window_test_library" to xla_test_library_deps, in order +# to refer to the cc_library compiled with the correct backend macros. The following test target +# "reduce_window_test" is an example. +xla_test_library( + name = "reduce_window_test_library", srcs = ["reduce_window_test.cc"], deps = [ + ":test_macros_header", "//tensorflow/compiler/xla:array2d", "//tensorflow/compiler/xla:array3d", "//tensorflow/compiler/xla:array4d", @@ -853,14 +898,23 @@ xla_test( "//tensorflow/compiler/xla/client:local_client", "//tensorflow/compiler/xla/client:padding", "//tensorflow/compiler/xla/client/lib:arithmetic", - "//tensorflow/compiler/xla/legacy_flags:cpu_compiler_flags", "//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", "//tensorflow/core:test", ], ) +xla_test( + name = "reduce_window_test", + timeout = "long", + srcs = [], + xla_test_library_deps = [":reduce_window_test_library"], + deps = [], +) + xla_test( name = "select_and_scatter_test", timeout = "long", @@ -877,9 +931,9 @@ xla_test( "//tensorflow/compiler/xla/client:local_client", "//tensorflow/compiler/xla/client:padding", "//tensorflow/compiler/xla/client/lib:arithmetic", - "//tensorflow/compiler/xla/legacy_flags:cpu_compiler_flags", "//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", ], @@ -889,14 +943,15 @@ xla_test( name = "copy_test", srcs = ["copy_test.cc"], deps = [ + ":client_library_test_base", "//tensorflow/compiler/xla:array2d", "//tensorflow/compiler/xla:literal_util", "//tensorflow/compiler/xla:util", "//tensorflow/compiler/xla:xla_data_proto", - "//tensorflow/compiler/xla/legacy_flags:cpu_compiler_flags", "//tensorflow/compiler/xla/service:hlo", "//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", "//tensorflow/core:test", ], @@ -908,12 +963,13 @@ xla_test( deps = [ "//tensorflow/compiler/xla:literal_util", "//tensorflow/compiler/xla:shape_util", + "//tensorflow/compiler/xla:test_helpers", "//tensorflow/compiler/xla:xla_data_proto", "//tensorflow/compiler/xla/client:computation", "//tensorflow/compiler/xla/client:computation_builder", - "//tensorflow/compiler/xla/legacy_flags:cpu_compiler_flags", "//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", ], @@ -928,10 +984,10 @@ xla_test( "//tensorflow/compiler/xla:shape_util", "//tensorflow/compiler/xla:util", "//tensorflow/compiler/xla:xla_data_proto", - "//tensorflow/compiler/xla/legacy_flags:cpu_compiler_flags", "//tensorflow/compiler/xla/service:hlo", "//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", "//tensorflow/core:test", ], @@ -946,9 +1002,9 @@ xla_test( "//tensorflow/compiler/xla:reference_util", "//tensorflow/compiler/xla/client:computation_builder", "//tensorflow/compiler/xla/client:local_client", - "//tensorflow/compiler/xla/legacy_flags:cpu_compiler_flags", "//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", ], ) @@ -961,13 +1017,12 @@ xla_test( "//tensorflow/compiler/xla:array4d", "//tensorflow/compiler/xla:literal_util", "//tensorflow/compiler/xla:statusor", - "//tensorflow/compiler/xla:test_helpers", + "//tensorflow/compiler/xla:test", "//tensorflow/compiler/xla/client:computation_builder", "//tensorflow/compiler/xla/client:local_client", - "//tensorflow/compiler/xla/legacy_flags:cpu_compiler_flags", "//tensorflow/compiler/xla/tests:client_library_test_base", "//tensorflow/compiler/xla/tests:literal_test_util", - "//tensorflow/core:test", + "//tensorflow/compiler/xla/tests:xla_internal_test_main", ], ) @@ -983,9 +1038,9 @@ xla_test( "//tensorflow/compiler/xla/client:computation_builder", "//tensorflow/compiler/xla/client:local_client", "//tensorflow/compiler/xla/client/lib:arithmetic", - "//tensorflow/compiler/xla/legacy_flags:cpu_compiler_flags", "//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", ], @@ -997,9 +1052,9 @@ xla_test( deps = [ "//tensorflow/compiler/xla/client:computation_builder", "//tensorflow/compiler/xla/client:local_client", - "//tensorflow/compiler/xla/legacy_flags:cpu_compiler_flags", "//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", ], ) @@ -1010,9 +1065,9 @@ xla_test( deps = [ "//tensorflow/compiler/xla/client:computation_builder", "//tensorflow/compiler/xla/client:local_client", - "//tensorflow/compiler/xla/legacy_flags:cpu_compiler_flags", "//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", ], ) @@ -1032,9 +1087,9 @@ xla_test( "//tensorflow/compiler/xla/client:computation", "//tensorflow/compiler/xla/client:computation_builder", "//tensorflow/compiler/xla/client:local_client", - "//tensorflow/compiler/xla/legacy_flags:cpu_compiler_flags", "//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", ], @@ -1046,12 +1101,13 @@ xla_test( deps = [ "//tensorflow/compiler/xla:literal_util", "//tensorflow/compiler/xla:shape_util", + "//tensorflow/compiler/xla:test", "//tensorflow/compiler/xla:util", "//tensorflow/compiler/xla:xla_data_proto", "//tensorflow/compiler/xla/client:computation_builder", "//tensorflow/compiler/xla/client:local_client", - "//tensorflow/compiler/xla/legacy_flags:cpu_compiler_flags", "//tensorflow/compiler/xla/tests:client_library_test_base", + "//tensorflow/compiler/xla/tests:xla_internal_test_main", "//tensorflow/core:lib", "//tensorflow/core:test", ], @@ -1069,15 +1125,16 @@ xla_test( "//tensorflow/compiler/xla:shape_util", "//tensorflow/compiler/xla:status_macros", "//tensorflow/compiler/xla:statusor", + "//tensorflow/compiler/xla:test", "//tensorflow/compiler/xla:test_helpers", "//tensorflow/compiler/xla:xla_data_proto", "//tensorflow/compiler/xla/client:computation", "//tensorflow/compiler/xla/client:computation_builder", "//tensorflow/compiler/xla/client:global_data", "//tensorflow/compiler/xla/client:local_client", - "//tensorflow/compiler/xla/legacy_flags:cpu_compiler_flags", "//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", ], @@ -1091,9 +1148,9 @@ xla_test( "//tensorflow/compiler/xla:array4d", "//tensorflow/compiler/xla/client:computation_builder", "//tensorflow/compiler/xla/client:local_client", - "//tensorflow/compiler/xla/legacy_flags:cpu_compiler_flags", "//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", ], @@ -1113,9 +1170,9 @@ xla_test( "//tensorflow/compiler/xla/client:global_data", "//tensorflow/compiler/xla/client:local_client", "//tensorflow/compiler/xla/client/lib:arithmetic", - "//tensorflow/compiler/xla/legacy_flags:cpu_compiler_flags", "//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:stream_executor_no_cuda", "//tensorflow/core:test", @@ -1130,13 +1187,14 @@ xla_test( "//tensorflow/compiler/xla:array3d", "//tensorflow/compiler/xla:reference_util", "//tensorflow/compiler/xla:statusor", + "//tensorflow/compiler/xla:test", "//tensorflow/compiler/xla:test_helpers", "//tensorflow/compiler/xla/client:computation", "//tensorflow/compiler/xla/client:computation_builder", "//tensorflow/compiler/xla/client:local_client", - "//tensorflow/compiler/xla/legacy_flags:cpu_compiler_flags", "//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", ], ) @@ -1149,9 +1207,9 @@ xla_test( "//tensorflow/compiler/xla:xla_data_proto", "//tensorflow/compiler/xla/client:computation_builder", "//tensorflow/compiler/xla/client:local_client", - "//tensorflow/compiler/xla/legacy_flags:cpu_compiler_flags", "//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:stream_executor_no_cuda", "//tensorflow/core:test", @@ -1171,10 +1229,10 @@ xla_test( "//tensorflow/compiler/xla/client:computation_builder", "//tensorflow/compiler/xla/client:global_data", "//tensorflow/compiler/xla/client:local_client", - "//tensorflow/compiler/xla/legacy_flags:cpu_compiler_flags", "//tensorflow/compiler/xla/tests:client_library_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", ], @@ -1186,9 +1244,9 @@ xla_test( deps = [ "//tensorflow/compiler/xla/client:computation_builder", "//tensorflow/compiler/xla/client:local_client", - "//tensorflow/compiler/xla/legacy_flags:cpu_compiler_flags", "//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", ], @@ -1202,16 +1260,16 @@ xla_test( "//tensorflow/compiler/xla:shape_util", "//tensorflow/compiler/xla:status_macros", "//tensorflow/compiler/xla:statusor", + "//tensorflow/compiler/xla:test", "//tensorflow/compiler/xla:test_helpers", "//tensorflow/compiler/xla:xla_data_proto", + "//tensorflow/compiler/xla/client:client_library", "//tensorflow/compiler/xla/client:computation", "//tensorflow/compiler/xla/client:computation_builder", "//tensorflow/compiler/xla/client:global_data", - "//tensorflow/compiler/xla/client:local_client", - "//tensorflow/compiler/xla/legacy_flags:cpu_compiler_flags", - "//tensorflow/compiler/xla/tests:client_library_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", ], @@ -1228,31 +1286,10 @@ xla_test( "//tensorflow/compiler/xla/client:computation_builder", "//tensorflow/compiler/xla/client:global_data", "//tensorflow/compiler/xla/client:local_client", - "//tensorflow/compiler/xla/legacy_flags:cpu_compiler_flags", "//tensorflow/compiler/xla/tests:client_library_test_base", "//tensorflow/compiler/xla/tests:literal_test_util", "//tensorflow/compiler/xla/tests:test_utils", - "//tensorflow/core:lib", - "//tensorflow/core:test", - ], -) - -xla_test( - name = "inprocess_service_test", - srcs = ["inprocess_service_test.cc"], - deps = [ - "//tensorflow/compiler/xla:array2d", - "//tensorflow/compiler/xla:shape_util", - "//tensorflow/compiler/xla:statusor", - "//tensorflow/compiler/xla:test_helpers", - "//tensorflow/compiler/xla:xla_data_proto", - "//tensorflow/compiler/xla/client:computation", - "//tensorflow/compiler/xla/client:computation_builder", - "//tensorflow/compiler/xla/client:global_data", - "//tensorflow/compiler/xla/client:local_client", - "//tensorflow/compiler/xla/legacy_flags:cpu_compiler_flags", - "//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", ], @@ -1271,10 +1308,10 @@ xla_test( "//tensorflow/compiler/xla/client:computation_builder", "//tensorflow/compiler/xla/client:global_data", "//tensorflow/compiler/xla/client:local_client", - "//tensorflow/compiler/xla/legacy_flags:cpu_compiler_flags", "//tensorflow/compiler/xla/service:session_proto", "//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", ], @@ -1288,14 +1325,34 @@ xla_test( "//tensorflow/compiler/xla:shape_util", "//tensorflow/compiler/xla:util", "//tensorflow/compiler/xla:xla_data_proto", - "//tensorflow/compiler/xla/legacy_flags:cpu_compiler_flags", "//tensorflow/compiler/xla/service:hlo", "//tensorflow/compiler/xla/tests:hlo_test_base", "//tensorflow/compiler/xla/tests:literal_test_util", + "//tensorflow/compiler/xla/tests:xla_internal_test_main", "//tensorflow/core:test", ], ) +xla_test( + name = "llvm_compiler_test", + srcs = ["llvm_compiler_test.cc"], + backends = [ + "cpu", + "gpu", + "cpu_parallel", + ], + deps = [ + "//tensorflow/compiler/xla:literal_util", + "//tensorflow/compiler/xla/service:hlo", + "//tensorflow/compiler/xla/service:llvm_compiler", + "//tensorflow/compiler/xla/tests:hlo_test_base", + "//tensorflow/compiler/xla/tests:literal_test_util", + "//tensorflow/compiler/xla/tests:xla_internal_test_main", + "//tensorflow/core:test", + "@llvm//:core", + ], +) + xla_test( name = "round_trip_packed_literal_test", srcs = ["round_trip_packed_literal_test.cc"], @@ -1307,9 +1364,9 @@ xla_test( "//tensorflow/compiler/xla:xla_data_proto", "//tensorflow/compiler/xla/client:global_data", "//tensorflow/compiler/xla/client:local_client", - "//tensorflow/compiler/xla/legacy_flags:cpu_compiler_flags", "//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", ], @@ -1324,10 +1381,41 @@ xla_test( "//tensorflow/compiler/xla:shape_util", "//tensorflow/compiler/xla:util", "//tensorflow/compiler/xla:xla_data_proto", - "//tensorflow/compiler/xla/legacy_flags:cpu_compiler_flags", + "//tensorflow/compiler/xla/client:client_library", + "//tensorflow/compiler/xla/client:computation", + "//tensorflow/compiler/xla/client:computation_builder", + "//tensorflow/compiler/xla/client:local_client", "//tensorflow/compiler/xla/service:hlo", + "//tensorflow/compiler/xla/service:platform_util", + "//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", + "//tensorflow/core:test", + ], +) + +xla_test( + name = "multioutput_fusion_test", + srcs = ["multioutput_fusion_test.cc"], + deps = [ + "//tensorflow/compiler/xla:array2d", + "//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", + "//tensorflow/compiler/xla/client:computation", + "//tensorflow/compiler/xla/client:computation_builder", + "//tensorflow/compiler/xla/client:local_client", + "//tensorflow/compiler/xla/service:hlo", + "//tensorflow/compiler/xla/service:platform_util", + "//tensorflow/compiler/xla/tests:client_library_test_base", + "//tensorflow/compiler/xla/tests:hlo_test_base", + "//tensorflow/compiler/xla/tests:literal_test_util", + "//tensorflow/compiler/xla/tests:test_utils", + "//tensorflow/compiler/xla/tests:xla_internal_test_main", "//tensorflow/core:lib", "//tensorflow/core:test", ], @@ -1342,6 +1430,7 @@ cc_test( linkstatic = 1, deps = [ "//tensorflow/compiler/xla:executable_run_options", + "//tensorflow/core:lib", "//tensorflow/core:test", "//tensorflow/core:test_main", ], @@ -1355,8 +1444,8 @@ cc_test( deps = [ ":local_client_test_base", "//tensorflow/compiler/xla:test_helpers", - "//tensorflow/compiler/xla/client:client_library", "//tensorflow/compiler/xla/client:computation_builder", + "//tensorflow/compiler/xla/client:local_client", "//tensorflow/compiler/xla/service:computation_tracker", "//tensorflow/compiler/xla/service:local_service", "//tensorflow/core:test_main", @@ -1374,10 +1463,10 @@ xla_test( "//tensorflow/compiler/xla:xla_data_proto", "//tensorflow/compiler/xla/client:global_data", "//tensorflow/compiler/xla/client:local_client", - "//tensorflow/compiler/xla/legacy_flags:cpu_compiler_flags", "//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", "//tensorflow/core:test", ], @@ -1390,10 +1479,10 @@ xla_test( "//tensorflow/compiler/xla:shape_util", "//tensorflow/compiler/xla/client:computation_builder", "//tensorflow/compiler/xla/client:local_client", - "//tensorflow/compiler/xla/legacy_flags:cpu_compiler_flags", "//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", "//tensorflow/core:test", ], @@ -1416,14 +1505,23 @@ xla_test( "//tensorflow/compiler/xla/client:computation_builder", "//tensorflow/compiler/xla/client:global_data", "//tensorflow/compiler/xla/client:local_client", - "//tensorflow/compiler/xla/legacy_flags:cpu_compiler_flags", "//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", ], ) +xla_test( + name = "deep_graph_test", + srcs = ["deep_graph_test.cc"], + deps = [ + "//tensorflow/compiler/xla/tests:client_library_test_base", + "//tensorflow/compiler/xla/tests:xla_internal_test_main", + ], +) + cc_test( name = "literal_test_util_test", srcs = ["literal_test_util_test.cc"], diff --git a/tensorflow/compiler/xla/tests/array_elementwise_ops_test.cc b/tensorflow/compiler/xla/tests/array_elementwise_ops_test.cc index d18511a6b4a98d42640ed22f6aa69c2e66465f8a..532e2394c0d727d77ec0e4ed23f81fdc34a950a6 100644 --- a/tensorflow/compiler/xla/tests/array_elementwise_ops_test.cc +++ b/tensorflow/compiler/xla/tests/array_elementwise_ops_test.cc @@ -26,16 +26,15 @@ limitations under the License. #include "tensorflow/compiler/xla/client/global_data.h" #include "tensorflow/compiler/xla/client/local_client.h" #include "tensorflow/compiler/xla/layout_util.h" -#include "tensorflow/compiler/xla/legacy_flags/cpu_compiler_flags.h" #include "tensorflow/compiler/xla/literal_util.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/compiler/xla/xla_data.pb.h" #include "tensorflow/core/lib/core/casts.h" -#include "tensorflow/core/platform/test.h" #include "tensorflow/core/platform/types.h" namespace xla { @@ -43,7 +42,7 @@ namespace { class ArrayElementwiseOpTest : public ClientLibraryTestBase { public: - ErrorSpec error_spec_{0.0001}; + ErrorSpec error_spec_{0.0001, 0.0001}; }; class ArrayElementwiseOpTestParamCount @@ -58,7 +57,7 @@ XLA_TEST_F(ArrayElementwiseOpTest, NegConstantZeroElementF32) { ComputeAndCompareR1(&builder, {}, {}, error_spec_); } -TEST_F(ArrayElementwiseOpTest, NegConstantF32) { +XLA_TEST_F(ArrayElementwiseOpTest, NegConstantF32) { ComputationBuilder builder(client_, TestName()); auto a = builder.ConstantR1({-2.5f, 3.14f, 2.25f, -10.0f, 6.0f}); auto result = builder.Neg(a); @@ -67,7 +66,7 @@ TEST_F(ArrayElementwiseOpTest, NegConstantF32) { error_spec_); } -TEST_F(ArrayElementwiseOpTest, NegConstantS32) { +XLA_TEST_F(ArrayElementwiseOpTest, NegConstantS32) { ComputationBuilder builder(client_, TestName()); auto a = builder.ConstantR1({-1, 0, 1, 324, std::numeric_limits::min(), @@ -127,7 +126,7 @@ XLA_TEST_F(ArrayElementwiseOpTest, IsFiniteR1F32s) { {}); } -TEST_F(ArrayElementwiseOpTest, AddTwoConstantF32s) { +XLA_TEST_F(ArrayElementwiseOpTest, AddTwoConstantF32s) { ComputationBuilder builder(client_, TestName()); auto a = builder.ConstantR1({-2.5f, 3.14f, 2.25f, -10.0f, 6.0f}); auto b = builder.ConstantR1({100.0f, 3.13f, 2.75f, 10.5f, -999.0f}); @@ -156,13 +155,13 @@ TEST_P(ArrayElementwiseOpTestParamCount, AddManyValues) { b_values.push_back(2 * i / static_cast(count + 2)); } - std::unique_ptr a_literal = LiteralUtil::CreateR1({a_values}); + std::unique_ptr a_literal = Literal::CreateR1({a_values}); std::unique_ptr a_data = client_->TransferToServer(*a_literal).ConsumeValueOrDie(); auto a_constant = builder.ConstantR1(a_values); auto a_param = builder.Parameter(0, a_literal->shape(), "a_param"); - std::unique_ptr b_literal = LiteralUtil::CreateR1({b_values}); + std::unique_ptr b_literal = Literal::CreateR1({b_values}); std::unique_ptr b_data = client_->TransferToServer(*b_literal).ConsumeValueOrDie(); auto b_constant = builder.Parameter(1, a_literal->shape(), "b_param"); @@ -186,7 +185,7 @@ TEST_P(ArrayElementwiseOpTestParamCount, AddManyValues) { error_spec_); } -TEST_F(ArrayElementwiseOpTest, SubTwoConstantF32s) { +XLA_TEST_F(ArrayElementwiseOpTest, SubTwoConstantF32s) { ComputationBuilder builder(client_, TestName()); auto a = builder.ConstantR1({-2.5f, 3.14f, 2.25f, -10.0f, 6.0f}); auto b = builder.ConstantR1({100.0f, 3.13f, 2.75f, 10.5f, -999.0f}); @@ -205,7 +204,7 @@ XLA_TEST_F(ArrayElementwiseOpTest, SubTwoConstantZeroElementF32s) { ComputeAndCompareR1(&builder, {}, {}, error_spec_); } -TEST_F(ArrayElementwiseOpTest, SubTwoConstantS32s) { +XLA_TEST_F(ArrayElementwiseOpTest, SubTwoConstantS32s) { ComputationBuilder builder(client_, TestName()); auto a = builder.ConstantR1({-1, 0, 2, 1000000000}); auto b = builder.ConstantR1({-1, 2, 1, -1}); @@ -223,7 +222,7 @@ XLA_TEST_F(ArrayElementwiseOpTest, SubTwoConstantZeroElementS32s) { ComputeAndCompareR1(&builder, {}, {}); } -TEST_F(ArrayElementwiseOpTest, DivTwoConstantF32s) { +XLA_TEST_F(ArrayElementwiseOpTest, DivTwoConstantF32s) { ComputationBuilder builder(client_, TestName()); auto a = builder.ConstantR1({-2.5f, 25.5f, 2.25f, -10.0f, 6.0f}); auto b = builder.ConstantR1({10.0f, 5.1f, 1.0f, 10.0f, -6.0f}); @@ -242,6 +241,150 @@ XLA_TEST_F(ArrayElementwiseOpTest, DivTwoConstantZeroElementF32s) { ComputeAndCompareR1(&builder, {}, {}, error_spec_); } +XLA_TEST_F(ArrayElementwiseOpTest, DivS32s) { + // clang-format off + // Some interesting values to test. + std::vector vals = { + INT32_MIN, INT32_MIN + 1, INT32_MIN + 2, -0x40000000, -0x3fffffff, + -271181, -1309, -17, -10, -5, -3, -2, -1, 0, 1, 2, 3, 5, 10, 17, 26, 101, + 7919, 0x40000000, INT32_MAX - 2, INT32_MAX - 1, INT32_MAX}; + // clang-format on + + std::vector dividends, divisors, quotients, remainders; + for (int32 divisor : vals) { + if (divisor != 0) { + for (int32 dividend : vals) { + // Avoid integer overflow. + if (dividend != INT32_MIN || divisor != -1) { + dividends.push_back(dividend); + divisors.push_back(divisor); + quotients.push_back(dividend / divisor); + remainders.push_back(dividend % divisor); + } + } + } + } + + { + ComputationBuilder builder(client_, TestName()); + ComputationDataHandle dividend; + ComputationDataHandle divisor; + auto dividend_data = + CreateR1Parameter(dividends, 0, "dividend", &builder, ÷nd); + auto divisor_data = + CreateR1Parameter(divisors, 1, "divisor", &builder, &divisor); + builder.Div(dividend, divisor); + + ComputeAndCompareR1(&builder, quotients, + {dividend_data.get(), divisor_data.get()}); + } + + // Test with a compile-time constant divisor. + { + ComputationBuilder builder(client_, TestName()); + ComputationDataHandle dividend; + auto dividend_data = + CreateR1Parameter(dividends, 0, "dividend", &builder, ÷nd); + builder.Div(dividend, builder.ConstantR1(divisors)); + + ComputeAndCompareR1(&builder, quotients, {dividend_data.get()}); + } + + { + ComputationBuilder builder(client_, TestName()); + ComputationDataHandle dividend; + ComputationDataHandle divisor; + auto dividend_data = + CreateR1Parameter(dividends, 0, "dividend", &builder, ÷nd); + auto divisor_data = + CreateR1Parameter(divisors, 1, "divisor", &builder, &divisor); + builder.Rem(dividend, divisor); + + ComputeAndCompareR1(&builder, remainders, + {dividend_data.get(), divisor_data.get()}); + } + + // Test with a compile-time constant divisor. + { + ComputationBuilder builder(client_, TestName()); + ComputationDataHandle dividend; + auto dividend_data = + CreateR1Parameter(dividends, 0, "dividend", &builder, ÷nd); + builder.Rem(dividend, builder.ConstantR1(divisors)); + + ComputeAndCompareR1(&builder, remainders, {dividend_data.get()}); + } +} + +XLA_TEST_F(ArrayElementwiseOpTest, DivU32s) { + // clang-format off + // Some interesting values to test. + std::vector vals = { + 0, 1, 2, 17, 101, 3333, 0x7FFFFFFF, 0xABCDEF12, 0xCAFEBEEF, 0x80000000, + 0x80000001, UINT32_MAX - 2, UINT32_MAX - 1, UINT32_MAX}; + // clang-format on + + std::vector dividends, divisors, quotients, remainders; + for (uint32 divisor : vals) { + if (divisor != 0) { + for (uint32 dividend : vals) { + dividends.push_back(dividend); + divisors.push_back(divisor); + quotients.push_back(dividend / divisor); + remainders.push_back(dividend % divisor); + } + } + } + + { + ComputationBuilder builder(client_, TestName()); + ComputationDataHandle dividend; + ComputationDataHandle divisor; + auto dividend_data = CreateR1Parameter(dividends, 0, "dividend", + &builder, ÷nd); + auto divisor_data = + CreateR1Parameter(divisors, 1, "divisor", &builder, &divisor); + builder.Div(dividend, divisor); + + ComputeAndCompareR1(&builder, quotients, + {dividend_data.get(), divisor_data.get()}); + } + + { + ComputationBuilder builder(client_, TestName()); + ComputationDataHandle dividend; + auto dividend_data = CreateR1Parameter(dividends, 0, "dividend", + &builder, ÷nd); + builder.Div(dividend, builder.ConstantR1(divisors)); + + ComputeAndCompareR1(&builder, quotients, {dividend_data.get()}); + } + + { + ComputationBuilder builder(client_, TestName()); + ComputationDataHandle dividend; + ComputationDataHandle divisor; + auto dividend_data = CreateR1Parameter(dividends, 0, "dividend", + &builder, ÷nd); + auto divisor_data = + CreateR1Parameter(divisors, 1, "divisor", &builder, &divisor); + builder.Rem(dividend, divisor); + + ComputeAndCompareR1(&builder, remainders, + {dividend_data.get(), divisor_data.get()}); + } + + { + ComputationBuilder builder(client_, TestName()); + ComputationDataHandle dividend; + auto dividend_data = CreateR1Parameter(dividends, 0, "dividend", + &builder, ÷nd); + builder.Rem(dividend, builder.ConstantR1(divisors)); + + ComputeAndCompareR1(&builder, remainders, {dividend_data.get()}); + } +} + XLA_TEST_F(ArrayElementwiseOpTest, RemF32s) { ComputationBuilder builder(client_, TestName()); auto a = builder.ConstantR1( @@ -277,7 +420,7 @@ XLA_TEST_F(ArrayElementwiseOpTest, RemF64s) { error_spec_); } -TEST_F(ArrayElementwiseOpTest, MulTwoConstantF32s) { +XLA_TEST_F(ArrayElementwiseOpTest, MulTwoConstantF32s) { ComputationBuilder builder(client_, TestName()); auto a = builder.ConstantR1({-2.5f, 25.5f, 2.25f, -10.0f, 6.0f}); auto b = builder.ConstantR1({10.0f, 5.0f, 1.0f, 10.0f, -6.0f}); @@ -296,7 +439,7 @@ XLA_TEST_F(ArrayElementwiseOpTest, MulTwoConstantZeroElementF32s) { ComputeAndCompareR1(&builder, {}, {}, error_spec_); } -TEST_F(ArrayElementwiseOpTest, MulTwoConstantS32s) { +XLA_TEST_F(ArrayElementwiseOpTest, MulTwoConstantS32s) { std::vector data = {0, 1, -1, @@ -331,7 +474,7 @@ XLA_TEST_F(ArrayElementwiseOpTest, MulTwoConstantZeroElementS32s) { ComputeAndCompareR1(&builder, {}, {}); } -TEST_F(ArrayElementwiseOpTest, MulTwoConstantU32s) { +XLA_TEST_F(ArrayElementwiseOpTest, MulTwoConstantU32s) { std::vector data = {0, 1, 0xDEADBEEF, 1234, 0x1a243514, 0xFFFFFFFF, 0x80808080}; @@ -353,7 +496,7 @@ TEST_F(ArrayElementwiseOpTest, MulTwoConstantU32s) { ComputeAndCompareR1(&builder, expected, {}); } -TEST_F(ArrayElementwiseOpTest, LogicalAnd) { +XLA_TEST_F(ArrayElementwiseOpTest, LogicalAnd) { ComputationBuilder builder(client_, TestName()); auto a = builder.ConstantR1({false, false, true, true}); auto b = builder.ConstantR1({false, true, false, true}); @@ -371,7 +514,7 @@ XLA_TEST_F(ArrayElementwiseOpTest, LogicalAndZeroElement) { ComputeAndCompareR1(&builder, {}, {}); } -TEST_F(ArrayElementwiseOpTest, LogicalOr) { +XLA_TEST_F(ArrayElementwiseOpTest, LogicalOr) { ComputationBuilder builder(client_, TestName()); auto a = builder.ConstantR1({false, false, true, true}); auto b = builder.ConstantR1({false, true, false, true}); @@ -389,7 +532,7 @@ XLA_TEST_F(ArrayElementwiseOpTest, LogicalOrZeroElement) { ComputeAndCompareR1(&builder, {}, {}); } -TEST_F(ArrayElementwiseOpTest, LogicalNot) { +XLA_TEST_F(ArrayElementwiseOpTest, LogicalNot) { ComputationBuilder builder(client_, TestName()); auto a = builder.ConstantR1({false, true, true, false}); auto out = builder.LogicalNot(a); @@ -405,7 +548,7 @@ XLA_TEST_F(ArrayElementwiseOpTest, LogicalNotZeroElement) { ComputeAndCompareR1(&builder, {}, {}); } -TEST_F(ArrayElementwiseOpTest, CompareEqF32s) { +XLA_TEST_F(ArrayElementwiseOpTest, CompareEqF32s) { SetFastMathDisabled(true); ComputationBuilder builder(client_, TestName()); auto lhs = builder.ConstantR1({-2.5f, 25.5f, 2.25f, NAN, 6.0f}); @@ -424,7 +567,7 @@ XLA_TEST_F(ArrayElementwiseOpTest, CompareEqZeroElementF32s) { ComputeAndCompareR1(&builder, {}, {}); } -TEST_F(ArrayElementwiseOpTest, CompareGeF32s) { +XLA_TEST_F(ArrayElementwiseOpTest, CompareGeF32s) { SetFastMathDisabled(true); ComputationBuilder builder(client_, TestName()); auto lhs = builder.ConstantR1({-2.5f, 25.5f, 2.25f, NAN, 6.0f}); @@ -434,7 +577,7 @@ TEST_F(ArrayElementwiseOpTest, CompareGeF32s) { ComputeAndCompareR1(&builder, {false, true, true, false, false}, {}); } -TEST_F(ArrayElementwiseOpTest, CompareGtF32s) { +XLA_TEST_F(ArrayElementwiseOpTest, CompareGtF32s) { SetFastMathDisabled(true); ComputationBuilder builder(client_, TestName()); auto lhs = builder.ConstantR1({-2.5f, 25.5f, 2.25f, NAN, 6.0f}); @@ -444,7 +587,7 @@ TEST_F(ArrayElementwiseOpTest, CompareGtF32s) { ComputeAndCompareR1(&builder, {false, true, true, false, false}, {}); } -TEST_F(ArrayElementwiseOpTest, CompareLeF32s) { +XLA_TEST_F(ArrayElementwiseOpTest, CompareLeF32s) { SetFastMathDisabled(true); ComputationBuilder builder(client_, TestName()); auto lhs = builder.ConstantR1({-2.5f, 5.0f, 2.25f, NAN, 6.0f}); @@ -454,7 +597,7 @@ TEST_F(ArrayElementwiseOpTest, CompareLeF32s) { ComputeAndCompareR1(&builder, {true, true, false, false, false}, {}); } -TEST_F(ArrayElementwiseOpTest, CompareLtF32s) { +XLA_TEST_F(ArrayElementwiseOpTest, CompareLtF32s) { SetFastMathDisabled(true); ComputationBuilder builder(client_, TestName()); auto lhs = builder.ConstantR1({-2.5f, 25.5f, 2.25f, NAN, 6.0f}); @@ -464,7 +607,7 @@ TEST_F(ArrayElementwiseOpTest, CompareLtF32s) { ComputeAndCompareR1(&builder, {true, false, false, false, false}, {}); } -TEST_F(ArrayElementwiseOpTest, CompareEqS32s) { +XLA_TEST_F(ArrayElementwiseOpTest, CompareEqS32s) { const int32 min = std::numeric_limits::min(); const int32 max = std::numeric_limits::max(); ComputationBuilder builder(client_, TestName()); @@ -486,7 +629,19 @@ XLA_TEST_F(ArrayElementwiseOpTest, CompareEqZeroElementS32s) { ComputeAndCompareR1(&builder, {}, {}); } -TEST_F(ArrayElementwiseOpTest, CompareNeS32s) { +XLA_TEST_F(ArrayElementwiseOpTest, CompareNeF32s) { + // Disable fast-math because we're operating on NaNs. + SetFastMathDisabled(true); + + ComputationBuilder builder(client_, TestName()); + auto lhs = builder.ConstantR1({-2.5f, 25.5f, 2.25f, NAN, 6.0f}); + auto rhs = builder.ConstantR1({10.0f, 25.5f, 1.0f, 10.0f, NAN}); + auto compare = builder.Ne(lhs, rhs); + + ComputeAndCompareR1(&builder, {true, false, true, true, true}, {}); +} + +XLA_TEST_F(ArrayElementwiseOpTest, CompareNeS32s) { const int32 min = std::numeric_limits::min(); const int32 max = std::numeric_limits::max(); ComputationBuilder builder(client_, TestName()); @@ -498,7 +653,7 @@ TEST_F(ArrayElementwiseOpTest, CompareNeS32s) { &builder, {false, true, true, true, false, true, true, true, false}, {}); } -TEST_F(ArrayElementwiseOpTest, CompareGeS32s) { +XLA_TEST_F(ArrayElementwiseOpTest, CompareGeS32s) { const int32 min = std::numeric_limits::min(); const int32 max = std::numeric_limits::max(); ComputationBuilder builder(client_, TestName()); @@ -510,7 +665,7 @@ TEST_F(ArrayElementwiseOpTest, CompareGeS32s) { &builder, {true, false, false, true, true, false, true, true, true}, {}); } -TEST_F(ArrayElementwiseOpTest, CompareGtS32s) { +XLA_TEST_F(ArrayElementwiseOpTest, CompareGtS32s) { const int32 min = std::numeric_limits::min(); const int32 max = std::numeric_limits::max(); ComputationBuilder builder(client_, TestName()); @@ -523,7 +678,7 @@ TEST_F(ArrayElementwiseOpTest, CompareGtS32s) { {}); } -TEST_F(ArrayElementwiseOpTest, CompareLeS32s) { +XLA_TEST_F(ArrayElementwiseOpTest, CompareLeS32s) { const int32 min = std::numeric_limits::min(); const int32 max = std::numeric_limits::max(); ComputationBuilder builder(client_, TestName()); @@ -535,7 +690,7 @@ TEST_F(ArrayElementwiseOpTest, CompareLeS32s) { &builder, {true, true, true, false, true, true, false, false, true}, {}); } -TEST_F(ArrayElementwiseOpTest, CompareLtS32s) { +XLA_TEST_F(ArrayElementwiseOpTest, CompareLtS32s) { const int32 min = std::numeric_limits::min(); const int32 max = std::numeric_limits::max(); ComputationBuilder builder(client_, TestName()); @@ -548,7 +703,7 @@ TEST_F(ArrayElementwiseOpTest, CompareLtS32s) { {}); } -TEST_F(ArrayElementwiseOpTest, CompareEqU32s) { +XLA_TEST_F(ArrayElementwiseOpTest, CompareEqU32s) { const uint32 max = std::numeric_limits::max(); ComputationBuilder builder(client_, TestName()); auto lhs = builder.ConstantR1({0, 0, 0, 5, 5, 5, max, max, max}); @@ -560,7 +715,7 @@ TEST_F(ArrayElementwiseOpTest, CompareEqU32s) { {}); } -TEST_F(ArrayElementwiseOpTest, CompareNeU32s) { +XLA_TEST_F(ArrayElementwiseOpTest, CompareNeU32s) { const uint32 max = std::numeric_limits::max(); ComputationBuilder builder(client_, TestName()); auto lhs = builder.ConstantR1({0, 0, 0, 5, 5, 5, max, max, max}); @@ -571,7 +726,7 @@ TEST_F(ArrayElementwiseOpTest, CompareNeU32s) { &builder, {false, true, true, true, false, true, true, true, false}, {}); } -TEST_F(ArrayElementwiseOpTest, CompareGeU32s) { +XLA_TEST_F(ArrayElementwiseOpTest, CompareGeU32s) { const uint32 max = std::numeric_limits::max(); ComputationBuilder builder(client_, TestName()); auto lhs = builder.ConstantR1({0, 0, 0, 5, 5, 5, max, max, max}); @@ -582,7 +737,7 @@ TEST_F(ArrayElementwiseOpTest, CompareGeU32s) { &builder, {true, false, false, true, true, false, true, true, true}, {}); } -TEST_F(ArrayElementwiseOpTest, CompareGtU32s) { +XLA_TEST_F(ArrayElementwiseOpTest, CompareGtU32s) { const uint32 max = std::numeric_limits::max(); ComputationBuilder builder(client_, TestName()); auto lhs = builder.ConstantR1({0, 0, 0, 5, 5, 5, max, max, max}); @@ -594,7 +749,7 @@ TEST_F(ArrayElementwiseOpTest, CompareGtU32s) { {}); } -TEST_F(ArrayElementwiseOpTest, CompareLeU32s) { +XLA_TEST_F(ArrayElementwiseOpTest, CompareLeU32s) { const uint32 max = std::numeric_limits::max(); ComputationBuilder builder(client_, TestName()); auto lhs = builder.ConstantR1({0, 0, 0, 5, 5, 5, max, max, max}); @@ -605,7 +760,7 @@ TEST_F(ArrayElementwiseOpTest, CompareLeU32s) { &builder, {true, true, true, false, true, true, false, false, true}, {}); } -TEST_F(ArrayElementwiseOpTest, CompareLtU32s) { +XLA_TEST_F(ArrayElementwiseOpTest, CompareLtU32s) { const uint32 max = std::numeric_limits::max(); ComputationBuilder builder(client_, TestName()); auto lhs = builder.ConstantR1({0, 0, 0, 5, 5, 5, max, max, max}); @@ -617,14 +772,27 @@ TEST_F(ArrayElementwiseOpTest, CompareLtU32s) { {}); } -TEST_F(ArrayElementwiseOpTest, PowF32s) { +XLA_TEST_F(ArrayElementwiseOpTest, PowF32s) { SetFastMathDisabled(true); ComputationBuilder builder(client_, TestName()); - auto lhs = builder.ConstantR1({4.0f, 2.0f, 2.0f, NAN, 6.0f}); - auto rhs = builder.ConstantR1({2.0f, -2.0f, 3.0f, 10.0f, NAN}); + auto lhs = + builder.ConstantR1({4.0f, 2.0f, 2.0f, NAN, 6.0f, -2.0f, -2.0f}); + auto rhs = + builder.ConstantR1({2.0f, -2.0f, 3.0f, 10.0f, NAN, 3.0f, 4.0f}); auto minimum = builder.Pow(lhs, rhs); - ComputeAndCompareR1(&builder, {16.0f, 0.25f, 8.0f, NAN, NAN}, {}, + ComputeAndCompareR1( + &builder, {16.0f, 0.25f, 8.0f, NAN, NAN, -8.0f, 16.0f}, {}, error_spec_); +} + +XLA_TEST_F(ArrayElementwiseOpTest, PowNonIntegerF32s) { + SetFastMathDisabled(true); + ComputationBuilder builder(client_, TestName()); + auto lhs = builder.ConstantR1({-2.0f, -0.6f, -0.6f, 0.0f}); + auto rhs = builder.ConstantR1({0.5f, 0.6f, -0.6f, -0.6f}); + auto minimum = builder.Pow(lhs, rhs); + + ComputeAndCompareR1(&builder, {NAN, NAN, NAN, INFINITY}, {}, error_spec_); } @@ -638,13 +806,13 @@ XLA_TEST_F(ArrayElementwiseOpTest, PowZeroElementF32s) { } // Some Pow cases that can be implemented more efficiently. -TEST_F(ArrayElementwiseOpTest, PowSpecialF32) { +XLA_TEST_F(ArrayElementwiseOpTest, PowSpecialF32) { ComputationBuilder b(client_, TestName()); std::vector values = {1.0f, 2.0f, 3.2f, -4.0f}; std::vector exponents = {0.0f, 1.0f, 2.0f, 0.5f, -1.0f, -0.5f}; - std::unique_ptr param_literal = LiteralUtil::CreateR1(values); + std::unique_ptr param_literal = Literal::CreateR1(values); std::unique_ptr param_data = client_->TransferToServer(*param_literal).ConsumeValueOrDie(); @@ -666,10 +834,249 @@ TEST_F(ArrayElementwiseOpTest, PowSpecialF32) { ComputeAndCompareR1(&b, expected, {param_data.get()}, error_spec_); } +XLA_TEST_F(ArrayElementwiseOpTest, PowOfExpF32) { + ComputationBuilder b(client_, TestName()); + + std::vector values0 = {1.0f, 2.0f, 3.2f, -4.0f, 0.0f, 5.7f}; + std::vector values1 = {0.0f, 1.0f, 2.0f, 0.5f, -1.0f, -0.5f}; + + std::unique_ptr literal0 = Literal::CreateR1(values0); + std::unique_ptr data0 = + client_->TransferToServer(*literal0).ConsumeValueOrDie(); + std::unique_ptr literal1 = Literal::CreateR1(values1); + std::unique_ptr data1 = + client_->TransferToServer(*literal1).ConsumeValueOrDie(); + auto param0 = b.Parameter(0, literal0->shape(), "param0"); + auto param1 = b.Parameter(1, literal1->shape(), "param1"); + b.Pow(b.Exp(param0), param1); + + std::vector expected(values0.size()); + for (int64 i = 0; i < values0.size(); ++i) { + expected[i] = std::pow(std::exp(values0[i]), values1[i]); + } + + ComputeAndCompareR1(&b, expected, {data0.get(), data1.get()}, + error_spec_); +} + +XLA_TEST_F(ArrayElementwiseOpTest, LogOfPowerF32) { + ComputationBuilder b(client_, TestName()); + + std::vector values0 = {1.0f, 2.0f, 3.2f, 4.0f, 0.5f, 5.7f}; + std::vector values1 = {0.0f, 1.0f, 2.0f, 0.5f, -1.0f, -0.5f}; + + std::unique_ptr literal0 = Literal::CreateR1(values0); + std::unique_ptr data0 = + client_->TransferToServer(*literal0).ConsumeValueOrDie(); + std::unique_ptr literal1 = Literal::CreateR1(values1); + std::unique_ptr data1 = + client_->TransferToServer(*literal1).ConsumeValueOrDie(); + auto param0 = b.Parameter(0, literal0->shape(), "param0"); + auto param1 = b.Parameter(1, literal1->shape(), "param1"); + b.Log(b.Pow(param0, param1)); + + std::vector expected(values0.size()); + for (int64 i = 0; i < values0.size(); ++i) { + expected[i] = std::log(std::pow(values0[i], values1[i])); + } + + ComputeAndCompareR1(&b, expected, {data0.get(), data1.get()}, + error_spec_); +} + +XLA_TEST_F(ArrayElementwiseOpTest, MulOfExpF32) { + ComputationBuilder b(client_, TestName()); + + std::vector values0 = {1.0f, 2.0f, 3.2f, -4.0f, 0.0f, 5.7f}; + std::vector values1 = {0.0f, 1.0f, 2.0f, 0.5f, -1.0f, -0.5f}; + + std::unique_ptr literal0 = Literal::CreateR1(values0); + std::unique_ptr data0 = + client_->TransferToServer(*literal0).ConsumeValueOrDie(); + std::unique_ptr literal1 = Literal::CreateR1(values1); + std::unique_ptr data1 = + client_->TransferToServer(*literal1).ConsumeValueOrDie(); + auto param0 = b.Parameter(0, literal0->shape(), "param0"); + auto param1 = b.Parameter(1, literal1->shape(), "param1"); + b.Mul(b.Exp(param0), b.Exp(param1)); + + std::vector expected(values0.size()); + for (int64 i = 0; i < values0.size(); ++i) { + expected[i] = std::exp(values0[i]) * std::exp(values1[i]); + } + + ComputeAndCompareR1(&b, expected, {data0.get(), data1.get()}, + error_spec_); +} + +XLA_TEST_F(ArrayElementwiseOpTest, DivOfExpF32) { + ComputationBuilder b(client_, TestName()); + + std::vector values0 = {1.0f, 2.0f, 3.2f, -4.0f, 0.0f, 5.7f}; + std::vector values1 = {0.0f, 1.0f, 2.0f, 0.5f, -1.0f, -0.5f}; + + std::unique_ptr literal0 = Literal::CreateR1(values0); + std::unique_ptr data0 = + client_->TransferToServer(*literal0).ConsumeValueOrDie(); + std::unique_ptr literal1 = Literal::CreateR1(values1); + std::unique_ptr data1 = + client_->TransferToServer(*literal1).ConsumeValueOrDie(); + auto param0 = b.Parameter(0, literal0->shape(), "param0"); + auto param1 = b.Parameter(1, literal1->shape(), "param1"); + b.Div(param0, b.Exp(param1)); + + std::vector expected(values0.size()); + for (int64 i = 0; i < values0.size(); ++i) { + expected[i] = values0[i] / std::exp(values1[i]); + } + + ComputeAndCompareR1(&b, expected, {data0.get(), data1.get()}, + error_spec_); +} + +XLA_TEST_F(ArrayElementwiseOpTest, Div3_lhs_F32) { + ComputationBuilder b(client_, TestName()); + + std::vector values0 = {1.0f, 2.0f, 3.2f, -4.0f, 0.45f, 5.7f}; + std::vector values1 = {0.1f, 1.0f, 2.0f, 0.5f, -1.0f, -0.5f}; + std::vector values2 = {0.1f, 1.1f, 6.9f, 12.5f, -15.0f, -0.5f}; + + std::unique_ptr literal0 = Literal::CreateR1(values0); + std::unique_ptr data0 = + client_->TransferToServer(*literal0).ConsumeValueOrDie(); + + std::unique_ptr literal1 = Literal::CreateR1(values1); + std::unique_ptr data1 = + client_->TransferToServer(*literal1).ConsumeValueOrDie(); + + std::unique_ptr literal2 = Literal::CreateR1(values2); + std::unique_ptr data2 = + client_->TransferToServer(*literal2).ConsumeValueOrDie(); + auto param0 = b.Parameter(0, literal0->shape(), "param0"); + auto param1 = b.Parameter(1, literal1->shape(), "param1"); + auto param2 = b.Parameter(2, literal2->shape(), "param2"); + b.Div(b.Div(param0, param1), param2); + + std::vector expected(values0.size()); + for (int64 i = 0; i < values0.size(); ++i) { + expected[i] = (values0[i] / values1[i]) / values2[i]; + } + + ComputeAndCompareR1( + &b, expected, {data0.get(), data1.get(), data2.get()}, error_spec_); +} + +XLA_TEST_F(ArrayElementwiseOpTest, Div3_rhs_F32) { + ComputationBuilder b(client_, TestName()); + + std::vector values0 = {1.0f, 2.0f, 3.2f, -4.0f, 0.45f, 5.7f}; + std::vector values1 = {0.1f, 1.0f, 2.0f, 0.5f, -1.0f, -0.5f}; + std::vector values2 = {0.1f, 1.1f, 6.9f, 12.5f, -15.0f, -0.5f}; + + std::unique_ptr literal0 = Literal::CreateR1(values0); + std::unique_ptr data0 = + client_->TransferToServer(*literal0).ConsumeValueOrDie(); + + std::unique_ptr literal1 = Literal::CreateR1(values1); + std::unique_ptr data1 = + client_->TransferToServer(*literal1).ConsumeValueOrDie(); + + std::unique_ptr literal2 = Literal::CreateR1(values2); + std::unique_ptr data2 = + client_->TransferToServer(*literal2).ConsumeValueOrDie(); + + auto param0 = b.Parameter(0, literal0->shape(), "param0"); + auto param1 = b.Parameter(1, literal1->shape(), "param1"); + auto param2 = b.Parameter(2, literal2->shape(), "param2"); + b.Div(param0, b.Div(param1, param2)); + + std::vector expected(values0.size()); + for (int64 i = 0; i < values0.size(); ++i) { + expected[i] = values0[i] / (values1[i] / values2[i]); + } + + ComputeAndCompareR1( + &b, expected, {data0.get(), data1.get(), data2.get()}, error_spec_); +} + +XLA_TEST_F(ArrayElementwiseOpTest, DivOfPowerF32) { + ComputationBuilder b(client_, TestName()); + + std::vector values0 = {1.0f, 2.0f, 3.2f, -4.0f, 0.45f, 5.7f}; + std::vector values1 = {0.1f, 1.0f, 2.0f, 0.5f, 1.0f, 0.5f}; + std::vector values2 = {0.1f, 1.1f, 6.9f, 9.5f, -11.0f, -0.5f}; + + std::unique_ptr literal0 = Literal::CreateR1(values0); + std::unique_ptr data0 = + client_->TransferToServer(*literal0).ConsumeValueOrDie(); + + std::unique_ptr literal1 = Literal::CreateR1(values1); + std::unique_ptr data1 = + client_->TransferToServer(*literal1).ConsumeValueOrDie(); + + std::unique_ptr literal2 = Literal::CreateR1(values2); + std::unique_ptr data2 = + client_->TransferToServer(*literal2).ConsumeValueOrDie(); + + auto param0 = b.Parameter(0, literal0->shape(), "param0"); + auto param1 = b.Parameter(1, literal1->shape(), "param1"); + auto param2 = b.Parameter(2, literal2->shape(), "param2"); + b.Div(param0, b.Pow(param1, param2)); + + std::vector expected(values0.size()); + for (int64 i = 0; i < values0.size(); ++i) { + expected[i] = values0[i] / std::pow(values1[i], values2[i]); + } + + ComputeAndCompareR1( + &b, expected, {data0.get(), data1.get(), data2.get()}, error_spec_); +} + +XLA_TEST_F(ArrayElementwiseOpTest, Div4F32) { + ComputationBuilder b(client_, TestName()); + + std::vector values0 = {1.0f, 2.0f, 3.2f, -4.0f, 0.45f, 5.7f}; + std::vector values1 = {0.1f, 1.0f, 2.0f, 0.5f, -1.0f, -0.5f}; + std::vector values2 = {0.1f, 1.1f, 6.9f, 12.5f, -15.0f, -0.5f}; + std::vector values3 = {2.1f, 3.1f, 9.9f, -4.5f, -11.0f, -21.5f}; + + std::unique_ptr literal0 = Literal::CreateR1(values0); + std::unique_ptr data0 = + client_->TransferToServer(*literal0).ConsumeValueOrDie(); + + std::unique_ptr literal1 = Literal::CreateR1(values1); + std::unique_ptr data1 = + client_->TransferToServer(*literal1).ConsumeValueOrDie(); + + std::unique_ptr literal2 = Literal::CreateR1(values2); + std::unique_ptr data2 = + client_->TransferToServer(*literal2).ConsumeValueOrDie(); + + std::unique_ptr literal3 = Literal::CreateR1(values3); + std::unique_ptr data3 = + client_->TransferToServer(*literal3).ConsumeValueOrDie(); + + auto param0 = b.Parameter(0, literal0->shape(), "param0"); + auto param1 = b.Parameter(1, literal1->shape(), "param1"); + auto param2 = b.Parameter(2, literal2->shape(), "param2"); + auto param3 = b.Parameter(3, literal3->shape(), "param2"); + b.Div(b.Div(param0, param1), b.Div(param2, param3)); + + std::vector expected(values0.size()); + for (int64 i = 0; i < values0.size(); ++i) { + expected[i] = (values0[i] / values1[i]) / (values2[i] / values3[i]); + } + + ComputeAndCompareR1( + &b, expected, {data0.get(), data1.get(), data2.get(), data3.get()}, + error_spec_); +} + TEST_P(ArrayElementwiseOpTestParamCount, SquareManyValues) { const int count = GetParam(); ComputationBuilder builder(client_, TestName()); std::vector values; + values.reserve(count); for (int i = 0; i < count; ++i) { values.push_back(i / static_cast(count)); } @@ -677,6 +1084,7 @@ TEST_P(ArrayElementwiseOpTestParamCount, SquareManyValues) { auto exp = builder.Pow(x, builder.ConstantR0(2.0f)); std::vector expected; + expected.reserve(values.size()); for (float value : values) { expected.push_back(value * value); } @@ -684,7 +1092,7 @@ TEST_P(ArrayElementwiseOpTestParamCount, SquareManyValues) { ComputeAndCompareR1(&builder, expected, {}, error_spec_); } -TEST_F(ArrayElementwiseOpTest, SquareIn4D) { +XLA_TEST_F(ArrayElementwiseOpTest, SquareIn4D) { ComputationBuilder builder(client_, TestName()); Array4D values(2, 2, 2, 2); @@ -723,7 +1131,7 @@ XLA_TEST_F(ArrayElementwiseOpTest, SquareIn4DZeroElements) { // // TODO(b/28180546): Make this compile in a way that is consistent // among backends. -TEST_F(ArrayElementwiseOpTest, MinF32s) { +XLA_TEST_F(ArrayElementwiseOpTest, MinF32s) { ComputationBuilder builder(client_, TestName()); #if !defined(XLA_TEST_BACKEND_CPU) auto lhs = builder.ConstantR1({1.0f, 1.0f, 2.25f}); @@ -777,7 +1185,7 @@ XLA_TEST_F(ArrayElementwiseOpTest, MinF64s) { // TODO(b/28180546): Make this compile in a way that is consistent // among backends. See comment on MinF32s test above. -TEST_F(ArrayElementwiseOpTest, MaxF32s) { +XLA_TEST_F(ArrayElementwiseOpTest, MaxF32s) { ComputationBuilder builder(client_, TestName()); #if !defined(XLA_TEST_BACKEND_CPU) auto lhs = builder.ConstantR1({1.0f, 1.0f, 2.25f}); @@ -829,7 +1237,7 @@ XLA_TEST_F(ArrayElementwiseOpTest, MaxF64s) { {}, error_spec_); } -TEST_F(ArrayElementwiseOpTest, MaxS32s) { +XLA_TEST_F(ArrayElementwiseOpTest, MaxS32s) { const int32 min = std::numeric_limits::min(); const int32 max = std::numeric_limits::max(); ComputationBuilder builder(client_, TestName()); @@ -844,7 +1252,7 @@ TEST_F(ArrayElementwiseOpTest, MaxS32s) { ComputeAndCompareR1(&builder, expected, {}); } -TEST_F(ArrayElementwiseOpTest, MinS32s) { +XLA_TEST_F(ArrayElementwiseOpTest, MinS32s) { const int32 min = std::numeric_limits::min(); const int32 max = std::numeric_limits::max(); ComputationBuilder builder(client_, TestName()); @@ -859,7 +1267,7 @@ TEST_F(ArrayElementwiseOpTest, MinS32s) { ComputeAndCompareR1(&builder, expected, {}); } -TEST_F(ArrayElementwiseOpTest, MaxU32s) { +XLA_TEST_F(ArrayElementwiseOpTest, MaxU32s) { const uint32 max = std::numeric_limits::max(); ComputationBuilder builder(client_, TestName()); auto x = builder.ConstantR1({0, 0, 1, 1, 1, max, max, max}); @@ -870,7 +1278,7 @@ TEST_F(ArrayElementwiseOpTest, MaxU32s) { ComputeAndCompareR1(&builder, expected, {}); } -TEST_F(ArrayElementwiseOpTest, MinU32s) { +XLA_TEST_F(ArrayElementwiseOpTest, MinU32s) { const uint32 max = std::numeric_limits::max(); ComputationBuilder builder(client_, TestName()); auto x = builder.ConstantR1({0, 0, 1, 1, 1, max, max, max}); @@ -881,7 +1289,7 @@ TEST_F(ArrayElementwiseOpTest, MinU32s) { ComputeAndCompareR1(&builder, expected, {}); } -TEST_F(ArrayElementwiseOpTest, MaxTenF32s) { +XLA_TEST_F(ArrayElementwiseOpTest, MaxTenF32s) { ComputationBuilder builder(client_, TestName()); auto x = builder.ConstantR1( {-0.0, 1.0, 2.0, -3.0, -4.0, 5.0, 6.0, -7.0, -8.0, 9.0}); @@ -914,7 +1322,7 @@ XLA_TEST_F(ArrayElementwiseOpTest, MaxR1S0AndR2S0x2F32s) { } } -TEST_F(ArrayElementwiseOpTest, Max1DAnd2DF32s) { +XLA_TEST_F(ArrayElementwiseOpTest, Max1DAnd2DF32s) { ComputationBuilder builder(client_, TestName()); auto v = builder.ConstantR1({2.0f, 3.0f, 4.0f}); auto m = @@ -957,7 +1365,7 @@ XLA_TEST_F(ArrayElementwiseOpTest, Max3DAndScalarZeroElementS32s) { ComputeAndCompareR3(&builder, expected, {}); } -TEST_F(ArrayElementwiseOpTest, Min2DTo1DF32s) { +XLA_TEST_F(ArrayElementwiseOpTest, Min2DTo1DF32s) { ComputationBuilder builder(client_, TestName()); auto m = builder.ConstantR2({{-10.4f, 64.0f, 6.0f}, {0.1f, 32.0f, 16.1f}}); @@ -1034,7 +1442,7 @@ XLA_TEST_F(ArrayElementwiseOpTest, RemTwoConstantS32s) { ComputeAndCompareR1(&builder, {-3, 1, 0, -1, 1}, {}); } -TEST_F(ArrayElementwiseOpTest, NonNanClampF32) { +XLA_TEST_F(ArrayElementwiseOpTest, NonNanClampF32) { ComputationBuilder builder(client_, TestName()); auto minimum = builder.ConstantR1({1.0f, -6.5f, 1.0f, 2.25f, 0.0f}); auto argument = builder.ConstantR1({2.0f, 10.0f, -5.0f, 1.0f, 10.0f}); @@ -1045,7 +1453,7 @@ TEST_F(ArrayElementwiseOpTest, NonNanClampF32) { error_spec_); } -TEST_F(ArrayElementwiseOpTest, ClampF32Scalar) { +XLA_TEST_F(ArrayElementwiseOpTest, ClampF32Scalar) { ComputationBuilder builder(client_, TestName()); auto minimum = builder.ConstantR0(0.0f); auto argument = builder.ConstantR1({2.0f, 10.0f, -5.0f, 1.0f, 4.0f}); @@ -1056,7 +1464,7 @@ TEST_F(ArrayElementwiseOpTest, ClampF32Scalar) { error_spec_); } -TEST_F(ArrayElementwiseOpTest, ClampF32ScalarVector) { +XLA_TEST_F(ArrayElementwiseOpTest, ClampF32ScalarVector) { ComputationBuilder builder(client_, TestName()); auto min_scalar = builder.ConstantR0(0.0f); auto min_vector = builder.ConstantR1({1.0f, -6.5f, 1.0f, 2.25f, 0.0f}); @@ -1075,16 +1483,16 @@ TEST_F(ArrayElementwiseOpTest, ClampF32ScalarVector) { error_spec_); } -TEST_F(ArrayElementwiseOpTest, AddTwoParametersF32s) { +XLA_TEST_F(ArrayElementwiseOpTest, AddTwoParametersF32s) { ComputationBuilder builder(client_, TestName()); std::unique_ptr param0_literal = - LiteralUtil::CreateR1({1.1f, 2.2f, 3.3f, 5.5f}); + Literal::CreateR1({1.1f, 2.2f, 3.3f, 5.5f}); std::unique_ptr param0_data = client_->TransferToServer(*param0_literal).ConsumeValueOrDie(); std::unique_ptr param1_literal = - LiteralUtil::CreateR1({7.2f, 2.3f, 3.4f, 5.6f}); + Literal::CreateR1({7.2f, 2.3f, 3.4f, 5.6f}); std::unique_ptr param1_data = client_->TransferToServer(*param1_literal).ConsumeValueOrDie(); @@ -1101,12 +1509,12 @@ XLA_TEST_F(ArrayElementwiseOpTest, AddTwoParametersZeroElementF32s) { ComputationBuilder builder(client_, TestName()); std::unique_ptr param0_literal = - LiteralUtil::CreateR3FromArray3D(Array3D(0, 7, 0)); + Literal::CreateR3FromArray3D(Array3D(0, 7, 0)); std::unique_ptr param0_data = client_->TransferToServer(*param0_literal).ConsumeValueOrDie(); std::unique_ptr param1_literal = - LiteralUtil::CreateR3FromArray3D(Array3D(0, 7, 0)); + Literal::CreateR3FromArray3D(Array3D(0, 7, 0)); std::unique_ptr param1_data = client_->TransferToServer(*param1_literal).ConsumeValueOrDie(); @@ -1119,11 +1527,11 @@ XLA_TEST_F(ArrayElementwiseOpTest, AddTwoParametersZeroElementF32s) { &builder, expected, {param0_data.get(), param1_data.get()}, error_spec_); } -TEST_F(ArrayElementwiseOpTest, AddParameterToConstantF32s) { +XLA_TEST_F(ArrayElementwiseOpTest, AddParameterToConstantF32s) { ComputationBuilder builder(client_, TestName()); std::unique_ptr param0_literal = - LiteralUtil::CreateR1({1.1f, 2.2f, 3.3f, 5.5f}); + Literal::CreateR1({1.1f, 2.2f, 3.3f, 5.5f}); std::unique_ptr param0_data = client_->TransferToServer(*param0_literal).ConsumeValueOrDie(); @@ -1135,7 +1543,25 @@ TEST_F(ArrayElementwiseOpTest, AddParameterToConstantF32s) { {param0_data.get()}, error_spec_); } -TEST_F(ArrayElementwiseOpTest, TanhF32s) { +XLA_TEST_F(ArrayElementwiseOpTest, CosF32s) { + ComputationBuilder builder(client_, TestName()); + auto a = builder.ConstantR1({3.14159f, 0.0f, 1.570796f, -0.78539f}); + auto result = builder.Cos(a); + + ComputeAndCompareR1(&builder, {-1.0f, 1.0f, 0.0f, 0.707107f}, {}, + error_spec_); +} + +XLA_TEST_F(ArrayElementwiseOpTest, SinF32s) { + ComputationBuilder builder(client_, TestName()); + auto a = builder.ConstantR1({3.14159f, 0.0f, 1.570796f, -0.78539f}); + auto result = builder.Sin(a); + + ComputeAndCompareR1(&builder, {0.0f, 0.0f, 1.0f, -0.707107f}, {}, + error_spec_); +} + +XLA_TEST_F(ArrayElementwiseOpTest, TanhF32s) { ComputationBuilder builder(client_, TestName()); auto a = builder.ConstantR1({-2.5f, 3.14f, 2.25f}); auto result = builder.Tanh(a); @@ -1144,7 +1570,51 @@ TEST_F(ArrayElementwiseOpTest, TanhF32s) { error_spec_); } -TEST_F(ArrayElementwiseOpTest, AddChainFoldLeft) { +XLA_TEST_F(ArrayElementwiseOpTest, TanhF32sVector) { + // This is like the test ArrayElementwiseOpTest.TanhF32s above, except that + // the input tensor is large enough to exercise the vectorized tanh + // implementation. + ComputationBuilder builder(client_, TestName()); + auto input_literal = Literal::CreateR2( + {{1.02, -0.32, 0.85, 0.90, 1.23, -0.91, -0.49, 0.80}, + {-0.67, 0.16, -0.07, 0.39, -0.41, 0.04, 1.36, 1.25}, + {0.41, 0.65, -1.08, 0.32, -1.45, -0.77, -1.09, 0.91}, + {-1.03, -0.30, -1.11, -1.17, 1.50, -0.85, 0.04, 1.02}, + {0.34, -0.61, 0.41, 0.07, -0.02, 1.42, -0.62, 0.81}, + {0.08, 0.81, -0.30, 1.17, -0.65, -0.44, 0.92, 1.26}, + {-1.29, 1.35, 0.08, -1.24, -0.92, 0.49, 1.17, -0.45}, + {-1.31, -1.44, -0.13, -1.31, -0.79, 1.41, 1.21, 1.05}}); + auto input_data = + client_->TransferToServer(*input_literal).ConsumeValueOrDie(); + + auto input = builder.Parameter(0, input_literal->shape(), "input"); + builder.Tanh(input); + + ComputeAndCompareR2( + &builder, + {{0.77009583, -0.30665702, 0.69070244, 0.71401149, 0.84400684, + -0.71985596, -0.45764771, 0.66664988}, + {-0.58278900, 0.16050975, -0.06770509, 0.36843640, -0.38476998, + 0.04018109, 0.87562293, 0.84788644}, + {0.38603750, 0.57294142, -0.79140943, 0.31032649, -0.89590985, + -0.64770776, -0.79625875, 0.72234446}, + {-0.77389336, -0.28871772, -0.80428445, -0.82541436, 0.90456349, + -0.68856895, 0.03877772, 0.76877952}, + {0.32561871, -0.54546672, 0.39072621, 0.07273290, -0.01924866, + 0.88924897, -0.55283129, 0.67183107}, + {0.08006320, 0.66944766, -0.29068485, 0.82573754, -0.57170743, + -0.41581789, 0.72739530, 0.85025692}, + {-0.85931867, 0.87357593, 0.07782833, -0.84597743, -0.72748238, + 0.45396307, 0.82449573, -0.42462519}, + {-0.86363792, -0.89368379, -0.12621804, -0.86445558, -0.65565848, + 0.88789743, 0.83566397, 0.78287679}}, + {input_data.get()}, + // The error spec is unusually high here to account for the fact that we + // use a rational interpolant to approximate tanh. + ErrorSpec(0.004, 0.004)); +} + +XLA_TEST_F(ArrayElementwiseOpTest, AddChainFoldLeft) { // a ------ (add) --------- (add) // / / // b -----/ / @@ -1162,7 +1632,7 @@ TEST_F(ArrayElementwiseOpTest, AddChainFoldLeft) { error_spec_); } -TEST_F(ArrayElementwiseOpTest, AddChainFoldRight) { +XLA_TEST_F(ArrayElementwiseOpTest, AddChainFoldRight) { // b ------ (add) --------- (add) // / / // c -----/ / @@ -1180,7 +1650,7 @@ TEST_F(ArrayElementwiseOpTest, AddChainFoldRight) { error_spec_); } -TEST_F(ArrayElementwiseOpTest, AddWithNeg) { +XLA_TEST_F(ArrayElementwiseOpTest, AddWithNeg) { // a ----- (neg) ----- (add) // / // b ----- (neg) ----/ @@ -1197,7 +1667,7 @@ TEST_F(ArrayElementwiseOpTest, AddWithNeg) { error_spec_); } -TEST_F(ArrayElementwiseOpTest, AddChainTwoSide) { +XLA_TEST_F(ArrayElementwiseOpTest, AddChainTwoSide) { // a ------ (add) ------------\ // / \ // b -----/ (add) @@ -1220,7 +1690,7 @@ TEST_F(ArrayElementwiseOpTest, AddChainTwoSide) { error_spec_); } -TEST_F(ArrayElementwiseOpTest, 2DBinaryOpF32s) { +XLA_TEST_F(ArrayElementwiseOpTest, 2DBinaryOpF32s) { ComputationBuilder builder(client_, TestName()); auto a = builder.ConstantR2({{-2.5f, 3.14f, 1.0f}, {2.25f, -10.0f, 3.33f}}); @@ -1245,7 +1715,7 @@ XLA_TEST_F(ArrayElementwiseOpTest, ScalarPlus2DF32) { ComputeAndCompareR2(&builder, expected_array, {}, error_spec_); } -TEST_F(ArrayElementwiseOpTest, 2DPlusScalarF32) { +XLA_TEST_F(ArrayElementwiseOpTest, 2DPlusScalarF32) { // Add a matrix + scalar. ComputationBuilder builder(client_, TestName()); auto a = @@ -1285,9 +1755,9 @@ XLA_TEST_F(ArrayElementwiseOpTest, Compare1DTo2DS32Eq) { auto cmp_dim_1 = builder.Eq(v, m, /*broadcast_dimensions=*/{0}); auto result = builder.Tuple({cmp_dim_0, cmp_dim_1}); - auto expected = LiteralUtil::MakeTuple( - {LiteralUtil::CreateR2({{true, true}, {true, false}}).get(), - LiteralUtil::CreateR2({{true, false}, {false, false}}).get()}); + auto expected = Literal::MakeTuple( + {Literal::CreateR2({{true, true}, {true, false}}).get(), + Literal::CreateR2({{true, false}, {false, false}}).get()}); ComputeAndCompareTuple(&builder, *expected, {}, error_spec_); } @@ -1361,7 +1831,7 @@ XLA_TEST_F(ArrayElementwiseOpTest, Compare1DTo2DS32Lt) { EXPECT_EQ(expected, ExecuteToString(&builder, {})); } -TEST_F(ArrayElementwiseOpTest, Mul2Dby1DF32) { +XLA_TEST_F(ArrayElementwiseOpTest, Mul2Dby1DF32) { // Test simple broadcasting of a R1F32 over R2F32 when the order of binary op // arguments is reversed. ComputationBuilder builder(client_, TestName()); @@ -1372,7 +1842,7 @@ TEST_F(ArrayElementwiseOpTest, Mul2Dby1DF32) { ComputeAndCompareR2(&builder, expected_array, {}, error_spec_); } -TEST_F(ArrayElementwiseOpTest, Add2DTo2DWithDegenerateDim1) { +XLA_TEST_F(ArrayElementwiseOpTest, Add2DTo2DWithDegenerateDim1) { // Tests broadcasting for arrays with degenerate (size == 1) dimensions. ComputationBuilder builder(client_, TestName()); // m's shape in XLA notation is {3, 2} @@ -1432,7 +1902,7 @@ XLA_TEST_F(ArrayElementwiseOpTest, Add1DTo2DF32TwoWaysOver1) { ComputeAndCompareR2(&builder, expected_array, {}, error_spec_); } -TEST_F(ArrayElementwiseOpTest, Add1DTo2DF32TwoWaysOver0) { +XLA_TEST_F(ArrayElementwiseOpTest, Add1DTo2DF32TwoWaysOver0) { // Add together a (2,2) array and a (2) array, using dimension 1 for // broadcasting (though there are two ways to broadcast these shapes). ComputationBuilder builder(client_, TestName()); @@ -1443,7 +1913,7 @@ TEST_F(ArrayElementwiseOpTest, Add1DTo2DF32TwoWaysOver0) { ComputeAndCompareR2(&builder, expected_array, {}, error_spec_); } -TEST_F(ArrayElementwiseOpTest, 3DBinaryOpF32s) { +XLA_TEST_F(ArrayElementwiseOpTest, 3DBinaryOpF32s) { // Binary add of two R3s together ComputationBuilder builder(client_, TestName()); Array3D a_3d({{{1.0f, 2.0f}, {3.0f, 4.0f}, {5.0f, 6.0f}}, @@ -1574,7 +2044,7 @@ XLA_TEST_F(ArrayElementwiseOpTest, CompareGtR3F32sWithDegenerateDim2) { EXPECT_EQ(expected, ExecuteToString(&builder, {})); } -TEST_F(ArrayElementwiseOpTest, 4DBinaryOpF32s) { +XLA_TEST_F(ArrayElementwiseOpTest, 4DBinaryOpF32s) { ComputationBuilder builder(client_, TestName()); std::unique_ptr> operand_a_4d(new Array4D(2, 3, 4, 5)); @@ -1601,7 +2071,7 @@ TEST_F(ArrayElementwiseOpTest, 4DBinaryOpF32s) { ComputeAndCompareR4(&builder, *expected_4d, {}, error_spec_); } -TEST_F(ArrayElementwiseOpTest, R4PlusR1InDim1) { +XLA_TEST_F(ArrayElementwiseOpTest, R4PlusR1InDim1) { ComputationBuilder builder(client_, TestName()); std::unique_ptr> operand_a_4d(new Array4D(2, 3, 4, 5)); @@ -1629,7 +2099,7 @@ TEST_F(ArrayElementwiseOpTest, R4PlusR1InDim1) { ComputeAndCompareR4(&builder, *expected_4d, {}, error_spec_); } -TEST_F(ArrayElementwiseOpTest, R4_32x64x2x2_Plus_R1_64) { +XLA_TEST_F(ArrayElementwiseOpTest, R4_16x16x2x2_Plus_R1_16) { constexpr int d0 = 16; constexpr int d1 = 16; constexpr int d2 = 2; @@ -1640,7 +2110,7 @@ TEST_F(ArrayElementwiseOpTest, R4_32x64x2x2_Plus_R1_64) { std::iota(r1.begin(), r1.end(), 1.0); ComputationBuilder builder(client_, TestName()); - std::unique_ptr a_literal = LiteralUtil::CreateR4FromArray4D(r4); + std::unique_ptr a_literal = Literal::CreateR4FromArray4D(r4); *a_literal->mutable_shape()->mutable_layout() = LayoutUtil::MakeLayout({0, 1, 2, 3}); auto a = builder.ConstantLiteral(*a_literal); @@ -1660,30 +2130,30 @@ TEST_F(ArrayElementwiseOpTest, R4_32x64x2x2_Plus_R1_64) { } // Show that we can't add two opaques. -TEST_F(ArrayElementwiseOpTest, CannotAddOpaques) { +XLA_TEST_F(ArrayElementwiseOpTest, CannotAddOpaques) { ComputationBuilder builder(client_, TestName()); auto shape = ShapeUtil::MakeOpaqueShape(); auto x = builder.Parameter(0, shape, "x"); auto concatenated = builder.Add(x, x); StatusOr computation_status = builder.Build(); ASSERT_FALSE(computation_status.ok()); - EXPECT_MATCH(computation_status.status().ToString(), - testing::ContainsRegex( - "Expected non-opaque argument for lhs of binary operation")); + EXPECT_THAT(computation_status.status().ToString(), + ::testing::ContainsRegex( + "Expected non-opaque argument for lhs of binary operation")); } // Regression test for b/31927799. "slice - y" is fused and requires implicit // broadcast. -TEST_F(ArrayElementwiseOpTest, ImplictBroadcastInFusedExpressions) { +XLA_TEST_F(ArrayElementwiseOpTest, ImplictBroadcastInFusedExpressions) { ComputationBuilder builder(client_, TestName()); - auto x_literal = LiteralUtil::CreateR1({1, 2, 3}); - auto y_literal = LiteralUtil::CreateR1({4, 5}); + auto x_literal = Literal::CreateR1({1, 2, 3}); + auto y_literal = Literal::CreateR1({4, 5}); auto x_data = client_->TransferToServer(*x_literal).ConsumeValueOrDie(); auto y_data = client_->TransferToServer(*y_literal).ConsumeValueOrDie(); auto x = builder.Parameter(0, x_literal->shape(), "x"); auto y = builder.Parameter(1, y_literal->shape(), "y"); - auto slice = builder.Slice(x, {1}, {2}); + auto slice = builder.Slice(x, {1}, {2}, {1}); builder.Sub(slice, y); ComputeAndCompareR1(&builder, {-2, -3}, {x_data.get(), y_data.get()}, @@ -1696,20 +2166,3 @@ INSTANTIATE_TEST_CASE_P(ArrayElementwiseOpTestParamCount, } // namespace } // namespace xla - -int main(int argc, char** argv) { - std::vector flag_list; - xla::legacy_flags::AppendCpuCompilerFlags(&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; - } - testing::InitGoogleTest(&argc, argv); - if (argc > 1) { - LOG(ERROR) << "Unknown argument " << argv[1] << "\n" << usage; - return 2; - } - return RUN_ALL_TESTS(); -} diff --git a/tensorflow/compiler/xla/tests/axpy_simple_test.cc b/tensorflow/compiler/xla/tests/axpy_simple_test.cc index adffac09e36d896d9bffa24f4423ab92d993e539..627a9c3e7d9f6eb8d360228362ea5adf12c6c798 100644 --- a/tensorflow/compiler/xla/tests/axpy_simple_test.cc +++ b/tensorflow/compiler/xla/tests/axpy_simple_test.cc @@ -17,7 +17,6 @@ limitations under the License. #include "tensorflow/compiler/xla/client/computation_builder.h" #include "tensorflow/compiler/xla/client/local_client.h" -#include "tensorflow/compiler/xla/legacy_flags/cpu_compiler_flags.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" @@ -71,20 +70,3 @@ TEST_F(AxpySimpleTest, AxpyTenValues) { } // namespace } // namespace xla - -int main(int argc, char** argv) { - std::vector flag_list; - xla::legacy_flags::AppendCpuCompilerFlags(&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; - } - testing::InitGoogleTest(&argc, argv); - if (argc > 1) { - LOG(ERROR) << "Unknown argument " << argv[1] << "\n" << usage; - return 2; - } - return RUN_ALL_TESTS(); -} diff --git a/tensorflow/compiler/xla/tests/bad_rng_shape_validation_test.cc b/tensorflow/compiler/xla/tests/bad_rng_shape_validation_test.cc index c7b533b80f1901a32324a15a8f6584e628a4ad30..e4bf1827acf24bcdbfe20fe39e794a0265ab89e3 100644 --- a/tensorflow/compiler/xla/tests/bad_rng_shape_validation_test.cc +++ b/tensorflow/compiler/xla/tests/bad_rng_shape_validation_test.cc @@ -21,14 +21,12 @@ limitations under the License. #include "tensorflow/compiler/xla/client/computation.h" #include "tensorflow/compiler/xla/client/computation_builder.h" #include "tensorflow/compiler/xla/client/local_client.h" -#include "tensorflow/compiler/xla/legacy_flags/cpu_compiler_flags.h" #include "tensorflow/compiler/xla/statusor.h" -#include "tensorflow/compiler/xla/test_helpers.h" +#include "tensorflow/compiler/xla/test.h" #include "tensorflow/compiler/xla/tests/client_library_test_base.h" #include "tensorflow/compiler/xla/types.h" #include "tensorflow/compiler/xla/xla_data.pb.h" #include "tensorflow/core/platform/logging.h" -#include "tensorflow/core/platform/test.h" namespace xla { namespace { @@ -45,8 +43,8 @@ TEST_F(BadRngShapeValidationTest, DefaultConstructedShapeCreatesError) { StatusOr computation = builder.Build(); EXPECT_FALSE(computation.ok()); LOG(INFO) << "status received: " << computation.status(); - EXPECT_MATCH(computation.status().error_message(), - testing::HasSubstr("shape has invalid")); + EXPECT_THAT(computation.status().error_message(), + ::testing::HasSubstr("shape has invalid")); } TEST_F(BadRngShapeValidationTest, ShapeWithoutLayoutIsOk) { @@ -66,20 +64,3 @@ TEST_F(BadRngShapeValidationTest, ShapeWithoutLayoutIsOk) { } // namespace } // namespace xla - -int main(int argc, char** argv) { - std::vector flag_list; - xla::legacy_flags::AppendCpuCompilerFlags(&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; - } - testing::InitGoogleTest(&argc, argv); - if (argc > 1) { - LOG(ERROR) << "Unknown argument " << argv[1] << "\n" << usage; - return 2; - } - return RUN_ALL_TESTS(); -} diff --git a/tensorflow/compiler/xla/tests/batch_normalization_test.cc b/tensorflow/compiler/xla/tests/batch_normalization_test.cc index 598fd69909b549a796a36cb8788b9b3adb06b165..028d1251b455b82a291c236f7866e52e27d3590e 100644 --- a/tensorflow/compiler/xla/tests/batch_normalization_test.cc +++ b/tensorflow/compiler/xla/tests/batch_normalization_test.cc @@ -23,12 +23,21 @@ limitations under the License. #include "tensorflow/compiler/xla/client/computation_builder.h" #include "tensorflow/compiler/xla/client/lib/arithmetic.h" #include "tensorflow/compiler/xla/client/local_client.h" -#include "tensorflow/compiler/xla/legacy_flags/cpu_compiler_flags.h" #include "tensorflow/compiler/xla/literal_util.h" +#include "tensorflow/compiler/xla/reference_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/shape_util.h" #include "tensorflow/compiler/xla/statusor.h" +#include "tensorflow/compiler/xla/test.h" +#include "tensorflow/compiler/xla/test_helpers.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/compiler/xla/util.h" #include "tensorflow/compiler/xla/xla_data.pb.h" #include "tensorflow/core/platform/logging.h" #include "tensorflow/core/platform/test.h" @@ -47,7 +56,7 @@ class BatchNormalizationTest : public ClientLibraryTestBase { {5.0f, 4.4f}, // p2 }); input_array_.FillWithPZ(pz); - input_literal_ = *LiteralUtil::CreateR4FromArray4D(input_array_); + input_literal_ = *Literal::CreateR4FromArray4D(input_array_); CHECK_EQ(kSamples, input_array_.planes()); CHECK_EQ(kZ, input_array_.depth()); CHECK_EQ(kY, input_array_.height()); @@ -189,22 +198,552 @@ TEST_F(BatchNormalizationTest, SpecComparisonForward) { ComputeAndCompareR4(&builder, expected, {}, error_spec_); } -} // namespace -} // namespace xla +struct BatchNormTestParam { + std::vector bounds; + int64 feature_index; + float random_value_mean; + float random_value_var; +}; + +// Tests to test the fused operation of BatchNorm. +class BatchNormTest : public ClientLibraryTestBase, + public ::testing::WithParamInterface { +}; + +XLA_TEST_P(BatchNormTest, RandomizedTests) { + float epsilon = 0.001; + ComputationBuilder builder(client_, TestName()); + const std::vector& bounds = GetParam().bounds; + Array4D input_array(bounds[0], bounds[1], bounds[2], bounds[3]); + input_array.FillRandom(GetParam().random_value_var, + GetParam().random_value_mean); + + const int64 feature_index = GetParam().feature_index; + const int64 num_elements_per_feature = + Product(bounds) / bounds[feature_index]; + const int64 feature_bound = bounds[feature_index]; + std::vector offset(feature_bound, 1); + std::vector scale(feature_bound, 2); + + auto input_squared = + ReferenceUtil::MapArray4D(input_array, [](float a) { return a * a; }); + std::vector reduce_dims; + for (int64 i = 0; i < static_cast(bounds.size()); ++i) { + if (i != feature_index) { + reduce_dims.push_back(i); + } + } + + auto sum = + ReferenceUtil::Reduce4DTo1D(input_array, /*init=*/0.0f, reduce_dims, + [](float a, float b) { return a + b; }); + + auto sum_squared = + ReferenceUtil::Reduce4DTo1D(*input_squared, /*init=*/0.0f, reduce_dims, + [](float a, float b) { return a + b; }); + + std::vector mean(feature_bound); + + for (int64 i = 0; i < feature_bound; ++i) { + mean[i] = sum[i] / num_elements_per_feature; + } + + std::vector mean_square(feature_bound); + for (int64 i = 0; i < feature_bound; ++i) { + mean_square[i] = mean[i] * mean[i]; + } + + std::vector square_mean(feature_bound); + for (int64 i = 0; i < feature_bound; ++i) { + square_mean[i] = sum_squared[i] / num_elements_per_feature; + } + + std::vector var(feature_bound); + for (int64 i = 0; i < feature_bound; ++i) { + var[i] = square_mean[i] - mean_square[i]; + } + + Array4D mean4D = + *ReferenceUtil::Broadcast1DTo4D(mean, bounds, feature_index); + auto var4D = *ReferenceUtil::Broadcast1DTo4D(var, bounds, feature_index); + auto scale4D = *ReferenceUtil::Broadcast1DTo4D(scale, bounds, feature_index); + auto offset4D = + *ReferenceUtil::Broadcast1DTo4D(offset, bounds, feature_index); + + auto normalized = *ReferenceUtil::BatchNorm4D(input_array, mean4D, var4D, + scale4D, offset4D, epsilon); + + auto expected_normalized = Literal::CreateR4FromArray4D(normalized); + + auto offset_literal = Literal::CreateR1(offset); + auto scale_literal = Literal::CreateR1(scale); + auto input_literal = Literal::CreateR4FromArray4D(input_array); + + auto input_activations = + builder.Parameter(0, input_literal->shape(), "input"); + auto scale_activations = + builder.Parameter(1, scale_literal->shape(), "offset"); + auto offset_activations = + builder.Parameter(2, offset_literal->shape(), "scale"); + + auto expected = *Literal::MakeTuple({expected_normalized.get(), + Literal::CreateR1(mean).get(), + Literal::CreateR1(var).get()}); + + std::unique_ptr input_data = + client_->TransferToServer(*input_literal).ConsumeValueOrDie(); + std::unique_ptr scale_data = + client_->TransferToServer(*scale_literal).ConsumeValueOrDie(); + std::unique_ptr offset_data = + client_->TransferToServer(*offset_literal).ConsumeValueOrDie(); + + builder.BatchNormTraining(input_activations, scale_activations, + offset_activations, epsilon, feature_index); + + ComputeAndCompareTuple( + &builder, expected, + {input_data.get(), scale_data.get(), offset_data.get()}, + ErrorSpec(0.01, 1)); +} + +XLA_TEST_P(BatchNormTest, RandomizedInferencingTests) { + float epsilon = 0.001; + ComputationBuilder builder(client_, TestName()); + const std::vector& bounds = GetParam().bounds; + Array4D input_array(bounds[0], bounds[1], bounds[2], bounds[3]); + input_array.FillRandom(GetParam().random_value_var, + GetParam().random_value_mean); + + const int64 feature_index = GetParam().feature_index; + const int64 num_elements_per_feature = + Product(bounds) / bounds[feature_index]; + const int64 feature_bound = bounds[feature_index]; + std::vector offset(feature_bound, 1); + std::vector scale(feature_bound, 2); + + auto input_squared = + ReferenceUtil::MapArray4D(input_array, [](float a) { return a * a; }); + std::vector reduce_dims; + for (int64 i = 0; i < static_cast(bounds.size()); ++i) { + if (i != feature_index) { + reduce_dims.push_back(i); + } + } + + auto sum = + ReferenceUtil::Reduce4DTo1D(input_array, /*init=*/0.0f, reduce_dims, + [](float a, float b) { return a + b; }); + + auto sum_squared = + ReferenceUtil::Reduce4DTo1D(*input_squared, /*init=*/0.0f, reduce_dims, + [](float a, float b) { return a + b; }); + + std::vector mean(feature_bound); + + for (int64 i = 0; i < feature_bound; ++i) { + mean[i] = sum[i] / num_elements_per_feature; + } + + std::vector mean_square(feature_bound); + for (int64 i = 0; i < feature_bound; ++i) { + mean_square[i] = mean[i] * mean[i]; + } + + std::vector square_mean(feature_bound); + for (int64 i = 0; i < feature_bound; ++i) { + square_mean[i] = sum_squared[i] / num_elements_per_feature; + } + + std::vector var(feature_bound); + for (int64 i = 0; i < feature_bound; ++i) { + var[i] = square_mean[i] - mean_square[i]; + } + + Array4D mean4D = + *ReferenceUtil::Broadcast1DTo4D(mean, bounds, feature_index); + auto var4D = *ReferenceUtil::Broadcast1DTo4D(var, bounds, feature_index); + auto scale4D = *ReferenceUtil::Broadcast1DTo4D(scale, bounds, feature_index); + auto offset4D = + *ReferenceUtil::Broadcast1DTo4D(offset, bounds, feature_index); + + auto normalized = *ReferenceUtil::BatchNorm4D(input_array, mean4D, var4D, + scale4D, offset4D, epsilon); + + auto offset_literal = Literal::CreateR1(offset); + auto scale_literal = Literal::CreateR1(scale); + auto mean_literal = Literal::CreateR1(mean); + auto var_literal = Literal::CreateR1(var); + auto input_literal = Literal::CreateR4FromArray4D(input_array); + + auto input_activations = + builder.Parameter(0, input_literal->shape(), "input"); + auto scale_activations = + builder.Parameter(1, scale_literal->shape(), "offset"); + auto offset_activations = + builder.Parameter(2, offset_literal->shape(), "scale"); + auto mean_activations = builder.Parameter(3, mean_literal->shape(), "mean"); + auto variance_activations = + builder.Parameter(4, var_literal->shape(), "variance"); + + Array4D expected = normalized; + + std::unique_ptr input_data = + client_->TransferToServer(*input_literal).ConsumeValueOrDie(); + std::unique_ptr scale_data = + client_->TransferToServer(*scale_literal).ConsumeValueOrDie(); + std::unique_ptr offset_data = + client_->TransferToServer(*offset_literal).ConsumeValueOrDie(); + std::unique_ptr mean_data = + client_->TransferToServer(*mean_literal).ConsumeValueOrDie(); + std::unique_ptr variance_data = + client_->TransferToServer(*var_literal).ConsumeValueOrDie(); + + builder.BatchNormInference(input_activations, scale_activations, + offset_activations, mean_activations, + variance_activations, epsilon, feature_index); + + ComputeAndCompareR4( + &builder, expected, + {input_data.get(), scale_data.get(), offset_data.get(), mean_data.get(), + variance_data.get()}, + ErrorSpec(0.01, 1)); +} + +XLA_TEST_P(BatchNormTest, RandomizedGradTests) { + float epsilon = 0.001; + ComputationBuilder builder(client_, TestName()); + const std::vector& bounds = GetParam().bounds; + Array4D input_array(bounds[0], bounds[1], bounds[2], bounds[3]); + input_array.FillRandom(GetParam().random_value_var, + GetParam().random_value_mean); + + Array4D grad_output_array(bounds[0], bounds[1], bounds[2], bounds[3]); + grad_output_array.FillRandom(GetParam().random_value_var, + GetParam().random_value_mean); + + const int64 feature_index = GetParam().feature_index; + const int64 num_elements_per_feature = + Product(bounds) / bounds[feature_index]; + const int64 feature_bound = bounds[feature_index]; + std::vector scale(feature_bound, 2); + + auto input_squared = + ReferenceUtil::MapArray4D(input_array, [](float a) { return a * a; }); + std::vector reduce_dims; + for (int64 i = 0; i < static_cast(bounds.size()); ++i) { + if (i != feature_index) { + reduce_dims.push_back(i); + } + } + + auto sum = + ReferenceUtil::Reduce4DTo1D(input_array, /*init=*/0.0f, reduce_dims, + [](float a, float b) { return a + b; }); + + auto sum_squared = + ReferenceUtil::Reduce4DTo1D(*input_squared, /*init=*/0.0f, reduce_dims, + [](float a, float b) { return a + b; }); + + std::vector mean(feature_bound); + + for (int64 i = 0; i < feature_bound; ++i) { + mean[i] = sum[i] / num_elements_per_feature; + } + + std::vector mean_square(feature_bound); + for (int64 i = 0; i < feature_bound; ++i) { + mean_square[i] = mean[i] * mean[i]; + } + + std::vector square_mean(feature_bound); + for (int64 i = 0; i < feature_bound; ++i) { + square_mean[i] = sum_squared[i] / num_elements_per_feature; + } + + std::vector var(feature_bound); + for (int64 i = 0; i < feature_bound; ++i) { + var[i] = square_mean[i] - mean_square[i]; + } + + Array4D mean4D = + *ReferenceUtil::Broadcast1DTo4D(mean, bounds, feature_index); + auto var4D = *ReferenceUtil::Broadcast1DTo4D(var, bounds, feature_index); + auto scale4D = *ReferenceUtil::Broadcast1DTo4D(scale, bounds, feature_index); + + auto var_add_epsilon = *ReferenceUtil::MapArray4D( + var4D, [epsilon](float a) { return a + epsilon; }); + + auto rsqrt_var_add_epsilon = *ReferenceUtil::MapArray4D( + var_add_epsilon, [epsilon](float a) { return 1 / std::sqrt(a); }); + + auto grad_output_times_var = + *ReferenceUtil::MapArray4D(grad_output_array, var_add_epsilon, + [](float a, float b) { return a * b; }); + + auto activation_shifted = *ReferenceUtil::MapArray4D( + input_array, mean4D, [](float a, float b) { return a - b; }); + + auto activation_shifted_times_grad_output = + *ReferenceUtil::MapArray4D(grad_output_array, activation_shifted, + [](float a, float b) { return a * b; }); + + auto grad_scale_before_reduction = *ReferenceUtil::MapArray4D( + activation_shifted_times_grad_output, rsqrt_var_add_epsilon, + [](float a, float b) { return a * b; }); + + auto grad_scale = ReferenceUtil::Reduce4DTo1D( + grad_scale_before_reduction, /*init=*/0.0f, reduce_dims, + [](float a, float b) { return a + b; }); + + auto grad_offset = + ReferenceUtil::Reduce4DTo1D(grad_output_array, /*init=*/0.0f, reduce_dims, + [](float a, float b) { return a + b; }); + + auto scale_times_rsqrt_var_add_epsilon = *ReferenceUtil::MapArray4D( + scale4D, rsqrt_var_add_epsilon, [](float a, float b) { return a * b; }); + + auto I1 = *ReferenceUtil::MapArray4D( + grad_output_array, [&](float a) { return num_elements_per_feature * a; }); + + auto I2 = *ReferenceUtil::Broadcast1DTo4D(grad_offset, bounds, feature_index); + + // I3 = sum(output_grad * (activation - mean(activation))) + auto I3 = *ReferenceUtil::Broadcast1DTo4D( + ReferenceUtil::Reduce4DTo1D(activation_shifted_times_grad_output, + /*init=*/0.0f, reduce_dims, + [](float a, float b) { return a + b; }), + bounds, feature_index); + + // I4 = (activation - mean(activation)) * + // sum(output_grad * (activation - mean(activation))) + auto I4 = *ReferenceUtil::MapArray4D(I3, activation_shifted, + [](float a, float b) { return a * b; }); + + // I5 = (activation - mean(activation)) * + // sum(output_grad * (activation - mean(activation))) / (variance + + // epsilon)) + auto I5 = *ReferenceUtil::MapArray4D(I4, var_add_epsilon, + [](float a, float b) { return a / b; }); + + auto grad_activation = *ReferenceUtil::MapArray4D( + I1, I2, [](float a, float b) { return a - b; }); + + grad_activation = *ReferenceUtil::MapArray4D( + grad_activation, I5, [](float a, float b) { return a - b; }); -int main(int argc, char** argv) { - std::vector flag_list; - xla::legacy_flags::AppendCpuCompilerFlags(&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; - } - testing::InitGoogleTest(&argc, argv); - if (argc > 1) { - LOG(ERROR) << "Unknown argument " << argv[1] << "\n" << usage; - return 2; - } - return RUN_ALL_TESTS(); + grad_activation = *ReferenceUtil::MapArray4D( + grad_activation, scale4D, [](float a, float b) { return a * b; }); + + grad_activation = *ReferenceUtil::MapArray4D( + grad_activation, rsqrt_var_add_epsilon, + [=](float a, float b) { return a * b / num_elements_per_feature; }); + + auto expected_grad_activation = + Literal::CreateR4FromArray4D(grad_activation); + + auto input_literal = Literal::CreateR4FromArray4D(input_array); + auto scale_literal = Literal::CreateR1(scale); + auto mean_literal = Literal::CreateR1(mean); + auto var_literal = Literal::CreateR1(var); + auto grad_output_literal = + Literal::CreateR4FromArray4D(grad_output_array); + + auto input_parameter = builder.Parameter(0, input_literal->shape(), "input"); + auto scale_parameter = builder.Parameter(1, scale_literal->shape(), "scale"); + auto mean_parameter = builder.Parameter(2, mean_literal->shape(), "mean"); + auto var_parameter = builder.Parameter(3, var_literal->shape(), "variance"); + auto grad_output_parameter = + builder.Parameter(4, grad_output_literal->shape(), "grad_output"); + + std::unique_ptr input_data = + client_->TransferToServer(*input_literal).ConsumeValueOrDie(); + std::unique_ptr scale_data = + client_->TransferToServer(*scale_literal).ConsumeValueOrDie(); + std::unique_ptr mean_data = + client_->TransferToServer(*mean_literal).ConsumeValueOrDie(); + std::unique_ptr var_data = + client_->TransferToServer(*var_literal).ConsumeValueOrDie(); + std::unique_ptr grad_output_data = + client_->TransferToServer(*grad_output_literal).ConsumeValueOrDie(); + + auto t = builder.BatchNormGrad(input_parameter, scale_parameter, + mean_parameter, var_parameter, + grad_output_parameter, epsilon, feature_index); + + auto expected = + *Literal::MakeTuple({expected_grad_activation.get(), + Literal::CreateR1(grad_scale).get(), + Literal::CreateR1(grad_offset).get()}); + + ComputeAndCompareTuple(&builder, expected, + {input_data.get(), scale_data.get(), mean_data.get(), + var_data.get(), grad_output_data.get()}, + ErrorSpec(0.01, 1)); } + +INSTANTIATE_TEST_CASE_P( + BatchNormTest_Instantiation, BatchNormTest, + ::testing::Values(BatchNormTestParam{{2, 2, 2, 2}, 0, 100.2f, 200.0f}, + BatchNormTestParam{{2, 2, 2, 2}, 3, 300.f, 400.0f}, + + BatchNormTestParam{{1, 10, 1, 1}, 0, 10.1f, 20.1f}, + BatchNormTestParam{{10, 10, 10, 10}, 1, 3.14f, 314.15f}, + BatchNormTestParam{{10, 10, 10, 10}, 2, 666.6f, 777.7f}, + BatchNormTestParam{{10, 10, 10, 10}, 1, -666.6f, 777.7f}, + BatchNormTestParam{{10, 10, 10, 10}, 2, 0.f, 777.7f}, + BatchNormTestParam{{1, 1, 10, 130}, 2, 0.f, 777.7f}, + BatchNormTestParam{{1, 1, 130, 11}, 2, 0.f, 777.7f}, + BatchNormTestParam{{1, 1, 10, 1}, 3, 888.8f, 9.9f}, + + BatchNormTestParam{{24, 129, 1, 2}, 2, 10000, 10000}, + BatchNormTestParam{{24, 129, 1, 2}, 3, 10000, 10000}, + + // Feature on low dimension to trigger relayout, test + // internal logical to physical dimension calculation + // is correct after relayout. + BatchNormTestParam{{1, 2, 3, 4}, 0, 100, 100})); + +XLA_TEST_F(BatchNormTest, BasicTraining) { + const int kFeatureIndex = 3; + ComputationBuilder builder(client_, TestName()); + + auto operand = builder.ConstantR4FromArray4D( + {{{{1.f, 2.f}}, {{3.f, 4.f}}}, {{{5.f, 6.f}}, {{7.f, 8.f}}}}); + + auto scale = builder.ConstantR1({2.0f, 3.0f}); + + auto offset = builder.ConstantR1({1.0f, 2.0f}); + + auto tuple = builder.BatchNormTraining(operand, scale, offset, + /*epsilon=*/0.001, kFeatureIndex); + + auto expected = *Literal::MakeTuple( + {Literal::CreateR4({{{{-1.6f, -2.0f}}, {{0.1f, 0.6f}}}, + {{{1.9f, 3.3f}}, {{3.7f, 6.0f}}}}) + .get(), + Literal::CreateR1({4, 5}).get(), + Literal::CreateR1({5, 5}).get()}); + + ComputeAndCompareTuple(&builder, expected, {}, ErrorSpec(0.1)); +} + +XLA_TEST_F(BatchNormTest, BasicTrainingOnSublane) { + const int kFeatureIndex = 2; + ComputationBuilder builder(client_, TestName()); + + auto operand = builder.ConstantR4FromArray4D( + {{{{1.f}, {2.f}}, {{3.f}, {4.f}}}, {{{5.f}, {6.f}}, {{7.f}, {8.f}}}}); + + auto scale = builder.ConstantR1({2.0f, 3.0f}); + + auto offset = builder.ConstantR1({1.0f, 2.0f}); + + auto tuple = builder.BatchNormTraining(operand, scale, offset, + /*epsilon=*/0.001, kFeatureIndex); + + auto expected = *Literal::MakeTuple( + {Literal::CreateR4({{{{-1.6f}, {-2.0f}}, {{0.1f}, {0.6f}}}, + {{{1.9f}, {3.3f}}, {{3.7f}, {6.0f}}}}) + .get(), + Literal::CreateR1({4, 5}).get(), + Literal::CreateR1({5, 5}).get()}); + + ComputeAndCompareTuple(&builder, expected, {}, ErrorSpec(0.1)); +} + +XLA_TEST_F(BatchNormTest, DISABLED_ON_GPU(TrainingWithFeatureOnLowDimension)) { + // Use 0 dimension as feature, tests layout analyzer. + const int kFeatureIndex = 0; + ComputationBuilder builder(client_, TestName()); + + ComputationDataHandle h0; + auto operand = CreateR3Parameter(Array3D(260, 2, 2, 1.0f), + /*parameter_number=*/0, "operand", + &builder, &h0); + ComputationDataHandle h1; + auto scale = + CreateR1Parameter(std::vector(260, 1.0f), + /*parameter_number=*/1, "scale", &builder, &h1); + ComputationDataHandle h2; + auto offset = + CreateR1Parameter(std::vector(260, 1.0f), + /*parameter_number=*/2, "offset", &builder, &h2); + + auto tuple = builder.BatchNormTraining(h0, h1, h2, + /*epsilon=*/1, kFeatureIndex); + + auto expected = *Literal::MakeTuple( + {Literal::CreateR3FromArray3D(Array3D(260, 2, 2, 1.0f)) + .get(), + Literal::CreateR1(std::vector(260, 1.0f)).get(), + Literal::CreateR1(std::vector(260, 0.0f)).get()}); + + ComputeAndCompareTuple(&builder, expected, + {operand.get(), scale.get(), offset.get()}, + ErrorSpec(0.1)); +} + +XLA_TEST_F(BatchNormTest, LargeEpsilonTest) { + // Test the correctness of choosing a large epsilon value. + const int kFeatureIndex = 2; + ComputationBuilder builder(client_, TestName()); + + ComputationDataHandle h0; + auto operand = CreateR3Parameter({{{0.0f}, {10.0f}, {20.0f}, {30.0f}}}, + /*parameter_number=*/0, "operand", + &builder, &h0); + ComputationDataHandle h1; + auto scale = + CreateR1Parameter(std::vector(1, 1.0f), + /*parameter_number=*/1, "scale", &builder, &h1); + ComputationDataHandle h2; + auto offset = + CreateR1Parameter(std::vector(1, 0.0f), + /*parameter_number=*/2, "offset", &builder, &h2); + + // var = 125, mean = 15, epsilon = -100 + auto tuple = builder.BatchNormTraining(h0, h1, h2, + /*epsilon=*/-100, kFeatureIndex); + + auto expected = *Literal::MakeTuple( + {Literal::CreateR3FromArray3D({{{-3.0f}, {-1.0f}, {1.0f}, {3.0f}}}) + .get(), + Literal::CreateR1(std::vector(1, 15.0f)).get(), + Literal::CreateR1(std::vector(1, 125.0f)).get()}); + + ComputeAndCompareTuple(&builder, expected, + {operand.get(), scale.get(), offset.get()}, + ErrorSpec(0.1)); +} + +XLA_TEST_F(BatchNormTest, BatchNormGradBasic) { + const int kFeatureIndex = 2; + ComputationBuilder builder(client_, TestName()); + + auto operand = + builder.ConstantR4FromArray4D(Array4D(2, 2, 2, 1, 0.0f)); + + auto scale = builder.ConstantR1({1.0f, 1.0f}); + + auto mean = builder.ConstantR1({0.0f, 0.0f}); + + auto var = builder.ConstantR1({1.0f, 1.0f}); + + auto grad_output = builder.ConstantR4FromArray4D( + {{{{1.f}, {2.f}}, {{3.f}, {4.f}}}, {{{5.f}, {6.f}}, {{7.f}, {8.f}}}}); + + builder.BatchNormGrad(operand, scale, mean, var, grad_output, + /*epsilon=*/0.0, kFeatureIndex); + + auto expected = *Literal::MakeTuple( + {Literal::CreateR4({{{{-3.f}, {-3.f}}, {{-1.f}, {-1.f}}}, + {{{1.f}, {1.f}}, {{3.f}, {3.f}}}}) + .get(), + Literal::CreateR1({0, 0}).get(), + Literal::CreateR1({16, 20}).get()}); + + ComputeAndCompareTuple(&builder, expected, {}, ErrorSpec(0.1)); +} + +} // namespace +} // namespace xla diff --git a/tensorflow/compiler/xla/tests/binop_scaling_test.cc b/tensorflow/compiler/xla/tests/binop_scaling_test.cc index e825bd435b677fbc4e6be278b5a243339a231873..97fec89b63fb8d3a4264275f3253a91e1ea2ce68 100644 --- a/tensorflow/compiler/xla/tests/binop_scaling_test.cc +++ b/tensorflow/compiler/xla/tests/binop_scaling_test.cc @@ -17,7 +17,6 @@ limitations under the License. #include "tensorflow/compiler/xla/array4d.h" #include "tensorflow/compiler/xla/client/computation_builder.h" #include "tensorflow/compiler/xla/client/local_client.h" -#include "tensorflow/compiler/xla/legacy_flags/cpu_compiler_flags.h" #include "tensorflow/compiler/xla/reference_util.h" #include "tensorflow/compiler/xla/tests/client_library_test_base.h" #include "tensorflow/compiler/xla/tests/literal_test_util.h" @@ -138,20 +137,3 @@ TEST_F(BinopScalingTest, R4PlusR0S32) { } // namespace } // namespace xla - -int main(int argc, char** argv) { - std::vector flag_list; - xla::legacy_flags::AppendCpuCompilerFlags(&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; - } - testing::InitGoogleTest(&argc, argv); - if (argc > 1) { - LOG(ERROR) << "Unknown argument " << argv[1] << "\n" << usage; - return 2; - } - return RUN_ALL_TESTS(); -} diff --git a/tensorflow/compiler/xla/tests/broadcast_simple_test.cc b/tensorflow/compiler/xla/tests/broadcast_simple_test.cc index 63744afb4ea72006262aad74e9b8d75a09b107e6..4f26bf47ae6d29f525351692612648d6432f9518 100644 --- a/tensorflow/compiler/xla/tests/broadcast_simple_test.cc +++ b/tensorflow/compiler/xla/tests/broadcast_simple_test.cc @@ -21,19 +21,88 @@ limitations under the License. #include "tensorflow/compiler/xla/array4d.h" #include "tensorflow/compiler/xla/client/computation_builder.h" #include "tensorflow/compiler/xla/client/local_client.h" -#include "tensorflow/compiler/xla/legacy_flags/cpu_compiler_flags.h" #include "tensorflow/compiler/xla/literal_util.h" #include "tensorflow/compiler/xla/statusor.h" -#include "tensorflow/compiler/xla/test_helpers.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/core/platform/test.h" namespace xla { namespace { -using BroadcastSimpleTest = ClientLibraryTestBase; +class BroadcastSimpleTest : public ClientLibraryTestBase { + public: + ComputationDataHandle BuildBinOp(HloOpcode op, + const ComputationDataHandle& lhs, + const ComputationDataHandle& rhs, + ComputationBuilder* builder) { + switch (op) { + case HloOpcode::kMinimum: { + return builder->Min(lhs, rhs); + } + case HloOpcode::kMaximum: { + return builder->Max(lhs, rhs); + } + case HloOpcode::kMultiply: { + return builder->Mul(lhs, rhs); + } + default: { + // Default to Add + return builder->Add(lhs, rhs); + } + } + } + + std::unique_ptr MakeR3Data( + tensorflow::gtl::ArraySlice bounds, + tensorflow::gtl::ArraySlice minor_to_major, Shape* r3_shape, + Array3D* r3_array, float start, float end, int seed) { + *r3_shape = ShapeUtil::MakeShapeWithLayout(F32, bounds, minor_to_major); + r3_array->FillRandom(start, end, seed); + auto r3_data = Literal::CreateR3FromArray3D(*r3_array)->Relayout( + LayoutUtil::MakeLayout(minor_to_major)); + std::unique_ptr r3_global_data = + client_->TransferToServer(*r3_data).ConsumeValueOrDie(); + return r3_global_data; + } + + std::unique_ptr MakeR2Data( + tensorflow::gtl::ArraySlice bounds, + tensorflow::gtl::ArraySlice minor_to_major, Shape* r2_shape, + Array2D* r2_array, float start, float end, int seed) { + *r2_shape = ShapeUtil::MakeShapeWithLayout(F32, bounds, minor_to_major); + r2_array->FillRandom(start, end, seed); + auto r2_data = Literal::CreateR2FromArray2D(*r2_array)->Relayout( + LayoutUtil::MakeLayout(minor_to_major)); + std::unique_ptr r2_global_data = + client_->TransferToServer(*r2_data).ConsumeValueOrDie(); + return r2_global_data; + } + + float ApplyOpToFloats(HloOpcode op, float lhs, float rhs) { + switch (op) { + case HloOpcode::kMinimum: { + return std::min(lhs, rhs); + } + case HloOpcode::kMaximum: { + return std::max(lhs, rhs); + } + case HloOpcode::kMultiply: { + return lhs * rhs; + } + case HloOpcode::kAdd: { + return lhs + rhs; + } + default: { + // Default to Add + CHECK(false); + } + } + } +}; + +using ::testing::HasSubstr; XLA_TEST_F(BroadcastSimpleTest, ScalarNoOpBroadcast) { ComputationBuilder b(client_, TestName()); @@ -89,6 +158,33 @@ XLA_TEST_F(BroadcastSimpleTest, 1DTo2D) { ComputeAndCompareR2(&b, expected, {}, ErrorSpec(0.0001)); } +// Tests implicit broadcasting of PREDs. +XLA_TEST_F(BroadcastSimpleTest, LogicalAnd2DTo3D_Pred) { + ComputationBuilder b(client_, TestName()); + + Array2D x_vals(2, 1); + x_vals(0, 0) = true; + x_vals(1, 0) = false; + Array3D y_vals(2, 2, 1); + y_vals(0, 0, 0) = false; + y_vals(0, 1, 0) = false; + y_vals(1, 0, 0) = true; + y_vals(1, 1, 0) = true; + + ComputationDataHandle x, y; + auto x_data = CreateR2Parameter(x_vals, 0, "x", &b, &x); + auto y_data = CreateR3Parameter(y_vals, 1, "y", &b, &y); + b.LogicalAnd(x, y, /*broadcast_dimensions=*/{1, 2}); + + Array3D expected(2, 2, 1); + expected(0, 0, 0) = false; + expected(0, 1, 0) = false; + expected(1, 0, 0) = true; + expected(1, 1, 0) = false; + + ComputeAndCompareR3(&b, expected, {x_data.get(), y_data.get()}); +} + XLA_TEST_F(BroadcastSimpleTest, ZeroElement_1DTo2D) { ComputationBuilder b(client_, TestName()); b.Broadcast(b.ConstantR1({}), {2}); @@ -116,13 +212,358 @@ XLA_TEST_F(BroadcastSimpleTest, InDimensionAndDegenerateBroadcasting) { ComputationBuilder b(client_, TestName()); b.Add(b.ConstantR2({{1.0, 5.0}}), - b.ConstantLiteral(*LiteralUtil::CreateR3( + b.ConstantLiteral(*Literal::CreateR3( {{{2.0}, {3.0}, {4.0}}, {{5.0}, {6.0}, {7.0}}})), /*broadcast_dimensions=*/{1, 2}); auto expected = - LiteralUtil::CreateR3({{{3.0, 7.0}, {4.0, 8.0}, {5.0, 9.0}}, - {{6.0, 10.0}, {7.0, 11.0}, {8.0, 12.0}}}); + Literal::CreateR3({{{3.0, 7.0}, {4.0, 8.0}, {5.0, 9.0}}, + {{6.0, 10.0}, {7.0, 11.0}, {8.0, 12.0}}}); + + ComputeAndCompareLiteral(&b, *expected, {}, ErrorSpec(0.0001)); +} + +struct R3ImplicitBroadcastSpec { + std::array output_bounds; + std::array minor2major_layout; + std::array input_bounds; + HloOpcode op; +} kR3ImplicitBroadcastTestCases[] = { + {{{1, 1, 1}}, {{2, 1, 0}}, {{1, 1, 1}}, HloOpcode::kAdd}, + {{{3, 4, 5}}, {{2, 1, 0}}, {{1, 1, 5}}, HloOpcode::kMaximum}, + {{{3, 4, 5}}, {{2, 1, 0}}, {{1, 4, 1}}, HloOpcode::kMinimum}, + {{{3, 4, 5}}, {{2, 1, 0}}, {{3, 1, 1}}, HloOpcode::kMultiply}, + {{{3, 4, 5}}, {{2, 1, 0}}, {{1, 1, 1}}, HloOpcode::kAdd}, + {{{3, 4, 5}}, {{2, 1, 0}}, {{1, 4, 5}}, HloOpcode::kAdd}, + {{{3, 4, 5}}, {{2, 1, 0}}, {{3, 4, 1}}, HloOpcode::kAdd}, + {{{3, 4, 5}}, {{2, 1, 0}}, {{3, 1, 5}}, HloOpcode::kAdd}, + {{{3, 199, 5}}, {{2, 1, 0}}, {{1, 199, 1}}, HloOpcode::kMinimum}, + {{{3, 4, 199}}, {{2, 1, 0}}, {{1, 1, 199}}, HloOpcode::kAdd}, +}; + +class BroadcastR3ImplicitTest + : public BroadcastSimpleTest, + public ::testing::WithParamInterface {}; + +XLA_TEST_P(BroadcastR3ImplicitTest, Doit) { + const R3ImplicitBroadcastSpec& spec = GetParam(); + ComputationBuilder builder(client_, TestName()); + + Shape r3_shape, r3_implicit_shape; + Array3D r3_array(spec.output_bounds[0], spec.output_bounds[1], + spec.output_bounds[2]); + Array3D r3_implicit_array(spec.input_bounds[0], spec.input_bounds[1], + spec.input_bounds[2]); + + std::unique_ptr r3_global_data = + MakeR3Data(spec.output_bounds, spec.minor2major_layout, &r3_shape, + &r3_array, 1.0, 2.5, 56789); + std::unique_ptr r3_implicit_global_data = + MakeR3Data(spec.input_bounds, spec.minor2major_layout, &r3_implicit_shape, + &r3_implicit_array, 1.0, 0.2, 56789); + + auto r3_implicit_parameter = builder.Parameter(0, r3_implicit_shape, "input"); + auto r3_parameter = builder.Parameter(1, r3_shape, "input"); + ComputationDataHandle op = + BuildBinOp(spec.op, r3_implicit_parameter, r3_parameter, &builder); + + Array3D expected_array(spec.output_bounds[0], spec.output_bounds[1], + spec.output_bounds[2]); + auto Each = ([&](tensorflow::gtl::ArraySlice indices, float* value) { + float r3_implicit = r3_implicit_array(indices[0] % spec.input_bounds[0], + indices[1] % spec.input_bounds[1], + indices[2] % spec.input_bounds[2]); + float r3 = r3_array(indices[0], indices[1], indices[2]); + *value = ApplyOpToFloats(spec.op, r3_implicit, r3); + }); + + int n1 = expected_array.n1(); + int n2 = expected_array.n2(); + int n3 = expected_array.n3(); + for (int64 i = 0; i < n1; i++) { + for (int64 j = 0; j < n2; j++) { + for (int64 k = 0; k < n3; k++) { + Each({i, j, k}, &expected_array(i, j, k)); + } + } + } + auto expected = Literal::CreateR3FromArray3D(expected_array); + ComputeAndCompareLiteral( + &builder, *expected, + {r3_implicit_global_data.get(), r3_global_data.get()}, + ErrorSpec(1e-7, 1e-7)); +} + +INSTANTIATE_TEST_CASE_P(BroadcastR3ImplicitTestInstances, + BroadcastR3ImplicitTest, + ::testing::ValuesIn(kR3ImplicitBroadcastTestCases)); + +// r1 and r3's dim0 matches, and r1's dim1 and dim2 have size 1: +XLA_TEST_F(BroadcastSimpleTest, Add3DTo3DDegenerate_1_2) { + ComputationBuilder b(client_, TestName()); + ComputationDataHandle r1h; + ComputationDataHandle r3h; + + Array3D r1d = {{{1}}, {{2}}}; + Array3D r3d = {{{1, 2}, {3, 4}}, {{5, 6}, {7, 8}}}; + auto r1 = CreateR3Parameter(r1d, 1, "r1", &b, &r1h); + auto r3 = CreateR3Parameter(r3d, 0, "r3", &b, &r3h); + + b.Add(r3h, r1h); + + auto expected = + Literal::CreateR3({{{2, 3}, {4, 5}}, {{7, 8}, {9, 10}}}); + + ComputeAndCompareLiteral(&b, *expected, {r3.get(), r1.get()}, + ErrorSpec(0.0001)); +} + +XLA_TEST_F(BroadcastSimpleTest, Add3DTo3DDegenerate_0_1) { + ComputationBuilder b(client_, TestName()); + auto r1 = b.ConstantLiteral(*Literal::CreateR3({{{1, 2}}})); + auto r3 = b.ConstantLiteral( + *Literal::CreateR3({{{1, 2}, {3, 4}}, {{5, 6}, {7, 8}}})); + b.Add(r3, r1); + + auto expected = + Literal::CreateR3({{{2, 4}, {4, 6}}, {{6, 8}, {8, 10}}}); + + ComputeAndCompareLiteral(&b, *expected, {}, ErrorSpec(0.0001)); +} + +XLA_TEST_F(BroadcastSimpleTest, Add3DTo3DDegenerate_0_2) { + ComputationBuilder b(client_, TestName()); + auto r1 = b.ConstantLiteral(*Literal::CreateR3({{{1}, {2}}})); + auto r3 = b.ConstantLiteral( + *Literal::CreateR3({{{1, 2}, {3, 4}}, {{5, 6}, {7, 8}}})); + b.Add(r3, r1); + + auto expected = + Literal::CreateR3({{{2, 3}, {5, 6}}, {{6, 7}, {9, 10}}}); + + ComputeAndCompareLiteral(&b, *expected, {}, ErrorSpec(0.0001)); +} + +XLA_TEST_F(BroadcastSimpleTest, Add3DTo3DDegenerate_0) { + ComputationBuilder b(client_, TestName()); + auto r1 = b.ConstantLiteral(*Literal::CreateR3({{{1, 2}, {3, 4}}})); + auto r3 = b.ConstantLiteral( + *Literal::CreateR3({{{1, 2}, {3, 4}}, {{5, 6}, {7, 8}}})); + b.Add(r3, r1); + + auto expected = + Literal::CreateR3({{{2, 4}, {6, 8}}, {{6, 8}, {10, 12}}}); + + ComputeAndCompareLiteral(&b, *expected, {}, ErrorSpec(0.0001)); +} + +XLA_TEST_F(BroadcastSimpleTest, Add3DTo3DDegenerate_1) { + ComputationBuilder b(client_, TestName()); + auto r1 = b.ConstantLiteral(*Literal::CreateR3({{{1, 2}}, {{3, 4}}})); + auto r3 = b.ConstantLiteral( + *Literal::CreateR3({{{1, 2}, {3, 4}}, {{5, 6}, {7, 8}}})); + b.Add(r3, r1); + + auto expected = + Literal::CreateR3({{{2, 4}, {4, 6}}, {{8, 10}, {10, 12}}}); + + ComputeAndCompareLiteral(&b, *expected, {}, ErrorSpec(0.0001)); +} + +XLA_TEST_F(BroadcastSimpleTest, Add3DTo3DDegenerate_2) { + ComputationBuilder b(client_, TestName()); + auto r1 = + b.ConstantLiteral(*Literal::CreateR3({{{1}, {2}}, {{3}, {4}}})); + auto r3 = b.ConstantLiteral( + *Literal::CreateR3({{{1, 2}, {3, 4}}, {{5, 6}, {7, 8}}})); + b.Add(r3, r1); + + auto expected = + Literal::CreateR3({{{2, 3}, {5, 6}}, {{8, 9}, {11, 12}}}); + + ComputeAndCompareLiteral(&b, *expected, {}, ErrorSpec(0.0001)); +} + +XLA_TEST_F(BroadcastSimpleTest, Add3DTo3DDegenerate_0_1_2) { + ComputationBuilder b(client_, TestName()); + auto r1 = b.ConstantLiteral(*Literal::CreateR3({{{1}}})); + auto r3 = b.ConstantLiteral( + *Literal::CreateR3({{{1, 2}, {3, 4}}, {{5, 6}, {7, 8}}})); + b.Add(r3, r1); + + auto expected = + Literal::CreateR3({{{2, 3}, {4, 5}}, {{6, 7}, {8, 9}}}); + + ComputeAndCompareLiteral(&b, *expected, {}, ErrorSpec(0.0001)); +} + +struct R2ImplicitBroadcastSpec { + std::array output_bounds; + std::array minor2major_layout; + std::array input_bounds1; + std::array input_bounds2; + HloOpcode op1; + HloOpcode op2; +} kR2ImplicitBroadcastTestCases[] = { + {{{2, 3}}, {{1, 0}}, {{2, 1}}, {{2, 1}}, HloOpcode::kAdd, HloOpcode::kAdd}, + {{{2, 3}}, {{1, 0}}, {{2, 1}}, {{1, 3}}, HloOpcode::kAdd, HloOpcode::kAdd}, + {{{2, 3}}, + {{1, 0}}, + {{2, 1}}, + {{1, 1}}, + HloOpcode::kAdd, + HloOpcode::kMinimum}, + {{{2, 3}}, + {{1, 0}}, + {{1, 3}}, + {{1, 1}}, + HloOpcode::kAdd, + HloOpcode::kMinimum}, + {{{2, 3}}, + {{1, 0}}, + {{1, 1}}, + {{1, 1}}, + HloOpcode::kAdd, + HloOpcode::kMinimum}, + {{{2, 3}}, {{0, 1}}, {{2, 1}}, {{2, 1}}, HloOpcode::kAdd, HloOpcode::kAdd}, + {{{150, 150}}, + {{1, 0}}, + {{150, 1}}, + {{150, 1}}, + HloOpcode::kAdd, + HloOpcode::kAdd}, + {{{150, 150}}, + {{1, 0}}, + {{150, 1}}, + {{1, 150}}, + HloOpcode::kAdd, + HloOpcode::kAdd}, + {{{150, 150}}, + {{1, 0}}, + {{150, 1}}, + {{1, 1}}, + HloOpcode::kAdd, + HloOpcode::kAdd}, + {{{50, 150}}, + {{1, 0}}, + {{50, 1}}, + {{50, 1}}, + HloOpcode::kAdd, + HloOpcode::kAdd}, + {{{50, 150}}, + {{1, 0}}, + {{50, 1}}, + {{1, 150}}, + HloOpcode::kAdd, + HloOpcode::kAdd}, + {{{50, 150}}, + {{1, 0}}, + {{50, 1}}, + {{1, 1}}, + HloOpcode::kAdd, + HloOpcode::kAdd}, + {{{150, 50}}, + {{1, 0}}, + {{150, 1}}, + {{150, 1}}, + HloOpcode::kAdd, + HloOpcode::kAdd}, + {{{150, 50}}, + {{1, 0}}, + {{150, 1}}, + {{1, 50}}, + HloOpcode::kAdd, + HloOpcode::kAdd}, + {{{150, 50}}, + {{1, 0}}, + {{150, 1}}, + {{1, 1}}, + HloOpcode::kAdd, + HloOpcode::kAdd}}; + +class BroadcastR2ImplicitTest + : public BroadcastSimpleTest, + public ::testing::WithParamInterface {}; + +// Test r2 op1 r2_implicit_1 op2 r2_implicit_2 +// where R2 is a rank-2 operand, and r2_implicit_2 are two +// rank-2 operands with degenerate dimensions: +XLA_TEST_P(BroadcastR2ImplicitTest, Doit) { + const R2ImplicitBroadcastSpec& spec = GetParam(); + + ComputationBuilder builder(client_, TestName()); + + // Operands with degenerate dimensions require implicit broadcasting: + Shape r2_shape, r2_implicit_shape1, r2_implicit_shape2; + Array2D r2_array(spec.output_bounds[0], spec.output_bounds[1]); + Array2D r2_implicit_array1(spec.input_bounds1[0], + spec.input_bounds1[1]); + Array2D r2_implicit_array2(spec.input_bounds2[0], + spec.input_bounds2[1]); + + std::unique_ptr r2_global_data = + MakeR2Data(spec.output_bounds, spec.minor2major_layout, &r2_shape, + &r2_array, 1.0, 2.5, 56789); + std::unique_ptr r2_implicit_global_data1 = + MakeR2Data(spec.input_bounds1, spec.minor2major_layout, + &r2_implicit_shape1, &r2_implicit_array1, 1.0, 0.2, 56789); + std::unique_ptr r2_implicit_global_data2 = + MakeR2Data(spec.input_bounds2, spec.minor2major_layout, + &r2_implicit_shape2, &r2_implicit_array2, 0.8, 0.4, 56789); + + auto r2_implicit_parameter1 = + builder.Parameter(0, r2_implicit_shape1, "input0"); + auto r2_parameter = builder.Parameter(1, r2_shape, "input1"); + auto r2_implicit_parameter2 = + builder.Parameter(2, r2_implicit_shape2, "input2"); + + ComputationDataHandle op1 = + BuildBinOp(spec.op1, r2_implicit_parameter1, r2_parameter, &builder); + ComputationDataHandle op2 = + BuildBinOp(spec.op2, op1, r2_implicit_parameter2, &builder); + + Array2D expected_array(spec.output_bounds[0], spec.output_bounds[1]); + + expected_array.Each([&](int64 i, int64 j, float* v) { + float v1 = r2_implicit_array1(i % spec.input_bounds1[0], + j % spec.input_bounds1[1]); + float v2 = r2_array(i, j); + float v3 = r2_implicit_array2(i % spec.input_bounds2[0], + j % spec.input_bounds2[1]); + float tmp = ApplyOpToFloats(spec.op1, v1, v2); + *v = ApplyOpToFloats(spec.op2, tmp, v3); + }); + + auto expected = Literal::CreateR2FromArray2D(expected_array); + ComputeAndCompareLiteral( + &builder, *expected, + {r2_implicit_global_data1.get(), r2_global_data.get(), + r2_implicit_global_data2.get()}, + ErrorSpec(1e-6, 1e-6)); +} + +INSTANTIATE_TEST_CASE_P(BroadcastR2ImplicitTestInstances, + BroadcastR2ImplicitTest, + ::testing::ValuesIn(kR2ImplicitBroadcastTestCases)); + +XLA_TEST_F(BroadcastSimpleTest, Add2DTo2DDegenerate_0) { + ComputationBuilder b(client_, TestName()); + auto r1 = b.ConstantLiteral(*Literal::CreateR2({{1, 2}})); + auto r2 = b.ConstantLiteral(*Literal::CreateR2({{1, 2}, {3, 4}})); + b.Add(r2, r1); + + auto expected = Literal::CreateR2({{2, 4}, {4, 6}}); + + ComputeAndCompareLiteral(&b, *expected, {}, ErrorSpec(0.0001)); +} + +XLA_TEST_F(BroadcastSimpleTest, Add2DTo2DDegenerate_1) { + ComputationBuilder b(client_, TestName()); + auto r1 = b.ConstantLiteral(*Literal::CreateR2({{1}, {2}})); + auto r2 = b.ConstantLiteral(*Literal::CreateR2({{1, 2}, {3, 4}})); + b.Add(r2, r1); + + auto expected = Literal::CreateR2({{2, 3}, {5, 6}}); ComputeAndCompareLiteral(&b, *expected, {}, ErrorSpec(0.0001)); } @@ -131,11 +572,11 @@ XLA_TEST_F(BroadcastSimpleTest, Add1DTo3DInDim0) { ComputationBuilder b(client_, TestName()); auto r1 = b.ConstantR1({10, 20}); auto r3 = b.ConstantLiteral( - *LiteralUtil::CreateR3({{{1, 2}, {3, 4}}, {{5, 6}, {7, 8}}})); + *Literal::CreateR3({{{1, 2}, {3, 4}}, {{5, 6}, {7, 8}}})); b.Add(r3, r1, {0}); - auto expected = LiteralUtil::CreateR3( - {{{11, 12}, {13, 14}}, {{25, 26}, {27, 28}}}); + auto expected = + Literal::CreateR3({{{11, 12}, {13, 14}}, {{25, 26}, {27, 28}}}); ComputeAndCompareLiteral(&b, *expected, {}, ErrorSpec(0.0001)); } @@ -144,11 +585,11 @@ XLA_TEST_F(BroadcastSimpleTest, Add1DTo3DInDim1) { ComputationBuilder b(client_, TestName()); auto r1 = b.ConstantR1({10, 20}); auto r3 = b.ConstantLiteral( - *LiteralUtil::CreateR3({{{1, 2}, {3, 4}}, {{5, 6}, {7, 8}}})); + *Literal::CreateR3({{{1, 2}, {3, 4}}, {{5, 6}, {7, 8}}})); b.Add(r1, r3, {1}); - auto expected = LiteralUtil::CreateR3( - {{{11, 12}, {23, 24}}, {{15, 16}, {27, 28}}}); + auto expected = + Literal::CreateR3({{{11, 12}, {23, 24}}, {{15, 16}, {27, 28}}}); ComputeAndCompareLiteral(&b, *expected, {}, ErrorSpec(0.0001)); } @@ -157,11 +598,11 @@ XLA_TEST_F(BroadcastSimpleTest, Add1DTo3DInDim2) { ComputationBuilder b(client_, TestName()); auto r1 = b.ConstantR1({10, 20}); auto r3 = b.ConstantLiteral( - *LiteralUtil::CreateR3({{{1, 2}, {3, 4}}, {{5, 6}, {7, 8}}})); + *Literal::CreateR3({{{1, 2}, {3, 4}}, {{5, 6}, {7, 8}}})); b.Add(r1, r3, {2}); - auto expected = LiteralUtil::CreateR3( - {{{11, 22}, {13, 24}}, {{15, 26}, {17, 28}}}); + auto expected = + Literal::CreateR3({{{11, 22}, {13, 24}}, {{15, 26}, {17, 28}}}); ComputeAndCompareLiteral(&b, *expected, {}, ErrorSpec(0.0001)); } @@ -172,7 +613,7 @@ XLA_TEST_F(BroadcastSimpleTest, Add1DTo3DInDimAll) { auto r1_1 = b.ConstantR1({100, 200}); auto r1_2 = b.ConstantR1({10, 20}); auto r3 = b.ConstantLiteral( - *LiteralUtil::CreateR3({{{1, 2}, {3, 4}}, {{5, 6}, {7, 8}}})); + *Literal::CreateR3({{{1, 2}, {3, 4}}, {{5, 6}, {7, 8}}})); for (int i = 0; i < 3; ++i) { r3 = b.Add(r1_0, r3, {0}); r3 = b.Add(r3, r1_1, {1}); @@ -180,7 +621,7 @@ XLA_TEST_F(BroadcastSimpleTest, Add1DTo3DInDimAll) { } r3 = b.Mul(r3, b.ConstantR0(-2)); - auto expected = LiteralUtil::CreateR3( + auto expected = Literal::CreateR3( {{{-6 * 1110 - 2, -6 * 1120 - 4}, {-6 * 1210 - 6, -6 * 1220 - 8}}, {{-6 * 2110 - 10, -6 * 2120 - 12}, {-6 * 2210 - 14, -6 * 2220 - 16}}}); @@ -201,7 +642,7 @@ XLA_TEST_F(BroadcastSimpleTest, Add1DTo3DInDimAllWithScalarBroadcast) { } r3 = b.Mul(r3, b.ConstantR0(-1)); - auto expected = LiteralUtil::CreateR3( + auto expected = Literal::CreateR3( {{{-3 * 1110 - 3, -3 * 1120 - 3}, {-3 * 1210 - 3, -3 * 1220 - 3}}, {{-3 * 2110 - 3, -3 * 2120 - 3}, {-3 * 2210 - 3, -3 * 2220 - 3}}}); @@ -214,14 +655,14 @@ XLA_TEST_F(BroadcastSimpleTest, InvalidBinaryAndDegenerateBroadcasting) { ComputationBuilder b(client_, TestName()); b.Add(b.ConstantR2({{1.0, 5.0}, {1.0, 5.0}}), - b.ConstantLiteral(*LiteralUtil::CreateR3( + b.ConstantLiteral(*Literal::CreateR3( {{{2.0}, {3.0}, {4.0}}, {{5.0}, {6.0}, {7.0}}})), /*broadcast_dimensions=*/{1, 2}); auto result_status = Execute(&b, {}); EXPECT_FALSE(result_status.ok()); - EXPECT_MATCH(result_status.status().error_message(), - testing::ContainsRegex("broadcast dimension 0 mismatch")); + EXPECT_THAT(result_status.status().error_message(), + HasSubstr("broadcast dimension 0 mismatch")); } XLA_TEST_F(BroadcastSimpleTest, InvalidInDimensionBroadcasting) { @@ -233,9 +674,8 @@ XLA_TEST_F(BroadcastSimpleTest, InvalidInDimensionBroadcasting) { auto result_status = Execute(&b, {}); EXPECT_FALSE(result_status.ok()); - EXPECT_MATCH( - result_status.status().error_message(), - testing::ContainsRegex("binary op BINOP_ADD with incompatible shapes")); + EXPECT_THAT(result_status.status().error_message(), + HasSubstr("binary op BINOP_ADD with incompatible shapes")); } XLA_TEST_F(BroadcastSimpleTest, InvalidDegenerateBroadcasting) { @@ -247,27 +687,9 @@ XLA_TEST_F(BroadcastSimpleTest, InvalidDegenerateBroadcasting) { auto result_status = Execute(&b, {}); EXPECT_FALSE(result_status.ok()); - EXPECT_MATCH( - result_status.status().error_message(), - testing::ContainsRegex("binary op BINOP_ADD with incompatible shapes")); + EXPECT_THAT(result_status.status().error_message(), + HasSubstr("binary op BINOP_ADD with incompatible shapes")); } } // namespace } // namespace xla - -int main(int argc, char** argv) { - std::vector flag_list; - xla::legacy_flags::AppendCpuCompilerFlags(&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; - } - testing::InitGoogleTest(&argc, argv); - if (argc > 1) { - LOG(ERROR) << "Unknown argument " << argv[1] << "\n" << usage; - return 2; - } - return RUN_ALL_TESTS(); -} diff --git a/tensorflow/compiler/xla/tests/broadcast_test.cc b/tensorflow/compiler/xla/tests/broadcast_test.cc index 1796a732e543b7f040adf6055e349d72cfcfad6e..0294628a127c9d506e6387d0b80f3da583c5a174 100644 --- a/tensorflow/compiler/xla/tests/broadcast_test.cc +++ b/tensorflow/compiler/xla/tests/broadcast_test.cc @@ -16,7 +16,6 @@ limitations under the License. #include #include -#include "tensorflow/compiler/xla/legacy_flags/cpu_compiler_flags.h" #include "tensorflow/compiler/xla/literal_util.h" #include "tensorflow/compiler/xla/ptr_util.h" #include "tensorflow/compiler/xla/service/hlo_computation.h" @@ -38,40 +37,40 @@ XLA_TEST_F(BroadcastTest, BroadcastScalarToScalar) { // Test degenerate case of broadcasting a scalar into a scalar. auto builder = HloComputation::Builder(TestName()); auto input = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(42.0))); + HloInstruction::CreateConstant(Literal::CreateR0(42.0))); builder.AddInstruction(HloInstruction::CreateBroadcast( ShapeUtil::MakeShape(F32, {}), input, {})); // Create HLO module, compile, and execute. - auto hlo_module = MakeUnique(TestName()); + auto hlo_module = CreateNewModule(); hlo_module->AddEntryComputation(builder.Build()); auto result = ExecuteAndTransfer(std::move(hlo_module), {}); - LiteralTestUtil::ExpectNear(*LiteralUtil::CreateR0(42.0), *result, + LiteralTestUtil::ExpectNear(*Literal::CreateR0(42.0), *result, error_spec_); } XLA_TEST_F(BroadcastTest, BroadcastScalarTo2D) { auto builder = HloComputation::Builder(TestName()); auto input = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(42.0))); + HloInstruction::CreateConstant(Literal::CreateR0(42.0))); builder.AddInstruction(HloInstruction::CreateBroadcast( ShapeUtil::MakeShape(F32, {2, 2}), input, {})); // Create HLO module, compile, and execute. - auto hlo_module = MakeUnique(TestName()); + auto hlo_module = CreateNewModule(); hlo_module->AddEntryComputation(builder.Build()); auto result = ExecuteAndTransfer(std::move(hlo_module), {}); LiteralTestUtil::ExpectNear( - *LiteralUtil::CreateR2({{42.0, 42.0}, {42.0, 42.0}}), *result, + *Literal::CreateR2({{42.0, 42.0}, {42.0, 42.0}}), *result, error_spec_); } XLA_TEST_F(BroadcastTest, BroadcastVectorTo2D) { auto builder = HloComputation::Builder(TestName()); auto input = builder.AddInstruction(HloInstruction::CreateConstant( - LiteralUtil::CreateR1({1.0, 2.0, 3.0}))); + Literal::CreateR1({1.0, 2.0, 3.0}))); // Broadcast vector in both dimension 0 and dimension 1. Join them in a tuple // to enable testing of the results. @@ -82,33 +81,33 @@ XLA_TEST_F(BroadcastTest, BroadcastVectorTo2D) { builder.AddInstruction(HloInstruction::CreateTuple({element1, element2})); // Create HLO module, compile, and execute. - auto hlo_module = MakeUnique(TestName()); + auto hlo_module = CreateNewModule(); hlo_module->AddEntryComputation(builder.Build()); auto result = ExecuteAndTransfer(std::move(hlo_module), {}); LiteralTestUtil::ExpectNear( - *LiteralUtil::CreateR2({{1.0, 1.0}, {2.0, 2.0}, {3.0, 3.0}}), + *Literal::CreateR2({{1.0, 1.0}, {2.0, 2.0}, {3.0, 3.0}}), result->tuple_literals(0), error_spec_); LiteralTestUtil::ExpectNear( - *LiteralUtil::CreateR2({{1.0, 2.0, 3.0}, {1.0, 2.0, 3.0}}), + *Literal::CreateR2({{1.0, 2.0, 3.0}, {1.0, 2.0, 3.0}}), result->tuple_literals(1), error_spec_); } XLA_TEST_F(BroadcastTest, Broadcast2DTo2D) { auto builder = HloComputation::Builder(TestName()); auto input = builder.AddInstruction(HloInstruction::CreateConstant( - LiteralUtil::CreateR2({{1.0, 2.0}, {3.0, 4.0}}))); + Literal::CreateR2({{1.0, 2.0}, {3.0, 4.0}}))); builder.AddInstruction(HloInstruction::CreateBroadcast( ShapeUtil::MakeShape(F32, {2, 2}), input, {0, 1})); // Create HLO module, compile, and execute. - auto hlo_module = MakeUnique(TestName()); + auto hlo_module = CreateNewModule(); hlo_module->AddEntryComputation(builder.Build()); auto result = ExecuteAndTransfer(std::move(hlo_module), {}); LiteralTestUtil::ExpectNear( - *LiteralUtil::CreateR2({{1.0, 2.0}, {3.0, 4.0}}), *result, + *Literal::CreateR2({{1.0, 2.0}, {3.0, 4.0}}), *result, error_spec_); } @@ -117,49 +116,49 @@ XLA_TEST_F(BroadcastTest, Broadcast2DTo2DTranspose) { // the dimensions, ie transpose. auto builder = HloComputation::Builder(TestName()); auto input = builder.AddInstruction(HloInstruction::CreateConstant( - LiteralUtil::CreateR2({{1.0, 2.0}, {3.0, 4.0}}))); + Literal::CreateR2({{1.0, 2.0}, {3.0, 4.0}}))); builder.AddInstruction(HloInstruction::CreateBroadcast( ShapeUtil::MakeShape(F32, {2, 2}), input, {1, 0})); // Create HLO module, compile, and execute. - auto hlo_module = MakeUnique(TestName()); + auto hlo_module = CreateNewModule(); hlo_module->AddEntryComputation(builder.Build()); auto result = ExecuteAndTransfer(std::move(hlo_module), {}); LiteralTestUtil::ExpectNear( - *LiteralUtil::CreateR2({{1.0, 3.0}, {2.0, 4.0}}), *result, + *Literal::CreateR2({{1.0, 3.0}, {2.0, 4.0}}), *result, error_spec_); } XLA_TEST_F(BroadcastTest, Broadcast2DTo3D) { auto builder = HloComputation::Builder(TestName()); auto input = builder.AddInstruction(HloInstruction::CreateConstant( - LiteralUtil::CreateR2({{1.0, 2.0}, {3.0, 4.0}}))); + Literal::CreateR2({{1.0, 2.0}, {3.0, 4.0}}))); builder.AddInstruction(HloInstruction::CreateBroadcast( ShapeUtil::MakeShape(F32, {2, 3, 2}), input, {0, 2})); // Create HLO module, compile, and execute. - auto hlo_module = MakeUnique(TestName()); + auto hlo_module = CreateNewModule(); hlo_module->AddEntryComputation(builder.Build()); auto result = ExecuteAndTransfer(std::move(hlo_module), {}); LiteralTestUtil::ExpectNear( - *LiteralUtil::CreateR3({{{1.0, 2.0}, {1.0, 2.0}, {1.0, 2.0}}, - {{3.0, 4.0}, {3.0, 4.0}, {3.0, 4.0}}}), + *Literal::CreateR3({{{1.0, 2.0}, {1.0, 2.0}, {1.0, 2.0}}, + {{3.0, 4.0}, {3.0, 4.0}, {3.0, 4.0}}}), *result, error_spec_); } TEST_F(BroadcastTest, Broadcast_R1_2_To_R4_2x2x3x3) { auto builder = HloComputation::Builder(TestName()); auto input = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR1({1.0, 2.0}))); + HloInstruction::CreateConstant(Literal::CreateR1({1.0, 2.0}))); // Broadcast vector in dimension 1. builder.AddInstruction(HloInstruction::CreateBroadcast( ShapeUtil::MakeShape(F32, {2, 2, 3, 3}), input, {1})); // Create HLO module, compile, and execute. - auto hlo_module = MakeUnique(TestName()); + auto hlo_module = CreateNewModule(); hlo_module->AddEntryComputation(builder.Build()); auto result = ExecuteAndTransfer(std::move(hlo_module), {}); @@ -167,8 +166,8 @@ TEST_F(BroadcastTest, Broadcast_R1_2_To_R4_2x2x3x3) { Array2D pz({{1, 2}, {1, 2}}); expected.FillWithPZ(pz); - LiteralTestUtil::ExpectNear( - *LiteralUtil::CreateR4FromArray4D(expected), *result, error_spec_); + LiteralTestUtil::ExpectNear(*Literal::CreateR4FromArray4D(expected), + *result, error_spec_); } TEST_F(BroadcastTest, Broadcast_R1_1025_To_R4_3x3x3x1025) { @@ -177,19 +176,19 @@ TEST_F(BroadcastTest, Broadcast_R1_1025_To_R4_3x3x3x1025) { int64 r1_size = input_data.size(); std::iota(input_data.begin(), input_data.end(), 0.0f); auto input = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR1(input_data))); + HloInstruction::CreateConstant(Literal::CreateR1(input_data))); // Broadcast vector in dimension 3. builder.AddInstruction(HloInstruction::CreateBroadcast( ShapeUtil::MakeShape(F32, {3, 3, 3, r1_size}), input, {3})); // Create HLO module, compile, and execute. - auto hlo_module = MakeUnique(TestName()); + auto hlo_module = CreateNewModule(); hlo_module->AddEntryComputation(builder.Build()); auto result = ExecuteAndTransfer(std::move(hlo_module), {}); Array4D expected(3, 3, 3, 1025); - Array2D yx(/*height=*/3, /*width=*/r1_size); + Array2D yx(3, r1_size); for (int64 y = 0; y < 3; ++y) { for (int64 x = 0; x < r1_size; ++x) { yx(y, x) = input_data[x]; @@ -197,8 +196,8 @@ TEST_F(BroadcastTest, Broadcast_R1_1025_To_R4_3x3x3x1025) { } expected.FillWithYX(yx); - LiteralTestUtil::ExpectNear( - *LiteralUtil::CreateR4FromArray4D(expected), *result, error_spec_); + LiteralTestUtil::ExpectNear(*Literal::CreateR4FromArray4D(expected), + *result, error_spec_); } XLA_TEST_F(BroadcastTest, Broadcast_R1_64_To_R4_32x64x7x7) { @@ -208,30 +207,30 @@ XLA_TEST_F(BroadcastTest, Broadcast_R1_64_To_R4_32x64x7x7) { std::vector r1_array(64, 42.0); auto input = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR1(r1_array))); + HloInstruction::CreateConstant(Literal::CreateR1(r1_array))); // Broadcast vector in dimension 1. builder.AddInstruction(HloInstruction::CreateBroadcast( ShapeUtil::MakeShape(F32, {32, 64, 7, 7}), input, {1})); // Create HLO module, compile, and execute. - auto hlo_module = MakeUnique(TestName()); + auto hlo_module = CreateNewModule(); hlo_module->AddEntryComputation(builder.Build()); auto result = ExecuteAndTransfer(std::move(hlo_module), {}); - LiteralTestUtil::ExpectNear(*LiteralUtil::CreateR4FromArray4D(r4_array), - *result, error_spec_); + LiteralTestUtil::ExpectNear(*Literal::CreateR4FromArray4D(r4_array), *result, + error_spec_); } TEST_F(BroadcastTest, Broadcast_R0_to_R4_64x64x3x3) { auto builder = HloComputation::Builder(TestName()); auto input = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(1.0f))); + HloInstruction::CreateConstant(Literal::CreateR0(1.0f))); builder.AddInstruction(HloInstruction::CreateBroadcast( ShapeUtil::MakeShape(F32, {64, 64, 3, 3}), input, {})); // Create HLO module, compile, and execute. - auto hlo_module = MakeUnique(TestName()); + auto hlo_module = CreateNewModule(); hlo_module->AddEntryComputation(builder.Build()); LOG(INFO) << hlo_module->ToString(); auto result = ExecuteAndTransfer(std::move(hlo_module), {}); @@ -239,48 +238,62 @@ TEST_F(BroadcastTest, Broadcast_R0_to_R4_64x64x3x3) { Array4D expected(64, 64, 3, 3); expected.Fill(1.0f); - LiteralTestUtil::ExpectNear( - *LiteralUtil::CreateR4FromArray4D(expected), *result, error_spec_); + LiteralTestUtil::ExpectNear(*Literal::CreateR4FromArray4D(expected), + *result, error_spec_); } TEST_F(BroadcastTest, Broadcast_R2_2x2_To_R4_3x3x2x2) { auto builder = HloComputation::Builder(TestName()); Array2D to_broadcast({{1.0f, 2.0f}, {3.0f, 4.0f}}); auto input = builder.AddInstruction(HloInstruction::CreateConstant( - LiteralUtil::CreateR2FromArray2D(to_broadcast))); + Literal::CreateR2FromArray2D(to_broadcast))); // Broadcast vector in dimensions 2 and 3. builder.AddInstruction(HloInstruction::CreateBroadcast( ShapeUtil::MakeShape(F32, {3, 3, 2, 2}), input, {2, 3})); // Create HLO module, compile, and execute. - auto hlo_module = MakeUnique(TestName()); + auto hlo_module = CreateNewModule(); hlo_module->AddEntryComputation(builder.Build()); auto result = ExecuteAndTransfer(std::move(hlo_module), {}); Array4D expected(3, 3, 2, 2); expected.FillWithYX(to_broadcast); - LiteralTestUtil::ExpectNear( - *LiteralUtil::CreateR4FromArray4D(expected), *result, error_spec_); + LiteralTestUtil::ExpectNear(*Literal::CreateR4FromArray4D(expected), + *result, error_spec_); } -} // namespace -} // namespace xla - -int main(int argc, char** argv) { - std::vector flag_list; - xla::legacy_flags::AppendCpuCompilerFlags(&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; - } - testing::InitGoogleTest(&argc, argv); - if (argc > 1) { - LOG(ERROR) << "Unknown argument " << argv[1] << "\n" << usage; - return 2; +TEST_F(BroadcastTest, Broadcast_R3_2x3x4_to_R4_2x3x4x5) { + auto builder = HloComputation::Builder(TestName()); + Array3D input_vals(2, 3, 4); + input_vals.FillRandom(1.0); + + Array4D expected(2, 3, 4, 5); + for (int i = 0; i < 2; ++i) { + for (int j = 0; j < 3; ++j) { + for (int k = 0; k < 4; ++k) { + for (int m = 0; m < 5; ++m) { + expected(i, j, k, m) = input_vals(i, j, k); + } + } + } } - return RUN_ALL_TESTS(); + auto input = builder.AddInstruction(HloInstruction::CreateConstant( + Literal::CreateR3FromArray3D(input_vals))); + + // Broadcast vector in dimensions 2 and 3. + builder.AddInstruction(HloInstruction::CreateBroadcast( + ShapeUtil::MakeShape(F32, {2, 3, 4, 5}), input, {0, 1, 2})); + + // Create HLO module, compile, and execute. + auto hlo_module = CreateNewModule(); + hlo_module->AddEntryComputation(builder.Build()); + auto result = ExecuteAndTransfer(std::move(hlo_module), {}); + + LiteralTestUtil::ExpectNear(*Literal::CreateR4FromArray4D(expected), + *result, error_spec_); } + +} // namespace +} // namespace xla diff --git a/tensorflow/compiler/xla/tests/build_defs.bzl b/tensorflow/compiler/xla/tests/build_defs.bzl index 2c7eeb820d3ea6fc38cf2132657d82ad8e9400d2..82935157f51a0948c88388c93852605c554a8640 100644 --- a/tensorflow/compiler/xla/tests/build_defs.bzl +++ b/tensorflow/compiler/xla/tests/build_defs.bzl @@ -1,20 +1,37 @@ """Build rules for XLA testing.""" load("@local_config_cuda//cuda:build_defs.bzl", "cuda_is_configured") +load("//tensorflow/compiler/xla/tests:plugin.bzl", "plugins") -def all_backends(): +all_backends = ["cpu", "cpu_parallel", "gpu"] + plugins.keys() + +def filter_backends(backends): + """Removes "gpu" from a backend list if CUDA is not enabled. + + This allows us to simply hardcode lists including "gpu" here and in the + BUILD file, without causing failures when CUDA isn't enabled.' + + Args: + backends: A list of backends to filter. + + Returns: + The filtered list of backends. + """ if cuda_is_configured(): - return ["cpu", "cpu_parallel", "gpu"] + return backends else: - return ["cpu", "cpu_parallel"] + return [backend for backend in backends if backend != "gpu"] + def xla_test(name, srcs, deps, + xla_test_library_deps=[], backends=[], args=[], tags=[], copts=[], + data=[], backend_tags={}, backend_args={}, **kwargs): @@ -69,6 +86,8 @@ def xla_test(name, name: Name of the target. srcs: Sources for the target. deps: Dependencies of the target. + xla_test_library_deps: If set, the generated test targets will depend on the + respective cc_libraries generated by the xla_test_library rule. backends: A list of backends to generate tests for. Supported values: "cpu", "cpu_parallel", "gpu". If this list is empty, the test will be generated for all supported backends. @@ -81,7 +100,7 @@ def xla_test(name, """ test_names = [] if not backends: - backends = all_backends() + backends = all_backends native.cc_library( name="%s_lib" % name, @@ -91,25 +110,36 @@ def xla_test(name, deps=deps + ["//tensorflow/compiler/xla/tests:test_macros_header"], ) - for backend in backends: + for backend in filter_backends(backends): test_name = "%s_%s" % (name, backend) this_backend_tags = ["xla_%s" % backend] this_backend_copts = [] this_backend_args = backend_args.get(backend, []) + this_backend_data = [] if backend == "cpu": backend_deps = ["//tensorflow/compiler/xla/service:cpu_plugin"] backend_deps += ["//tensorflow/compiler/xla/tests:test_macros_cpu"] elif backend == "cpu_parallel": backend_deps = ["//tensorflow/compiler/xla/service:cpu_plugin"] backend_deps += ["//tensorflow/compiler/xla/tests:test_macros_cpu"] - this_backend_args += ["--xla_cpu_parallel=true"] + this_backend_args += ["--xla_backend_extra_options=\"xla_cpu_parallel\""] elif backend == "gpu": backend_deps = ["//tensorflow/compiler/xla/service:gpu_plugin"] backend_deps += ["//tensorflow/compiler/xla/tests:test_macros_gpu"] this_backend_tags += ["requires-gpu-sm35"] + elif backend in plugins: + backend_deps = plugins[backend]["deps"] + this_backend_copts += plugins[backend]["copts"] + this_backend_tags += plugins[backend]["tags"] + this_backend_args += plugins[backend]["args"] + this_backend_data += plugins[backend]["data"] else: fail("Unknown backend %s" % backend) + if xla_test_library_deps: + for lib_dep in xla_test_library_deps: + backend_deps += ["%s_%s" % (lib_dep, backend)] + native.cc_test( name=test_name, srcs=srcs, @@ -118,30 +148,102 @@ def xla_test(name, this_backend_copts, args=args + this_backend_args, deps=deps + backend_deps, + data=data + this_backend_data, **kwargs) test_names.append(test_name) native.test_suite(name=name, tests=test_names) +def xla_test_library(name, + srcs, + hdrs=[], + deps=[], + backends=[]): + """Generates cc_library targets for the given XLA backends. + + This rule forces the sources to be compiled for each backend so that the + backend specific macros could expand correctly. It's useful when test targets + in different directories referring to the same sources but test with different + arguments. + + Examples: + + # Generates the targets: foo_test_library_cpu and foo_test_gpu. + xla_test_library( + name = "foo_test_library", + srcs = ["foo_test.cc"], + backends = ["cpu", "gpu"], + deps = [...], + ) + # Then use the xla_test rule to generate test targets: + xla_test( + name = "foo_test", + srcs = [], + backends = ["cpu", "gpu"], + deps = [...], + xla_test_library_deps = [":foo_test_library"], + ) + + Args: + name: Name of the target. + srcs: Sources for the target. + hdrs: Headers for the target. + deps: Dependencies of the target. + backends: A list of backends to generate libraries for. + Supported values: "cpu", "cpu_parallel", "gpu". If this list is empty, the + library will be generated for all supported backends. + """ + + if not backends: + backends = all_backends + + for backend in filter_backends(backends): + this_backend_copts = [] + if backend in ["cpu", "cpu_parallel", "gpu"]: + backend_deps = ["//tensorflow/compiler/xla/tests:test_macros_%s" % backend] + elif backend in plugins: + backend_deps = plugins[backend]["deps"] + this_backend_copts += plugins[backend]["copts"] + else: + fail("Unknown backend %s" % backend) + + native.cc_library( + name = "%s_%s" % (name, backend), + srcs = srcs, + testonly = True, + hdrs = hdrs, + copts = ["-DXLA_TEST_BACKEND_%s=1" % backend.upper()] + + this_backend_copts, + deps = deps + backend_deps, + ) + def generate_backend_suites(backends=[]): if not backends: - backends = all_backends() - for backend in backends: + backends = all_backends + for backend in filter_backends(backends): native.test_suite(name="%s_tests" % backend, tags = ["xla_%s" % backend]) def generate_backend_test_macros(backends=[]): if not backends: - backends = all_backends() - for backend in backends: + backends = all_backends + for backend in filter_backends(backends): + manifest = "" + if backend in plugins: + manifest = plugins[backend]["disabled_manifest"] + native.cc_library( name="test_macros_%s" % backend, testonly = True, + srcs = ["test_macros.cc"], hdrs = ["test_macros.h"], - copts = ["-DXLA_PLATFORM=\\\"%s\\\"" % backend.upper()], + copts = [ + "-DXLA_PLATFORM=\\\"%s\\\"" % backend.upper(), + "-DXLA_DISABLED_MANIFEST=\\\"%s\\\"" % manifest, + ], deps = [ "//tensorflow/compiler/xla:types", "//tensorflow/core:lib", diff --git a/tensorflow/compiler/xla/tests/call_test.cc b/tensorflow/compiler/xla/tests/call_test.cc index 0b5e6d512771cc6aebfd92af81bfdfa56d176088..214bc79198d56b4eb2b66823de2d8317bf041987 100644 --- a/tensorflow/compiler/xla/tests/call_test.cc +++ b/tensorflow/compiler/xla/tests/call_test.cc @@ -18,9 +18,9 @@ limitations under the License. #include "tensorflow/compiler/xla/client/computation.h" #include "tensorflow/compiler/xla/client/computation_builder.h" -#include "tensorflow/compiler/xla/legacy_flags/cpu_compiler_flags.h" #include "tensorflow/compiler/xla/literal_util.h" #include "tensorflow/compiler/xla/shape_util.h" +#include "tensorflow/compiler/xla/test_helpers.h" #include "tensorflow/compiler/xla/tests/client_library_test_base.h" #include "tensorflow/compiler/xla/tests/literal_test_util.h" #include "tensorflow/compiler/xla/tests/test_macros.h" @@ -73,40 +73,76 @@ class CallOpTest : public ClientLibraryTestBase { Shape r1s2f32_ = ShapeUtil::MakeShape(F32, {2}); }; +// TODO(b/64094172) Failing on GPU as of 2017-07-26. XLA_TEST_F(CallOpTest, DISABLED_ON_GPU(CallR0F32IdentityScalar)) { ComputationBuilder builder(client_, TestName()); Computation callee = CreateR0F32IdentityComputation(); - auto constant = builder.ConstantLiteral(*LiteralUtil::CreateR0(42.0)); + auto constant = builder.ConstantLiteral(*Literal::CreateR0(42.0)); builder.Call(callee, {constant}); ComputeAndCompareR0(&builder, 42.0, {}, ErrorSpec(0.01f)); } +// TODO(b/64094172) Failing on GPU as of 2017-07-26. XLA_TEST_F(CallOpTest, DISABLED_ON_GPU(CallR1S0F32AddArray)) { ComputationBuilder builder(client_, TestName()); Computation callee = CreateR1S0F32AdditionComputation(); - auto x = builder.ConstantLiteral(*LiteralUtil::CreateR1({})); - auto y = builder.ConstantLiteral(*LiteralUtil::CreateR1({})); + auto x = builder.ConstantLiteral(*Literal::CreateR1({})); + auto y = builder.ConstantLiteral(*Literal::CreateR1({})); builder.Call(callee, {x, y}); ComputeAndCompareR1(&builder, {}, {}, ErrorSpec(0.01f)); } +// TODO(b/64094172) Failing on GPU as of 2017-07-26. XLA_TEST_F(CallOpTest, DISABLED_ON_GPU(CallR1S2F32AddArray)) { ComputationBuilder builder(client_, TestName()); Computation callee = CreateR1S2F32AdditionComputation(); - auto x = builder.ConstantLiteral(*LiteralUtil::CreateR1({1.0f, 2.0f})); - auto y = builder.ConstantLiteral(*LiteralUtil::CreateR1({2.0f, 3.0f})); + auto x = builder.ConstantLiteral(*Literal::CreateR1({1.0f, 2.0f})); + auto y = builder.ConstantLiteral(*Literal::CreateR1({2.0f, 3.0f})); builder.Call(callee, {x, y}); ComputeAndCompareR1(&builder, {3.0f, 5.0f}, {}, ErrorSpec(0.01f)); } +// TODO(b/64094172) Failing on GPU as of 2017-07-26. +XLA_TEST_F(CallOpTest, DISABLED_ON_GPU(CallTreeTwoDeepBranchFactorThree)) { + ComputationBuilder builder(client_, "inner"); + { + auto x = builder.Parameter(0, r0f32_, "x"); + builder.Add(x, builder.ConstantR0(1.0)); + } + TF_ASSERT_OK_AND_ASSIGN(Computation inner, builder.Build()); + + ComputationBuilder builder2(client_, "outer"); + { + auto x = builder2.Parameter(0, r0f32_, "x"); + x = builder2.Call(inner, {x}); + x = builder2.Call(inner, {x}); + x = builder2.Call(inner, {x}); + } + TF_ASSERT_OK_AND_ASSIGN(Computation outer, builder2.Build()); + + ComputationBuilder builder3(client_, "outermost"); + { + auto x = builder3.Parameter(0, r0f32_, "x"); + x = builder3.Call(outer, {x}); + x = builder3.Call(outer, {x}); + x = builder3.Call(outer, {x}); + } + + TF_ASSERT_OK_AND_ASSIGN( + std::unique_ptr start, + client_->TransferToServer(*Literal::CreateR0(1.0f))); + ComputeAndCompareR0(&builder3, 10.0f, {start.get()}, ErrorSpec(0.0f)); +} + +// TODO(b/64094172) Failing on GPU as of 2017-07-26. XLA_TEST_F(CallOpTest, DISABLED_ON_GPU(CallR0F32Tuple)) { ComputationBuilder builder(client_, TestName()); Computation callee = CreateR0F32TupleComputation(); - auto elem = LiteralUtil::CreateR0(42.0); - auto tuple = LiteralUtil::MakeTuple({elem.get()}); + auto elem = Literal::CreateR0(42.0); + auto tuple = Literal::MakeTuple({elem.get()}); builder.Call(callee, {builder.ConstantLiteral(*elem)}); ComputeAndCompareTuple(&builder, *tuple, {}, ErrorSpec(0.01f)); @@ -114,20 +150,3 @@ XLA_TEST_F(CallOpTest, DISABLED_ON_GPU(CallR0F32Tuple)) { } // namespace } // namespace xla - -int main(int argc, char** argv) { - std::vector flag_list; - xla::legacy_flags::AppendCpuCompilerFlags(&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; - } - testing::InitGoogleTest(&argc, argv); - if (argc > 1) { - LOG(ERROR) << "Unknown argument " << argv[1] << "\n" << usage; - return 2; - } - return RUN_ALL_TESTS(); -} diff --git a/tensorflow/compiler/xla/tests/check_execution_arity_test.cc b/tensorflow/compiler/xla/tests/check_execution_arity_test.cc index 675c9fccb007f5a0a16b50618e849d3740877403..659660d91e519b428d28ced8591d05b4e4d45f53 100644 --- a/tensorflow/compiler/xla/tests/check_execution_arity_test.cc +++ b/tensorflow/compiler/xla/tests/check_execution_arity_test.cc @@ -18,24 +18,25 @@ limitations under the License. #include "tensorflow/compiler/xla/client/computation_builder.h" #include "tensorflow/compiler/xla/client/global_data.h" #include "tensorflow/compiler/xla/client/local_client.h" -#include "tensorflow/compiler/xla/legacy_flags/cpu_compiler_flags.h" #include "tensorflow/compiler/xla/literal_util.h" #include "tensorflow/compiler/xla/shape_util.h" #include "tensorflow/compiler/xla/statusor.h" +#include "tensorflow/compiler/xla/test.h" #include "tensorflow/compiler/xla/test_helpers.h" #include "tensorflow/compiler/xla/tests/client_library_test_base.h" #include "tensorflow/compiler/xla/tests/test_macros.h" #include "tensorflow/compiler/xla/xla_data.pb.h" -#include "tensorflow/core/platform/test.h" namespace xla { namespace { +using ::testing::ContainsRegex; + class CheckExecutionArityTest : public ClientLibraryTestBase {}; TEST_F(CheckExecutionArityTest, TwoParamComputationNumArguments) { ComputationBuilder builder(client_, "add_two_params"); - auto param_literal = LiteralUtil::CreateR1({1.1f, 2.2f}); + auto param_literal = Literal::CreateR1({1.1f, 2.2f}); auto p0 = builder.Parameter(0, param_literal->shape(), "param0"); auto p1 = builder.Parameter(1, param_literal->shape(), "param1"); @@ -52,23 +53,25 @@ TEST_F(CheckExecutionArityTest, TwoParamComputationNumArguments) { // The arity of the UserComputation is 2 arguments. Execution will succeed // with 2 arguments, but fail with a different number. - auto result_two_args = - client_->Execute(computation, {param0_data.get(), param1_data.get()}); + auto result_two_args = client_->Execute( + computation, {param0_data.get(), param1_data.get()}, &execution_options_); ASSERT_IS_OK(result_two_args.status()); - auto result_one_arg = client_->Execute(computation, {param0_data.get()}); + auto result_one_arg = + client_->Execute(computation, {param0_data.get()}, &execution_options_); ASSERT_FALSE(result_one_arg.ok()); ASSERT_EQ(result_one_arg.status().code(), tensorflow::error::INVALID_ARGUMENT); - ASSERT_MATCH(result_one_arg.status().error_message(), - testing::ContainsRegex("takes 2")); + ASSERT_THAT(result_one_arg.status().error_message(), + ContainsRegex("takes 2")); - auto result_zero_args = client_->Execute(computation, {}); + auto result_zero_args = + client_->Execute(computation, {}, &execution_options_); ASSERT_FALSE(result_zero_args.ok()); ASSERT_EQ(result_zero_args.status().code(), tensorflow::error::INVALID_ARGUMENT); - ASSERT_MATCH(result_zero_args.status().error_message(), - testing::ContainsRegex("takes 2")); + ASSERT_THAT(result_zero_args.status().error_message(), + ContainsRegex("takes 2")); } XLA_TEST_F(CheckExecutionArityTest, CheckArgumentShapes) { @@ -82,57 +85,43 @@ XLA_TEST_F(CheckExecutionArityTest, CheckArgumentShapes) { ASSERT_IS_OK(computation_status.status()); auto computation = computation_status.ConsumeValueOrDie(); - auto f32_literal = LiteralUtil::CreateR0(1.1f); + auto f32_literal = Literal::CreateR0(1.1f); auto f32_data = client_->TransferToServer(*f32_literal).ConsumeValueOrDie(); - auto f32_4_literal = LiteralUtil::CreateR1({1.0f, 2.0f, 3.0f, 4.0f}); + auto f32_4_literal = Literal::CreateR1({1.0f, 2.0f, 3.0f, 4.0f}); auto f32_4_data = client_->TransferToServer(*f32_4_literal).ConsumeValueOrDie(); - auto u8_4_literal = LiteralUtil::CreateR1U8("hola"); + auto u8_4_literal = Literal::CreateR1U8("hola"); auto u8_4_data = client_->TransferToServer(*u8_4_literal).ConsumeValueOrDie(); // Match - auto status = - client_->Execute(computation, {f32_data.get(), f32_4_data.get()}); + auto status = client_->Execute( + computation, {f32_data.get(), f32_4_data.get()}, &execution_options_); ASSERT_IS_OK(status.status()); // Shape mismatch in parameter 0 - status = client_->Execute(computation, {f32_4_data.get(), f32_4_data.get()}); + status = client_->Execute(computation, {f32_4_data.get(), f32_4_data.get()}, + &execution_options_); ASSERT_FALSE(status.ok()); ASSERT_EQ(status.status().code(), tensorflow::error::INVALID_ARGUMENT); - ASSERT_MATCH(status.status().error_message(), - testing::ContainsRegex("expects parameter 0")); + ASSERT_THAT(status.status().error_message(), + ContainsRegex("expects parameter 0")); // Shape mismatch in parameter 1 (rank) - status = client_->Execute(computation, {f32_data.get(), f32_data.get()}); + status = client_->Execute(computation, {f32_data.get(), f32_data.get()}, + &execution_options_); ASSERT_FALSE(status.ok()); ASSERT_EQ(status.status().code(), tensorflow::error::INVALID_ARGUMENT); - ASSERT_MATCH(status.status().error_message(), - testing::ContainsRegex("expects parameter 1")); + ASSERT_THAT(status.status().error_message(), + ContainsRegex("expects parameter 1")); // Shape mismatch in parameter 1 (element type) - status = client_->Execute(computation, {f32_data.get(), u8_4_data.get()}); + status = client_->Execute(computation, {f32_data.get(), u8_4_data.get()}, + &execution_options_); ASSERT_FALSE(status.ok()); ASSERT_EQ(status.status().code(), tensorflow::error::INVALID_ARGUMENT); - ASSERT_MATCH(status.status().error_message(), - testing::ContainsRegex("expects parameter 1")); + ASSERT_THAT(status.status().error_message(), + ContainsRegex("expects parameter 1")); } } // namespace } // namespace xla - -int main(int argc, char** argv) { - std::vector flag_list; - xla::legacy_flags::AppendCpuCompilerFlags(&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; - } - testing::InitGoogleTest(&argc, argv); - if (argc > 1) { - LOG(ERROR) << "Unknown argument " << argv[1] << "\n" << usage; - return 2; - } - return RUN_ALL_TESTS(); -} diff --git a/tensorflow/compiler/xla/tests/client_library_test_base.cc b/tensorflow/compiler/xla/tests/client_library_test_base.cc index 7bf1168dc3911825c9d8539e33a2bfacd0e81d7f..3001813dd4c0f990854070219b49e9f97d97d925 100644 --- a/tensorflow/compiler/xla/tests/client_library_test_base.cc +++ b/tensorflow/compiler/xla/tests/client_library_test_base.cc @@ -20,7 +20,7 @@ limitations under the License. #include "tensorflow/compiler/xla/client/client_library.h" #include "tensorflow/compiler/xla/client/computation.h" #include "tensorflow/compiler/xla/client/local_client.h" -#include "tensorflow/compiler/xla/legacy_flags/hlo_pass_pipeline_flags.h" +#include "tensorflow/compiler/xla/execution_options_util.h" #include "tensorflow/compiler/xla/literal_util.h" #include "tensorflow/compiler/xla/ptr_util.h" #include "tensorflow/compiler/xla/shape_util.h" @@ -37,21 +37,37 @@ namespace xla { namespace { // Wrapper function that creates a nicer error message (than a bare // ValueOrDie()) if the platform we intend to test is not available. -Client* GetOrCreateLocalClientOrDie(se::Platform* platform) { - StatusOr result = ClientLibrary::GetOrCreateLocalClient(platform); +Client* GetOrCreateLocalClientOrDie(const LocalClientOptions& client_options) { + StatusOr result = + ClientLibrary::GetOrCreateLocalClient(client_options); TF_CHECK_OK(result.status()) << "could not create local client for testing"; return result.ValueOrDie(); } } // namespace ClientLibraryTestBase::ClientLibraryTestBase( - se::Platform* platform, - tensorflow::gtl::ArraySlice disabled_pass_names) - : client_(GetOrCreateLocalClientOrDie(platform)) { - legacy_flags::HloPassPipelineFlags* flags = - legacy_flags::GetHloPassPipelineFlags(); - flags->xla_disable_hlo_passes = - tensorflow::str_util::Join(disabled_pass_names, ","); + perftools::gputools::Platform* platform, + const LocalClientOptions& client_options) + : client_(GetOrCreateLocalClientOrDie(client_options)), + execution_options_(CreateDefaultExecutionOptions()) { + CHECK_EQ(platform, client_options.platform()); + // Disabling constant_folding so that tests (usually written using Constants) + // will exercise the intended code paths, instead of being constant folded. + // + // TODO(b/38354253): Constant folding is currently disabled. Change tests to + // use Parameters instead of Constants, and re-enable constant folding by + // default. + execution_options_.mutable_debug_options()->add_xla_disable_hlo_passes( + "constant_folding"); +} + +ClientLibraryTestBase::ClientLibraryTestBase(se::Platform* platform) + : execution_options_(CreateDefaultExecutionOptions()) { + LocalClientOptions default_options; + default_options.set_platform(platform); + client_ = GetOrCreateLocalClientOrDie(default_options); + execution_options_.mutable_debug_options()->add_xla_disable_hlo_passes( + "constant_folding"); } string ClientLibraryTestBase::TestName() const { @@ -67,12 +83,9 @@ StatusOr> ClientLibraryTestBase::Execute( } StatusOr> ClientLibraryTestBase::ExecuteAndTransfer( - ComputationBuilder* builder, + const Computation& computation, tensorflow::gtl::ArraySlice arguments, const Shape* shape_with_output_layout) { - // Build the computation, as a convenience. - TF_ASSIGN_OR_RETURN(auto computation, builder->Build()); - ExecutionOptions execution_options = execution_options_; if (shape_with_output_layout != nullptr) { *execution_options.mutable_shape_with_output_layout() = @@ -82,6 +95,15 @@ StatusOr> ClientLibraryTestBase::ExecuteAndTransfer( &execution_options); } +StatusOr> ClientLibraryTestBase::ExecuteAndTransfer( + ComputationBuilder* builder, + tensorflow::gtl::ArraySlice arguments, + const Shape* shape_with_output_layout) { + // Build the computation, as a convenience. + TF_ASSIGN_OR_RETURN(auto computation, builder->Build()); + return ExecuteAndTransfer(computation, arguments, shape_with_output_layout); +} + std::unique_ptr ClientLibraryTestBase::ExecuteOrDie( ComputationBuilder* builder, tensorflow::gtl::ArraySlice arguments) { @@ -108,14 +130,14 @@ string ClientLibraryTestBase::ExecuteToString( if (!result.ok()) { return result.status().ToString(); } else { - return LiteralUtil::ToString(*result.ValueOrDie()); + return result.ValueOrDie()->ToString(); } } void ClientLibraryTestBase::ComputeAndCompareR1( ComputationBuilder* builder, const tensorflow::core::Bitmap& expected, tensorflow::gtl::ArraySlice arguments) { - std::unique_ptr expected_literal = LiteralUtil::CreateR1(expected); + std::unique_ptr expected_literal = Literal::CreateR1(expected); ClientLibraryTestBase::ComputeAndCompareLiteral(builder, *expected_literal, arguments); } @@ -136,18 +158,121 @@ void ClientLibraryTestBase::ComputeAndCompareLiteral( error, shape_with_layout)); } +tensorflow::Status +ClientLibraryTestBase::ComputeAndCompareLiteralWithAllOutputLayouts( + const xla::Computation& computation, const Literal& expected, + tensorflow::gtl::ArraySlice arguments, + const std::function& verify_output) { + // Try with no layout requirement. + TF_ASSIGN_OR_RETURN(auto actual, ExecuteAndTransfer(computation, arguments)); + verify_output(*actual, ""); + + // Try with all output layouts. + std::vector minor_to_major(ShapeUtil::Rank(expected.shape())); + std::iota(minor_to_major.begin(), minor_to_major.end(), 0); + do { + auto layout = ShapeUtil::MakeShapeWithLayout( + expected.shape().element_type(), + AsInt64Slice(expected.shape().dimensions()), minor_to_major); + TF_ASSIGN_OR_RETURN(auto actual, + ExecuteAndTransfer(computation, arguments, &layout)); + verify_output(*actual, tensorflow::strings::StrCat( + "Test with output layout: ", + ShapeUtil::HumanStringWithLayout(layout))); + } while (std::next_permutation(minor_to_major.begin(), minor_to_major.end())); + return tensorflow::Status::OK(); +} + +tensorflow::Status +ClientLibraryTestBase::ComputeAndCompareLiteralWithAllInputLayouts( + const xla::Computation& computation, const Literal& expected, + tensorflow::gtl::ArraySlice arguments, + const std::function& verify_output, + const Shape* output_with_layout) { + std::vector arguments_with_layout; + std::vector layout_strings; + // This is a recursive function. It's an std::function instead of a lambda + // because it needs to capture itself. The index is the index of the argument + // to try all layouts for. + std::function choose; + choose = [&, this](int64 index) -> tensorflow::Status { + if (index < arguments.size()) { + // Try out all layouts for the operand. + 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())) { + layout_strings.push_back( + ShapeUtil::HumanStringWithLayout(literal->shape())); + arguments_with_layout.push_back(arguments[index]); + TF_RETURN_IF_ERROR(choose(index + 1)); + arguments_with_layout.pop_back(); + layout_strings.pop_back(); + return tensorflow::Status::OK(); + } + + std::vector minor_to_major(ShapeUtil::Rank(literal->shape())); + std::iota(minor_to_major.begin(), minor_to_major.end(), 0); + do { + auto literal_relayout = + literal->Relayout(LayoutUtil::MakeLayout(minor_to_major)); + layout_strings.push_back( + ShapeUtil::HumanStringWithLayout(literal_relayout->shape())); + TF_ASSIGN_OR_RETURN(auto data, + client_->TransferToServer(*literal_relayout)); + arguments_with_layout.push_back(data.get()); + TF_RETURN_IF_ERROR(choose(index + 1)); + arguments_with_layout.pop_back(); + layout_strings.pop_back(); + } while ( + std::next_permutation(minor_to_major.begin(), minor_to_major.end())); + return tensorflow::Status::OK(); + } + + // Every argument has an assigned layout. + TF_ASSIGN_OR_RETURN( + auto actual, + ExecuteAndTransfer( + computation, + tensorflow::gtl::ArraySlice(arguments_with_layout), + output_with_layout)); + string error_message = "Test with input layouts: "; + for (const auto& str : layout_strings) { + tensorflow::strings::StrAppend(&error_message, str, " "); + } + verify_output(*actual, error_message); + return tensorflow::Status::OK(); + }; + + return choose(0); +} + tensorflow::Status ClientLibraryTestBase::ComputeAndCompareLiteralWithStatus( ComputationBuilder* builder, const Literal& expected, tensorflow::gtl::ArraySlice arguments, const Shape* shape_with_layout) { - TF_ASSIGN_OR_RETURN( - auto actual, ExecuteAndTransfer(builder, arguments, shape_with_layout)); + TF_ASSIGN_OR_RETURN(auto computation, builder->Build()); if (ShapeUtil::ElementIsFloating(expected.shape())) { LOG(WARNING) << "performing exact comparison of floating point numbers"; } else { TF_RET_CHECK(ShapeUtil::ElementIsIntegral(expected.shape()) || expected.shape().element_type() == PRED); } + auto expect_equal = [&](const Literal& actual, const string& error_message) { + LiteralTestUtil::ExpectEqual(expected, actual, error_message); + }; + if (execution_options_.debug_options().xla_test_all_output_layouts()) { + return ComputeAndCompareLiteralWithAllOutputLayouts( + computation, expected, arguments, expect_equal); + } + if (execution_options_.debug_options().xla_test_all_input_layouts()) { + return ComputeAndCompareLiteralWithAllInputLayouts( + computation, expected, arguments, expect_equal, shape_with_layout); + } + TF_ASSIGN_OR_RETURN(auto actual, ExecuteAndTransfer(computation, arguments, + shape_with_layout)); LiteralTestUtil::ExpectEqual(expected, *actual); return tensorflow::Status::OK(); } @@ -156,9 +281,21 @@ tensorflow::Status ClientLibraryTestBase::ComputeAndCompareLiteralWithStatus( ComputationBuilder* builder, const Literal& expected, tensorflow::gtl::ArraySlice arguments, ErrorSpec error, const Shape* shape_with_layout) { - TF_ASSIGN_OR_RETURN( - auto actual, ExecuteAndTransfer(builder, arguments, shape_with_layout)); TF_RET_CHECK(ShapeUtil::ElementIsFloating(expected.shape())); + TF_ASSIGN_OR_RETURN(auto computation, builder->Build()); + auto expect_near = [&](const Literal& actual, const string& error_message) { + LiteralTestUtil::ExpectNear(expected, actual, error, error_message); + }; + if (execution_options_.debug_options().xla_test_all_output_layouts()) { + return ComputeAndCompareLiteralWithAllOutputLayouts(computation, expected, + arguments, expect_near); + } + if (execution_options_.debug_options().xla_test_all_input_layouts()) { + return ComputeAndCompareLiteralWithAllInputLayouts( + computation, expected, arguments, expect_near, shape_with_layout); + } + TF_ASSIGN_OR_RETURN(auto actual, ExecuteAndTransfer(computation, arguments, + shape_with_layout)); LiteralTestUtil::ExpectNear(expected, *actual, error); return tensorflow::Status::OK(); } @@ -174,12 +311,12 @@ void ClientLibraryTestBase::ComputeAndCompareR1U8( auto actual = actual_status.ConsumeValueOrDie(); // Turn the expected value into a literal. - std::unique_ptr expected_literal = LiteralUtil::CreateR1U8(expected); + std::unique_ptr expected_literal = Literal::CreateR1U8(expected); - VLOG(1) << "expected: " << LiteralUtil::ToString(*expected_literal); - VLOG(1) << "actual: " << LiteralUtil::ToString(*actual); + VLOG(1) << "expected: " << expected_literal->ToString(); + VLOG(1) << "actual: " << actual->ToString(); - EXPECT_EQ(expected, actual->u8s()); + EXPECT_EQ(expected, actual->u8s_string()); } void ClientLibraryTestBase::ComputeAndCompareTuple( diff --git a/tensorflow/compiler/xla/tests/client_library_test_base.h b/tensorflow/compiler/xla/tests/client_library_test_base.h index 34f82603e897748ee4ff8521cab8451b65e4488e..7fe1445b94097f762b777fc6936a0a1ab5a726c8 100644 --- a/tensorflow/compiler/xla/tests/client_library_test_base.h +++ b/tensorflow/compiler/xla/tests/client_library_test_base.h @@ -24,6 +24,7 @@ limitations under the License. #include "tensorflow/compiler/xla/array2d.h" #include "tensorflow/compiler/xla/array3d.h" #include "tensorflow/compiler/xla/array4d.h" +#include "tensorflow/compiler/xla/client/client_library.h" #include "tensorflow/compiler/xla/client/computation.h" #include "tensorflow/compiler/xla/client/computation_builder.h" #include "tensorflow/compiler/xla/client/global_data.h" @@ -46,21 +47,33 @@ namespace xla { class ClientLibraryTestBase : public ::testing::Test { protected: explicit ClientLibraryTestBase( - perftools::gputools::Platform* platform = nullptr, - tensorflow::gtl::ArraySlice disabled_pass_names = {}); + perftools::gputools::Platform* platform = nullptr); + + // Creates a new ClientLibraryTestBase with custom client options. + ClientLibraryTestBase(perftools::gputools::Platform* platform, + const LocalClientOptions& client_options); // Returns the name of the test currently being run. string TestName() const; void SetFastMathDisabled(bool disabled) { - execution_options_.set_disable_fast_math(disabled); + execution_options_.mutable_debug_options()->set_xla_enable_fast_math( + !disabled); } void SetSeed(uint64 seed) { execution_options_.set_seed(seed); } + // Provides mutable access to the execution DebugOptions field; this lets + // tests tweak the options that will be used to compile/run the graph. + DebugOptions* mutable_debug_options() { + return execution_options_.mutable_debug_options(); + } + // TODO(b/25566808): Add helper that populates a literal from a testdata file. - // Convenience methods for building and running a computation from a builder. + // Convenience methods for building and running a computation with the member + // execution options. Modify execution_options_ in your test if you want to + // customize the options. StatusOr> Execute( ComputationBuilder* builder, tensorflow::gtl::ArraySlice arguments); @@ -68,6 +81,10 @@ class ClientLibraryTestBase : public ::testing::Test { ComputationBuilder* builder, tensorflow::gtl::ArraySlice arguments, const Shape* shape_with_output_layout = nullptr); + StatusOr> ExecuteAndTransfer( + const Computation& computation, + tensorflow::gtl::ArraySlice arguments, + const Shape* shape_with_output_layout = nullptr); // Convenience OrDie variants of above methods. std::unique_ptr ExecuteOrDie( @@ -265,6 +282,22 @@ class ClientLibraryTestBase : public ::testing::Test { Client* client_; ExecutionOptions execution_options_; + + private: + // Build and run the computation with all permutations of output layouts. + tensorflow::Status ComputeAndCompareLiteralWithAllOutputLayouts( + const xla::Computation& computation, const Literal& expected, + tensorflow::gtl::ArraySlice arguments, + const std::function& verify_output); + // Build and run the computation with all permutations of layouts of all input + // arguments. + tensorflow::Status ComputeAndCompareLiteralWithAllInputLayouts( + const xla::Computation& computation, const Literal& expected, + tensorflow::gtl::ArraySlice arguments, + const std::function& verify_output, + const Shape* output_with_layout = nullptr); }; template @@ -272,7 +305,7 @@ void ClientLibraryTestBase::ComputeAndCompareR0( ComputationBuilder* builder, NativeT expected, tensorflow::gtl::ArraySlice arguments) { std::unique_ptr expected_literal = - LiteralUtil::CreateR0(expected); + Literal::CreateR0(expected); ClientLibraryTestBase::ComputeAndCompareLiteral(builder, *expected_literal, arguments); } @@ -285,7 +318,7 @@ void ClientLibraryTestBase::ComputeAndCompareR0( std::is_same::value, "Floating point type required when specifying an ErrorSpec"); std::unique_ptr expected_literal = - LiteralUtil::CreateR0(expected); + Literal::CreateR0(expected); ClientLibraryTestBase::ComputeAndCompareLiteral(builder, *expected_literal, arguments, error); } @@ -295,7 +328,7 @@ void ClientLibraryTestBase::ComputeAndCompareR1( ComputationBuilder* builder, tensorflow::gtl::ArraySlice expected, tensorflow::gtl::ArraySlice arguments) { std::unique_ptr expected_literal = - LiteralUtil::CreateR1(expected); + Literal::CreateR1(expected); ClientLibraryTestBase::ComputeAndCompareLiteral(builder, *expected_literal, arguments); } @@ -308,7 +341,7 @@ void ClientLibraryTestBase::ComputeAndCompareR1( std::is_same::value, "Floating point type required when specifying an ErrorSpec"); std::unique_ptr expected_literal = - LiteralUtil::CreateR1(expected); + Literal::CreateR1(expected); ClientLibraryTestBase::ComputeAndCompareLiteral(builder, *expected_literal, arguments, error); } @@ -318,7 +351,7 @@ void ClientLibraryTestBase::ComputeAndCompareR2( ComputationBuilder* builder, const Array2D& expected, tensorflow::gtl::ArraySlice arguments) { std::unique_ptr expected_literal = - LiteralUtil::CreateR2FromArray2D(expected); + Literal::CreateR2FromArray2D(expected); ClientLibraryTestBase::ComputeAndCompareLiteral(builder, *expected_literal, arguments); } @@ -331,7 +364,7 @@ void ClientLibraryTestBase::ComputeAndCompareR2( std::is_same::value, "Floating point type required when specifying an ErrorSpec"); std::unique_ptr expected_literal = - LiteralUtil::CreateR2FromArray2D(expected); + Literal::CreateR2FromArray2D(expected); ClientLibraryTestBase::ComputeAndCompareLiteral(builder, *expected_literal, arguments, error); } @@ -341,7 +374,7 @@ void ClientLibraryTestBase::ComputeAndCompareR3( ComputationBuilder* builder, const Array3D& expected, tensorflow::gtl::ArraySlice arguments) { std::unique_ptr expected_literal = - LiteralUtil::CreateR3FromArray3D(expected); + Literal::CreateR3FromArray3D(expected); ClientLibraryTestBase::ComputeAndCompareLiteral(builder, *expected_literal, arguments); } @@ -354,7 +387,7 @@ void ClientLibraryTestBase::ComputeAndCompareR3( std::is_same::value, "Floating point type required when specifying an ErrorSpec"); std::unique_ptr expected_literal = - LiteralUtil::CreateR3FromArray3D(expected); + Literal::CreateR3FromArray3D(expected); ClientLibraryTestBase::ComputeAndCompareLiteral(builder, *expected_literal, arguments, error); } @@ -364,7 +397,7 @@ void ClientLibraryTestBase::ComputeAndCompareR4( ComputationBuilder* builder, const Array4D& expected, tensorflow::gtl::ArraySlice arguments) { std::unique_ptr expected_literal = - LiteralUtil::CreateR4FromArray4D(expected); + Literal::CreateR4FromArray4D(expected); ClientLibraryTestBase::ComputeAndCompareLiteral(builder, *expected_literal, arguments); } @@ -377,7 +410,7 @@ void ClientLibraryTestBase::ComputeAndCompareR4( std::is_same::value, "Floating point type required when specifying an ErrorSpec"); std::unique_ptr expected_literal = - LiteralUtil::CreateR4FromArray4D(expected); + Literal::CreateR4FromArray4D(expected); ClientLibraryTestBase::ComputeAndCompareLiteral(builder, *expected_literal, arguments, error); } @@ -386,7 +419,7 @@ template std::unique_ptr ClientLibraryTestBase::CreateR0Parameter( NativeT value, int64 parameter_number, const string& name, ComputationBuilder* builder, ComputationDataHandle* data_handle) { - std::unique_ptr literal = LiteralUtil::CreateR0(value); + std::unique_ptr literal = Literal::CreateR0(value); std::unique_ptr data = client_->TransferToServer(*literal).ConsumeValueOrDie(); *data_handle = builder->Parameter(parameter_number, literal->shape(), name); @@ -398,7 +431,7 @@ std::unique_ptr ClientLibraryTestBase::CreateR1Parameter( tensorflow::gtl::ArraySlice values, int64 parameter_number, const string& name, ComputationBuilder* builder, ComputationDataHandle* data_handle) { - std::unique_ptr literal = LiteralUtil::CreateR1(values); + std::unique_ptr literal = Literal::CreateR1(values); std::unique_ptr data = client_->TransferToServer(*literal).ConsumeValueOrDie(); *data_handle = builder->Parameter(parameter_number, literal->shape(), name); @@ -410,7 +443,7 @@ std::unique_ptr ClientLibraryTestBase::CreateR2Parameter( const Array2D& array_2d, int64 parameter_number, const string& name, ComputationBuilder* builder, ComputationDataHandle* data_handle) { - std::unique_ptr literal = LiteralUtil::CreateR2FromArray2D(array_2d); + std::unique_ptr literal = Literal::CreateR2FromArray2D(array_2d); std::unique_ptr data = client_->TransferToServer(*literal).ConsumeValueOrDie(); *data_handle = builder->Parameter(parameter_number, literal->shape(), name); @@ -422,7 +455,7 @@ std::unique_ptr ClientLibraryTestBase::CreateR3Parameter( const Array3D& array_3d, int64 parameter_number, const string& name, ComputationBuilder* builder, ComputationDataHandle* data_handle) { - std::unique_ptr literal = LiteralUtil::CreateR3FromArray3D(array_3d); + std::unique_ptr literal = Literal::CreateR3FromArray3D(array_3d); std::unique_ptr data = client_->TransferToServer(*literal).ConsumeValueOrDie(); *data_handle = builder->Parameter(parameter_number, literal->shape(), name); diff --git a/tensorflow/compiler/xla/tests/client_test.cc b/tensorflow/compiler/xla/tests/client_test.cc index 86ce636ee566ab0ed5535e0f57b6d34776549528..0853feeebd6f7a249cf767e1f8a63675d4bddd27 100644 --- a/tensorflow/compiler/xla/tests/client_test.cc +++ b/tensorflow/compiler/xla/tests/client_test.cc @@ -19,7 +19,6 @@ limitations under the License. #include "tensorflow/compiler/xla/client/computation_builder.h" #include "tensorflow/compiler/xla/client/global_data.h" #include "tensorflow/compiler/xla/client/local_client.h" -#include "tensorflow/compiler/xla/legacy_flags/cpu_compiler_flags.h" #include "tensorflow/compiler/xla/shape_util.h" #include "tensorflow/compiler/xla/statusor.h" #include "tensorflow/compiler/xla/test_helpers.h" @@ -46,7 +45,7 @@ TEST_F(ClientTest, ExecuteWithLayout) { auto computation = b.Build(); ASSERT_TRUE(computation.ok()) << computation.status(); - ExecutionOptions execution_options; + ExecutionOptions execution_options = execution_options_; *execution_options.mutable_shape_with_output_layout() = ShapeUtil::MakeShapeWithLayout(S32, /*dimensions=*/{2, 2}, execute_layout); @@ -76,7 +75,7 @@ TEST_F(ClientTest, ExecuteWithTupleLayout) { auto computation = b.Build(); ASSERT_TRUE(computation.ok()) << computation.status(); - ExecutionOptions execution_options; + ExecutionOptions execution_options = execution_options_; // Create a result shape with one element column major and the other row // major. *execution_options.mutable_shape_with_output_layout() = @@ -110,20 +109,3 @@ TEST_F(ClientTest, ExecuteWithTupleLayout) { } // namespace } // namespace xla - -int main(int argc, char** argv) { - std::vector flag_list; - xla::legacy_flags::AppendCpuCompilerFlags(&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; - } - testing::InitGoogleTest(&argc, argv); - if (argc > 1) { - LOG(ERROR) << "Unknown argument " << argv[1] << "\n" << usage; - return 2; - } - return RUN_ALL_TESTS(); -} diff --git a/tensorflow/compiler/xla/tests/codegen_test_base.cc b/tensorflow/compiler/xla/tests/codegen_test_base.cc index e6f3225bb79fca99f189d1e7ae7913715c5c2246..a52be3ffd1e4767ee711ff4f7bb826da3015e38c 100644 --- a/tensorflow/compiler/xla/tests/codegen_test_base.cc +++ b/tensorflow/compiler/xla/tests/codegen_test_base.cc @@ -15,80 +15,25 @@ limitations under the License. #include "tensorflow/compiler/xla/tests/codegen_test_base.h" -#include -#include - -#include "tensorflow/compiler/xla/ptr_util.h" -#include "tensorflow/compiler/xla/service/backend.h" -#include "tensorflow/compiler/xla/service/compiler.h" -#include "tensorflow/compiler/xla/service/hlo_module_config.h" -#include "tensorflow/compiler/xla/statusor.h" -#include "tensorflow/core/lib/core/status.h" -#include "tensorflow/core/lib/io/path.h" -#include "tensorflow/core/lib/strings/strcat.h" -#include "tensorflow/core/platform/env.h" -#include "tensorflow/core/platform/logging.h" -#include "tensorflow/core/platform/subprocess.h" -#include "tensorflow/core/platform/test.h" - namespace xla { -void CodegenTestBase::CompileAndVerifyIr(std::unique_ptr hlo_module, - const string& pattern) { - std::unique_ptr executable = - CompileToExecutable(std::move(hlo_module)); - string ir_module_string = GetIrFromExecutable(*executable); - RunFileCheck(ir_module_string, pattern); -} - -std::unique_ptr CodegenTestBase::CompileToExecutable( +StatusOr> CodegenTestBase::CompileToExecutable( std::unique_ptr hlo_module) { - auto module_config = MakeUnique( - hlo_module->entry_computation()->ComputeProgramShape()); - module_config->set_fast_math_disabled(fast_math_disabled_); - return backend_->compiler() - ->Compile(std::move(hlo_module), std::move(module_config), - test_hlo_dumper_, backend_->default_stream_executor()) - .ConsumeValueOrDie(); + return backend_->compiler()->Compile(std::move(hlo_module), + backend_->default_stream_executor()); } -void CodegenTestBase::RunFileCheck(const string& input, const string& pattern) { - using tensorflow::io::JoinPath; - - // Write input to a temporary file. - char tempdir_template[] = "/tmp/ir_testXXXXXX"; - char* tempdir_name = mkdtemp(tempdir_template); - CHECK_NOTNULL(tempdir_name); - string pattern_path = JoinPath(tempdir_name, "xla_hlo_test_ir_pattern"); - TF_CHECK_OK(tensorflow::WriteStringToFile(tensorflow::Env::Default(), - pattern_path, pattern)); - - // Invoke FileCheck to check whether input matches `pattern`. - const char* file_check_path_suffix = "external/llvm/FileCheck"; - string file_check_path; - if (const char* test_srcdir = getenv("TEST_SRCDIR")) { - file_check_path = JoinPath(test_srcdir, file_check_path_suffix); - } else { - file_check_path = file_check_path_suffix; - } - - tensorflow::SubProcess file_check_process; - file_check_process.SetProgram(file_check_path, - {file_check_path, pattern_path}); - file_check_process.SetChannelAction(tensorflow::CHAN_STDIN, - tensorflow::ACTION_PIPE); - file_check_process.SetChannelAction(tensorflow::CHAN_STDERR, - tensorflow::ACTION_PIPE); - CHECK(file_check_process.Start()); - string standard_error; - int exit_status = file_check_process.Communicate( - /*stdin_input=*/&input, /*stdout_output=*/nullptr, - /*stderr_output=*/&standard_error); - - // FileCheck returns 0 when the inputs match. If matching failed, we output - // the error message generated by FileCheck. - SCOPED_TRACE(tensorflow::strings::StrCat("Input to FileCheck:\n", input)); - EXPECT_EQ(0, exit_status) << standard_error; +StatusOr> +CodegenTestBase::CompileToAotCompilationResult( + std::unique_ptr hlo_module, + const AotCompilationOptions& options) { + std::vector> hlo_modules; + hlo_modules.push_back(std::move(hlo_module)); + TF_ASSIGN_OR_RETURN( + std::vector> results, + backend_->compiler()->CompileAheadOfTime(std::move(hlo_modules), + options)); + return std::move(results.front()); } } // namespace xla diff --git a/tensorflow/compiler/xla/tests/codegen_test_base.h b/tensorflow/compiler/xla/tests/codegen_test_base.h index ba32aac8e4b0d189d0699ea4d659524ef9f39b99..441fcd6890ebf703962dbd3eaaac765f81b08303 100644 --- a/tensorflow/compiler/xla/tests/codegen_test_base.h +++ b/tensorflow/compiler/xla/tests/codegen_test_base.h @@ -17,43 +17,25 @@ limitations under the License. #define TENSORFLOW_COMPILER_XLA_TESTS_CODEGEN_TEST_BASE_H_ #include -#include +#include "tensorflow/compiler/xla/service/compiler.h" #include "tensorflow/compiler/xla/service/executable.h" #include "tensorflow/compiler/xla/service/hlo_module.h" #include "tensorflow/compiler/xla/tests/hlo_test_base.h" namespace xla { -// Tests that verify IR emitted by the CPU/GPU backend is as expected. +// Provides access to both the JIT and the AOT compiler for testing. class CodegenTestBase : public HloTestBase { protected: - CodegenTestBase() {} - - // Returns the embedded LLVM IR from the given executable. Codegen tests must - // override this method, but execution tests do not have to because they do - // not examine the embedded IR. - virtual string GetIrFromExecutable(const Executable& executable) = 0; - - // Compiles the given HLO module to LLVM IR and verifies the IR matches the - // given pattern. `pattern` is in the FileCheck pattern matching syntax - // (http://llvm.org/docs/CommandGuide/FileCheck.html). - void CompileAndVerifyIr(std::unique_ptr hlo_module, - const string& pattern); - - // Sets the fast-math-disabled flag on the config we use when compiling. - void set_fast_math_disabled(bool disabled) { fast_math_disabled_ = disabled; } - - protected: - // Compiles hlo_module to an executable, CHECK-failing if this fails. - std::unique_ptr CompileToExecutable( + // Compiles hlo_module with the JIT compiler. + StatusOr> CompileToExecutable( std::unique_ptr hlo_module); - // Runs FileCheck with the given pattern over the given string and EXPECTs - // that FileCheck succeeded in matching the input. - void RunFileCheck(const string& input, const string& pattern); - - bool fast_math_disabled_ = false; + // Compiles hlo_module with the AOT compiler. + StatusOr> CompileToAotCompilationResult( + std::unique_ptr hlo_module, + const AotCompilationOptions& options); }; } // namespace xla diff --git a/tensorflow/compiler/xla/tests/compilation_cache_test.cc b/tensorflow/compiler/xla/tests/compilation_cache_test.cc index 1d0df615824b034b06e5265befa69104fd6312c5..707e439245c29a1ddf80bfd9205aa14b0d4765f6 100644 --- a/tensorflow/compiler/xla/tests/compilation_cache_test.cc +++ b/tensorflow/compiler/xla/tests/compilation_cache_test.cc @@ -21,7 +21,6 @@ limitations under the License. #include "tensorflow/compiler/xla/client/computation_builder.h" #include "tensorflow/compiler/xla/client/global_data.h" #include "tensorflow/compiler/xla/client/local_client.h" -#include "tensorflow/compiler/xla/legacy_flags/cpu_compiler_flags.h" #include "tensorflow/compiler/xla/literal_util.h" #include "tensorflow/compiler/xla/shape_util.h" #include "tensorflow/compiler/xla/statusor.h" @@ -47,10 +46,10 @@ class CompilationCacheTest : public ClientLibraryTestBase { std::unique_ptr result = client_ ->ExecuteAndTransfer(computation, arguments, - /*execution_options=*/nullptr, + /*execution_options=*/&execution_options_, &execution_profile) .ConsumeValueOrDie(); - LiteralTestUtil::ExpectNear(*LiteralUtil::CreateR0(expected_result), + LiteralTestUtil::ExpectNear(*Literal::CreateR0(expected_result), *result, error_spec_); EXPECT_EQ(expect_cache_hit, execution_profile.compilation_cache_hit()); } @@ -61,14 +60,13 @@ class CompilationCacheTest : public ClientLibraryTestBase { std::initializer_list> expected_result, bool expect_cache_hit) { ExecutionProfile execution_profile; - auto data_handle = - client_ - ->Execute(computation, arguments, /*execution_options=*/nullptr, - &execution_profile) - .ConsumeValueOrDie(); + auto data_handle = client_ + ->Execute(computation, arguments, + &execution_options_, &execution_profile) + .ConsumeValueOrDie(); std::unique_ptr result = client_->Transfer(*data_handle).ConsumeValueOrDie(); - LiteralTestUtil::ExpectNear(*LiteralUtil::CreateR2(expected_result), + LiteralTestUtil::ExpectNear(*Literal::CreateR2(expected_result), *result, error_spec_); EXPECT_EQ(expect_cache_hit, execution_profile.compilation_cache_hit()); } @@ -88,13 +86,13 @@ XLA_TEST_F(CompilationCacheTest, ComputationCalledMultipleTimes) { XLA_TEST_F(CompilationCacheTest, ComputationCalledWithDifferentParameters) { std::unique_ptr data_42 = - client_->TransferToServer(*LiteralUtil::CreateR0(42.0f)) + client_->TransferToServer(*Literal::CreateR0(42.0f)) .ConsumeValueOrDie(); std::unique_ptr data_123 = - client_->TransferToServer(*LiteralUtil::CreateR0(123.0f)) + client_->TransferToServer(*Literal::CreateR0(123.0f)) .ConsumeValueOrDie(); std::unique_ptr data_456 = - client_->TransferToServer(*LiteralUtil::CreateR0(456.0f)) + client_->TransferToServer(*Literal::CreateR0(456.0f)) .ConsumeValueOrDie(); ComputationBuilder builder(client_, TestName()); @@ -200,20 +198,3 @@ XLA_TEST_F(CompilationCacheTest, MutatedComputation) { } // namespace } // namespace xla - -int main(int argc, char** argv) { - std::vector flag_list; - xla::legacy_flags::AppendCpuCompilerFlags(&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; - } - testing::InitGoogleTest(&argc, argv); - if (argc > 1) { - LOG(ERROR) << "Unknown argument " << argv[1] << "\n" << usage; - return 2; - } - return RUN_ALL_TESTS(); -} diff --git a/tensorflow/compiler/xla/tests/compute_constant_test.cc b/tensorflow/compiler/xla/tests/compute_constant_test.cc index 709ce5029c82d52fe7a577d1e4cf7ea6ec07cecb..b2e9743af79d0e4658451e7a9522c338036851ba 100644 --- a/tensorflow/compiler/xla/tests/compute_constant_test.cc +++ b/tensorflow/compiler/xla/tests/compute_constant_test.cc @@ -17,44 +17,73 @@ limitations under the License. #include #include +#include "tensorflow/compiler/xla/client/client_library.h" #include "tensorflow/compiler/xla/client/computation.h" #include "tensorflow/compiler/xla/client/computation_builder.h" #include "tensorflow/compiler/xla/client/global_data.h" -#include "tensorflow/compiler/xla/client/local_client.h" #include "tensorflow/compiler/xla/layout_util.h" -#include "tensorflow/compiler/xla/legacy_flags/cpu_compiler_flags.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/compiler/xla/test_helpers.h" -#include "tensorflow/compiler/xla/tests/client_library_test_base.h" +#include "tensorflow/compiler/xla/test.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/compiler/xla/xla_data.pb.h" -#include "tensorflow/core/platform/test.h" #include "tensorflow/core/platform/types.h" namespace xla { namespace { -class ComputeConstantTest : public ClientLibraryTestBase { +// An enumerator for the client types that we want to iterate over in +// the various tests. +enum class ClientType { kLocal, kCompileOnly }; +ClientType client_types[] = {ClientType::kLocal, ClientType::kCompileOnly}; + +class ComputeConstantTest : public ::testing::Test { public: + explicit ComputeConstantTest( + perftools::gputools::Platform* platform = nullptr) + : platform_(platform) {} + + string TestName() const { + return ::testing::UnitTest::GetInstance()->current_test_info()->name(); + } + + Client* ClientOrDie(::perftools::gputools::Platform* platform, + ClientType client_type) { + if (client_type == ClientType::kLocal) { + StatusOr result = + ClientLibrary::GetOrCreateLocalClient(platform); + TF_CHECK_OK(result.status()) + << "could not create LocalClient for testing"; + return result.ValueOrDie(); + } else if (client_type == ClientType::kCompileOnly) { + StatusOr result = + ClientLibrary::GetOrCreateCompileOnlyClient(platform); + TF_CHECK_OK(result.status()) + << "could not create CompileOnlyClient for testing"; + return result.ValueOrDie(); + } + LOG(FATAL) << "invalid client_type value"; + } + StatusOr> ComputeConstantLiteral( - ComputationDataHandle operand, ComputationBuilder* builder, - Layout* output_layout = nullptr) { - TF_ASSIGN_OR_RETURN(auto remote_computed, + Client* client, const ComputationDataHandle& operand, + ComputationBuilder* builder, Layout* output_layout = nullptr) { + TF_ASSIGN_OR_RETURN(auto computed, builder->ComputeConstant(operand, output_layout)); - TF_ASSIGN_OR_RETURN(auto computed, client_->Transfer(*remote_computed)); return std::move(computed); } template - StatusOr ComputeConstantScalar(ComputationDataHandle operand, + StatusOr ComputeConstantScalar(Client* client, + const ComputationDataHandle& operand, ComputationBuilder* builder) { - TF_ASSIGN_OR_RETURN(auto literal, ComputeConstantLiteral(operand, builder)); - return LiteralUtil::Get(*literal, {}); + TF_ASSIGN_OR_RETURN(auto literal, + ComputeConstantLiteral(client, operand, builder)); + return literal->Get({}); } bool IsConstant(const ComputationDataHandle& operand, @@ -64,186 +93,164 @@ class ComputeConstantTest : public ClientLibraryTestBase { return result.ok() ? result.ValueOrDie() : false; } - template - void ExpectConstantComputedScalar(ComputationDataHandle operand, - Scalar expected, - ComputationBuilder* builder) { - Scalar computed = ComputeConstantScalar(operand, builder); - ASSERT_TRUE(computed.ok()) << computed.status(); - std::unique_ptr expected_literal = LiteralUtil::CreateR0(expected); - LiteralTestUtil::ExpectEqual(*expected_literal, *computed); - } + perftools::gputools::Platform* platform_; }; TEST_F(ComputeConstantTest, ScalarInt32Literal) { - ComputationBuilder b(client_, TestName()); - auto computation = b.ConstantR0(42); - EXPECT_TRUE(IsConstant(computation, &b)); - - auto value = ComputeConstantScalar(computation, &b); - ASSERT_TRUE(value.ok()) << value.status(); - EXPECT_EQ(value.ValueOrDie(), 42); + for (ClientType client_type : client_types) { + Client* client = ClientOrDie(platform_, client_type); + ComputationBuilder b(client, TestName()); + auto computation = b.ConstantR0(42); + EXPECT_TRUE(IsConstant(computation, &b)); + + auto value = ComputeConstantScalar(client, computation, &b); + ASSERT_TRUE(value.ok()) << value.status(); + EXPECT_EQ(value.ValueOrDie(), 42); + } } TEST_F(ComputeConstantTest, ScalarFloatAdd) { - ComputationBuilder b(client_, TestName()); - auto computation = - b.Add(b.ConstantR0(42.5f), b.ConstantR0(1.5f)); - EXPECT_TRUE(IsConstant(computation, &b)); - - auto value = ComputeConstantScalar(computation, &b); - ASSERT_TRUE(value.ok()) << value.status(); - EXPECT_EQ(value.ValueOrDie(), 44.0f); + for (ClientType client_type : client_types) { + Client* client = ClientOrDie(platform_, client_type); + ComputationBuilder b(client, TestName()); + auto computation = + b.Add(b.ConstantR0(42.5f), b.ConstantR0(1.5f)); + EXPECT_TRUE(IsConstant(computation, &b)); + + auto value = ComputeConstantScalar(client, computation, &b); + ASSERT_TRUE(value.ok()) << value.status(); + EXPECT_EQ(value.ValueOrDie(), 44.0f); + } } TEST_F(ComputeConstantTest, ScalarRng) { - ComputationBuilder b(client_, TestName()); - auto computation = - b.RngUniform(b.ConstantR0(1.1f), b.ConstantR0(2.1f), - ShapeUtil::MakeShape(F32, {})); - EXPECT_FALSE(IsConstant(computation, &b)); - - auto value = ComputeConstantScalar(computation, &b); - ASSERT_FALSE(value.ok()) - << "computing a RNG value should not be considered a constant"; + for (ClientType client_type : client_types) { + Client* client = ClientOrDie(platform_, client_type); + ComputationBuilder b(client, TestName()); + auto computation = + b.RngUniform(b.ConstantR0(1.1f), b.ConstantR0(2.1f), + ShapeUtil::MakeShape(F32, {})); + EXPECT_FALSE(IsConstant(computation, &b)); + + auto value = ComputeConstantScalar(client, computation, &b); + ASSERT_FALSE(value.ok()) + << "computing a RNG value should not be considered a constant"; + } } TEST_F(ComputeConstantTest, DirectParam) { - ComputationBuilder b(client_, TestName()); - auto computation = b.Parameter(0, ShapeUtil::MakeShape(F32, {}), "param"); - EXPECT_FALSE(IsConstant(computation, &b)); - - auto value = ComputeConstantScalar(computation, &b); - EXPECT_TRUE(tensorflow::StringPiece(value.status().ToString()) - .contains("depends on parameter")) - << value.status(); + for (ClientType client_type : client_types) { + Client* client = ClientOrDie(platform_, client_type); + ComputationBuilder b(client, TestName()); + auto computation = b.Parameter(0, ShapeUtil::MakeShape(F32, {}), "param"); + EXPECT_FALSE(IsConstant(computation, &b)); + + auto value = ComputeConstantScalar(client, computation, &b); + EXPECT_TRUE(tensorflow::StringPiece(value.status().ToString()) + .contains("depends on parameter")) + << value.status(); + } } TEST_F(ComputeConstantTest, IndirectParam) { - ComputationBuilder b(client_, TestName()); - auto computation = - b.Add(b.ConstantR0(1.0f), - b.Parameter(0, ShapeUtil::MakeShape(F32, {}), "param")); - EXPECT_FALSE(IsConstant(computation, &b)); - - auto value = ComputeConstantScalar(computation, &b); - EXPECT_TRUE(tensorflow::StringPiece(value.status().ToString()) - .contains("depends on parameter")) - << value.status(); + for (ClientType client_type : client_types) { + Client* client = ClientOrDie(platform_, client_type); + ComputationBuilder b(client, TestName()); + auto computation = + b.Add(b.ConstantR0(1.0f), + b.Parameter(0, ShapeUtil::MakeShape(F32, {}), "param")); + EXPECT_FALSE(IsConstant(computation, &b)); + + auto value = ComputeConstantScalar(client, computation, &b); + EXPECT_TRUE(tensorflow::StringPiece(value.status().ToString()) + .contains("depends on parameter")) + << value.status(); + } } // Test computation of an expression interspersed with param nodes but // the expression does not depend on the param nodes. TEST_F(ComputeConstantTest, UnrelatedParam) { - ComputationBuilder b(client_, TestName()); + for (ClientType client_type : client_types) { + Client* client = ClientOrDie(platform_, client_type); + ComputationBuilder b(client, TestName()); - auto param_a = b.Parameter(10, ShapeUtil::MakeShape(F32, {}), "param0"); - auto constant_4 = b.Add(b.ConstantR0(2.5f), b.ConstantR0(1.5f)); - auto not_constant_a = b.Add(constant_4, param_a); + auto param_a = b.Parameter(10, ShapeUtil::MakeShape(F32, {}), "param0"); + auto constant_4 = + b.Add(b.ConstantR0(2.5f), b.ConstantR0(1.5f)); + auto not_constant_a = b.Add(constant_4, param_a); - auto param_b = b.Parameter(1, ShapeUtil::MakeShape(F32, {}), "param1"); - auto constant_9 = b.Mul(b.ConstantR0(2.0f), b.ConstantR0(4.5f)); - auto not_constant_b = b.Add(param_b, constant_9); + auto param_b = b.Parameter(1, ShapeUtil::MakeShape(F32, {}), "param1"); + auto constant_9 = + b.Mul(b.ConstantR0(2.0f), b.ConstantR0(4.5f)); + auto not_constant_b = b.Add(param_b, constant_9); - auto constant_13 = b.Add(constant_4, constant_9); - b.Add(not_constant_b, b.Add(constant_13, not_constant_a)); + auto constant_13 = b.Add(constant_4, constant_9); + b.Add(not_constant_b, b.Add(constant_13, not_constant_a)); - EXPECT_TRUE(IsConstant(constant_13, &b)); + EXPECT_TRUE(IsConstant(constant_13, &b)); - auto value = ComputeConstantScalar(constant_13, &b); - ASSERT_TRUE(value.ok()) << value.status(); - EXPECT_EQ(value.ValueOrDie(), 13.0f); + auto value = ComputeConstantScalar(client, constant_13, &b); + ASSERT_TRUE(value.ok()) << value.status(); + EXPECT_EQ(value.ValueOrDie(), 13.0f); + } } TEST_F(ComputeConstantTest, NonScalarAdd) { - ComputationBuilder b(client_, TestName()); + for (ClientType client_type : client_types) { + Client* client = ClientOrDie(platform_, client_type); + ComputationBuilder b(client, TestName()); - auto computation = - b.Add(b.ConstantR1({1, 2}), b.ConstantR1({3, 4})); - EXPECT_TRUE(IsConstant(computation, &b)); + auto computation = + b.Add(b.ConstantR1({1, 2}), b.ConstantR1({3, 4})); + EXPECT_TRUE(IsConstant(computation, &b)); - auto computed = ComputeConstantLiteral(computation, &b); - ASSERT_TRUE(computed.ok()) << computed.status(); - std::unique_ptr expected_literal = - LiteralUtil::CreateR1({4, 6}); - LiteralTestUtil::ExpectEqual(*expected_literal, *computed.ValueOrDie()); + auto computed = ComputeConstantLiteral(client, computation, &b); + ASSERT_TRUE(computed.ok()) << computed.status(); + std::unique_ptr expected_literal = + Literal::CreateR1({4, 6}); + LiteralTestUtil::ExpectEqual(*expected_literal, *computed.ValueOrDie()); + } } TEST_F(ComputeConstantTest, IntegerDivide) { - ComputationBuilder b(client_, TestName()); - auto computation = b.Div(b.ConstantR0(15), b.ConstantR0(3)); - EXPECT_TRUE(IsConstant(computation, &b)); - - auto computed = ComputeConstantLiteral(computation, &b); - ASSERT_TRUE(computed.ok()) << computed.status(); - std::unique_ptr expected_literal = LiteralUtil::CreateR0(5); - LiteralTestUtil::ExpectEqual(*expected_literal, *computed.ValueOrDie()); -} + for (ClientType client_type : client_types) { + Client* client = ClientOrDie(platform_, client_type); + ComputationBuilder b(client, TestName()); + auto computation = b.Div(b.ConstantR0(15), b.ConstantR0(3)); + EXPECT_TRUE(IsConstant(computation, &b)); -XLA_TEST_F(ComputeConstantTest, Layout) { - ComputationBuilder b(client_, TestName()); - - std::vector> layouts = {{0, 1}, {1, 0}}; - for (const std::vector& layout : layouts) { - auto layout_proto = LayoutUtil::MakeLayout(layout); - auto computed = - ComputeConstantLiteral(b.Add(b.ConstantR2({{1, 2}, {3, 4}}), - b.ConstantR2({{10, 20}, {30, 40}})), - &b, &layout_proto); + auto computed = ComputeConstantLiteral(client, computation, &b); ASSERT_TRUE(computed.ok()) << computed.status(); - - std::unique_ptr expected_literal = - test_utils::CreateR2LiteralWithLayout({{11, 22}, {33, 44}}, - layout); - LiteralTestUtil::AssertEqualShapesAndLayouts( - expected_literal->shape(), computed.ValueOrDie()->shape()); + std::unique_ptr expected_literal = Literal::CreateR0(5); LiteralTestUtil::ExpectEqual(*expected_literal, *computed.ValueOrDie()); } } -// This test is permanently disabled on CPU because it requires that the -// backend used for execution is different than the backend used for -// ComputeConstant which is always cpu. -TEST_F(ComputeConstantTest, DISABLED_ON_CPU(ReuseComputedConstant)) { - // Compute a trivial constant, then try to use the value in an Execute - // call. This should fail because the constant resides on the CPU and the - // Execute call is executed on a different backend. - ComputationBuilder constant_b(client_, TestName()); - auto constant = constant_b.ConstantR0(42); - auto handle = constant_b.ComputeConstant(constant).ConsumeValueOrDie(); - auto literal = client_->Transfer(*handle).ConsumeValueOrDie(); - LiteralTestUtil::ExpectR0Equal(42, *literal); - - // Build trivial computation which takes one parameter. - ComputationBuilder b(client_, TestName()); - b.Neg(b.Parameter(0, ShapeUtil::MakeShape(S32, {}), "param0")); - auto computation = b.Build().ConsumeValueOrDie(); - - // Try to use value from ComputeConstant in Execute. - auto execute_status = client_->Execute(computation, {handle.get()}); - EXPECT_FALSE(execute_status.ok()); - EXPECT_MATCH( - execute_status.status().error_message(), - testing::ContainsRegex("argument 0 is on device Host:0 but computation " - "will be executed on device")); +XLA_TEST_F(ComputeConstantTest, Layout) { + for (ClientType client_type : client_types) { + Client* client = ClientOrDie(platform_, client_type); + ComputationBuilder b(client, TestName()); + + std::vector> layouts = {{0, 1}, {1, 0}}; + for (const std::vector& layout : layouts) { + auto layout_proto = LayoutUtil::MakeLayout(layout); + auto computed = ComputeConstantLiteral( + client, + b.Add(b.ConstantR2({{1, 2}, {3, 4}}), + b.ConstantR2({{10, 20}, {30, 40}})), + &b, &layout_proto); + ASSERT_TRUE(computed.ok()) << computed.status(); + + std::unique_ptr expected_literal = + test_utils::CreateR2LiteralWithLayout({{11, 22}, {33, 44}}, + layout); + LiteralTestUtil::AssertEqualShapesAndLayouts( + expected_literal->shape(), computed.ValueOrDie()->shape()); + LiteralTestUtil::ExpectEqual(*expected_literal, *computed.ValueOrDie()); + } + } } } // namespace } // namespace xla - -int main(int argc, char** argv) { - std::vector flag_list; - xla::legacy_flags::AppendCpuCompilerFlags(&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; - } - testing::InitGoogleTest(&argc, argv); - if (argc > 1) { - LOG(ERROR) << "Unknown argument " << argv[1] << "\n" << usage; - return 2; - } - return RUN_ALL_TESTS(); -} diff --git a/tensorflow/compiler/xla/tests/concat_test.cc b/tensorflow/compiler/xla/tests/concat_test.cc index 9a48b19b96aea829ded626ddb4ac64c0fa42b64c..1bcad5a3f37a37c9d482f3a5a899ac527666cca3 100644 --- a/tensorflow/compiler/xla/tests/concat_test.cc +++ b/tensorflow/compiler/xla/tests/concat_test.cc @@ -21,9 +21,9 @@ limitations under the License. #include "tensorflow/compiler/xla/client/computation.h" #include "tensorflow/compiler/xla/client/computation_builder.h" #include "tensorflow/compiler/xla/client/local_client.h" -#include "tensorflow/compiler/xla/legacy_flags/cpu_compiler_flags.h" #include "tensorflow/compiler/xla/reference_util.h" #include "tensorflow/compiler/xla/statusor.h" +#include "tensorflow/compiler/xla/test.h" #include "tensorflow/compiler/xla/test_helpers.h" #include "tensorflow/compiler/xla/tests/client_library_test_base.h" #include "tensorflow/compiler/xla/tests/literal_test_util.h" @@ -34,6 +34,7 @@ namespace xla { namespace { using ConcatTest = ClientLibraryTestBase; +using ::testing::HasSubstr; // Concatenate expects at least one argument. XLA_TEST_F(ConcatTest, Concat_Nothing) { @@ -41,9 +42,8 @@ XLA_TEST_F(ConcatTest, Concat_Nothing) { auto concatenated = builder.ConcatInDim({}, 0); StatusOr computation_status = builder.Build(); ASSERT_FALSE(computation_status.ok()); - EXPECT_MATCH( - computation_status.status().ToString(), - testing::ContainsRegex("Concatenate expects at least one argument")); + EXPECT_THAT(computation_status.status().ToString(), + HasSubstr("Concatenate expects at least one argument")); } // Concatenate with one argument works. @@ -56,6 +56,15 @@ XLA_TEST_F(ConcatTest, Concat_R1_With_Nothing) { ComputeAndCompareR1(&builder, expected, {}, ErrorSpec(0.0001)); } +XLA_TEST_F(ConcatTest, Concat_R1_L0_With_Nothing) { + ComputationBuilder builder(client_, TestName()); + auto a = builder.ConstantR1({}); + auto concatenated = builder.ConcatInDim({a}, 0); + + std::vector expected = {}; + ComputeAndCompareR1(&builder, expected, {}, ErrorSpec(0.0001)); +} + // Show that we can't concatenate R0 with R0 because we can't name the dimension // to concatenate on. XLA_TEST_F(ConcatTest, CannotConcatR0WithR0) { @@ -65,9 +74,8 @@ XLA_TEST_F(ConcatTest, CannotConcatR0WithR0) { auto concatenated = builder.ConcatInDim({a, b}, 0); StatusOr computation_status = builder.Build(); ASSERT_FALSE(computation_status.ok()); - EXPECT_MATCH(computation_status.status().ToString(), - testing::ContainsRegex( - "dimension to concatenate along out of bounds: 0")); + EXPECT_THAT(computation_status.status().ToString(), + HasSubstr("dimension to concatenate along out of bounds: 0")); } XLA_TEST_F(ConcatTest, Concat_R1_L0_With_R1_L0) { @@ -404,10 +412,9 @@ XLA_TEST_F(ConcatTest, CannotConcatOpaques) { auto concatenated = builder.ConcatInDim({x, y}, 0); StatusOr computation_status = builder.Build(); ASSERT_FALSE(computation_status.ok()); - EXPECT_MATCH( + EXPECT_THAT( computation_status.status().ToString(), - testing::ContainsRegex( - "Expected non-opaque argument for operand of concatenation")); + HasSubstr("Expected non-opaque argument for operand of concatenation")); } XLA_TEST_F(ConcatTest, ConcatSeveralBoxedPredicates) { @@ -434,6 +441,39 @@ XLA_TEST_F(ConcatTest, ConcatSeveralR1S32s) { ComputeAndCompareR1(&builder, expected, {}); } +XLA_TEST_F(ConcatTest, ConcatR3WeirdDims) { + ComputationBuilder builder(client_, TestName()); + + Array3D arr0(9, 17, 1); + arr0.Fill(1); + + Array3D arr1(9, 17, 256); + arr1.Fill(2); + + Array3D expected(9, 17, arr0.n3() + arr1.n3()); + for (int64 i = 0; i < expected.n1(); ++i) { + for (int64 j = 0; j < expected.n2(); ++j) { + int64 kk = 0; + for (const Array3D& arr : {arr0, arr1}) { + for (int64 k = 0; k < arr.n3(); ++k, ++kk) { + expected(i, j, kk) = arr(i, j, k); + } + } + } + } + + ComputationDataHandle h0; + auto p0 = CreateR3Parameter(arr0, /*parameter_number=*/0, "p0", + &builder, &h0); + ComputationDataHandle h1; + auto p1 = CreateR3Parameter(arr1, /*parameter_number=*/1, "p1", + &builder, &h1); + + auto concatenated = builder.ConcatInDim({h0, h1}, 2); + + ComputeAndCompareR3(&builder, expected, {p0.get(), p1.get()}); +} + // Describes a binary rank-2 concatenation test. struct R2BinarySpec { int64 lhs_dim0; @@ -476,8 +516,8 @@ TEST_P(ConcatR2BinaryTest, DoIt) { // concat XLA_TEST_F(ConcatTest, ConcatOperandsOfSameOperand) { auto f32_scalar = ShapeUtil::MakeShape(xla::F32, {}); - auto x_literal = LiteralUtil::CreateR0(2.f); - auto y_literal = LiteralUtil::CreateR0(3.f); + auto x_literal = Literal::CreateR0(2.f); + auto y_literal = Literal::CreateR0(3.f); auto x_data = client_->TransferToServer(*x_literal).ConsumeValueOrDie(); auto y_data = client_->TransferToServer(*y_literal).ConsumeValueOrDie(); @@ -494,6 +534,63 @@ XLA_TEST_F(ConcatTest, ConcatOperandsOfSameOperand) { {x_data.get(), y_data.get()}, ErrorSpec(1e-4)); } +// Test that the HLO optimization to replace a concat of a bradcasted scalar +// produces the correct result in rank 1. +XLA_TEST_F(ConcatTest, ConcatBroadcastArgument) { + auto f32_scalar = ShapeUtil::MakeShape(xla::F32, {}); + auto x_literal = Literal::CreateR1({2.0f, 3.0f, 5.0f, 6.0f}); + auto y_literal = Literal::CreateR0(1.5f); + auto z_literal = Literal::CreateR0(5.5f); + auto x_data = client_->TransferToServer(*x_literal).ConsumeValueOrDie(); + auto y_data = client_->TransferToServer(*y_literal).ConsumeValueOrDie(); + auto z_data = client_->TransferToServer(*z_literal).ConsumeValueOrDie(); + + ComputationBuilder builder(client_, TestName()); + auto x = builder.Parameter(0, x_literal->shape(), "x"); + auto y = builder.Parameter(1, f32_scalar, "y"); + auto z = builder.Parameter(2, f32_scalar, "z"); + auto bcast = builder.Broadcast(y, {5}); + auto bcast2 = builder.Broadcast(z, {3}); + auto concat = builder.ConcatInDim({bcast, x}, /*dimension=*/0); + builder.ConcatInDim({concat, bcast2}, /*dimension=*/0); + + ComputeAndCompareR1( + &builder, + {1.5f, 1.5f, 1.5f, 1.5f, 1.5f, 2.0f, 3.0f, 5.0f, 6.0f, 5.5f, 5.5f, 5.5f}, + {x_data.get(), y_data.get(), z_data.get()}, ErrorSpec(1e-4)); +} + +// Test that the HLO optimization to replace a concat of a bradcasted scalar +// produces the correct result in rank 3 with both high and low padding in +// different dimensions. +XLA_TEST_F(ConcatTest, ConcatBroadcastArgumentR3) { + auto f32_scalar = ShapeUtil::MakeShape(xla::F32, {}); + Array3D x3d(3, 5, 7, 3.14f); + auto x_literal = Literal::CreateR3FromArray3D(x3d); + auto y_literal = Literal::CreateR0(1.5f); + auto z_literal = Literal::CreateR0(5.5f); + auto x_data = client_->TransferToServer(*x_literal).ConsumeValueOrDie(); + auto y_data = client_->TransferToServer(*y_literal).ConsumeValueOrDie(); + auto z_data = client_->TransferToServer(*z_literal).ConsumeValueOrDie(); + + ComputationBuilder builder(client_, TestName()); + auto x = builder.Parameter(0, x_literal->shape(), "x"); + auto y = builder.Parameter(1, f32_scalar, "y"); + auto z = builder.Parameter(2, f32_scalar, "y"); + auto y_bcast = builder.Broadcast(y, {1, 5, 7}); + auto z_bcast = builder.Broadcast(z, {4, 1, 7}); + auto concat = builder.ConcatInDim({y_bcast, x}, /*dimension=*/0); + builder.ConcatInDim({concat, z_bcast}, /*dimension=*/1); + Array3D y_bcast3d(1, 5, 7, 1.5f); + Array3D z_bcast3d(4, 1, 7, 5.5f); + auto concat0 = ReferenceUtil::Concat3D(y_bcast3d, x3d, 0); + auto concat1 = ReferenceUtil::Concat3D(*concat0, z_bcast3d, 1); + + ComputeAndCompareR3(&builder, *concat1, + {x_data.get(), y_data.get(), z_data.get()}, + ErrorSpec(1e-4)); +} + INSTANTIATE_TEST_CASE_P(ConcatR2BinaryTestInstantiation, ConcatR2BinaryTest, ::testing::Values(R2BinarySpec{1, 1, 1, 1, 0}, R2BinarySpec{1, 1, 1, 1, 1}, @@ -504,20 +601,3 @@ INSTANTIATE_TEST_CASE_P(ConcatR2BinaryTestInstantiation, ConcatR2BinaryTest, } // namespace } // namespace xla - -int main(int argc, char** argv) { - std::vector flag_list; - xla::legacy_flags::AppendCpuCompilerFlags(&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; - } - testing::InitGoogleTest(&argc, argv); - if (argc > 1) { - LOG(ERROR) << "Unknown argument " << argv[1] << "\n" << usage; - return 2; - } - return RUN_ALL_TESTS(); -} diff --git a/tensorflow/compiler/xla/tests/constants_test.cc b/tensorflow/compiler/xla/tests/constants_test.cc index 58d52ac116841fc158046b2858f44ade3b488379..97bd1553664a6c0fcb097b441ec42efb4eaa9cc2 100644 --- a/tensorflow/compiler/xla/tests/constants_test.cc +++ b/tensorflow/compiler/xla/tests/constants_test.cc @@ -23,7 +23,6 @@ limitations under the License. #include "tensorflow/compiler/xla/array4d.h" #include "tensorflow/compiler/xla/client/computation_builder.h" #include "tensorflow/compiler/xla/client/local_client.h" -#include "tensorflow/compiler/xla/legacy_flags/cpu_compiler_flags.h" #include "tensorflow/compiler/xla/literal_util.h" #include "tensorflow/compiler/xla/tests/client_library_test_base.h" #include "tensorflow/compiler/xla/tests/literal_test_util.h" @@ -112,7 +111,7 @@ TEST_F(ConstantsTest, Small_2x2) { TEST_F(ConstantsTest, Empty_3x0x2) { ComputationBuilder builder(client_, TestName()); auto constant = builder.ConstantLiteral( - *LiteralUtil::CreateR3FromArray3D(Array3D(3, 0, 2))); + *Literal::CreateR3FromArray3D(Array3D(3, 0, 2))); ComputeAndCompareR3(&builder, Array3D(3, 0, 2), {}); } @@ -127,8 +126,8 @@ TEST_F(ConstantsTest, Small_2x2x2) { {{5.f, 6.f}, // y0 {7.f, 8.f}}, // y1 }); - auto constant = builder.ConstantLiteral( - *LiteralUtil::CreateR3FromArray3D(array3d)); + auto constant = + builder.ConstantLiteral(*Literal::CreateR3FromArray3D(array3d)); ComputeAndCompareR3(&builder, array3d, {}); } @@ -142,7 +141,7 @@ TEST_F(ConstantsTest, Small_3x2x1x1) { {5.0f, 4.4f}, // p2 }); input_array.FillWithPZ(pz); - Literal input_literal = *LiteralUtil::CreateR4FromArray4D(input_array); + Literal input_literal = *Literal::CreateR4FromArray4D(input_array); { ComputationBuilder builder(client_, TestName()); @@ -160,9 +159,9 @@ TEST_F(ConstantsTest, Small_3x2x1x1) { // TODO(b/29263943): Support tuple constants. TEST_F(ConstantsTest, DISABLED_TupleConstant) { ComputationBuilder builder(client_, TestName()); - builder.ConstantLiteral(*LiteralUtil::MakeTuple( - {LiteralUtil::CreateR2({{1.0}, {2.0}}).get(), - LiteralUtil::CreateR1({2.0, 42}).get()})); + builder.ConstantLiteral( + *Literal::MakeTuple({Literal::CreateR2({{1.0}, {2.0}}).get(), + Literal::CreateR1({2.0, 42}).get()})); std::unique_ptr result = ExecuteAndTransferOrDie(&builder, {}); @@ -174,20 +173,3 @@ TEST_F(ConstantsTest, DISABLED_TupleConstant) { } // namespace } // namespace xla - -int main(int argc, char** argv) { - std::vector flag_list; - xla::legacy_flags::AppendCpuCompilerFlags(&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; - } - testing::InitGoogleTest(&argc, argv); - if (argc > 1) { - LOG(ERROR) << "Unknown argument " << argv[1] << "\n" << usage; - return 2; - } - return RUN_ALL_TESTS(); -} diff --git a/tensorflow/compiler/xla/tests/convert_test.cc b/tensorflow/compiler/xla/tests/convert_test.cc index 9f8c3a9aeb7666a519ef4f06a900f66271785dd1..12b5e8426a78dc3a00794abdb892c0dcc1d15927 100644 --- a/tensorflow/compiler/xla/tests/convert_test.cc +++ b/tensorflow/compiler/xla/tests/convert_test.cc @@ -20,7 +20,6 @@ limitations under the License. #include "tensorflow/compiler/xla/client/computation_builder.h" #include "tensorflow/compiler/xla/client/local_client.h" -#include "tensorflow/compiler/xla/legacy_flags/cpu_compiler_flags.h" #include "tensorflow/compiler/xla/shape_util.h" #include "tensorflow/compiler/xla/tests/client_library_test_base.h" #include "tensorflow/compiler/xla/tests/literal_test_util.h" @@ -36,8 +35,10 @@ namespace { class ConvertTest : public ClientLibraryTestBase { public: explicit ConvertTest(perftools::gputools::Platform* platform = nullptr) - : ClientLibraryTestBase(platform, - /*disabled_pass_names=*/{"algsimp", "inline"}) {} + : ClientLibraryTestBase(platform) { + mutable_debug_options()->add_xla_disable_hlo_passes("algsimp"); + mutable_debug_options()->add_xla_disable_hlo_passes("inline"); + } }; TEST_F(ConvertTest, ConvertR1S32ToR1S32) { @@ -67,6 +68,24 @@ TEST_F(ConvertTest, ConvertR1S32ToR1F32) { ComputeAndCompareR1(&builder, expected, {}, ErrorSpec(0.0001)); } +TEST_F(ConvertTest, ConvertR1PREDToR1S32) { + ComputationBuilder builder(client_, TestName()); + auto a = builder.ConstantR1({true, false, true}); + builder.ConvertElementType(a, S32); + + std::vector expected = {1, 0, 1}; + ComputeAndCompareR1(&builder, expected, {}); +} + +TEST_F(ConvertTest, ConvertR1PREDToR1F32) { + ComputationBuilder builder(client_, TestName()); + auto a = builder.ConstantR1({true, false, true}); + builder.ConvertElementType(a, F32); + + std::vector expected = {1., 0., 1.}; + ComputeAndCompareR1(&builder, expected, {}); +} + XLA_TEST_F(ConvertTest, ConvertR1S0S32ToR1S0F32) { ComputationBuilder builder(client_, TestName()); auto a = builder.ConstantR1({}); @@ -191,20 +210,3 @@ TEST_F(ConvertTest, ConvertReshape) { } // namespace } // namespace xla - -int main(int argc, char** argv) { - std::vector flag_list; - xla::legacy_flags::AppendCpuCompilerFlags(&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; - } - testing::InitGoogleTest(&argc, argv); - if (argc > 1) { - LOG(ERROR) << "Unknown argument " << argv[1] << "\n" << usage; - return 2; - } - return RUN_ALL_TESTS(); -} diff --git a/tensorflow/compiler/xla/tests/convolution_dimension_numbers_test.cc b/tensorflow/compiler/xla/tests/convolution_dimension_numbers_test.cc index 9f38dc4b365672733ed773043f77bc4a3e8405ef..83882ca75e93ee9edec8e292991b53f1af57bb62 100644 --- a/tensorflow/compiler/xla/tests/convolution_dimension_numbers_test.cc +++ b/tensorflow/compiler/xla/tests/convolution_dimension_numbers_test.cc @@ -21,16 +21,14 @@ limitations under the License. #include "tensorflow/compiler/xla/client/computation_builder.h" #include "tensorflow/compiler/xla/client/local_client.h" #include "tensorflow/compiler/xla/client/padding.h" -#include "tensorflow/compiler/xla/legacy_flags/cpu_compiler_flags.h" #include "tensorflow/compiler/xla/ptr_util.h" #include "tensorflow/compiler/xla/reference_util.h" #include "tensorflow/compiler/xla/statusor.h" -#include "tensorflow/compiler/xla/test_helpers.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/core/platform/logging.h" -#include "tensorflow/core/platform/test.h" #include "tensorflow/core/platform/types.h" namespace xla { @@ -43,8 +41,8 @@ TEST_F(ConvolutionDimensionNumbersTest, InvalidInputDimensionNumbers) { auto dimension_numbers_status = ComputationBuilder::CreateConvDimensionNumbers(0, 2, 2, 3, 0, 1, 2, 3); ASSERT_FALSE(dimension_numbers_status.ok()); - ASSERT_MATCH(dimension_numbers_status.status().error_message(), - testing::ContainsRegex("input are not unique")); + ASSERT_THAT(dimension_numbers_status.status().error_message(), + ::testing::HasSubstr("input are not unique")); } // Tests the convolution operation with invalid weight dimension numbers. @@ -52,8 +50,8 @@ TEST_F(ConvolutionDimensionNumbersTest, InvalidWeightDimensionNumbers) { auto dimension_numbers_status = ComputationBuilder::CreateConvDimensionNumbers(0, 1, 2, 3, 2, 3, 2, 3); ASSERT_FALSE(dimension_numbers_status.ok()); - ASSERT_MATCH(dimension_numbers_status.status().error_message(), - testing::ContainsRegex("weight are not unique")); + ASSERT_THAT(dimension_numbers_status.status().error_message(), + ::testing::HasSubstr("weight are not unique")); } XLA_TEST_F(ConvolutionDimensionNumbersTest, @@ -63,8 +61,7 @@ XLA_TEST_F(ConvolutionDimensionNumbersTest, auto weight_array = MakeUnique>(4, 3, 1, 1); weight_array->FillWithMultiples(0.2); auto weight_data = - client_ - ->TransferToServer(*LiteralUtil::CreateR4FromArray4D(*weight_array)) + client_->TransferToServer(*Literal::CreateR4FromArray4D(*weight_array)) .ConsumeValueOrDie(); ComputationBuilder builder(client_, TestName()); @@ -98,20 +95,3 @@ XLA_TEST_F(ConvolutionDimensionNumbersTest, } // namespace } // namespace xla - -int main(int argc, char** argv) { - std::vector flag_list; - xla::legacy_flags::AppendCpuCompilerFlags(&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; - } - testing::InitGoogleTest(&argc, argv); - if (argc > 1) { - LOG(ERROR) << "Unknown argument " << argv[1] << "\n" << usage; - return 2; - } - return RUN_ALL_TESTS(); -} diff --git a/tensorflow/compiler/xla/tests/convolution_test.cc b/tensorflow/compiler/xla/tests/convolution_test.cc index 230c72f4bfa9ed40400fafc72328b38037f2f231..66b522b681d503473e02193764e649cac329ab6d 100644 --- a/tensorflow/compiler/xla/tests/convolution_test.cc +++ b/tensorflow/compiler/xla/tests/convolution_test.cc @@ -25,7 +25,6 @@ limitations under the License. #include "tensorflow/compiler/xla/client/local_client.h" #include "tensorflow/compiler/xla/client/padding.h" #include "tensorflow/compiler/xla/layout_util.h" -#include "tensorflow/compiler/xla/legacy_flags/cpu_compiler_flags.h" #include "tensorflow/compiler/xla/literal_util.h" #include "tensorflow/compiler/xla/ptr_util.h" #include "tensorflow/compiler/xla/reference_util.h" @@ -114,10 +113,10 @@ TEST_F(ConvolutionTest, Convolve_1x1x1x2_1x1x1x2_Valid) { ReferenceUtil::ConvArray4D(input, filter, {1, 1}, Padding::kValid); auto input_literal = - client_->TransferToServer(*LiteralUtil::CreateR4FromArray4D(input)) + client_->TransferToServer(*Literal::CreateR4FromArray4D(input)) .ConsumeValueOrDie(); auto filter_literal = - client_->TransferToServer(*LiteralUtil::CreateR4FromArray4D(filter)) + client_->TransferToServer(*Literal::CreateR4FromArray4D(filter)) .ConsumeValueOrDie(); ComputeAndCompareR4(&builder, *aexpected, @@ -157,10 +156,10 @@ TEST_F(ConvolutionTest, Convolve_1x1x4x4_1x1x2x2_Valid) { ReferenceUtil::ConvArray4D(input, filter, {1, 1}, Padding::kValid); auto input_literal = - client_->TransferToServer(*LiteralUtil::CreateR4FromArray4D(input)) + client_->TransferToServer(*Literal::CreateR4FromArray4D(input)) .ConsumeValueOrDie(); auto filter_literal = - client_->TransferToServer(*LiteralUtil::CreateR4FromArray4D(filter)) + client_->TransferToServer(*Literal::CreateR4FromArray4D(filter)) .ConsumeValueOrDie(); ComputeAndCompareR4(&builder, *aexpected, @@ -200,10 +199,10 @@ TEST_F(ConvolutionTest, Convolve_1x1x4x4_1x1x2x2_Same) { ReferenceUtil::ConvArray4D(input, filter, {1, 1}, Padding::kSame); auto input_literal = - client_->TransferToServer(*LiteralUtil::CreateR4FromArray4D(input)) + client_->TransferToServer(*Literal::CreateR4FromArray4D(input)) .ConsumeValueOrDie(); auto filter_literal = - client_->TransferToServer(*LiteralUtil::CreateR4FromArray4D(filter)) + client_->TransferToServer(*Literal::CreateR4FromArray4D(filter)) .ConsumeValueOrDie(); ComputeAndCompareR4(&builder, *aexpected, @@ -245,10 +244,10 @@ TEST_F(ConvolutionTest, Convolve_1x1x4x4_1x1x3x3_Same) { ReferenceUtil::ConvArray4D(input, filter, {1, 1}, Padding::kSame); auto input_literal = - client_->TransferToServer(*LiteralUtil::CreateR4FromArray4D(input)) + client_->TransferToServer(*Literal::CreateR4FromArray4D(input)) .ConsumeValueOrDie(); auto filter_literal = - client_->TransferToServer(*LiteralUtil::CreateR4FromArray4D(filter)) + client_->TransferToServer(*Literal::CreateR4FromArray4D(filter)) .ConsumeValueOrDie(); ComputeAndCompareR4(&builder, *aexpected, @@ -272,10 +271,10 @@ XLA_TEST_F(ConvolutionTest, Convolve1D_1x2x5_1x2x2_Valid) { Array3D expected({{{510, 610, 710, 810}}}); auto input_literal = - client_->TransferToServer(*LiteralUtil::CreateR3FromArray3D(input)) + client_->TransferToServer(*Literal::CreateR3FromArray3D(input)) .ConsumeValueOrDie(); auto filter_literal = - client_->TransferToServer(*LiteralUtil::CreateR3FromArray3D(filter)) + client_->TransferToServer(*Literal::CreateR3FromArray3D(filter)) .ConsumeValueOrDie(); ComputeAndCompareR3(&builder, expected, @@ -312,21 +311,18 @@ XLA_TEST_F(ConvolutionTest, Convolve3D_1x4x2x3x3_2x2x2x3x3_Valid) { std::vector input_elems(ShapeUtil::ElementsIn(input_shape)); std::iota(input_elems.begin(), input_elems.end(), 1.0f); - auto input_r1 = LiteralUtil::CreateR1(input_elems); - auto input_r5 = - LiteralUtil::Reshape(*input_r1, input_dims).ConsumeValueOrDie(); + auto input_r1 = Literal::CreateR1(input_elems); + auto input_r5 = input_r1->Reshape(input_dims).ConsumeValueOrDie(); std::vector filter_elems(ShapeUtil::ElementsIn(filter_shape)); std::iota(filter_elems.begin(), filter_elems.end(), 1.0f); - auto filter_r1 = LiteralUtil::CreateR1(filter_elems); - auto filter_r5 = - LiteralUtil::Reshape(*filter_r1, filter_dims).ConsumeValueOrDie(); + auto filter_r1 = Literal::CreateR1(filter_elems); + auto filter_r5 = filter_r1->Reshape(filter_dims).ConsumeValueOrDie(); - auto expected_r1 = LiteralUtil::CreateR1( + auto expected_r1 = Literal::CreateR1( {19554, 19962, 20370, 22110, 22590, 23070, 34890, 35730, 36570, 37446, 38358, 39270, 50226, 51498, 52770, 52782, 54126, 55470}); - auto expected_r5 = - LiteralUtil::Reshape(*expected_r1, {1, 3, 1, 2, 3}).ConsumeValueOrDie(); + auto expected_r5 = expected_r1->Reshape({1, 3, 1, 2, 3}).ConsumeValueOrDie(); auto input_literal = client_->TransferToServer(*input_r5).ConsumeValueOrDie(); auto filter_literal = @@ -337,22 +333,103 @@ XLA_TEST_F(ConvolutionTest, Convolve3D_1x4x2x3x3_2x2x2x3x3_Valid) { error_spec_); } -} // namespace -} // namespace xla +XLA_TEST_F(ConvolutionTest, Convolve2D_1x3x3x5_3x3x5x5_Valid) { + ComputationBuilder builder(client_, TestName()); + std::vector input_dims = {1, 3, 3, 5}; + std::vector filter_dims = {3, 3, 5, 3}; + Shape input_shape = ShapeUtil::MakeShape(F32, input_dims); + Shape filter_shape = ShapeUtil::MakeShape(F32, filter_dims); + { + auto input = builder.Parameter(0, input_shape, "input"); + auto filter = builder.Parameter(1, filter_shape, "filter"); -int main(int argc, char** argv) { - std::vector flag_list; - xla::legacy_flags::AppendCpuCompilerFlags(&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 dimension numbers for 2D convolution. + ConvolutionDimensionNumbers dnums; + dnums.set_batch_dimension(0); + dnums.add_spatial_dimensions(1); + dnums.add_spatial_dimensions(2); + dnums.set_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); + + builder.ConvWithGeneralDimensions(input, filter, {1, 1}, Padding::kValid, + dnums); } - testing::InitGoogleTest(&argc, argv); - if (argc > 1) { - LOG(ERROR) << "Unknown argument " << argv[1] << "\n" << usage; - return 2; + + std::vector input_elems(ShapeUtil::ElementsIn(input_shape)); + std::iota(input_elems.begin(), input_elems.end(), 1.0f); + auto input_r1 = Literal::CreateR1(input_elems); + auto input_r4 = input_r1->Reshape(input_dims).ConsumeValueOrDie(); + + std::vector filter_elems(ShapeUtil::ElementsIn(filter_shape)); + std::iota(filter_elems.begin(), filter_elems.end(), 1.0f); + auto filter_r1 = Literal::CreateR1(filter_elems); + auto filter_r4 = filter_r1->Reshape(filter_dims).ConsumeValueOrDie(); + + auto expected_r1 = Literal::CreateR1({92115, 93150, 94185}); + auto expected_r4 = expected_r1->Reshape({1, 1, 1, 3}).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_); +} + +XLA_TEST_F(ConvolutionTest, Convolve1D_Valid) { + ComputationBuilder builder(client_, TestName()); + int64 output_feature = 1; + int64 input_feature = 64; + int64 batch = 1; + int64 length = 1; + std::vector input_dims = {batch, 4 + length - 1, input_feature}; + std::vector filter_dims = {4, input_feature, output_feature}; + Shape input_shape = ShapeUtil::MakeShape(F32, input_dims); + Shape filter_shape = ShapeUtil::MakeShape(F32, filter_dims); + { + auto input = builder.Parameter(0, input_shape, "input"); + auto filter = builder.Parameter(1, filter_shape, "filter"); + + // Tensorflow dimension numbers for 2D convolution. + ConvolutionDimensionNumbers dnums; + dnums.set_batch_dimension(0); + dnums.add_spatial_dimensions(1); + dnums.set_feature_dimension(2); + dnums.add_kernel_spatial_dimensions(0); + dnums.set_kernel_input_feature_dimension(1); + dnums.set_kernel_output_feature_dimension(2); + + builder.ConvWithGeneralDimensions(input, filter, {1}, Padding::kValid, + dnums); } - return RUN_ALL_TESTS(); + + std::vector input_elems(ShapeUtil::ElementsIn(input_shape), 1.0); + // std::iota(input_elems.begin(), input_elems.end(), 1.0f); + auto input_r1 = Literal::CreateR1(input_elems); + auto input_r4 = input_r1->Reshape(input_dims).ConsumeValueOrDie(); + + std::vector filter_elems(ShapeUtil::ElementsIn(filter_shape), 1.0); + // std::iota(filter_elems.begin(), filter_elems.end(), 1.0f); + + auto filter_r1 = Literal::CreateR1(filter_elems); + auto filter_r4 = filter_r1->Reshape(filter_dims).ConsumeValueOrDie(); + + std::vector expect_elems(batch * output_feature * length, 256); + auto expected_r1 = Literal::CreateR1(expect_elems); + auto expected_r4 = + expected_r1->Reshape({batch, length, output_feature}).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_); } + +} // namespace +} // namespace xla diff --git a/tensorflow/compiler/xla/tests/convolution_variants_test.cc b/tensorflow/compiler/xla/tests/convolution_variants_test.cc index 86be451c26b3483e4390d9347ca8649427c0367d..145918db3e5e57c39054706d53bbfb7648af3143 100644 --- a/tensorflow/compiler/xla/tests/convolution_variants_test.cc +++ b/tensorflow/compiler/xla/tests/convolution_variants_test.cc @@ -28,7 +28,6 @@ limitations under the License. #include "tensorflow/compiler/xla/client/computation_builder.h" #include "tensorflow/compiler/xla/client/local_client.h" #include "tensorflow/compiler/xla/client/padding.h" -#include "tensorflow/compiler/xla/legacy_flags/cpu_compiler_flags.h" #include "tensorflow/compiler/xla/literal_util.h" #include "tensorflow/compiler/xla/reference_util.h" #include "tensorflow/compiler/xla/tests/client_library_test_base.h" @@ -52,7 +51,7 @@ class ConvolutionVariantsTest : public ClientLibraryTestBase { #endif }; -TEST_F(ConvolutionVariantsTest, Minimal) { +XLA_TEST_F(ConvolutionVariantsTest, Minimal) { ComputationBuilder builder(client_, TestName()); const Array4D input_array(1, 1, 1, 1, {2}); @@ -67,7 +66,7 @@ TEST_F(ConvolutionVariantsTest, Minimal) { ComputeAndCompareR4(&builder, expected, {}, error_spec_); } -TEST_F(ConvolutionVariantsTest, MinimalWithBatch) { +XLA_TEST_F(ConvolutionVariantsTest, MinimalWithBatch) { ComputationBuilder builder(client_, TestName()); const Array4D input_array(5, 1, 1, 1, {1, 2, 3, 4, 5}); @@ -82,7 +81,7 @@ TEST_F(ConvolutionVariantsTest, MinimalWithBatch) { ComputeAndCompareR4(&builder, expected, {}, error_spec_); } -TEST_F(ConvolutionVariantsTest, Flat1x1) { +XLA_TEST_F(ConvolutionVariantsTest, Flat1x1) { ComputationBuilder builder(client_, TestName()); Array4D input_array(2, 1, 3, 4); @@ -99,7 +98,7 @@ TEST_F(ConvolutionVariantsTest, Flat1x1) { ComputeAndCompareR4(&builder, expected, {}, error_spec_); } -TEST_F(ConvolutionVariantsTest, Deep1x1) { +XLA_TEST_F(ConvolutionVariantsTest, Deep1x1) { ComputationBuilder builder(client_, TestName()); Array4D input_array(1, 2, 1, 1, {10, 1}); @@ -114,7 +113,7 @@ TEST_F(ConvolutionVariantsTest, Deep1x1) { ComputeAndCompareR4(&builder, expected, {}, error_spec_); } -TEST_F(ConvolutionVariantsTest, Filter1x2in1x2) { +XLA_TEST_F(ConvolutionVariantsTest, Filter1x2in1x2) { ComputationBuilder builder(client_, TestName()); Array4D input_array(1, 1, 1, 2, {1, 2}); @@ -129,7 +128,7 @@ TEST_F(ConvolutionVariantsTest, Filter1x2in1x2) { ComputeAndCompareR4(&builder, expected, {}, error_spec_); } -TEST_F(ConvolutionVariantsTest, Filter1x2in1x3) { +XLA_TEST_F(ConvolutionVariantsTest, Filter1x2in1x3) { ComputationBuilder builder(client_, TestName()); Array4D input_array(1, 1, 1, 3, {1, 2, 3}); @@ -144,7 +143,7 @@ TEST_F(ConvolutionVariantsTest, Filter1x2in1x3) { ComputeAndCompareR4(&builder, expected, {}, error_spec_); } -TEST_F(ConvolutionVariantsTest, Filter1x2in2x2) { +XLA_TEST_F(ConvolutionVariantsTest, Filter1x2in2x2) { ComputationBuilder builder(client_, TestName()); Array4D input_array(1, 1, 2, 2, {1, 2, 3, 4}); @@ -159,7 +158,7 @@ TEST_F(ConvolutionVariantsTest, Filter1x2in2x2) { ComputeAndCompareR4(&builder, expected, {}, error_spec_); } -TEST_F(ConvolutionVariantsTest, Filter2x1in2x2) { +XLA_TEST_F(ConvolutionVariantsTest, Filter2x1in2x2) { ComputationBuilder builder(client_, TestName()); Array4D input_array(1, 1, 2, 2, {1, 2, 3, 4}); @@ -174,7 +173,7 @@ TEST_F(ConvolutionVariantsTest, Filter2x1in2x2) { ComputeAndCompareR4(&builder, expected, {}, error_spec_); } -TEST_F(ConvolutionVariantsTest, Filter2x2in2x2) { +XLA_TEST_F(ConvolutionVariantsTest, Filter2x2in2x2) { ComputationBuilder builder(client_, TestName()); Array4D input_array(1, 1, 2, 2, {1, 2, 3, 4}); @@ -189,7 +188,7 @@ TEST_F(ConvolutionVariantsTest, Filter2x2in2x2) { ComputeAndCompareR4(&builder, expected, {}, error_spec_); } -TEST_F(ConvolutionVariantsTest, Filter1x2in2x3WithDepthAndBatch) { +XLA_TEST_F(ConvolutionVariantsTest, Filter1x2in2x3WithDepthAndBatch) { ComputationBuilder builder(client_, TestName()); Array4D input_array( @@ -210,7 +209,7 @@ TEST_F(ConvolutionVariantsTest, Filter1x2in2x3WithDepthAndBatch) { ComputeAndCompareR4(&builder, expected, {}, error_spec_); } -TEST_F(ConvolutionVariantsTest, Filter1x1stride1x2in1x4) { +XLA_TEST_F(ConvolutionVariantsTest, Filter1x1stride1x2in1x4) { ComputationBuilder builder(client_, TestName()); Array4D input_array(1, 1, 1, 4, {1, 2, 3, 4}); @@ -225,7 +224,7 @@ TEST_F(ConvolutionVariantsTest, Filter1x1stride1x2in1x4) { ComputeAndCompareR4(&builder, expected, {}, error_spec_); } -TEST_F(ConvolutionVariantsTest, Filter1x1stride1x2in1x5) { +XLA_TEST_F(ConvolutionVariantsTest, Filter1x1stride1x2in1x5) { ComputationBuilder builder(client_, TestName()); Array4D input_array(1, 1, 1, 5, {1, 2, 3, 4, 5}); @@ -240,7 +239,7 @@ TEST_F(ConvolutionVariantsTest, Filter1x1stride1x2in1x5) { ComputeAndCompareR4(&builder, expected, {}, error_spec_); } -TEST_F(ConvolutionVariantsTest, Filter1x3stride1x2in1x4) { +XLA_TEST_F(ConvolutionVariantsTest, Filter1x3stride1x2in1x4) { ComputationBuilder builder(client_, TestName()); Array4D input_array(1, 1, 1, 4, {1, 2, 3, 4}); @@ -255,7 +254,7 @@ TEST_F(ConvolutionVariantsTest, Filter1x3stride1x2in1x4) { ComputeAndCompareR4(&builder, expected, {}, error_spec_); } -TEST_F(ConvolutionVariantsTest, Filter1x3stride1x2in1x5) { +XLA_TEST_F(ConvolutionVariantsTest, Filter1x3stride1x2in1x5) { ComputationBuilder builder(client_, TestName()); Array4D input_array(1, 1, 1, 5, {1, 2, 3, 4, 5}); @@ -270,7 +269,7 @@ TEST_F(ConvolutionVariantsTest, Filter1x3stride1x2in1x5) { ComputeAndCompareR4(&builder, expected, {}, error_spec_); } -TEST_F(ConvolutionVariantsTest, Filter1x1stride2x2in3x3) { +XLA_TEST_F(ConvolutionVariantsTest, Filter1x1stride2x2in3x3) { ComputationBuilder builder(client_, TestName()); Array4D input_array(1, 1, 3, 3, {1, 2, 3, 4, 5, 6, 7, 8, 9}); @@ -285,7 +284,7 @@ TEST_F(ConvolutionVariantsTest, Filter1x1stride2x2in3x3) { ComputeAndCompareR4(&builder, expected, {}, error_spec_); } -TEST_F(ConvolutionVariantsTest, Filter3x1in1x1Padded) { +XLA_TEST_F(ConvolutionVariantsTest, Filter3x1in1x1Padded) { ComputationBuilder builder(client_, TestName()); Array4D input_array(1, 1, 1, 1, {1}); @@ -300,7 +299,7 @@ TEST_F(ConvolutionVariantsTest, Filter3x1in1x1Padded) { ComputeAndCompareR4(&builder, expected, {}, error_spec_); } -TEST_F(ConvolutionVariantsTest, Filter5x1in3x1Padded) { +XLA_TEST_F(ConvolutionVariantsTest, Filter5x1in3x1Padded) { ComputationBuilder builder(client_, TestName()); Array4D input_array(1, 1, 1, 3, {1, 2, 3}); @@ -315,7 +314,7 @@ TEST_F(ConvolutionVariantsTest, Filter5x1in3x1Padded) { ComputeAndCompareR4(&builder, expected, {}, error_spec_); } -TEST_F(ConvolutionVariantsTest, Filter3x3in2x2Padded) { +XLA_TEST_F(ConvolutionVariantsTest, Filter3x3in2x2Padded) { ComputationBuilder builder(client_, TestName()); Array4D input_array(1, 1, 2, 2, {1, 2, 3, 4}); @@ -332,7 +331,7 @@ TEST_F(ConvolutionVariantsTest, Filter3x3in2x2Padded) { ComputeAndCompareR4(&builder, expected, {}, error_spec_); } -TEST_F(ConvolutionVariantsTest, Filter1x1in2x1WithPaddingAndDepth) { +XLA_TEST_F(ConvolutionVariantsTest, Filter1x1in2x1WithPaddingAndDepth) { ComputationBuilder builder(client_, TestName()); Array4D input_array(1, 2, 1, 2, {1, 2, 3, 4}); @@ -347,7 +346,7 @@ TEST_F(ConvolutionVariantsTest, Filter1x1in2x1WithPaddingAndDepth) { ComputeAndCompareR4(&builder, expected, {}, error_spec_); } -TEST_F(ConvolutionVariantsTest, Filter2x2Stride1x1Input3x3) { +XLA_TEST_F(ConvolutionVariantsTest, Filter2x2Stride1x1Input3x3) { ComputationBuilder builder(client_, TestName()); Array4D input_array(1, 1, 3, 3, {1, 2, 3, 4, 5, 6, 7, 8, 9}); @@ -362,7 +361,7 @@ TEST_F(ConvolutionVariantsTest, Filter2x2Stride1x1Input3x3) { ComputeAndCompareR4(&builder, expected, {}, error_spec_); } -TEST_F(ConvolutionVariantsTest, Filter1x2Stride1x1Input1x3) { +XLA_TEST_F(ConvolutionVariantsTest, Filter1x2Stride1x1Input1x3) { ComputationBuilder builder(client_, TestName()); Array4D input_array(1, 1, 1, 3, {1, 2, 3}); @@ -377,7 +376,7 @@ TEST_F(ConvolutionVariantsTest, Filter1x2Stride1x1Input1x3) { ComputeAndCompareR4(&builder, expected, {}, error_spec_); } -TEST_F(ConvolutionVariantsTest, Filter2x1x8x8Input1x1x8x8) { +XLA_TEST_F(ConvolutionVariantsTest, Filter2x1x8x8Input1x1x8x8) { ComputationBuilder builder(client_, TestName()); std::vector input_data(64); @@ -397,7 +396,7 @@ TEST_F(ConvolutionVariantsTest, Filter2x1x8x8Input1x1x8x8) { ComputeAndCompareR4(&builder, expected, {}, error_spec_); } -TEST_F(ConvolutionVariantsTest, Filter1x1x1x1Input16x1x1x1) { +XLA_TEST_F(ConvolutionVariantsTest, Filter1x1x1x1Input16x1x1x1) { ComputationBuilder builder(client_, TestName()); std::vector input_data(16 * 1 * 1 * 1); @@ -418,7 +417,7 @@ TEST_F(ConvolutionVariantsTest, Filter1x1x1x1Input16x1x1x1) { ComputeAndCompareR4(&builder, expected, {}, error_spec_); } -TEST_F(ConvolutionVariantsTest, Filter1x1x2x2Input16x1x2x2) { +XLA_TEST_F(ConvolutionVariantsTest, Filter1x1x2x2Input16x1x2x2) { ComputationBuilder builder(client_, TestName()); constexpr int bs = 16; @@ -449,7 +448,7 @@ TEST_F(ConvolutionVariantsTest, Filter1x1x2x2Input16x1x2x2) { ComputeAndCompareR4(&builder, expected, {}, error_spec_); } -TEST_F(ConvolutionVariantsTest, Filter1x1x2x2Input3x1x2x2) { +XLA_TEST_F(ConvolutionVariantsTest, Filter1x1x2x2Input3x1x2x2) { ComputationBuilder builder(client_, TestName()); constexpr int kx = 2; @@ -479,7 +478,7 @@ TEST_F(ConvolutionVariantsTest, Filter1x1x2x2Input3x1x2x2) { ComputeAndCompareR4(&builder, expected, {}, error_spec_); } -TEST_F(ConvolutionVariantsTest, Filter1x1x8x8Input16x1x8x8) { +XLA_TEST_F(ConvolutionVariantsTest, Filter1x1x8x8Input16x1x8x8) { ComputationBuilder builder(client_, TestName()); Array4D input_array(16, 1, 8, 8); @@ -507,7 +506,7 @@ TEST_F(ConvolutionVariantsTest, Filter1x1x8x8Input16x1x8x8) { ComputeAndCompareR4(&builder, expected, {}, error_spec_); } -TEST_F(ConvolutionVariantsTest, Filter2x2x8x8Input1x2x8x8) { +XLA_TEST_F(ConvolutionVariantsTest, Filter2x2x8x8Input1x2x8x8) { ComputationBuilder builder(client_, TestName()); std::vector input_data(2 * 8 * 8); @@ -533,7 +532,7 @@ TEST_F(ConvolutionVariantsTest, Filter2x2x8x8Input1x2x8x8) { ComputeAndCompareR4(&builder, expected, {}, error_spec_); } -TEST_F(ConvolutionVariantsTest, Filter2x2x8x8Input2x2x8x8) { +XLA_TEST_F(ConvolutionVariantsTest, Filter2x2x8x8Input2x2x8x8) { ComputationBuilder builder(client_, TestName()); std::vector input_data(2 * 2 * 8 * 8); @@ -559,7 +558,7 @@ TEST_F(ConvolutionVariantsTest, Filter2x2x8x8Input2x2x8x8) { ComputeAndCompareR4(&builder, expected, {}, error_spec_); } -TEST_F(ConvolutionVariantsTest, Filter2x2x8x8Input32x2x8x8) { +XLA_TEST_F(ConvolutionVariantsTest, Filter2x2x8x8Input32x2x8x8) { ComputationBuilder builder(client_, TestName()); std::vector input_data(32 * 2 * 8 * 8); @@ -599,7 +598,7 @@ TEST_F(ConvolutionVariantsTest, Filter2x2x8x8Input32x2x8x8) { ComputeAndCompareR4(&builder, expected, {}, error_spec_); } -TEST_F(ConvolutionVariantsTest, Filter16x16x1x1Input16x16x1x1) { +XLA_TEST_F(ConvolutionVariantsTest, Filter16x16x1x1Input16x16x1x1) { ComputationBuilder builder(client_, TestName()); Array4D input_array(16, 16, 1, 1); @@ -795,7 +794,7 @@ XLA_TEST_F(ConvolutionVariantsTest, NegativePaddingAndDilation) { ComputeAndCompareR4(&builder, expected, {}, error_spec_); } -TEST_F(ConvolutionVariantsTest, RandomData_Input1x1x2x3_Filter2x1x1x2) { +XLA_TEST_F(ConvolutionVariantsTest, RandomData_Input1x1x2x3_Filter2x1x1x2) { constexpr int bs = 1; constexpr int iz = 1; constexpr int oz = 2; @@ -828,7 +827,7 @@ TEST_F(ConvolutionVariantsTest, RandomData_Input1x1x2x3_Filter2x1x1x2) { ComputeAndCompareR4(&builder, *expected, {}, error_spec_); } -TEST_F(ConvolutionVariantsTest, RandomData_Input1x16x1x1_Filter1x16x1x1) { +XLA_TEST_F(ConvolutionVariantsTest, RandomData_Input1x16x1x1_Filter1x16x1x1) { constexpr int bs = 1; constexpr int iz = 16; constexpr int oz = 1; @@ -861,7 +860,7 @@ TEST_F(ConvolutionVariantsTest, RandomData_Input1x16x1x1_Filter1x16x1x1) { ComputeAndCompareR4(&builder, *expected, {}, error_spec_); } -TEST_F(ConvolutionVariantsTest, RandomData_Input16x16x1x1_Filter1x16x1x1) { +XLA_TEST_F(ConvolutionVariantsTest, RandomData_Input16x16x1x1_Filter1x16x1x1) { constexpr int bs = 16; constexpr int iz = 16; constexpr int oz = 1; @@ -894,7 +893,7 @@ TEST_F(ConvolutionVariantsTest, RandomData_Input16x16x1x1_Filter1x16x1x1) { ComputeAndCompareR4(&builder, *expected, {}, error_spec_); } -TEST_F(ConvolutionVariantsTest, RandomData_Input16x16x1x1_Filter16x16x1x1) { +XLA_TEST_F(ConvolutionVariantsTest, RandomData_Input16x16x1x1_Filter16x16x1x1) { constexpr int bs = 16; constexpr int iz = 16; constexpr int oz = 16; @@ -927,7 +926,7 @@ TEST_F(ConvolutionVariantsTest, RandomData_Input16x16x1x1_Filter16x16x1x1) { ComputeAndCompareR4(&builder, *expected, {}, error_spec_); } -TEST_F(ConvolutionVariantsTest, RandomData_Input16x16x16x16_Filter16x16x16x16) { +XLA_TEST_F(ConvolutionVariantsTest, RandomData_Input16x16x16x16_Filter16x16x16x16) { constexpr int bs = 16; constexpr int iz = 16; constexpr int oz = 16; @@ -960,7 +959,7 @@ TEST_F(ConvolutionVariantsTest, RandomData_Input16x16x16x16_Filter16x16x16x16) { ComputeAndCompareR4(&builder, *expected, {}, error_spec_); } -TEST_F(ConvolutionVariantsTest, Filter1x2x1x1Input1x2x3x1GeneralPadding) { +XLA_TEST_F(ConvolutionVariantsTest, Filter1x2x1x1Input1x2x3x1GeneralPadding) { ComputationBuilder builder(client_, TestName()); std::vector input_data(1 * 2 * 3 * 1); @@ -1000,7 +999,7 @@ TEST_F(ConvolutionVariantsTest, Filter1x2x1x1Input1x2x3x1GeneralPadding) { ComputeAndCompareR4(&builder, expected, {}, error_spec_); } -TEST_F(ConvolutionVariantsTest, Filter1x1x1x1Input1x2x3x1GeneralPadding) { +XLA_TEST_F(ConvolutionVariantsTest, Filter1x1x1x1Input1x2x3x1GeneralPadding) { ComputationBuilder builder(client_, TestName()); std::vector input_data(1 * 2 * 3 * 1); @@ -1040,7 +1039,7 @@ TEST_F(ConvolutionVariantsTest, Filter1x1x1x1Input1x2x3x1GeneralPadding) { ComputeAndCompareR4(&builder, expected, {}, error_spec_); } -TEST_F(ConvolutionVariantsTest, Filter1x1x1x1Input1x2x3x1NoPadding) { +XLA_TEST_F(ConvolutionVariantsTest, Filter1x1x1x1Input1x2x3x1NoPadding) { ComputationBuilder builder(client_, TestName()); std::vector input_data(1 * 2 * 3 * 1); @@ -1077,7 +1076,7 @@ TEST_F(ConvolutionVariantsTest, Filter1x1x1x1Input1x2x3x1NoPadding) { ComputeAndCompareR4(&builder, expected, {}, error_spec_); } -TEST_F(ConvolutionVariantsTest, Filter1x1x2x3Input1x2x3x2NoPadding) { +XLA_TEST_F(ConvolutionVariantsTest, Filter1x1x2x3Input1x2x3x2NoPadding) { ComputationBuilder builder(client_, TestName()); std::vector input_data(1 * 2 * 3 * 2); @@ -1124,7 +1123,7 @@ TEST_F(ConvolutionVariantsTest, Filter1x1x2x3Input1x2x3x2NoPadding) { // Conv([1,2,3], Reverse([5,6]), padding_low=1) // into // BackwardInputConv([1,2,3], [5,6], padding_low=0, padding_high=1) -TEST_F(ConvolutionVariantsTest, BackwardInputLowPaddingLessThanHighPadding) { +XLA_TEST_F(ConvolutionVariantsTest, BackwardInputLowPaddingLessThanHighPadding) { ComputationBuilder builder(client_, TestName()); auto gradients = builder.ConstantR4FromArray4D( @@ -1142,7 +1141,7 @@ TEST_F(ConvolutionVariantsTest, BackwardInputLowPaddingLessThanHighPadding) { // Conv([1], Reverse([1,10,100]), padding_high=3, base_dilation=3) // into // BackwardInputConv([1], [1,10,100], stride=3, padding=(2,1)) -TEST_F(ConvolutionVariantsTest, BackwardInputLowPaddingGreaterThanHighPadding) { +XLA_TEST_F(ConvolutionVariantsTest, BackwardInputLowPaddingGreaterThanHighPadding) { ComputationBuilder builder(client_, TestName()); auto gradients = builder.ConstantR4FromArray4D( @@ -1163,7 +1162,7 @@ TEST_F(ConvolutionVariantsTest, BackwardInputLowPaddingGreaterThanHighPadding) { // Conv([1], Reverse([1,10,100]), padding=(1,1)) // into // BackwardInputConv([1], [1,10,100], padding=(1,1)) -TEST_F(ConvolutionVariantsTest, BackwardInputEvenPadding) { +XLA_TEST_F(ConvolutionVariantsTest, BackwardInputEvenPadding) { ComputationBuilder builder(client_, TestName()); auto gradients = builder.ConstantR4FromArray4D( @@ -1184,7 +1183,7 @@ TEST_F(ConvolutionVariantsTest, BackwardInputEvenPadding) { // // However, XLA:GPU doesn't actually fuse it because PadInsertion doesn't // support negative padding on backward convolution yet (b/32744257). -TEST_F(ConvolutionVariantsTest, BackwardInputWithNegativePaddingHigh) { +XLA_TEST_F(ConvolutionVariantsTest, BackwardInputWithNegativePaddingHigh) { ComputationBuilder builder(client_, TestName()); auto gradients = builder.ConstantR4FromArray4D( @@ -1199,7 +1198,7 @@ TEST_F(ConvolutionVariantsTest, BackwardInputWithNegativePaddingHigh) { ComputeAndCompareR4(&builder, {{{{12, 23, 30, 0}}}}, {}, error_spec_); } -TEST_F(ConvolutionVariantsTest, BackwardFilterLowPaddingLessThanHighPadding) { +XLA_TEST_F(ConvolutionVariantsTest, BackwardFilterLowPaddingLessThanHighPadding) { ComputationBuilder builder(client_, TestName()); // activations: 1,2,3,4 ---pad--> 0,1,2,3,4,0,0 @@ -1222,7 +1221,7 @@ TEST_F(ConvolutionVariantsTest, BackwardFilterLowPaddingLessThanHighPadding) { ComputeAndCompareR4(&builder, {{{{24, 130, 240}}}}, {}, error_spec_); } -TEST_F(ConvolutionVariantsTest, +XLA_TEST_F(ConvolutionVariantsTest, BackwardFilterLowPaddingGreaterThanHighPadding) { ComputationBuilder builder(client_, TestName()); @@ -1248,7 +1247,7 @@ TEST_F(ConvolutionVariantsTest, ComputeAndCompareR4(&builder, {{{{13, 24}}}}, {}, error_spec_); } -TEST_F(ConvolutionVariantsTest, BackwardFilterEvenPadding) { +XLA_TEST_F(ConvolutionVariantsTest, BackwardFilterEvenPadding) { ComputationBuilder builder(client_, TestName()); // activations: 1,2,3,4 ---pad--> 0,0,1,2,3,4,0 @@ -1275,7 +1274,7 @@ TEST_F(ConvolutionVariantsTest, BackwardFilterEvenPadding) { ComputeAndCompareR4(&builder, {{{{13, 24, 130}}}}, {}, error_spec_); } -TEST_F(ConvolutionVariantsTest, BackwardInputEvenPadding1D) { +XLA_TEST_F(ConvolutionVariantsTest, BackwardInputEvenPadding1D) { ComputationBuilder builder(client_, TestName()); auto gradients = builder.ConstantR3FromArray3D( @@ -1289,7 +1288,7 @@ TEST_F(ConvolutionVariantsTest, BackwardInputEvenPadding1D) { ComputeAndCompareR3(&builder, {{{10}}}, {}, error_spec_); } -TEST_F(ConvolutionVariantsTest, BackwardFilterEvenPadding1D) { +XLA_TEST_F(ConvolutionVariantsTest, BackwardFilterEvenPadding1D) { ComputationBuilder builder(client_, TestName()); auto activations = @@ -1308,23 +1307,22 @@ TEST_F(ConvolutionVariantsTest, BackwardFilterEvenPadding1D) { ComputeAndCompareR3(&builder, {{{13, 24, 130}}}, {}, error_spec_); } -TEST_F(ConvolutionVariantsTest, BackwardInputEvenPadding3D) { +XLA_TEST_F(ConvolutionVariantsTest, BackwardInputEvenPadding3D) { ComputationBuilder builder(client_, TestName()); - auto gradients_flat = LiteralUtil::CreateR1({1}); + auto gradients_flat = Literal::CreateR1({1}); auto gradients_literal = - LiteralUtil::Reshape(*gradients_flat, {1, 1, 1, 1, 1}) - .ConsumeValueOrDie(); + gradients_flat->Reshape({1, 1, 1, 1, 1}).ConsumeValueOrDie(); auto gradients = builder.ConstantLiteral(*gradients_literal); - auto weights_flat = LiteralUtil::CreateR1({1, 10, 100}); + auto weights_flat = Literal::CreateR1({1, 10, 100}); auto weights_literal = - LiteralUtil::Reshape(*weights_flat, {1, 1, 1, 1, 3}).ConsumeValueOrDie(); + weights_flat->Reshape({1, 1, 1, 1, 3}).ConsumeValueOrDie(); auto weights = builder.ConstantLiteral(*weights_literal); - auto expected_flat = LiteralUtil::CreateR1({10}); + auto expected_flat = Literal::CreateR1({10}); auto expected_literal = - LiteralUtil::Reshape(*expected_flat, {1, 1, 1, 1, 1}).ConsumeValueOrDie(); + expected_flat->Reshape({1, 1, 1, 1, 1}).ConsumeValueOrDie(); auto mirrored_weights = builder.Rev(weights, {2, 3, 4}); builder.ConvWithGeneralPadding(gradients, mirrored_weights, @@ -1333,24 +1331,22 @@ TEST_F(ConvolutionVariantsTest, BackwardInputEvenPadding3D) { ComputeAndCompareLiteral(&builder, *expected_literal, {}, error_spec_); } -TEST_F(ConvolutionVariantsTest, BackwardFilterEvenPadding3D) { +XLA_TEST_F(ConvolutionVariantsTest, BackwardFilterEvenPadding3D) { ComputationBuilder builder(client_, TestName()); - auto activations_flat = LiteralUtil::CreateR1({1, 2, 3, 4}); + auto activations_flat = Literal::CreateR1({1, 2, 3, 4}); auto activations_literal = - LiteralUtil::Reshape(*activations_flat, {1, 1, 1, 1, 4}) - .ConsumeValueOrDie(); + activations_flat->Reshape({1, 1, 1, 1, 4}).ConsumeValueOrDie(); auto activations = builder.ConstantLiteral(*activations_literal); - auto gradients_flat = LiteralUtil::CreateR1({100, 10, 1}); + auto gradients_flat = Literal::CreateR1({100, 10, 1}); auto gradients_literal = - LiteralUtil::Reshape(*gradients_flat, {1, 1, 1, 1, 3}) - .ConsumeValueOrDie(); + gradients_flat->Reshape({1, 1, 1, 1, 3}).ConsumeValueOrDie(); auto gradients = builder.ConstantLiteral(*gradients_literal); - auto expected_flat = LiteralUtil::CreateR1({13, 24, 130}); + auto expected_flat = Literal::CreateR1({13, 24, 130}); auto expected_literal = - LiteralUtil::Reshape(*expected_flat, {1, 1, 1, 1, 3}).ConsumeValueOrDie(); + expected_flat->Reshape({1, 1, 1, 1, 3}).ConsumeValueOrDie(); auto forward_conv = builder.ConvGeneralDilated( activations, gradients, @@ -1365,20 +1361,3 @@ TEST_F(ConvolutionVariantsTest, BackwardFilterEvenPadding3D) { } // namespace } // namespace xla - -int main(int argc, char** argv) { - std::vector flag_list; - xla::legacy_flags::AppendCpuCompilerFlags(&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; - } - testing::InitGoogleTest(&argc, argv); - if (argc > 1) { - LOG(ERROR) << "Unknown argument " << argv[1] << "\n" << usage; - return 2; - } - return RUN_ALL_TESTS(); -} diff --git a/tensorflow/compiler/xla/tests/copy_test.cc b/tensorflow/compiler/xla/tests/copy_test.cc index 29e29505333b64926cdd0b3e9fe7ef3407eaaec2..bcb85b04eefa349df1c055e010d584b85b55a4a8 100644 --- a/tensorflow/compiler/xla/tests/copy_test.cc +++ b/tensorflow/compiler/xla/tests/copy_test.cc @@ -17,13 +17,13 @@ limitations under the License. #include #include "tensorflow/compiler/xla/array2d.h" -#include "tensorflow/compiler/xla/legacy_flags/cpu_compiler_flags.h" #include "tensorflow/compiler/xla/literal_util.h" #include "tensorflow/compiler/xla/ptr_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/service/hlo_opcode.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" @@ -44,11 +44,10 @@ class CopyOpTest : public HloTestBase { builder.AddInstruction(HloInstruction::CreateUnary( constant->shape(), HloOpcode::kCopy, constant)); auto computation = builder.Build(); - auto hlo_module = MakeUnique("test_module"); - hlo_module->AddEntryComputation(std::move(computation)); + auto module = CreateNewModule(); + module->AddEntryComputation(std::move(computation)); - std::unique_ptr result = - ExecuteAndTransfer(std::move(hlo_module), {}); + std::unique_ptr result = ExecuteAndTransfer(std::move(module), {}); LiteralTestUtil::ExpectEqual(literal, *result); } @@ -57,39 +56,34 @@ class CopyOpTest : public HloTestBase { tensorflow::gtl::ArraySlice permutation); }; -TEST_F(CopyOpTest, CopyR0Bool) { - TestCopyOp(*LiteralUtil::CreateR0(true)); -} +XLA_TEST_F(CopyOpTest, CopyR0Bool) { TestCopyOp(*Literal::CreateR0(true)); } -TEST_F(CopyOpTest, CopyR1S0U32) { - TestCopyOp(*LiteralUtil::CreateR1({})); -} +XLA_TEST_F(CopyOpTest, CopyR1S0U32) { TestCopyOp(*Literal::CreateR1({})); } -TEST_F(CopyOpTest, CopyR1S3U32) { - TestCopyOp(*LiteralUtil::CreateR1({1, 2, 3})); +XLA_TEST_F(CopyOpTest, CopyR1S3U32) { + TestCopyOp(*Literal::CreateR1({1, 2, 3})); } -TEST_F(CopyOpTest, CopyR3F32_2x2x3) { - TestCopyOp( - *LiteralUtil::CreateR3({{{1.0f, 2.0f, 3.0f}, {4.0f, 5.0f, 6.0f}}, - {{1.1f, 2.1f, 3.1f}, {6.1f, 3.5f, 2.8f}}})); +XLA_TEST_F(CopyOpTest, CopyR3F32_2x2x3) { + TestCopyOp(*Literal::CreateR3({{{1.0f, 2.0f, 3.0f}, {4.0f, 5.0f, 6.0f}}, + {{1.1f, 2.1f, 3.1f}, {6.1f, 3.5f, 2.8f}}})); } -TEST_F(CopyOpTest, CopyR4S32_2x2x3x2) { - TestCopyOp(*LiteralUtil::CreateR4( +XLA_TEST_F(CopyOpTest, CopyR4S32_2x2x3x2) { + TestCopyOp(*Literal::CreateR4( {{{{1, -2}, {-4, 5}, {6, 7}}, {{8, 9}, {10, 11}, {12, 13}}}, {{{10, 3}, {7, -2}, {3, 6}}, {{2, 5}, {-11, 5}, {-2, -5}}}})); } -TEST_F(CopyOpTest, CopyR4S32_0x2x3x2) { - TestCopyOp(*LiteralUtil::CreateR4FromArray4D(Array4D(0, 2, 3, 2))); +XLA_TEST_F(CopyOpTest, CopyR4S32_0x2x3x2) { + TestCopyOp(*Literal::CreateR4FromArray4D(Array4D(0, 2, 3, 2))); } -TEST_F(CopyOpTest, CopyParameterScalar) { +XLA_TEST_F(CopyOpTest, CopyParameterScalar) { auto builder = HloComputation::Builder(TestName()); // Copy literal to device to use as parameter. - auto literal = LiteralUtil::CreateR0(42.0); + auto literal = Literal::CreateR0(42.0); Shape shape = literal->shape(); auto constant_device_base = TransferToDevice(*literal); @@ -100,18 +94,18 @@ TEST_F(CopyOpTest, CopyParameterScalar) { auto computation = builder.Build(); - auto hlo_module = MakeUnique("test_module"); - hlo_module->AddEntryComputation(std::move(computation)); + auto module = CreateNewModule(); + module->AddEntryComputation(std::move(computation)); std::unique_ptr result = - ExecuteAndTransfer(std::move(hlo_module), {constant_device_base}); + ExecuteAndTransfer(std::move(module), {constant_device_base}); LiteralTestUtil::ExpectR0Near(42.0f, *result, error_spec_); } -TEST_F(CopyOpTest, CopyConstantR2Twice) { +XLA_TEST_F(CopyOpTest, CopyConstantR2Twice) { auto builder = HloComputation::Builder(TestName()); - auto literal = LiteralUtil::CreateR2({{1.0, 2.0}, {3.0, 4.0}}); + auto literal = Literal::CreateR2({{1.0, 2.0}, {3.0, 4.0}}); auto constant = builder.AddInstruction( HloInstruction::CreateConstant(std::move(literal))); @@ -122,19 +116,18 @@ TEST_F(CopyOpTest, CopyConstantR2Twice) { auto computation = builder.Build(); - auto hlo_module = MakeUnique("test_module"); - hlo_module->AddEntryComputation(std::move(computation)); - std::unique_ptr result = - ExecuteAndTransfer(std::move(hlo_module), {}); + auto module = CreateNewModule(); + module->AddEntryComputation(std::move(computation)); + std::unique_ptr result = ExecuteAndTransfer(std::move(module), {}); LiteralTestUtil::ExpectR2Near({{1.0, 2.0}, {3.0, 4.0}}, *result, error_spec_); } -TEST_F(CopyOpTest, CopyConstantR2DifferentLayouts) { +XLA_TEST_F(CopyOpTest, CopyConstantR2DifferentLayouts) { HloComputation::Builder builder(TestName()); std::unique_ptr literal = - LiteralUtil::CreateR2({{1.0, 2.0}, {3.0, 4.0}}); + Literal::CreateR2({{1.0, 2.0}, {3.0, 4.0}}); // Reverse the minor-to-major order of the literal. Layout* literal_layout = literal->mutable_shape()->mutable_layout(); ASSERT_EQ(2, literal_layout->minor_to_major_size()); @@ -148,10 +141,9 @@ TEST_F(CopyOpTest, CopyConstantR2DifferentLayouts) { std::unique_ptr computation = builder.Build(); - auto hlo_module = MakeUnique("test_module"); - hlo_module->AddEntryComputation(std::move(computation)); - std::unique_ptr result = - ExecuteAndTransfer(std::move(hlo_module), {}); + auto module = CreateNewModule(); + module->AddEntryComputation(std::move(computation)); + std::unique_ptr result = ExecuteAndTransfer(std::move(module), {}); // The result of the computation has the default layout, which is the inverse // of the layout of the source literal. @@ -171,7 +163,7 @@ void CopyOpTest::TestCopyConstantLayout021(size_t n1, size_t n2, size_t n3) { HloComputation::Builder builder(TestName()); - std::unique_ptr literal = LiteralUtil::CreateR3FromArray3D(a); + std::unique_ptr literal = Literal::CreateR3FromArray3D(a); HloInstruction* constant = builder.AddInstruction( HloInstruction::CreateConstant(std::move(literal))); @@ -181,15 +173,10 @@ void CopyOpTest::TestCopyConstantLayout021(size_t n1, size_t n2, size_t n3) { std::unique_ptr computation = builder.Build(); - auto hlo_module = MakeUnique("test_module"); - auto config = MakeUnique(computation->ComputeProgramShape()); - *config->mutable_entry_computation_layout()->mutable_result_layout() = - ShapeLayout(ShapeUtil::MakeShapeWithLayout( - constant->shape().element_type(), - AsInt64Slice(constant->shape().dimensions()), {1, 2, 0})); - hlo_module->AddEntryComputation(std::move(computation)); - std::unique_ptr result = - ExecuteAndTransfer(std::move(hlo_module), std::move(config), {}); + auto module = CreateNewModule(); + module->AddEntryComputation(std::move(computation)); + ForceResultLayout(module.get(), LayoutUtil::MakeLayout({1, 2, 0})); + std::unique_ptr result = ExecuteAndTransfer(std::move(module), {}); LiteralTestUtil::ExpectR3EqualArray3D(a, *result); } @@ -210,7 +197,7 @@ void CopyOpTest::TestCopyConstantLayoutR4( HloComputation::Builder builder(TestName()); - std::unique_ptr literal = LiteralUtil::CreateR4FromArray4D(a); + std::unique_ptr literal = Literal::CreateR4FromArray4D(a); HloInstruction* constant = builder.AddInstruction( HloInstruction::CreateConstant(std::move(literal))); @@ -220,18 +207,10 @@ void CopyOpTest::TestCopyConstantLayoutR4( std::unique_ptr computation = builder.Build(); - auto hlo_module = MakeUnique("test_module"); - auto config = MakeUnique(computation->ComputeProgramShape()); - *config->mutable_entry_computation_layout()->mutable_result_layout() = - ShapeLayout(ShapeUtil::MakeShapeWithLayout( - constant->shape().element_type(), - AsInt64Slice(constant->shape().dimensions()), ({ - std::vector p(permutation.rbegin(), permutation.rend()); - p; - }))); - hlo_module->AddEntryComputation(std::move(computation)); - std::unique_ptr result = - ExecuteAndTransfer(std::move(hlo_module), std::move(config), {}); + auto module = CreateNewModule(); + module->AddEntryComputation(std::move(computation)); + ForceResultLayout(module.get(), LayoutUtil::MakeLayout(permutation)); + std::unique_ptr result = ExecuteAndTransfer(std::move(module), {}); LiteralTestUtil::ExpectR4EqualArray4D(a, *result); } @@ -256,22 +235,21 @@ XLA_TEST_F(CopyOpTest, CopyConstantR4Layout0312_MultipleTilesPerLayer) { TestCopyConstantLayoutR4(2, 14, 5, 35, {0, 3, 1, 2}); } -} // namespace -} // namespace xla +using CopyOpClientTest = ClientLibraryTestBase; -int main(int argc, char** argv) { - std::vector flag_list; - xla::legacy_flags::AppendCpuCompilerFlags(&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; - } - testing::InitGoogleTest(&argc, argv); - if (argc > 1) { - LOG(ERROR) << "Unknown argument " << argv[1] << "\n" << usage; - return 2; - } - return RUN_ALL_TESTS(); +XLA_TEST_F(CopyOpClientTest, Copy0x0) { + Shape in_shape = ShapeUtil::MakeShapeWithLayout(F32, {0, 0}, {0, 1}); + Shape out_shape = ShapeUtil::MakeShapeWithLayout(F32, {0, 0}, {1, 0}); + auto empty = Literal::CreateFromShape(in_shape); + + ComputationBuilder builder(client_, TestName()); + auto param0 = builder.Parameter(0, in_shape, "input"); + auto input_data = client_->TransferToServer(*empty).ConsumeValueOrDie(); + + auto actual = ExecuteAndTransfer(&builder, {input_data.get()}, &out_shape) + .ConsumeValueOrDie(); + LiteralTestUtil::ExpectEqual(*empty, *actual); } + +} // namespace +} // namespace xla diff --git a/tensorflow/compiler/xla/tests/custom_call_test.cc b/tensorflow/compiler/xla/tests/custom_call_test.cc index dc54c9defec2394049c38781a8d02fc8bd05158a..342478bc744273be9deb8b750b5a6a47b7d9f91b 100644 --- a/tensorflow/compiler/xla/tests/custom_call_test.cc +++ b/tensorflow/compiler/xla/tests/custom_call_test.cc @@ -16,7 +16,6 @@ limitations under the License. #include #include -#include "tensorflow/compiler/xla/legacy_flags/cpu_compiler_flags.h" #include "tensorflow/compiler/xla/literal_util.h" #include "tensorflow/compiler/xla/ptr_util.h" #include "tensorflow/compiler/xla/service/hlo_computation.h" @@ -29,23 +28,22 @@ limitations under the License. #include "tensorflow/compiler/xla/tests/test_macros.h" #include "tensorflow/compiler/xla/xla_data.pb.h" #include "tensorflow/core/platform/dynamic_annotations.h" +#include "tensorflow/core/platform/macros.h" #include "tensorflow/core/platform/test.h" -extern "C" void __attribute__((visibility("default"))) -R0F32Add2(float* out, float** in) { + +extern "C" void TF_EXPORT R0F32Add2(float* out, float** in) { TF_ANNOTATE_MEMORY_IS_INITIALIZED(in, sizeof(float*)); *out = **in + 2.0f; } -extern "C" void __attribute__((visibility("default"))) -R2F32ReduceSum(float* out, float** in) { +extern "C" void TF_EXPORT R2F32ReduceSum(float* out, float** in) { TF_ANNOTATE_MEMORY_IS_INITIALIZED(in, sizeof(float) * 4); float* array = in[0]; *out = array[0] + array[1] + array[2] + array[3]; } -extern "C" void __attribute__((visibility("default"))) -Add1ToValues(float* out, float** in) { +extern "C" void TF_EXPORT Add1ToValues(float* out, float** in) { TF_ANNOTATE_MEMORY_IS_INITIALIZED(in, sizeof(float) * 4); float* array = in[0]; out[0] = array[0] + 1; @@ -64,23 +62,22 @@ class CustomCallTest : public HloTestBase { }; XLA_TEST_F(CustomCallTest, DISABLED_ON_GPU(CustomCallR0F32Add2)) { - auto hlo_module = MakeUnique("test_module"); + auto module = CreateNewModule(); auto builder = HloComputation::Builder(TestName()); auto constant = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(42.0f))); + HloInstruction::CreateConstant(Literal::CreateR0(42.0f))); builder.AddInstruction( HloInstruction::CreateCustomCall(r0f32_, {constant}, "R0F32Add2")); - hlo_module->AddEntryComputation(builder.Build()); + module->AddEntryComputation(builder.Build()); - std::unique_ptr result = - ExecuteAndTransfer(std::move(hlo_module), {}); + std::unique_ptr result = ExecuteAndTransfer(std::move(module), {}); LiteralTestUtil::ExpectR0Near(44.0f, *result, error_spec_); } XLA_TEST_F(CustomCallTest, DISABLED_ON_GPU(CustomCallR2F32Reduce)) { - auto hlo_module = MakeUnique("test_module"); + auto module = CreateNewModule(); auto builder = HloComputation::Builder(TestName()); Array2D array(2, 2); @@ -90,24 +87,23 @@ XLA_TEST_F(CustomCallTest, DISABLED_ON_GPU(CustomCallR2F32Reduce)) { array(1, 1) = 4.0f; auto constant = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR2FromArray2D(array))); + HloInstruction::CreateConstant(Literal::CreateR2FromArray2D(array))); builder.AddInstruction( HloInstruction::CreateCustomCall(r0f32_, {constant}, "R2F32ReduceSum")); - hlo_module->AddEntryComputation(builder.Build()); + module->AddEntryComputation(builder.Build()); - std::unique_ptr result = - ExecuteAndTransfer(std::move(hlo_module), {}); + std::unique_ptr result = ExecuteAndTransfer(std::move(module), {}); LiteralTestUtil::ExpectR0Near(10.0f, *result, error_spec_); } XLA_TEST_F(CustomCallTest, DISABLED_ON_GPU(CustomCall_UsedInOtherComputations)) { - auto hlo_module = MakeUnique("test_module"); + auto module = CreateNewModule(); auto b = HloComputation::Builder(TestName()); auto input = b.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR2FromArray2D( + HloInstruction::CreateConstant(Literal::CreateR2FromArray2D( Array2D{{1.0f, 2.0f}, {3.0f, 4.0f}}))); auto incremented = b.AddInstruction(HloInstruction::CreateCustomCall( ShapeUtil::MakeShape(F32, {1, 2, 2}), {input}, "Add1ToValues")); @@ -119,30 +115,12 @@ XLA_TEST_F(CustomCallTest, HloInstruction::CreateConcatenate(ShapeUtil::MakeShape(F32, {2, 2, 2}), {incremented, incremented_again}, 0)); - hlo_module->AddEntryComputation(b.Build()); + module->AddEntryComputation(b.Build()); - std::unique_ptr result = - ExecuteAndTransfer(std::move(hlo_module), {}); + std::unique_ptr result = ExecuteAndTransfer(std::move(module), {}); LiteralTestUtil::ExpectR3EqualArray3D( Array3D{{{2, 3}, {4, 5}}, {{3, 4}, {5, 6}}}, *result); } } // namespace } // namespace xla - -int main(int argc, char** argv) { - std::vector flag_list; - xla::legacy_flags::AppendCpuCompilerFlags(&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; - } - testing::InitGoogleTest(&argc, argv); - if (argc > 1) { - LOG(ERROR) << "Unknown argument " << argv[1] << "\n" << usage; - return 2; - } - return RUN_ALL_TESTS(); -} diff --git a/tensorflow/compiler/xla/tests/deallocation_test.cc b/tensorflow/compiler/xla/tests/deallocation_test.cc index 528efd2942b0ebbba16faba2a0543a2694cd5c2a..fe5621e8dc209d6113e74030444c198716d355dc 100644 --- a/tensorflow/compiler/xla/tests/deallocation_test.cc +++ b/tensorflow/compiler/xla/tests/deallocation_test.cc @@ -19,17 +19,18 @@ limitations under the License. #include "tensorflow/compiler/xla/client/computation_builder.h" #include "tensorflow/compiler/xla/client/global_data.h" #include "tensorflow/compiler/xla/client/local_client.h" -#include "tensorflow/compiler/xla/legacy_flags/cpu_compiler_flags.h" #include "tensorflow/compiler/xla/statusor.h" +#include "tensorflow/compiler/xla/test.h" #include "tensorflow/compiler/xla/test_helpers.h" #include "tensorflow/compiler/xla/tests/client_library_test_base.h" #include "tensorflow/compiler/xla/tests/test_macros.h" #include "tensorflow/core/lib/gtl/array_slice.h" -#include "tensorflow/core/platform/test.h" namespace xla { namespace { +using ::testing::HasSubstr; + class DeallocationTest : public ClientLibraryTestBase { protected: // Build and execute the given computation then verify the results can be @@ -39,7 +40,8 @@ class DeallocationTest : public ClientLibraryTestBase { tensorflow::gtl::ArraySlice arguments) { Computation computation = builder->Build().ConsumeValueOrDie(); auto global_data = - client_->Execute(computation, arguments).ConsumeValueOrDie(); + client_->Execute(computation, arguments, &execution_options_) + .ConsumeValueOrDie(); TF_CHECK_OK(client_->Transfer(*global_data).status()); return global_data; } @@ -50,7 +52,7 @@ TEST_F(DeallocationTest, DeallocateScalar) { builder.ConstantR0(42.0); auto global_data = ExecuteAndCheckTransfer(&builder, {}); - // A result can be transfered an arbitrary number of times. Add an extra + // A result can be transferred an arbitrary number of times. Add an extra // transfer here so we're not just testing that a second call to Transfer // fails. ASSERT_IS_OK(client_->Transfer(*global_data).status()); @@ -59,8 +61,8 @@ TEST_F(DeallocationTest, DeallocateScalar) { auto transfer_status = client_->Transfer(*global_data); ASSERT_FALSE(transfer_status.ok()); - ASSERT_MATCH(transfer_status.status().error_message(), - testing::HasSubstr("was previously deallocated")); + ASSERT_THAT(transfer_status.status().error_message(), + HasSubstr("was previously deallocated")); } TEST_F(DeallocationTest, DeallocateVector) { @@ -72,8 +74,8 @@ TEST_F(DeallocationTest, DeallocateVector) { auto transfer_status = client_->Transfer(*global_data); ASSERT_FALSE(transfer_status.ok()); - ASSERT_MATCH(transfer_status.status().error_message(), - testing::HasSubstr("was previously deallocated")); + ASSERT_THAT(transfer_status.status().error_message(), + HasSubstr("was previously deallocated")); } TEST_F(DeallocationTest, DeallocateEmptyVector) { @@ -85,8 +87,8 @@ TEST_F(DeallocationTest, DeallocateEmptyVector) { auto transfer_status = client_->Transfer(*global_data); ASSERT_FALSE(transfer_status.ok()); - ASSERT_MATCH(transfer_status.status().error_message(), - testing::HasSubstr("was previously deallocated")); + ASSERT_THAT(transfer_status.status().error_message(), + HasSubstr("was previously deallocated")); } XLA_TEST_F(DeallocationTest, DeallocateTuple) { @@ -99,8 +101,8 @@ XLA_TEST_F(DeallocationTest, DeallocateTuple) { auto transfer_status = client_->Transfer(*global_data); ASSERT_FALSE(transfer_status.ok()); - ASSERT_MATCH(transfer_status.status().error_message(), - testing::HasSubstr("was previously deallocated")); + ASSERT_THAT(transfer_status.status().error_message(), + HasSubstr("was previously deallocated")); } XLA_TEST_F(DeallocationTest, DeallocateTupleWithRepeatedElements) { @@ -114,8 +116,8 @@ XLA_TEST_F(DeallocationTest, DeallocateTupleWithRepeatedElements) { auto transfer_status = client_->Transfer(*global_data); ASSERT_FALSE(transfer_status.ok()); - ASSERT_MATCH(transfer_status.status().error_message(), - testing::HasSubstr("was previously deallocated")); + ASSERT_THAT(transfer_status.status().error_message(), + HasSubstr("was previously deallocated")); } XLA_TEST_F(DeallocationTest, DeallocateNestedTuple) { @@ -130,26 +132,9 @@ XLA_TEST_F(DeallocationTest, DeallocateNestedTuple) { auto transfer_status = client_->Transfer(*global_data); ASSERT_FALSE(transfer_status.ok()); - ASSERT_MATCH(transfer_status.status().error_message(), - testing::HasSubstr("was previously deallocated")); + ASSERT_THAT(transfer_status.status().error_message(), + HasSubstr("was previously deallocated")); } } // namespace } // namespace xla - -int main(int argc, char** argv) { - std::vector flag_list; - xla::legacy_flags::AppendCpuCompilerFlags(&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; - } - testing::InitGoogleTest(&argc, argv); - if (argc > 1) { - LOG(ERROR) << "Unknown argument " << argv[1] << "\n" << usage; - return 2; - } - return RUN_ALL_TESTS(); -} diff --git a/tensorflow/compiler/xla/tests/deconstruct_tuple_test.cc b/tensorflow/compiler/xla/tests/deconstruct_tuple_test.cc index 57a7c61b141f3e8c5cf3ecc7e34043a79129c01b..032c06cd3c9f872f57674d3d7b5adc201c91ea77 100644 --- a/tensorflow/compiler/xla/tests/deconstruct_tuple_test.cc +++ b/tensorflow/compiler/xla/tests/deconstruct_tuple_test.cc @@ -20,10 +20,10 @@ limitations under the License. #include "tensorflow/compiler/xla/client/computation_builder.h" #include "tensorflow/compiler/xla/client/global_data.h" #include "tensorflow/compiler/xla/client/local_client.h" -#include "tensorflow/compiler/xla/legacy_flags/cpu_compiler_flags.h" #include "tensorflow/compiler/xla/literal_util.h" #include "tensorflow/compiler/xla/shape_util.h" #include "tensorflow/compiler/xla/statusor.h" +#include "tensorflow/compiler/xla/test.h" #include "tensorflow/compiler/xla/test_helpers.h" #include "tensorflow/compiler/xla/tests/client_library_test_base.h" #include "tensorflow/compiler/xla/tests/test_macros.h" @@ -34,6 +34,9 @@ limitations under the License. namespace xla { namespace { +using ::testing::ContainsRegex; +using ::testing::HasSubstr; + class DeconstructTupleTest : public ClientLibraryTestBase { protected: // Build and execute the given computation then verify the results can be @@ -43,7 +46,8 @@ class DeconstructTupleTest : public ClientLibraryTestBase { tensorflow::gtl::ArraySlice arguments) { Computation computation = builder->Build().ConsumeValueOrDie(); auto global_data = - client_->Execute(computation, arguments).ConsumeValueOrDie(); + client_->Execute(computation, arguments, &execution_options_) + .ConsumeValueOrDie(); TF_CHECK_OK(client_->Transfer(*global_data).status()); return global_data; } @@ -61,11 +65,11 @@ TEST_F(DeconstructTupleTest, DeconstructTuple) { // Try copying the elements back and comparing it auto handles = result_status.ConsumeValueOrDie(); - std::vector copy(4); - ASSERT_IS_OK(client_->TransferInProcess(*handles[0], ©[0])); - EXPECT_MATCH(copy, testing::VectorMatcher({1.0, 2.0, 3.0, 4.0})); - ASSERT_IS_OK(client_->TransferInProcess(*handles[1], ©[0])); - EXPECT_MATCH(copy, testing::VectorMatcher({2.0, 4.0, 6.0, 8.0})); + std::unique_ptr literal; + TF_ASSERT_OK_AND_ASSIGN(literal, client_->Transfer(*handles[0])); + LiteralTestUtil::ExpectR1Equal({1.0, 2.0, 3.0, 4.0}, *literal); + TF_ASSERT_OK_AND_ASSIGN(literal, client_->Transfer(*handles[1])); + LiteralTestUtil::ExpectR1Equal({2.0, 4.0, 6.0, 8.0}, *literal); } TEST_F(DeconstructTupleTest, DeconstructTupleTwice) { @@ -82,19 +86,20 @@ TEST_F(DeconstructTupleTest, DeconstructTupleTwice) { auto handles1 = result_status1.ConsumeValueOrDie(); auto handles2 = result_status2.ConsumeValueOrDie(); - std::vector copy(4); - ASSERT_IS_OK(client_->TransferInProcess(*handles1[0], ©[0])); - EXPECT_MATCH(copy, testing::VectorMatcher({1.0, 2.0, 3.0, 4.0})); - ASSERT_IS_OK(client_->TransferInProcess(*handles1[1], ©[0])); - EXPECT_MATCH(copy, testing::VectorMatcher({2.0, 4.0, 6.0, 8.0})); + std::unique_ptr literal; + TF_ASSERT_OK_AND_ASSIGN(literal, client_->Transfer(*handles1[0])); + LiteralTestUtil::ExpectR1Equal({1.0, 2.0, 3.0, 4.0}, *literal); + TF_ASSERT_OK_AND_ASSIGN(literal, client_->Transfer(*handles1[1])); + LiteralTestUtil::ExpectR1Equal({2.0, 4.0, 6.0, 8.0}, *literal); + handles1[0].reset(); handles1[1].reset(); - ASSERT_IS_OK(client_->TransferInProcess(*handles2[0], ©[0])); - EXPECT_MATCH(copy, testing::VectorMatcher({1.0, 2.0, 3.0, 4.0})); - ASSERT_IS_OK(client_->TransferInProcess(*handles2[1], ©[0])); - EXPECT_MATCH(copy, testing::VectorMatcher({2.0, 4.0, 6.0, 8.0})); + TF_ASSERT_OK_AND_ASSIGN(literal, client_->Transfer(*handles2[0])); + LiteralTestUtil::ExpectR1Equal({1.0, 2.0, 3.0, 4.0}, *literal); + TF_ASSERT_OK_AND_ASSIGN(literal, client_->Transfer(*handles2[1])); + LiteralTestUtil::ExpectR1Equal({2.0, 4.0, 6.0, 8.0}, *literal); } XLA_TEST_F(DeconstructTupleTest, DeconstructTupleRepeatedElement) { @@ -112,15 +117,15 @@ XLA_TEST_F(DeconstructTupleTest, DeconstructTupleRepeatedElement) { // the same as handle[3] and handle[1] should be the same as handle[2]. auto handles = result_status.ConsumeValueOrDie(); - std::vector copy(4); - ASSERT_IS_OK(client_->TransferInProcess(*handles[0], ©[0])); - EXPECT_MATCH(copy, testing::VectorMatcher({1.0, 2.0, 3.0, 4.0})); - ASSERT_IS_OK(client_->TransferInProcess(*handles[1], ©[0])); - EXPECT_MATCH(copy, testing::VectorMatcher({2.0, 4.0, 6.0, 8.0})); - ASSERT_IS_OK(client_->TransferInProcess(*handles[2], ©[0])); - EXPECT_MATCH(copy, testing::VectorMatcher({2.0, 4.0, 6.0, 8.0})); - ASSERT_IS_OK(client_->TransferInProcess(*handles[3], ©[0])); - EXPECT_MATCH(copy, testing::VectorMatcher({1.0, 2.0, 3.0, 4.0})); + std::unique_ptr literal; + TF_ASSERT_OK_AND_ASSIGN(literal, client_->Transfer(*handles[0])); + LiteralTestUtil::ExpectR1Equal({1.0, 2.0, 3.0, 4.0}, *literal); + TF_ASSERT_OK_AND_ASSIGN(literal, client_->Transfer(*handles[1])); + LiteralTestUtil::ExpectR1Equal({2.0, 4.0, 6.0, 8.0}, *literal); + TF_ASSERT_OK_AND_ASSIGN(literal, client_->Transfer(*handles[2])); + LiteralTestUtil::ExpectR1Equal({2.0, 4.0, 6.0, 8.0}, *literal); + TF_ASSERT_OK_AND_ASSIGN(literal, client_->Transfer(*handles[3])); + LiteralTestUtil::ExpectR1Equal({1.0, 2.0, 3.0, 4.0}, *literal); } TEST_F(DeconstructTupleTest, DeconstructTupleThenDeallocate) { @@ -138,19 +143,19 @@ TEST_F(DeconstructTupleTest, DeconstructTupleThenDeallocate) { // should not have been deallocated because of reference counting. global_data.reset(); - std::vector copy(4); - ASSERT_IS_OK(client_->TransferInProcess(*handles[0], ©[0])); - EXPECT_MATCH(copy, testing::VectorMatcher({1.0, 2.0, 3.0, 4.0})); - ASSERT_IS_OK(client_->TransferInProcess(*handles[1], ©[0])); - EXPECT_MATCH(copy, testing::VectorMatcher({2.0, 4.0, 6.0, 8.0})); - ASSERT_IS_OK(client_->TransferInProcess(*handles[2], ©[0])); - EXPECT_MATCH(copy, testing::VectorMatcher({1.0, 2.0, 3.0, 4.0})); + std::unique_ptr literal; + TF_ASSERT_OK_AND_ASSIGN(literal, client_->Transfer(*handles[0])); + LiteralTestUtil::ExpectR1Equal({1.0, 2.0, 3.0, 4.0}, *literal); + TF_ASSERT_OK_AND_ASSIGN(literal, client_->Transfer(*handles[1])); + LiteralTestUtil::ExpectR1Equal({2.0, 4.0, 6.0, 8.0}, *literal); + TF_ASSERT_OK_AND_ASSIGN(literal, client_->Transfer(*handles[2])); + LiteralTestUtil::ExpectR1Equal({1.0, 2.0, 3.0, 4.0}, *literal); /// Try deallocating one of the repeated elements, then copy handles[0].reset(); - ASSERT_IS_OK(client_->TransferInProcess(*handles[2], ©[0])); - EXPECT_MATCH(copy, testing::VectorMatcher({1.0, 2.0, 3.0, 4.0})); + TF_ASSERT_OK_AND_ASSIGN(literal, client_->Transfer(*handles[2])); + LiteralTestUtil::ExpectR1Equal({1.0, 2.0, 3.0, 4.0}, *literal); } TEST_F(DeconstructTupleTest, DeconstructNonTuple) { @@ -160,14 +165,14 @@ TEST_F(DeconstructTupleTest, DeconstructNonTuple) { auto result_status = client_->DeconstructTuple(*global_data); EXPECT_FALSE(result_status.ok()); - EXPECT_MATCH(result_status.status().error_message(), - testing::ContainsRegex("global data handle .* is not a tuple")); + EXPECT_THAT(result_status.status().error_message(), + ContainsRegex("global data handle .* is not a tuple")); } XLA_TEST_F(DeconstructTupleTest, DeconstructTupleFromParam) { ComputationBuilder builder(client_, TestName()); std::unique_ptr param0_literal = - LiteralUtil::CreateR1({3.14f, -100.25f}); + Literal::CreateR1({3.14f, -100.25f}); std::unique_ptr param0_data = client_->TransferToServer(*param0_literal).ConsumeValueOrDie(); auto p = builder.Parameter(0, ShapeUtil::MakeShape(F32, {2}), "param0"); @@ -189,27 +194,9 @@ XLA_TEST_F(DeconstructTupleTest, DeconstructNestedTuple) { auto result_status = client_->DeconstructTuple(*global_data); EXPECT_FALSE(result_status.ok()); - EXPECT_MATCH( - result_status.status().error_message(), - testing::ContainsRegex("deconstructing nested tuples not yet supported")); + EXPECT_THAT(result_status.status().error_message(), + HasSubstr("deconstructing nested tuples not yet supported")); } } // namespace } // namespace xla - -int main(int argc, char** argv) { - std::vector flag_list; - xla::legacy_flags::AppendCpuCompilerFlags(&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; - } - testing::InitGoogleTest(&argc, argv); - if (argc > 1) { - LOG(ERROR) << "Unknown argument " << argv[1] << "\n" << usage; - return 2; - } - return RUN_ALL_TESTS(); -} diff --git a/tensorflow/compiler/xla/tests/deep_graph_test.cc b/tensorflow/compiler/xla/tests/deep_graph_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..1da7a96fe2388eabd647a72aac81bdf2ef5bb6c6 --- /dev/null +++ b/tensorflow/compiler/xla/tests/deep_graph_test.cc @@ -0,0 +1,38 @@ +/* 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/client_library_test_base.h" + +namespace xla { +namespace { +TEST_F(ClientLibraryTestBase, DeepGraph) { + // TODO(b/62624812): To trigger the stack overflow this test is + // intended to track, we need to set kDepth to 20000. + // Unfortunately, setting it that high causes the test to time out. + const int kDepth = 200; + ComputationBuilder b(client_, TestName()); + ComputationDataHandle x; + ComputationDataHandle y; + auto x_data = CreateR0Parameter(3, 0, "x", &b, &x); + auto y_data = CreateR0Parameter(1, 1, "y", &b, &y); + ComputationDataHandle z = x; + for (int i = 0; i < kDepth; ++i) { + z = b.Add(z, y); + } + ComputeAndCompareR0(&b, /*expected=*/kDepth + 3, + {x_data.get(), y_data.get()}); +} +} // namespace +} // namespace xla diff --git a/tensorflow/compiler/xla/tests/dot_operation_test.cc b/tensorflow/compiler/xla/tests/dot_operation_test.cc index 64d2af3c393aedb5b03d2c9ede8c88ef6a00db5e..224aa57899d04eb8309b2337bb8fc936a81d350f 100644 --- a/tensorflow/compiler/xla/tests/dot_operation_test.cc +++ b/tensorflow/compiler/xla/tests/dot_operation_test.cc @@ -20,9 +20,6 @@ limitations under the License. #include "tensorflow/compiler/xla/array3d.h" #include "tensorflow/compiler/xla/client/computation_builder.h" #include "tensorflow/compiler/xla/client/local_client.h" -#include "tensorflow/compiler/xla/legacy_flags/cpu_compiler_flags.h" -#include "tensorflow/compiler/xla/legacy_flags/cpu_runtime_flags.h" -#include "tensorflow/compiler/xla/legacy_flags/layout_util_flags.h" #include "tensorflow/compiler/xla/primitive_util.h" #include "tensorflow/compiler/xla/reference_util.h" #include "tensorflow/compiler/xla/shape_util.h" @@ -154,6 +151,27 @@ XLA_TEST_F(DotOperationTest, Dot_2x0_0x2) { error_spec_); } +XLA_TEST_F(DotOperationTest, FusedDot) { + ComputationBuilder builder(client_, TestName()); + auto param0 = builder.Parameter(0, ShapeUtil::MakeShape(F32, {2, 4}), "arg0"); + auto param1 = builder.Parameter(1, ShapeUtil::MakeShape(F32, {4, 1}), "arg1"); + auto exp0 = builder.Exp(param0); + auto result = builder.Dot(exp0, param1); + + auto lhs_handle = client_ + ->TransferToServer(*Literal::CreateR2( + {{1.0, 2.0, 3.0, 4.0}, {-1.0, -2.0, -3.0, -4.0}})) + .ConsumeValueOrDie(); + auto rhs_handle = client_ + ->TransferToServer(*Literal::CreateR2( + {{1.0}, {2.0}, {3.0}, {4.0}})) + .ConsumeValueOrDie(); + + ComputeAndCompareR2( + &builder, Array2D({{296.14560492846033}, {0.8611737683031964}}), + {lhs_handle.get(), rhs_handle.get()}, error_spec_); +} + template void DotOperationTest::TestSquareMatrixDot(bool lhs_row_major, bool rhs_row_major) { @@ -185,14 +203,14 @@ void DotOperationTest::TestMatrixDot(int M, int K, int N, bool lhs_row_major, bool rhs_row_major) { std::unique_ptr> lhs_data = MakeLinspaceArray2D(0.0, 1.0, M, K); - std::unique_ptr lhs_lit = LiteralUtil::CreateR2FromArray2DWithLayout( + std::unique_ptr lhs_lit = Literal::CreateR2FromArray2DWithLayout( *lhs_data, LayoutUtil::MakeLayout(MinorToMajorForIsRowMajor(lhs_row_major))); auto lhs_handle = client_->TransferToServer(*lhs_lit).ConsumeValueOrDie(); std::unique_ptr> rhs_data = MakeLinspaceArray2D(0.0, 1.0, K, N); - std::unique_ptr rhs_lit = LiteralUtil::CreateR2FromArray2DWithLayout( + std::unique_ptr rhs_lit = Literal::CreateR2FromArray2DWithLayout( *rhs_data, LayoutUtil::MakeLayout(MinorToMajorForIsRowMajor(rhs_row_major))); auto rhs_handle = client_->TransferToServer(*rhs_lit).ConsumeValueOrDie(); @@ -366,9 +384,9 @@ XLA_TEST_F(DotOperationTest, BatchMatMul) { std::vector out_slices; for (int i = 0; i < 4; ++i) { // Slice off individual matrices and reshape to 2D tensors. - auto x_slice = builder.Slice(x_flat, {i, 0, 0}, {i + 1, 2, 2}); + auto x_slice = builder.Slice(x_flat, {i, 0, 0}, {i + 1, 2, 2}, {1, 1, 1}); x_slice = builder.Reshape(x_slice, {0, 1, 2}, {2, 2}); - auto y_slice = builder.Slice(y_flat, {i, 0, 0}, {i + 1, 2, 2}); + auto y_slice = builder.Slice(y_flat, {i, 0, 0}, {i + 1, 2, 2}, {1, 1, 1}); y_slice = builder.Reshape(y_slice, {0, 1, 2}, {2, 2}); auto out = builder.Dot(x_slice, y_slice); @@ -379,12 +397,12 @@ XLA_TEST_F(DotOperationTest, BatchMatMul) { builder.Reshape(out_flat, {0, 1, 2}, {2, 2, 2, 2}); auto x_data = client_ - ->TransferToServer(*LiteralUtil::CreateR4( + ->TransferToServer(*Literal::CreateR4( {{{{1000, 100}, {10, 1}}, {{2000, 200}, {20, 2}}}, {{{3000, 300}, {30, 3}}, {{4000, 400}, {40, 4}}}})) .ConsumeValueOrDie(); auto y_data = client_ - ->TransferToServer(*LiteralUtil::CreateR4( + ->TransferToServer(*Literal::CreateR4( {{{{1, 2}, {3, 4}}, {{5, 6}, {7, 8}}}, {{{11, 22}, {33, 44}}, {{55, 66}, {77, 88}}}})) .ConsumeValueOrDie(); @@ -415,14 +433,14 @@ TEST_F(DotOperationTest, TransposeFolding) { auto lhs_handle = client_ ->TransferToServer( - *LiteralUtil::CreateR2FromArray2DWithLayout( + *Literal::CreateR2FromArray2DWithLayout( *lhs, LayoutUtil::MakeLayout( MinorToMajorForIsRowMajor(row_major)))) .ConsumeValueOrDie(); auto rhs_handle = client_ ->TransferToServer( - *LiteralUtil::CreateR2FromArray2DWithLayout( + *Literal::CreateR2FromArray2DWithLayout( *rhs, LayoutUtil::MakeLayout( MinorToMajorForIsRowMajor(row_major)))) .ConsumeValueOrDie(); @@ -456,22 +474,3 @@ TEST_F(DotOperationTest, TransposeFolding) { } // namespace } // namespace xla - -int main(int argc, char** argv) { - std::vector flag_list; - xla::legacy_flags::AppendLayoutUtilFlags(&flag_list); - xla::legacy_flags::AppendCpuRuntimeFlags(&flag_list); - xla::legacy_flags::AppendCpuCompilerFlags(&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; - } - testing::InitGoogleTest(&argc, argv); - if (argc > 1) { - LOG(ERROR) << "Unknown argument " << argv[1] << "\n" << usage; - return 2; - } - return RUN_ALL_TESTS(); -} diff --git a/tensorflow/compiler/xla/tests/dynamic_ops_test.cc b/tensorflow/compiler/xla/tests/dynamic_ops_test.cc index 1d1fb337075855372ae54ac3c7e9abf55a6c32f1..b32c9e160408d28ee679bd445db9a03aec86ffff 100644 --- a/tensorflow/compiler/xla/tests/dynamic_ops_test.cc +++ b/tensorflow/compiler/xla/tests/dynamic_ops_test.cc @@ -21,13 +21,13 @@ limitations under the License. #include "tensorflow/compiler/xla/client/computation.h" #include "tensorflow/compiler/xla/client/computation_builder.h" #include "tensorflow/compiler/xla/client/local_client.h" -#include "tensorflow/compiler/xla/legacy_flags/cpu_compiler_flags.h" #include "tensorflow/compiler/xla/reference_util.h" #include "tensorflow/compiler/xla/service/device_memory_allocator.h" #include "tensorflow/compiler/xla/service/local_service.h" #include "tensorflow/compiler/xla/service/platform_util.h" #include "tensorflow/compiler/xla/service/shaped_buffer.h" #include "tensorflow/compiler/xla/service/transfer_manager.h" +#include "tensorflow/compiler/xla/test_helpers.h" #include "tensorflow/compiler/xla/tests/client_library_test_base.h" #include "tensorflow/compiler/xla/tests/literal_test_util.h" #include "tensorflow/compiler/xla/tests/test_macros.h" @@ -43,278 +43,310 @@ namespace { class DynamicSliceTest : public ClientLibraryTestBase { protected: - template + template void TestR1() { // Slice at dimension start. - RunR1({0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0}, {0}, {5}, - {0.0, 1.0, 2.0, 3.0, 4.0}); + RunR1({0, 1, 2, 3, 4, 5, 6, 7}, {0}, {5}, {0, 1, 2, 3, 4}); // Slice in the middle. - RunR1({0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0}, {2}, {3}, - {2.0, 3.0, 4.0}); + RunR1({0, 1, 2, 3, 4, 5, 6, 7}, {2}, {3}, {2, 3, 4}); // Slice at dimension boundaries. - RunR1({0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0}, {5}, {3}, - {5.0, 6.0, 7.0}); + RunR1({0, 1, 2, 3, 4, 5, 6, 7}, {5}, {3}, {5, 6, 7}); // Slice at dimension boundaries, but with sizes that cause indices to wrap. - RunR1({0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0}, {6}, {4}, - {6.0, 7.0, 0.0, 1.0}); + RunR1({0, 1, 2, 3, 4, 5, 6, 7}, {6}, {4}, {6, 7, 0, 1}); + // Zero element slice. + RunR1({0, 1, 2, 3, 4, 5, 6, 7}, {2}, {0}, {}); } - template + template void TestR2() { // Slice at dimension start. - RunR2({{1.0f, 2.0f, 3.0f}, {4.0f, 5.0f, 6.0f}, {7.0f, 8.0f, 9.0f}}, - {0, 0}, {2, 2}, {{1.0f, 2.0f}, {4.0f, 5.0f}}); + RunR2({{1, 2, 3}, {4, 5, 6}, {7, 8, 9}}, {0, 0}, {2, 2}, + {{1, 2}, {4, 5}}); // Slice in the middle. - RunR2({{1.0f, 2.0f, 3.0f}, {4.0f, 5.0f, 6.0f}, {7.0f, 8.0f, 9.0f}}, - {1, 1}, {2, 1}, {{5.0f}, {8.0f}}); + RunR2({{1, 2, 3}, {4, 5, 6}, {7, 8, 9}}, {1, 1}, {2, 1}, + {{5}, {8}}); // Slice at dimension boundaries. - RunR2({{1.0f, 2.0f, 3.0f}, {4.0f, 5.0f, 6.0f}, {7.0f, 8.0f, 9.0f}}, - {1, 1}, {2, 1}, {{5.0f}, {8.0f}}); + RunR2({{1, 2, 3}, {4, 5, 6}, {7, 8, 9}}, {1, 1}, {2, 1}, + {{5}, {8}}); // Slice at dimension boundaries, but with sizes that cause indices to wrap. - RunR2({{1.0f, 2.0f, 3.0f}, {4.0f, 5.0f, 6.0f}, {7.0f, 8.0f, 9.0f}}, - {1, 1}, {3, 3}, - {{5.0f, 6.0f, 4.0f}, {8.0f, 9.0f, 7.0f}, {2.0f, 3.0f, 1.0f}}); + RunR2({{1, 2, 3}, {4, 5, 6}, {7, 8, 9}}, {1, 1}, {3, 3}, + {{5, 6, 4}, {8, 9, 7}, {2, 3, 1}}); + // Zero element slice: 2x0. + RunR2({{1, 2, 3}, {4, 5, 6}, {7, 8, 9}}, {0, 0}, {2, 0}, + {{}, {}}); + // Zero element slice: 0x2. + RunR2({{1, 2, 3}, {4, 5, 6}, {7, 8, 9}}, {0, 0}, {0, 2}, + Array2D(0, 2)); } - template + template void TestR3() { // R3 Shape: [2, 3, 2] // clang-format off // Slice at dimension start. - RunR3( - {{{1.0f, 2.0f}, {3.0f, 4.0f}, {5.0f, 6.0f}}, - {{7.0f, 8.0f}, {9.0f, 10.0f}, {11.0f, 12.0f}}}, - {0, 0, 0}, {2, 1, 2}, - {{{1.0f, 2.0f}}, {{7.0f, 8.0f}}}); + RunR3( + {{{1, 2}, {3, 4}, {5, 6}}, + {{7, 8}, {9, 10}, {11, 12}}}, + {0, 0, 0}, {2, 1, 2}, + {{{1, 2}}, {{7, 8}}}); // Slice in the middle. - RunR3( - {{{1.0f, 2.0f}, {3.0f, 4.0f}, {5.0f, 6.0f}}, - {{7.0f, 8.0f}, {9.0f, 10.0f}, {11.0f, 12.0f}}}, - {0, 1, 1}, {2, 2, 1}, - {{{4.0f}, {6.0f}}, {{10.0f}, {12.0f}}}); + RunR3( + {{{1, 2}, {3, 4}, {5, 6}}, + {{7, 8}, {9, 10}, {11, 12}}}, + {0, 1, 1}, {2, 2, 1}, + {{{4}, {6}}, {{10}, {12}}}); // Slice at dimension boundaries, but with sizes that cause indices to wrap. - RunR3( - {{{1.0f, 2.0f}, {3.0f, 4.0f}, {5.0f, 6.0f}}, - {{7.0f, 8.0f}, {9.0f, 10.0f}, {11.0f, 12.0f}}}, - {0, 2, 1}, {2, 2, 1}, - {{{6.0f}, {2.0f}}, {{12.0f}, {8.0f}}}); + RunR3( + {{{1, 2}, {3, 4}, {5, 6}}, + {{7, 8}, {9, 10}, {11, 12}}}, + {0, 2, 1}, {2, 1, 2}, + {{{6, 5}}, {{12, 11}}}); // clang-format on } - template - void RunR1(const std::vector& input_values, + template + void RunR1(tensorflow::gtl::ArraySlice input_values, const std::vector slice_starts, - const std::vector slice_sizes, - const std::vector& expected_values) { + const std::vector& slice_sizes, + tensorflow::gtl::ArraySlice expected_values) { ComputationBuilder builder(client_, TestName()); // Initialize and transfer dynamic slice start indices parameter. ComputationDataHandle starts; std::unique_ptr start_data = CreateR1Parameter( slice_starts, 0, "slice_starts", &builder, &starts); // Build dynamic slice computation. - auto input = builder.ConstantR1(input_values); + auto input = builder.ConstantR1(input_values); builder.DynamicSlice(input, starts, slice_sizes); // Run computation and compare against expected values. - ComputeAndCompareR1(&builder, expected_values, {start_data.get()}, - ErrorSpec(0.000001)); + ComputeAndCompareR1(&builder, expected_values, {start_data.get()}); } - template - void RunR2(const Array2D& input_values, + template + void RunR2(const Array2D& input_values, const std::vector slice_starts, - const std::vector slice_sizes, - const Array2D& expected_values) { + const std::vector& slice_sizes, + const Array2D& expected_values) { ComputationBuilder builder(client_, TestName()); // Initialize and transfer dynamic slice start indices parameter. ComputationDataHandle starts; std::unique_ptr start_data = CreateR1Parameter( slice_starts, 0, "slice_starts", &builder, &starts); // Build dynamic slice computation. - auto input = builder.ConstantR2FromArray2D(input_values); + auto input = builder.ConstantR2FromArray2D(input_values); builder.DynamicSlice(input, starts, slice_sizes); // Run computation and compare against expected values. - ComputeAndCompareR2(&builder, expected_values, {start_data.get()}, - ErrorSpec(0.000001)); + ComputeAndCompareR2(&builder, expected_values, {start_data.get()}); } - template - void RunR3(const Array3D& input_values, + template + void RunR3(const Array3D& input_values, const std::vector slice_starts, - const std::vector slice_sizes, - const Array3D& expected_values) { + const std::vector& slice_sizes, + const Array3D& expected_values) { ComputationBuilder builder(client_, TestName()); // Initialize and transfer dynamic slice start indices parameter. ComputationDataHandle starts; std::unique_ptr start_data = CreateR1Parameter( slice_starts, 0, "slice_starts", &builder, &starts); // Build dynamic slice computation. - auto input = builder.ConstantR3FromArray3D(input_values); + auto input = builder.ConstantR3FromArray3D(input_values); builder.DynamicSlice(input, starts, slice_sizes); // Run computation and compare against expected values. - ComputeAndCompareR3(&builder, expected_values, {start_data.get()}, - ErrorSpec(0.000001)); + ComputeAndCompareR3(&builder, expected_values, {start_data.get()}); } }; -XLA_TEST_F(DynamicSliceTest, Int32R1) { TestR1(); } +XLA_TEST_F(DynamicSliceTest, Int32R1) { TestR1(); } + +XLA_TEST_F(DynamicSliceTest, Int64R1) { TestR1(); } -XLA_TEST_F(DynamicSliceTest, Int64R1) { TestR1(); } +XLA_TEST_F(DynamicSliceTest, UInt64R1) { TestR1(); } -XLA_TEST_F(DynamicSliceTest, UInt64R1) { TestR1(); } +XLA_TEST_F(DynamicSliceTest, Int32R2) { TestR2(); } -XLA_TEST_F(DynamicSliceTest, Int32R2) { TestR2(); } +XLA_TEST_F(DynamicSliceTest, Int64R2) { TestR2(); } -XLA_TEST_F(DynamicSliceTest, Int64R2) { TestR2(); } +XLA_TEST_F(DynamicSliceTest, UInt64R2) { TestR2(); } -XLA_TEST_F(DynamicSliceTest, UInt64R2) { TestR2(); } +XLA_TEST_F(DynamicSliceTest, Int32R3) { TestR3(); } -XLA_TEST_F(DynamicSliceTest, Int32R3) { TestR3(); } +XLA_TEST_F(DynamicSliceTest, Int64R3) { TestR3(); } + +XLA_TEST_F(DynamicSliceTest, UInt64R3) { TestR3(); } + +XLA_TEST_F(DynamicSliceTest, Int32R1Pred) { + // Slice at dimension start. + RunR1({true, false, false, true, false, true, true, false}, {0}, + {5}, {true, false, false, true, false}); + // Slice in the middle. + RunR1({true, false, false, true, false, true, true, false}, {2}, + {3}, {false, true, false}); + // Slice at dimension boundaries. + RunR1({true, false, false, true, false, true, true, false}, {5}, + {3}, {true, true, false}); + // Zero element slice. + RunR1({true, false, false, true, false, true, true, false}, {2}, + {0}, {}); +} -XLA_TEST_F(DynamicSliceTest, Int64R3) { TestR3(); } +XLA_TEST_F(DynamicSliceTest, Int32R2Pred) { + // Slice at dimension start. + RunR2( + {{true, false, true}, {false, false, true}, {true, true, false}}, {0, 0}, + {2, 2}, {{true, false}, {false, false}}); + // Slice in the middle. + RunR2( + {{true, false, true}, {false, false, true}, {true, true, false}}, {1, 1}, + {2, 1}, {{false}, {true}}); + // Slice at dimension boundaries. + RunR2( + {{true, false, true}, {false, false, true}, {true, true, false}}, {1, 1}, + {2, 1}, {{false}, {true}}); + // Zero element slice: 2x0. + RunR2( + {{true, false, true}, {false, false, true}, {true, true, false}}, {0, 0}, + {2, 0}, {{}, {}}); + // Zero element slice: 0x2. + RunR2( + {{true, false, true}, {false, false, true}, {true, true, false}}, {0, 0}, + {0, 2}, Array2D(0, 2)); +} -XLA_TEST_F(DynamicSliceTest, UInt64R3) { TestR3(); } +XLA_TEST_F(DynamicSliceTest, Int32R3Pred) { + // R3 Shape: [2, 3, 2] + // clang-format off + + // Slice at dimension start. + RunR3( + {{{true, false}, {false, true}, {true, true}}, + {{false, true}, {true, false}, {false, false}}}, + {0, 0, 0}, {2, 1, 2}, + {{{true, false}}, {{false, true}}}); + + // Slice in the middle. + RunR3( + {{{true, false}, {false, true}, {true, true}}, + {{false, true}, {true, false}, {false, false}}}, + {0, 1, 1}, {2, 2, 1}, + {{{true}, {true}}, {{false}, {false}}}); + + // clang-format on +} class DynamicUpdateSliceTest : public ClientLibraryTestBase { protected: - template + template void TestR1() { - // clang-format off // Slice at dimension start. - RunR1({0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0}, - {8.0, 9.0, 10.0}, {0}, - {8.0, 9.0, 10.0, 3.0, 4.0, 5.0, 6.0, 7.0}); + RunR1({0, 1, 2, 3, 4, 5, 6, 7}, {8, 9, 10}, {0}, + {8, 9, 10, 3, 4, 5, 6, 7}); // Slice in the middle. - RunR1({0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0}, - {8.0, 9.0, 10.0}, {2}, - {0.0, 1.0, 8.0, 9.0, 10.0, 5.0, 6.0, 7.0}); + RunR1({0, 1, 2, 3, 4, 5, 6, 7}, {8, 9, 10}, {2}, + {0, 1, 8, 9, 10, 5, 6, 7}); // Slice at dimension boundaries. - RunR1({0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0}, - {8.0, 9.0, 10.0}, {5}, - {0.0, 1.0, 2.0, 3.0, 4.0, 8.0, 9.0, 10.0}); + RunR1({0, 1, 2, 3, 4, 5, 6, 7}, {8, 9, 10}, {5}, + {0, 1, 2, 3, 4, 8, 9, 10}); // Slice at dimension boundaries, but with sizes that cause indices to wrap. - RunR1({0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0}, - {8.0, 9.0, 10.0}, {6}, - {0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 8.0, 9.0}); - // clang-format on + RunR1({0, 1, 2, 3, 4, 5, 6, 7}, {8, 9, 10}, {6}, + {0, 1, 2, 3, 4, 5, 8, 9}); + // Zero-sized update. + RunR1({0, 1, 2, 3, 4, 5, 6, 7}, {}, {2}, + {0, 1, 2, 3, 4, 5, 6, 7}); } - template + template void TestR2() { - // clang-format off // Slice at dimension start. - RunR2( - {{1.0f, 2.0f, 3.0f}, {4.0f, 5.0f, 6.0f}, {7.0f, 8.0f, 9.0f}}, - {{10.0f, 11.0f}}, {0, 0}, - {{10.0f, 11.0f, 3.0f}, {4.0f, 5.0f, 6.0f}, {7.0f, 8.0f, 9.0f}}); + RunR2({{1, 2, 3}, {4, 5, 6}, {7, 8, 9}}, {{10, 11}}, {0, 0}, + {{10, 11, 3}, {4, 5, 6}, {7, 8, 9}}); // Slice in the middle. - RunR2( - {{1.0f, 2.0f, 3.0f}, {4.0f, 5.0f, 6.0f}, {7.0f, 8.0f, 9.0f}}, - {{10.0f, 11.0f}}, {1, 1}, - {{1.0f, 2.0f, 3.0f}, {4.0f, 10.0f, 11.0f}, {7.0f, 8.0f, 9.0f}}); + RunR2({{1, 2, 3}, {4, 5, 6}, {7, 8, 9}}, {{10, 11}}, {1, 1}, + {{1, 2, 3}, {4, 10, 11}, {7, 8, 9}}); // Slice at dimension boundaries. - RunR2( - {{1.0f, 2.0f, 3.0f}, {4.0f, 5.0f, 6.0f}, {7.0f, 8.0f, 9.0f}}, - {{10.0f, 11.0f}}, {2, 1}, - {{1.0f, 2.0f, 3.0f}, {4.0f, 5.0f, 6.0f}, {7.0f, 10.0f, 11.0f}}); + RunR2({{1, 2, 3}, {4, 5, 6}, {7, 8, 9}}, {{10, 11}}, {2, 1}, + {{1, 2, 3}, {4, 5, 6}, {7, 10, 11}}); // Slice at dimension boundaries, but with sizes that cause indices to wrap. - RunR2( - {{1.0f, 2.0f, 3.0f}, {4.0f, 5.0f, 6.0f}, {7.0f, 8.0f, 9.0f}}, - {{10.0f, 11.0f}}, {2, 2}, - {{1.0f, 2.0f, 3.0f}, {4.0f, 5.0f, 6.0f}, {7.0f, 8.0f, 10.0f}}); - // clang-format on + RunR2({{1, 2, 3}, {4, 5, 6}, {7, 8, 9}}, {{10, 11}}, {2, 2}, + {{1, 2, 3}, {4, 5, 6}, {7, 8, 10}}); + // Zero-sized update. + RunR2({{1, 2, 3}, {4, 5, 6}, {7, 8, 9}}, {{}}, {2, 1}, + {{1, 2, 3}, {4, 5, 6}, {7, 8, 9}}); } - template + template void TestR3() { // R3 Shape: [2, 3, 2] - // clang-format off // Slice at dimension start. - RunR3( - {{{1.0f, 2.0f}, {3.0f, 4.0f}, {5.0f, 6.0f}}, - {{7.0f, 8.0f}, {9.0f, 10.0f}, {11.0f, 12.0f}}}, - {{{13.0f, 14.0f}, {15.0f, 16.0f}}, - {{17.0f, 18.0f}, {19.0f, 20.0f}}}, - {0, 0, 0}, - {{{13.0f, 14.0f}, {15.0f, 16.0f}, {5.0f, 6.0f}}, - {{17.0f, 18.0f}, {19.0f, 20.0f}, {11.0f, 12.0f}}}); + RunR3( + {{{1, 2}, {3, 4}, {5, 6}}, {{7, 8}, {9, 10}, {11, 12}}}, + {{{13, 14}, {15, 16}}, {{17, 18}, {19, 20}}}, {0, 0, 0}, + {{{13, 14}, {15, 16}, {5, 6}}, {{17, 18}, {19, 20}, {11, 12}}}); // Slice in the middle. - RunR3( - {{{1.0f, 2.0f}, {3.0f, 4.0f}, {5.0f, 6.0f}}, - {{7.0f, 8.0f}, {9.0f, 10.0f}, {11.0f, 12.0f}}}, - {{{13.0f}, {15.0f}}}, - {1, 1, 1}, - {{{1.0f, 2.0f}, {3.0f, 4.0f}, {5.0f, 6.0f}}, - {{7.0f, 8.0f}, {9.0f, 13.0f}, {11.0f, 15.0f}}}); + RunR3( + {{{1, 2}, {3, 4}, {5, 6}}, {{7, 8}, {9, 10}, {11, 12}}}, {{{13}, {15}}}, + {1, 1, 1}, {{{1, 2}, {3, 4}, {5, 6}}, {{7, 8}, {9, 13}, {11, 15}}}); // Slice at dimension boundaries, but with sizes that cause indices to wrap. - RunR3( - {{{1.0f, 2.0f}, {3.0f, 4.0f}, {5.0f, 6.0f}}, - {{7.0f, 8.0f}, {9.0f, 10.0f}, {11.0f, 12.0f}}}, - {{{13.0f}, {15.0f}}}, - {1, 2, 1}, - {{{1.0f, 2.0f}, {3.0f, 4.0f}, {5.0f, 6.0f}}, - {{7.0f, 8.0f}, {9.0f, 10.0f}, {11.0f, 13.0f}}}); - // clang-format on + RunR3( + {{{1, 2}, {3, 4}, {5, 6}}, {{7, 8}, {9, 10}, {11, 12}}}, {{{13}, {15}}}, + {1, 2, 1}, {{{1, 2}, {3, 4}, {5, 6}}, {{7, 8}, {9, 10}, {11, 13}}}); } - template - void RunR1(const std::vector& input_values, - const std::vector& update_values, + template + void RunR1(tensorflow::gtl::ArraySlice input_values, + tensorflow::gtl::ArraySlice update_values, const std::vector slice_starts, - const std::vector& expected_values) { + tensorflow::gtl::ArraySlice expected_values) { ComputationBuilder builder(client_, TestName()); // Initialize and transfer dynamic slice start indices parameter. ComputationDataHandle starts; std::unique_ptr start_data = CreateR1Parameter( slice_starts, 0, "slice_starts", &builder, &starts); // Build dynamic slice computation. - auto input = builder.ConstantR1(input_values); - auto update = builder.ConstantR1(update_values); + auto input = builder.ConstantR1(input_values); + auto update = builder.ConstantR1(update_values); builder.DynamicUpdateSlice(input, update, starts); // Run computation and compare against expected values. - ComputeAndCompareR1(&builder, expected_values, {start_data.get()}, - ErrorSpec(0.000001)); + ComputeAndCompareR1(&builder, expected_values, {start_data.get()}); } - template - void RunR2(const Array2D& input_values, - const Array2D& update_values, + template + void RunR2(const Array2D& input_values, + const Array2D& update_values, const std::vector slice_starts, - const Array2D& expected_values) { + const Array2D& expected_values) { ComputationBuilder builder(client_, TestName()); // Initialize and transfer dynamic slice start indices parameter. ComputationDataHandle starts; std::unique_ptr start_data = CreateR1Parameter( slice_starts, 0, "slice_starts", &builder, &starts); // Build dynamic slice computation. - auto input = builder.ConstantR2FromArray2D(input_values); - auto update = builder.ConstantR2FromArray2D(update_values); + auto input = builder.ConstantR2FromArray2D(input_values); + auto update = builder.ConstantR2FromArray2D(update_values); builder.DynamicUpdateSlice(input, update, starts); // Run computation and compare against expected values. - ComputeAndCompareR2(&builder, expected_values, {start_data.get()}, - ErrorSpec(0.000001)); + ComputeAndCompareR2(&builder, expected_values, {start_data.get()}); } - template - void RunR3(const Array3D& input_values, - const Array3D& update_values, + template + void RunR3(const Array3D& input_values, + const Array3D& update_values, const std::vector slice_starts, - const Array3D& expected_values) { + const Array3D& expected_values) { ComputationBuilder builder(client_, TestName()); // Initialize and transfer dynamic slice start indices parameter. ComputationDataHandle starts; std::unique_ptr start_data = CreateR1Parameter( slice_starts, 0, "slice_starts", &builder, &starts); // Build dynamic slice computation. - auto input = builder.ConstantR3FromArray3D(input_values); - auto update = builder.ConstantR3FromArray3D(update_values); + auto input = builder.ConstantR3FromArray3D(input_values); + auto update = builder.ConstantR3FromArray3D(update_values); builder.DynamicUpdateSlice(input, update, starts); // Run computation and compare against expected values. - ComputeAndCompareR3(&builder, expected_values, {start_data.get()}, - ErrorSpec(0.000001)); + ComputeAndCompareR3(&builder, expected_values, {start_data.get()}); } void RunR3Contiguous(std::vector operand_shape, int32 index, @@ -370,28 +402,86 @@ class DynamicUpdateSliceTest : public ClientLibraryTestBase { template void DumpArray(const string& name, const Array3D values) { std::unique_ptr literal = - LiteralUtil::CreateR3FromArray3D(values); - LOG(INFO) << name << ":" << LiteralUtil::ToString(*literal); + Literal::CreateR3FromArray3D(values); + LOG(INFO) << name << ":" << literal->ToString(); } }; -XLA_TEST_F(DynamicUpdateSliceTest, Int32R1) { TestR1(); } +XLA_TEST_F(DynamicUpdateSliceTest, Int32R1) { TestR1(); } + +XLA_TEST_F(DynamicUpdateSliceTest, Int64R1) { TestR1(); } + +XLA_TEST_F(DynamicUpdateSliceTest, UInt64R1) { TestR1(); } + +XLA_TEST_F(DynamicUpdateSliceTest, Int32R2) { TestR2(); } -XLA_TEST_F(DynamicUpdateSliceTest, Int64R1) { TestR1(); } +XLA_TEST_F(DynamicUpdateSliceTest, Int64R2) { TestR2(); } -XLA_TEST_F(DynamicUpdateSliceTest, UInt64R1) { TestR1(); } +XLA_TEST_F(DynamicUpdateSliceTest, UInt64R2) { TestR2(); } -XLA_TEST_F(DynamicUpdateSliceTest, Int32R2) { TestR2(); } +XLA_TEST_F(DynamicUpdateSliceTest, Int32R3) { TestR3(); } -XLA_TEST_F(DynamicUpdateSliceTest, Int64R2) { TestR2(); } +XLA_TEST_F(DynamicUpdateSliceTest, Int64R3) { TestR3(); } -XLA_TEST_F(DynamicUpdateSliceTest, UInt64R2) { TestR2(); } +XLA_TEST_F(DynamicUpdateSliceTest, UInt64R3) { TestR3(); } -XLA_TEST_F(DynamicUpdateSliceTest, Int32R3) { TestR3(); } +XLA_TEST_F(DynamicUpdateSliceTest, Int32R1Pred) { + // Slice at dimension start. + RunR1({false, false, true, true, false, true, true, false}, + {true, true, false}, {0}, + {true, true, false, true, false, true, true, false}); + // Slice in the middle. + RunR1({false, false, true, true, false, true, true, false}, + {false, true, true}, {2}, + {false, false, false, true, true, true, true, false}); + // Slice at dimension boundaries. + RunR1({false, false, true, true, false, true, true, false}, + {false, true, true}, {5}, + {false, false, true, true, false, false, true, true}); + // Zero-sized update. + RunR1({false, false, true, true, false, true, true, false}, {}, + {2}, {false, false, true, true, false, true, true, false}); +} -XLA_TEST_F(DynamicUpdateSliceTest, Int64R3) { TestR3(); } +XLA_TEST_F(DynamicUpdateSliceTest, Int32R2Pred) { + // Slice at dimension start. + RunR2( + {{false, true, false}, {true, false, true}, {false, true, true}}, + {{true, false}}, {0, 0}, + {{true, false, false}, {true, false, true}, {false, true, true}}); + // Slice in the middle. + RunR2( + {{false, true, false}, {true, false, true}, {false, true, true}}, + {{true, false}}, {1, 1}, + {{false, true, false}, {true, true, false}, {false, true, true}}); + // Slice at dimension boundaries. + RunR2( + {{false, true, false}, {true, false, true}, {false, true, true}}, + {{true, false}}, {2, 1}, + {{false, true, false}, {true, false, true}, {false, true, false}}); + // Zero-sized update. + RunR2( + {{false, true, false}, {true, false, true}, {false, true, true}}, {{}}, + {2, 1}, {{false, true, false}, {true, false, true}, {false, true, true}}); +} -XLA_TEST_F(DynamicUpdateSliceTest, UInt64R3) { TestR3(); } +XLA_TEST_F(DynamicUpdateSliceTest, Int32R3Pred) { + // R3 Shape: [2, 3, 2] + // Slice at dimension start. + RunR3( + {{{true, false}, {false, true}, {true, true}}, + {{false, false}, {false, true}, {true, false}}}, + {{{false, true}, {true, false}}, {{true, true}, {false, true}}}, + {0, 0, 0}, + {{{false, true}, {true, false}, {true, true}}, + {{true, true}, {false, true}, {true, false}}}); + // Slice in the middle. + RunR3({{{true, false}, {false, true}, {true, true}}, + {{false, false}, {false, true}, {true, false}}}, + {{{false}, {true}}}, {1, 1, 1}, + {{{true, false}, {false, true}, {true, true}}, + {{false, false}, {false, false}, {true, true}}}); +} // Tests for simple R3 case where the update is contiguous (i.e. the minor // two dimensions are not sliced). @@ -451,7 +541,7 @@ void BM_DynamicSlice(int num_iters) { ComputationBuilder builder(client, "DynamicSlice"); // Create input as a constant: shape [1, 2, 3, 4] - auto input_literal = LiteralUtil::CreateR4( + auto input_literal = Literal::CreateR4( {{{{1, 2, 3, 4}, {5, 6, 7, 8}, {9, 10, 11, 12}}, {{13, 14, 15, 16}, {17, 18, 19, 20}, {21, 22, 23, 24}}}}); auto input = builder.ConstantLiteral(*input_literal); @@ -469,7 +559,7 @@ void BM_DynamicSlice(int num_iters) { &allocator, 0) .ConsumeValueOrDie(); - auto start_indices_literal = LiteralUtil::CreateR1({0, 1, 2, 3}); + auto start_indices_literal = Literal::CreateR1({0, 1, 2, 3}); ASSERT_IS_OK(transfer_manager->TransferLiteralToDevice( executors[device_ordinal], *start_indices_literal, buffer->mutable_buffer({}))); @@ -498,20 +588,3 @@ BENCHMARK(BM_DynamicSlice); } // namespace } // namespace xla - -int main(int argc, char** argv) { - std::vector flag_list; - xla::legacy_flags::AppendCpuCompilerFlags(&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; - } - testing::InitGoogleTest(&argc, argv); - if (argc > 1) { - LOG(ERROR) << "Unknown argument " << argv[1] << "\n" << usage; - return 2; - } - return RUN_ALL_TESTS(); -} diff --git a/tensorflow/compiler/xla/tests/filecheck.cc b/tensorflow/compiler/xla/tests/filecheck.cc new file mode 100644 index 0000000000000000000000000000000000000000..407b5f4ada517d54af2a44742376348a1625b9b9 --- /dev/null +++ b/tensorflow/compiler/xla/tests/filecheck.cc @@ -0,0 +1,77 @@ +/* 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/filecheck.h" + +#include + +#include "tensorflow/compiler/xla/util.h" +#include "tensorflow/core/lib/core/errors.h" +#include "tensorflow/core/lib/io/path.h" +#include "tensorflow/core/platform/env.h" +#include "tensorflow/core/platform/subprocess.h" + +namespace xla { + +StatusOr RunFileCheck(const string& input, const string& pattern) { + using tensorflow::io::JoinPath; + + // Generate an input file for the FileCheck pattern. + string pattern_path; + auto env = tensorflow::Env::Default(); + if (!env->LocalTempFilename(&pattern_path)) { + return tensorflow::errors::Internal("couldn't get a pattern file name"); + } + TF_RETURN_IF_ERROR(tensorflow::WriteStringToFile(env, pattern_path, pattern)); + + // Invoke FileCheck to check whether input matches `pattern`. + const char* file_check_path_suffix = "external/llvm/FileCheck"; + string file_check_path; + if (const char* test_srcdir = getenv("TEST_SRCDIR")) { + file_check_path = JoinPath(test_srcdir, file_check_path_suffix); + } else { + file_check_path = file_check_path_suffix; + } + + tensorflow::SubProcess file_check_process; + file_check_process.SetProgram(file_check_path, + {file_check_path, pattern_path}); + file_check_process.SetChannelAction(tensorflow::CHAN_STDIN, + tensorflow::ACTION_PIPE); + file_check_process.SetChannelAction(tensorflow::CHAN_STDERR, + tensorflow::ACTION_PIPE); + if (!file_check_process.Start()) { + return tensorflow::errors::Internal("couldn't start FileCheck"); + } + + string standard_error; + int exit_status = file_check_process.Communicate( + /*stdin_input=*/&input, /*stdout_output=*/nullptr, + /*stderr_output=*/&standard_error); + + // FileCheck returns 0 when the inputs match. If matching failed, log + // the error message generated by FileCheck and the inputs. + bool succeeded = (exit_status == 0); + if (!succeeded) { + VLOG(1) << "FileCheck error: " << standard_error; + VLOG(1) << "FileCheck input was:"; + XLA_VLOG_LINES(1, input); + VLOG(1) << "FileCheck pattern was:"; + XLA_VLOG_LINES(1, pattern); + } + return succeeded; +} + +} // namespace xla diff --git a/tensorflow/compiler/xla/tests/filecheck.h b/tensorflow/compiler/xla/tests/filecheck.h new file mode 100644 index 0000000000000000000000000000000000000000..599bf57ad327fe0ef3b4972395eb4e0c883f763b --- /dev/null +++ b/tensorflow/compiler/xla/tests/filecheck.h @@ -0,0 +1,32 @@ +/* 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 THIRD_PARTY_TENSORFLOW_COMPILER_XLA_TESTS_FILECHECK_H_ +#define THIRD_PARTY_TENSORFLOW_COMPILER_XLA_TESTS_FILECHECK_H_ + +#include + +#include "tensorflow/compiler/xla/statusor.h" + +namespace xla { + +// Runs FileCheck with the given pattern over given input string. Provided that +// FileCheck can execute, returns true if and only if FileCheck succeeded in +// matching the input. +StatusOr RunFileCheck(const string& input, const string& pattern); + +} // namespace xla + +#endif // THIRD_PARTY_TENSORFLOW_COMPILER_XLA_TESTS_FILECHECK_H_ diff --git a/tensorflow/compiler/xla/tests/floor_ceil_test.cc b/tensorflow/compiler/xla/tests/floor_ceil_test.cc index 8e300630858ab13c6960aee78bf54bd1566b5818..e75a41acacc3aaad770f8bba78b43d8bf99b911b 100644 --- a/tensorflow/compiler/xla/tests/floor_ceil_test.cc +++ b/tensorflow/compiler/xla/tests/floor_ceil_test.cc @@ -18,7 +18,6 @@ limitations under the License. #include "tensorflow/compiler/xla/client/computation_builder.h" #include "tensorflow/compiler/xla/client/local_client.h" -#include "tensorflow/compiler/xla/legacy_flags/cpu_compiler_flags.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" @@ -109,20 +108,3 @@ TEST_F(FloorCeilTest, R0Ceil) { } // namespace } // namespace xla - -int main(int argc, char** argv) { - std::vector flag_list; - xla::legacy_flags::AppendCpuCompilerFlags(&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; - } - testing::InitGoogleTest(&argc, argv); - if (argc > 1) { - LOG(ERROR) << "Unknown argument " << argv[1] << "\n" << usage; - return 2; - } - return RUN_ALL_TESTS(); -} diff --git a/tensorflow/compiler/xla/tests/fmax_test.cc b/tensorflow/compiler/xla/tests/fmax_test.cc index 2835038c90c49fd9d961e292327d09163cbb59de..f2aaf6621c1f0d7a7d1bc29b845859579d8e8d9d 100644 --- a/tensorflow/compiler/xla/tests/fmax_test.cc +++ b/tensorflow/compiler/xla/tests/fmax_test.cc @@ -17,7 +17,6 @@ limitations under the License. #include "tensorflow/compiler/xla/client/computation_builder.h" #include "tensorflow/compiler/xla/client/local_client.h" -#include "tensorflow/compiler/xla/legacy_flags/cpu_compiler_flags.h" #include "tensorflow/compiler/xla/tests/client_library_test_base.h" #include "tensorflow/compiler/xla/tests/literal_test_util.h" #include "tensorflow/core/platform/test.h" @@ -42,20 +41,3 @@ TEST_F(FmaxSimpleTest, FmaxTenValues) { } // namespace } // namespace xla - -int main(int argc, char** argv) { - std::vector flag_list; - xla::legacy_flags::AppendCpuCompilerFlags(&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; - } - testing::InitGoogleTest(&argc, argv); - if (argc > 1) { - LOG(ERROR) << "Unknown argument " << argv[1] << "\n" << usage; - return 2; - } - return RUN_ALL_TESTS(); -} diff --git a/tensorflow/compiler/xla/tests/fusion_test.cc b/tensorflow/compiler/xla/tests/fusion_test.cc index 7bddbfa894c82012306dde83df74cfd624c6ea48..2be409561ab3e23d9ea2e49aac381a90395380d0 100644 --- a/tensorflow/compiler/xla/tests/fusion_test.cc +++ b/tensorflow/compiler/xla/tests/fusion_test.cc @@ -20,7 +20,9 @@ limitations under the License. #include #include "tensorflow/compiler/xla/array2d.h" -#include "tensorflow/compiler/xla/legacy_flags/cpu_compiler_flags.h" +#include "tensorflow/compiler/xla/client/client_library.h" +#include "tensorflow/compiler/xla/client/computation.h" +#include "tensorflow/compiler/xla/client/computation_builder.h" #include "tensorflow/compiler/xla/literal_util.h" #include "tensorflow/compiler/xla/primitive_util.h" #include "tensorflow/compiler/xla/ptr_util.h" @@ -28,7 +30,9 @@ limitations under the License. #include "tensorflow/compiler/xla/service/hlo_instruction.h" #include "tensorflow/compiler/xla/service/hlo_module.h" #include "tensorflow/compiler/xla/service/hlo_opcode.h" +#include "tensorflow/compiler/xla/service/platform_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" @@ -36,11 +40,13 @@ limitations under the License. #include "tensorflow/core/lib/gtl/array_slice.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" #include "tensorflow/core/platform/types.h" using tensorflow::gtl::ArraySlice; +namespace se = ::perftools::gputools; + namespace xla { namespace { @@ -74,14 +80,14 @@ class FusionTest : public HloTestBase { } auto builder = HloComputation::Builder(TestName()); - auto hlo_module = MakeUnique(TestName()); + auto hlo_module = CreateNewModule(); auto prim_type = primitive_util::NativeToPrimitiveType(); HloInstruction* hlos[4]; for (int i = 0; i < Arity; ++i) { hlos[i + 1] = builder.AddInstruction(HloInstruction::CreateConstant( - LiteralUtil::CreateR2FromArray2D(operand_data[i]))); + Literal::CreateR2FromArray2D(operand_data[i]))); } auto answer_shape = ShapeUtil::MakeShape(prim_type, {test_width, test_height}); @@ -107,7 +113,7 @@ class FusionTest : public HloTestBase { ArraySlice(hlos, 0, Arity + 1), HloInstruction::FusionKind::kLoop); - auto expected = LiteralUtil::CreateR2FromArray2D(answer_data); + auto expected = Literal::CreateR2FromArray2D(answer_data); auto actual = ExecuteAndTransfer(std::move(hlo_module), {}); if (primitive_util::IsFloatingPointType(prim_type)) { LiteralTestUtil::ExpectNear(*expected, *actual, ErrorSpec(1e-4)); @@ -147,8 +153,8 @@ float FusionTest::ComputeElementwiseAnswer(HloOpcode opcode, } template <> -uint8 FusionTest::ComputeElementwiseAnswer(HloOpcode opcode, - ArraySlice xs) { +bool FusionTest::ComputeElementwiseAnswer(HloOpcode opcode, + ArraySlice xs) { switch (opcode) { case HloOpcode::kEq: return xs[0] == xs[1]; @@ -176,35 +182,34 @@ 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 = MakeUnique(TestName()); + auto hlo_module = CreateNewModule(); auto const0 = builder.AddInstruction(HloInstruction::CreateConstant( - LiteralUtil::CreateR2({{1.0}, {2.0}, {3.0}}))); + Literal::CreateR2({{1.0}, {2.0}, {3.0}}))); auto const1 = builder.AddInstruction(HloInstruction::CreateConstant( - LiteralUtil::CreateR2({{-1.0}, {-1.0}, {-1.0}}))); + Literal::CreateR2({{-1.0}, {-1.0}, {-1.0}}))); auto add2 = builder.AddInstruction(HloInstruction::CreateBinary( ShapeUtil::MakeShape(F32, {3, 1}), HloOpcode::kAdd, const0, const1)); auto reshape3 = builder.AddInstruction(HloInstruction::CreateTranspose( ShapeUtil::MakeShape(F32, {1, 3}), add2, {1, 0})); auto const4 = builder.AddInstruction(HloInstruction::CreateConstant( - LiteralUtil::CreateR2({{1.62, 2.72, 3.14}}))); + Literal::CreateR2({{1.62, 2.72, 3.14}}))); auto concat5 = builder.AddInstruction(HloInstruction::CreateConcatenate( ShapeUtil::MakeShape(F32, {2, 3}), {reshape3, const4}, 0)); auto const6 = builder.AddInstruction(HloInstruction::CreateConstant( - LiteralUtil::CreateR2({{1.0, 1.0, 1.0}, {0.0, 0.0, 0.0}}))); + Literal::CreateR2({{1.0, 1.0, 1.0}, {0.0, 0.0, 0.0}}))); auto negate7 = builder.AddInstruction(HloInstruction::CreateUnary( ShapeUtil::MakeShape(F32, {2, 3}), HloOpcode::kNegate, const6)); auto add8 = builder.AddInstruction(HloInstruction::CreateBinary( ShapeUtil::MakeShape(F32, {2, 3}), HloOpcode::kAdd, concat5, negate7)); auto const9 = builder.AddInstruction(HloInstruction::CreateConstant( - LiteralUtil::CreateR2({{0.5, 0.5, 0.5}, {0.5, 0.5, 0.5}}))); - auto const10 = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR2( - {{true, false, true}, {false, true, false}}))); + Literal::CreateR2({{0.5, 0.5, 0.5}, {0.5, 0.5, 0.5}}))); + auto const10 = builder.AddInstruction(HloInstruction::CreateConstant( + Literal::CreateR2({{true, false, true}, {false, true, false}}))); auto select11 = builder.AddInstruction( HloInstruction::CreateTernary(ShapeUtil::MakeShape(F32, {2, 3}), HloOpcode::kSelect, const10, add8, const9)); auto slice12 = builder.AddInstruction(HloInstruction::CreateSlice( - ShapeUtil::MakeShape(F32, {2, 1}), select11, {0, 1}, {2, 2})); + ShapeUtil::MakeShape(F32, {2, 1}), select11, {0, 1}, {2, 2}, {1, 1})); // CreateFusionInstruction needs the `instructions_to_fuse` argument in // reverse topological order, so the first element in `instructions_to_fuse` // must be the root. @@ -214,7 +219,7 @@ XLA_TEST_F(FusionTest, Test) { const4, reshape3, add2, const1, const0}, HloInstruction::FusionKind::kLoop); - LiteralTestUtil::ExpectNear(*LiteralUtil::CreateR2({{0.5}, {2.72}}), + LiteralTestUtil::ExpectNear(*Literal::CreateR2({{0.5}, {2.72}}), *ExecuteAndTransfer(std::move(hlo_module), {}), ErrorSpec(1e-4)); } @@ -224,13 +229,13 @@ 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 = MakeUnique(TestName()); + auto hlo_module = CreateNewModule(); auto const0 = builder.AddInstruction(HloInstruction::CreateConstant( - LiteralUtil::CreateR2({{1.0, 2.0, 3.0}}))); + Literal::CreateR2({{1.0, 2.0, 3.0}}))); auto copy1 = builder.AddInstruction(HloInstruction::CreateUnary( ShapeUtil::MakeShape(F32, {1, 3}), HloOpcode::kCopy, const0)); auto const2 = builder.AddInstruction(HloInstruction::CreateConstant( - LiteralUtil::CreateR2({{-2.0, -2.0, -2.0}}))); + Literal::CreateR2({{-2.0, -2.0, -2.0}}))); // add3 = copy1 + const2 = const0 + const2 = {1,2,3} + {-2,-2,-2} = {-1,0,+1} auto add3 = builder.AddInstruction(HloInstruction::CreateBinary( ShapeUtil::MakeShape(F32, {1, 3}), HloOpcode::kAdd, copy1, const2)); @@ -240,18 +245,18 @@ XLA_TEST_F(FusionTest, Parameter) { ->CreateFusionInstruction(/*instructions_to_fuse=*/{add3, const2}, HloInstruction::FusionKind::kLoop); - LiteralTestUtil::ExpectNear(*LiteralUtil::CreateR2({{-1.0, 0.0, 1.0}}), + LiteralTestUtil::ExpectNear(*Literal::CreateR2({{-1.0, 0.0, 1.0}}), *ExecuteAndTransfer(std::move(hlo_module), {}), ErrorSpec(1e-4)); } XLA_TEST_F(FusionTest, BroadcastIntoBinaryOp) { auto builder = HloComputation::Builder(TestName()); - auto hlo_module = MakeUnique(TestName()); + auto hlo_module = CreateNewModule(); auto const_vector = builder.AddInstruction(HloInstruction::CreateConstant( - LiteralUtil::CreateR1({1.0, 2.0, 3.0}))); + Literal::CreateR1({1.0, 2.0, 3.0}))); auto const_array = builder.AddInstruction(HloInstruction::CreateConstant( - LiteralUtil::CreateR2({{-1.0, -2.0, -4.0}, {10.0, 20.0, 30.0}}))); + Literal::CreateR2({{-1.0, -2.0, -4.0}, {10.0, 20.0, 30.0}}))); auto broadcast = builder.AddInstruction( HloInstruction::CreateBroadcast(const_array->shape(), const_vector, {1})); // add2 = broadcast(const_vector) + const_array @@ -265,153 +270,260 @@ XLA_TEST_F(FusionTest, BroadcastIntoBinaryOp) { HloInstruction::FusionKind::kLoop); LiteralTestUtil::ExpectNear( - *LiteralUtil::CreateR2({{0.0, 0.0, -1.0}, {11.0, 22.0, 33.0}}), + *Literal::CreateR2({{0.0, 0.0, -1.0}, {11.0, 22.0, 33.0}}), *ExecuteAndTransfer(std::move(hlo_module), {}), ErrorSpec(1e-4)); } XLA_TEST_F(FusionTest, ReshapeToScalar) { auto builder = HloComputation::Builder(TestName()); - auto hlo_module = MakeUnique(TestName()); + auto hlo_module = CreateNewModule(); auto single_element_array = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR2({{5}}))); + HloInstruction::CreateConstant(Literal::CreateR2({{5}}))); auto reshape = builder.AddInstruction(HloInstruction::CreateReshape( ShapeUtil::MakeShape(S32, {}), single_element_array)); hlo_module->AddEntryComputation(builder.Build()) ->CreateFusionInstruction(/*instructions_to_fuse=*/{reshape}, HloInstruction::FusionKind::kLoop); - LiteralTestUtil::ExpectEqual(*LiteralUtil::CreateR0(5), + LiteralTestUtil::ExpectEqual(*Literal::CreateR0(5), *ExecuteAndTransfer(std::move(hlo_module), {})); } XLA_TEST_F(FusionTest, Reshape_3by2_1by2by3) { auto builder = HloComputation::Builder(TestName()); - auto hlo_module = MakeUnique(TestName()); + auto hlo_module = CreateNewModule(); auto const0 = builder.AddInstruction(HloInstruction::CreateConstant( - LiteralUtil::CreateR2({{1, 2}, {3, 4}, {5, 6}}))); + Literal::CreateR2({{1, 2}, {3, 4}, {5, 6}}))); auto reshape1 = builder.AddInstruction(HloInstruction::CreateReshape( ShapeUtil::MakeShape(S32, {1, 2, 3}), const0)); hlo_module->AddEntryComputation(builder.Build()) ->CreateFusionInstruction(/*instructions_to_fuse=*/{reshape1}, HloInstruction::FusionKind::kLoop); LiteralTestUtil::ExpectEqual( - *LiteralUtil::CreateR3({{{1, 2, 3}, {4, 5, 6}}}), + *Literal::CreateR3({{{1, 2, 3}, {4, 5, 6}}}), *ExecuteAndTransfer(std::move(hlo_module), {})); } XLA_TEST_F(FusionTest, Reshape_1by2by3_3by2) { auto builder = HloComputation::Builder(TestName()); - auto hlo_module = MakeUnique(TestName()); + auto hlo_module = CreateNewModule(); auto const0 = builder.AddInstruction(HloInstruction::CreateConstant( - LiteralUtil::CreateR3({{{1, 2, 3}, {4, 5, 6}}}))); + Literal::CreateR3({{{1, 2, 3}, {4, 5, 6}}}))); auto reshape1 = builder.AddInstruction( HloInstruction::CreateReshape(ShapeUtil::MakeShape(S32, {3, 2}), const0)); hlo_module->AddEntryComputation(builder.Build()) ->CreateFusionInstruction(/*instructions_to_fuse=*/{reshape1}, HloInstruction::FusionKind::kLoop); LiteralTestUtil::ExpectEqual( - *LiteralUtil::CreateR2({{1, 2}, {3, 4}, {5, 6}}), + *Literal::CreateR2({{1, 2}, {3, 4}, {5, 6}}), *ExecuteAndTransfer(std::move(hlo_module), {})); } XLA_TEST_F(FusionTest, Reshape_1by1by1_) { auto builder = HloComputation::Builder(TestName()); - auto hlo_module = MakeUnique(TestName()); + auto hlo_module = CreateNewModule(); auto const0 = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR3({{{7}}}))); + HloInstruction::CreateConstant(Literal::CreateR3({{{7}}}))); auto reshape1 = builder.AddInstruction( HloInstruction::CreateReshape(ShapeUtil::MakeShape(S32, {}), const0)); hlo_module->AddEntryComputation(builder.Build()) ->CreateFusionInstruction(/*instructions_to_fuse=*/{reshape1}, HloInstruction::FusionKind::kLoop); - LiteralTestUtil::ExpectEqual(*LiteralUtil::CreateR0(7), + LiteralTestUtil::ExpectEqual(*Literal::CreateR0(7), *ExecuteAndTransfer(std::move(hlo_module), {})); } XLA_TEST_F(FusionTest, Reshape__1by1by1) { auto builder = HloComputation::Builder(TestName()); - auto hlo_module = MakeUnique(TestName()); + auto hlo_module = CreateNewModule(); auto const0 = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(7))); + HloInstruction::CreateConstant(Literal::CreateR0(7))); auto reshape1 = builder.AddInstruction(HloInstruction::CreateReshape( ShapeUtil::MakeShape(S32, {1, 1, 1}), const0)); hlo_module->AddEntryComputation(builder.Build()) ->CreateFusionInstruction(/*instructions_to_fuse=*/{reshape1}, HloInstruction::FusionKind::kLoop); - LiteralTestUtil::ExpectEqual(*LiteralUtil::CreateR3({{{7}}}), + LiteralTestUtil::ExpectEqual(*Literal::CreateR3({{{7}}}), *ExecuteAndTransfer(std::move(hlo_module), {})); } XLA_TEST_F(FusionTest, Reshape__) { auto builder = HloComputation::Builder(TestName()); - auto hlo_module = MakeUnique(TestName()); + auto hlo_module = CreateNewModule(); auto const0 = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(7))); + HloInstruction::CreateConstant(Literal::CreateR0(7))); auto reshape1 = builder.AddInstruction( HloInstruction::CreateReshape(ShapeUtil::MakeShape(S32, {}), const0)); hlo_module->AddEntryComputation(builder.Build()) ->CreateFusionInstruction(/*instructions_to_fuse=*/{reshape1}, HloInstruction::FusionKind::kLoop); - LiteralTestUtil::ExpectEqual(*LiteralUtil::CreateR0(7), + LiteralTestUtil::ExpectEqual(*Literal::CreateR0(7), *ExecuteAndTransfer(std::move(hlo_module), {})); } XLA_TEST_F(FusionTest, Reshape_3by3_3by3) { auto builder = HloComputation::Builder(TestName()); - auto hlo_module = MakeUnique(TestName()); + auto hlo_module = CreateNewModule(); auto const0 = builder.AddInstruction(HloInstruction::CreateConstant( - LiteralUtil::CreateR2({{1, 2, 3}, {4, 5, 6}, {7, 8, 9}}))); + Literal::CreateR2({{1, 2, 3}, {4, 5, 6}, {7, 8, 9}}))); auto reshape1 = builder.AddInstruction( HloInstruction::CreateReshape(ShapeUtil::MakeShape(S32, {3, 3}), const0)); hlo_module->AddEntryComputation(builder.Build()) ->CreateFusionInstruction(/*instructions_to_fuse=*/{reshape1}, HloInstruction::FusionKind::kLoop); LiteralTestUtil::ExpectEqual( - *LiteralUtil::CreateR2({{1, 2, 3}, {4, 5, 6}, {7, 8, 9}}), + *Literal::CreateR2({{1, 2, 3}, {4, 5, 6}, {7, 8, 9}}), *ExecuteAndTransfer(std::move(hlo_module), {})); } XLA_TEST_F(FusionTest, Transpose_2by3) { auto builder = HloComputation::Builder(TestName()); - auto hlo_module = MakeUnique(TestName()); + auto hlo_module = CreateNewModule(); auto const0 = builder.AddInstruction(HloInstruction::CreateConstant( - LiteralUtil::CreateR2({{1, 2, 3}, {4, 5, 6}}))); + Literal::CreateR2({{1, 2, 3}, {4, 5, 6}}))); auto reshape1 = builder.AddInstruction(HloInstruction::CreateTranspose( ShapeUtil::MakeShape(S32, {3, 2}), const0, {1, 0})); hlo_module->AddEntryComputation(builder.Build()) ->CreateFusionInstruction(/*instructions_to_fuse=*/{reshape1}, HloInstruction::FusionKind::kLoop); LiteralTestUtil::ExpectEqual( - *LiteralUtil::CreateR2({{1, 4}, {2, 5}, {3, 6}}), + *Literal::CreateR2({{1, 4}, {2, 5}, {3, 6}}), *ExecuteAndTransfer(std::move(hlo_module), {})); } XLA_TEST_F(FusionTest, Transpose_3by3) { auto builder = HloComputation::Builder(TestName()); - auto hlo_module = MakeUnique(TestName()); + auto hlo_module = CreateNewModule(); auto const0 = builder.AddInstruction(HloInstruction::CreateConstant( - LiteralUtil::CreateR2({{1, 2, 3}, {4, 5, 6}, {7, 8, 9}}))); + Literal::CreateR2({{1, 2, 3}, {4, 5, 6}, {7, 8, 9}}))); auto reshape1 = builder.AddInstruction(HloInstruction::CreateTranspose( ShapeUtil::MakeShape(S32, {3, 3}), const0, {1, 0})); hlo_module->AddEntryComputation(builder.Build()) ->CreateFusionInstruction(/*instructions_to_fuse=*/{reshape1}, HloInstruction::FusionKind::kLoop); LiteralTestUtil::ExpectEqual( - *LiteralUtil::CreateR2({{1, 4, 7}, {2, 5, 8}, {3, 6, 9}}), + *Literal::CreateR2({{1, 4, 7}, {2, 5, 8}, {3, 6, 9}}), *ExecuteAndTransfer(std::move(hlo_module), {})); } XLA_TEST_F(FusionTest, Reverse) { auto builder = HloComputation::Builder(TestName()); - auto hlo_module = MakeUnique(TestName()); + auto hlo_module = CreateNewModule(); auto const0 = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR1({1, 2, 3}))); + HloInstruction::CreateConstant(Literal::CreateR1({1, 2, 3}))); auto reverse1 = builder.AddInstruction(HloInstruction::CreateReverse( ShapeUtil::MakeShape(S32, {3}), const0, {0})); hlo_module->AddEntryComputation(builder.Build()) ->CreateFusionInstruction(/*instructions_to_fuse=*/{reverse1}, HloInstruction::FusionKind::kLoop); - LiteralTestUtil::ExpectEqual(*LiteralUtil::CreateR1({3, 2, 1}), + LiteralTestUtil::ExpectEqual(*Literal::CreateR1({3, 2, 1}), + *ExecuteAndTransfer(std::move(hlo_module), {})); +} + +XLA_TEST_F(FusionTest, ReverseNegate) { + auto builder = HloComputation::Builder(TestName()); + auto hlo_module = CreateNewModule(); + auto const0 = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR1({1, 2, 3}))); + auto reverse1 = builder.AddInstruction(HloInstruction::CreateReverse( + ShapeUtil::MakeShape(S32, {3}), const0, {0})); + auto negate2 = builder.AddInstruction(HloInstruction::CreateUnary( + ShapeUtil::MakeShape(S32, {3}), HloOpcode::kNegate, reverse1)); + hlo_module->AddEntryComputation(builder.Build()) + ->CreateFusionInstruction(/*instructions_to_fuse=*/{negate2, reverse1}, + HloInstruction::FusionKind::kLoop); + + LiteralTestUtil::ExpectEqual(*Literal::CreateR1({-3, -2, -1}), + *ExecuteAndTransfer(std::move(hlo_module), {})); +} + +XLA_TEST_F(FusionTest, BroadcastNegate) { + auto builder = HloComputation::Builder(TestName()); + auto hlo_module = CreateNewModule(); + auto const0 = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(1))); + auto broadcast1 = builder.AddInstruction(HloInstruction::CreateBroadcast( + ShapeUtil::MakeShape(S32, {2}), const0, {})); + auto negate2 = builder.AddInstruction(HloInstruction::CreateUnary( + ShapeUtil::MakeShape(S32, {2}), HloOpcode::kNegate, broadcast1)); + hlo_module->AddEntryComputation(builder.Build()) + ->CreateFusionInstruction(/*instructions_to_fuse=*/{negate2, broadcast1}, + HloInstruction::FusionKind::kLoop); + + LiteralTestUtil::ExpectEqual(*Literal::CreateR1({-1, -1}), + *ExecuteAndTransfer(std::move(hlo_module), {})); +} + +XLA_TEST_F(FusionTest, SliceNegate) { + auto builder = HloComputation::Builder(TestName()); + auto hlo_module = CreateNewModule(); + auto const0 = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR1({1, 2, 3, 4}))); + auto slice1 = builder.AddInstruction(HloInstruction::CreateSlice( + ShapeUtil::MakeShape(S32, {2}), const0, {0}, {4}, {2})); + auto negate2 = builder.AddInstruction(HloInstruction::CreateUnary( + ShapeUtil::MakeShape(S32, {2}), HloOpcode::kNegate, slice1)); + hlo_module->AddEntryComputation(builder.Build()) + ->CreateFusionInstruction(/*instructions_to_fuse=*/{negate2, slice1}, + HloInstruction::FusionKind::kLoop); + + LiteralTestUtil::ExpectEqual(*Literal::CreateR1({-1, -3}), + *ExecuteAndTransfer(std::move(hlo_module), {})); +} + +XLA_TEST_F(FusionTest, DynamicSliceNegate) { + auto builder = HloComputation::Builder(TestName()); + auto hlo_module = CreateNewModule(); + auto const0 = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR1({1, 2, 3, 4}))); + auto const1 = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR1({1}))); + auto dynamic_slice2 = + builder.AddInstruction(HloInstruction::CreateDynamicSlice( + 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()) + ->CreateFusionInstruction( + /*instructions_to_fuse=*/{negate3, dynamic_slice2}, + HloInstruction::FusionKind::kLoop); + + LiteralTestUtil::ExpectEqual(*Literal::CreateR1({-2, -3}), + *ExecuteAndTransfer(std::move(hlo_module), {})); +} + +XLA_TEST_F(FusionTest, ReshapeNegate) { + auto builder = HloComputation::Builder(TestName()); + auto hlo_module = CreateNewModule(); + auto const0 = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR1({1, 2, 3, 4}))); + auto reshape1 = builder.AddInstruction( + HloInstruction::CreateReshape(ShapeUtil::MakeShape(S32, {2, 2}), const0)); + auto negate2 = builder.AddInstruction(HloInstruction::CreateUnary( + ShapeUtil::MakeShape(S32, {2, 2}), HloOpcode::kNegate, reshape1)); + hlo_module->AddEntryComputation(builder.Build()) + ->CreateFusionInstruction(/*instructions_to_fuse=*/{negate2, reshape1}, + HloInstruction::FusionKind::kLoop); + + LiteralTestUtil::ExpectEqual(*Literal::CreateR2({{-1, -2}, {-3, -4}}), + *ExecuteAndTransfer(std::move(hlo_module), {})); +} + +// TODO(b/64070202): Investigate failure. +XLA_TEST_F(FusionTest, DISABLED_ON_GPU(TransposeNegate)) { + auto builder = HloComputation::Builder(TestName()); + auto hlo_module = CreateNewModule(); + auto const0 = builder.AddInstruction(HloInstruction::CreateConstant( + Literal::CreateR2({{1, 2}, {3, 4}}))); + auto transpose1 = builder.AddInstruction(HloInstruction::CreateTranspose( + ShapeUtil::MakeShape(S32, {2, 2}), const0, {1, 0})); + auto negate2 = builder.AddInstruction(HloInstruction::CreateUnary( + ShapeUtil::MakeShape(S32, {2, 2}), HloOpcode::kNegate, transpose1)); + hlo_module->AddEntryComputation(builder.Build()) + ->CreateFusionInstruction(/*instructions_to_fuse=*/{negate2, transpose1}, + HloInstruction::FusionKind::kLoop); + + LiteralTestUtil::ExpectEqual(*Literal::CreateR2({{-1, -3}, {-2, -4}}), *ExecuteAndTransfer(std::move(hlo_module), {})); } @@ -427,13 +539,13 @@ std::unique_ptr MakeReduceTestComputation() { } XLA_TEST_F(FusionTest, DISABLED_ON_CPU(Reduce)) { - auto hlo_module = MakeUnique(TestName()); + auto hlo_module = CreateNewModule(); auto builder = HloComputation::Builder(TestName()); - auto const0 = builder.AddInstruction(HloInstruction::CreateConstant( - LiteralUtil::CreateR1({1, 2, 4, 8}))); + auto const0 = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR1({1, 2, 4, 8}))); auto const1 = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(0))); + HloInstruction::CreateConstant(Literal::CreateR0(0))); auto reduce2 = builder.AddInstruction(HloInstruction::CreateReduce( ShapeUtil::MakeShape(S32, {}), const0, const1, {0}, hlo_module->AddEmbeddedComputation(MakeReduceTestComputation()))); @@ -441,38 +553,38 @@ XLA_TEST_F(FusionTest, DISABLED_ON_CPU(Reduce)) { ->CreateFusionInstruction(/*instructions_to_fuse=*/{reduce2}, HloInstruction::FusionKind::kLoop); - LiteralTestUtil::ExpectEqual(*LiteralUtil::CreateR0(15), + LiteralTestUtil::ExpectEqual(*Literal::CreateR0(15), *ExecuteAndTransfer(std::move(hlo_module), {})); } XLA_TEST_F(FusionTest, DISABLED_ON_CPU(ReduceImplicitBroadcast)) { - auto hlo_module = MakeUnique(TestName()); + auto hlo_module = CreateNewModule(); auto builder = HloComputation::Builder(TestName()); - auto const0 = builder.AddInstruction(HloInstruction::CreateConstant( - LiteralUtil::CreateR1({1, 2, 4, 8}))); + auto const0 = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR1({1, 2, 4, 8}))); auto const1 = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(0))); + HloInstruction::CreateConstant(Literal::CreateR0(0))); auto reduce2 = builder.AddInstruction(HloInstruction::CreateReduce( ShapeUtil::MakeShape(S32, {}), const0, const1, {0}, hlo_module->AddEmbeddedComputation(MakeReduceTestComputation()))); auto negate3 = builder.AddInstruction(HloInstruction::CreateUnary( - ShapeUtil::MakeShape(S32, {1}), HloOpcode::kNegate, reduce2)); + ShapeUtil::MakeShape(S32, {}), HloOpcode::kNegate, reduce2)); hlo_module->AddEntryComputation(builder.Build()) ->CreateFusionInstruction(/*instructions_to_fuse=*/{negate3, reduce2}, HloInstruction::FusionKind::kLoop); - LiteralTestUtil::ExpectEqual(*LiteralUtil::CreateR1({-15}), + LiteralTestUtil::ExpectEqual(*Literal::CreateR0(-15), *ExecuteAndTransfer(std::move(hlo_module), {})); } XLA_TEST_F(FusionTest, DISABLED_ON_CPU(ReduceWindow)) { auto builder = HloComputation::Builder(TestName()); - auto hlo_module = MakeUnique(TestName()); + auto hlo_module = CreateNewModule(); auto const0 = builder.AddInstruction(HloInstruction::CreateConstant( - LiteralUtil::CreateR2({{2, 3, 5}, {7, 11, 13}, {17, 19, 23}}))); + Literal::CreateR2({{2, 3, 5}, {7, 11, 13}, {17, 19, 23}}))); auto const1 = builder.AddInstruction( - HloInstruction::CreateConstant(LiteralUtil::CreateR0(1))); + HloInstruction::CreateConstant(Literal::CreateR0(1))); Window window; ASSERT_TRUE( tensorflow::protobuf::TextFormat::ParseFromString("dimensions:{\n" @@ -512,10 +624,46 @@ XLA_TEST_F(FusionTest, DISABLED_ON_CPU(ReduceWindow)) { HloInstruction::FusionKind::kLoop); LiteralTestUtil::ExpectEqual( - *LiteralUtil::CreateR2({{462, 2145}, {24871, 62491}}), + *Literal::CreateR2({{462, 2145}, {24871, 62491}}), *ExecuteAndTransfer(std::move(hlo_module), {})); } +// When a constant (or other op) which has multiple users is imported +// into a fusion, it should remain shared, rather than being duplicated +// within the fusion. +XLA_TEST_F(FusionTest, SharedConstant) { + auto hlo_module = CreateNewModule(); + + auto builder = HloComputation::Builder(TestName()); + auto const0 = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR1({0}))); + auto const1 = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR1({2}))); + auto add1 = builder.AddInstruction(HloInstruction::CreateBinary( + ShapeUtil::MakeShape(S32, {1}), HloOpcode::kAdd, const1, const0)); + auto add2 = builder.AddInstruction(HloInstruction::CreateBinary( + ShapeUtil::MakeShape(S32, {1}), HloOpcode::kAdd, const1, add1)); + auto add3 = builder.AddInstruction(HloInstruction::CreateBinary( + ShapeUtil::MakeShape(S32, {1}), HloOpcode::kAdd, const1, add2)); + auto add4 = builder.AddInstruction(HloInstruction::CreateBinary( + ShapeUtil::MakeShape(S32, {1}), HloOpcode::kAdd, const1, add3)); + hlo_module->AddEntryComputation(builder.Build()) + ->CreateFusionInstruction( + {add4, add3, add2, add1, const1}, + HloInstruction::FusionKind::kLoop); + + HloComputation* entry_comp = hlo_module->entry_computation(); + + // entry computation contains the constant(0) and the fusion + EXPECT_EQ(entry_comp->instructions().size(), 2); + + // fused instruction contains the constant(2), the parameter, and 4 adds + EXPECT_EQ(entry_comp->root_instruction()->fused_instructions().size(), 6); + + LiteralTestUtil::ExpectEqual(*Literal::CreateR1({8}), + *ExecuteAndTransfer(std::move(hlo_module), {})); +} + XLA_TEST_F(FusionTest, Add2D) { TestElementwise2D(HloOpcode::kAdd); } XLA_TEST_F(FusionTest, Subtract2D) { @@ -542,48 +690,84 @@ XLA_TEST_F(FusionTest, Maximum2D) { TestElementwise2D(HloOpcode::kMaximum); } -XLA_TEST_F(FusionTest, Equal2D) { TestElementwise2D(HloOpcode::kEq); } +XLA_TEST_F(FusionTest, Equal2D) { TestElementwise2D(HloOpcode::kEq); } XLA_TEST_F(FusionTest, Inequal2D) { - TestElementwise2D(HloOpcode::kNe); + TestElementwise2D(HloOpcode::kNe); } XLA_TEST_F(FusionTest, Greater2D) { - TestElementwise2D(HloOpcode::kGt); + TestElementwise2D(HloOpcode::kGt); } -XLA_TEST_F(FusionTest, Lesser2D) { - TestElementwise2D(HloOpcode::kLt); -} +XLA_TEST_F(FusionTest, Lesser2D) { TestElementwise2D(HloOpcode::kLt); } XLA_TEST_F(FusionTest, GreaterOrEqual2D) { - TestElementwise2D(HloOpcode::kGe); + TestElementwise2D(HloOpcode::kGe); } XLA_TEST_F(FusionTest, LesserOrEqual2D) { - TestElementwise2D(HloOpcode::kLe); + TestElementwise2D(HloOpcode::kLe); } XLA_TEST_F(FusionTest, Clamp2D) { TestElementwise2D(HloOpcode::kClamp); } -} // namespace -} // namespace xla - -int main(int argc, char** argv) { - std::vector flag_list; - xla::legacy_flags::AppendCpuCompilerFlags(&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; +void BM_ParallelFusion(int num_iters) { + // Simple element-wise computation to benchmark parallel task partitioning. + tensorflow::testing::StopTiming(); + + se::Platform* platform = PlatformUtil::GetDefaultPlatform().ValueOrDie(); + auto executors = PlatformUtil::GetStreamExecutors(platform).ValueOrDie(); + StreamExecutorMemoryAllocator allocator(platform, executors); + + const int64 intra_op_parallelism_threads = 16; + xla::LocalClientOptions client_options; + client_options.set_platform(platform); + client_options.set_intra_op_parallelism_threads(intra_op_parallelism_threads); + auto client = + ClientLibrary::GetOrCreateLocalClient(client_options).ValueOrDie(); + + const int64 dim_size = 1024; + // Create a simple fusable elementwise computation. + ComputationBuilder builder(client, "ParallelFusion"); + Shape input_shape = ShapeUtil::MakeShape(F32, {dim_size, dim_size}); + auto input0 = builder.Broadcast(builder.ConstantR0(1.5f), + AsInt64Slice(input_shape.dimensions())); + auto input1 = builder.Broadcast(builder.ConstantR0(2.0f), + AsInt64Slice(input_shape.dimensions())); + auto input2 = builder.Broadcast(builder.ConstantR0(3.0f), + AsInt64Slice(input_shape.dimensions())); + auto x = builder.Mul(input0, input1); + auto y = builder.Add(x, input2); + auto computation = builder.Build().ConsumeValueOrDie(); + + std::unique_ptr executable = + client->Compile(computation, {}, ExecutableBuildOptions()) + .ConsumeValueOrDie(); + + // Run some warm-up executions. + ExecutableRunOptions options; + options.set_allocator(&allocator); + const int kWarmups = 2; + for (int i = 0; i < kWarmups; ++i) { + auto result = executable->Run({}, options); + ASSERT_TRUE(result.ok()); } - testing::InitGoogleTest(&argc, argv); - if (argc > 1) { - LOG(ERROR) << "Unknown argument " << argv[1] << "\n" << usage; - return 2; + + // Run benchmark. + tensorflow::testing::BytesProcessed(static_cast(num_iters) * dim_size * + dim_size * sizeof(float)); + tensorflow::testing::UseRealTime(); + tensorflow::testing::StartTiming(); + for (int i = 0; i < num_iters; ++i) { + auto result = executable->Run({}, options); + ASSERT_TRUE(result.ok()); } - return RUN_ALL_TESTS(); } + +BENCHMARK(BM_ParallelFusion); + +} // namespace +} // namespace xla diff --git a/tensorflow/compiler/xla/tests/hlo_metadata_test.cc b/tensorflow/compiler/xla/tests/hlo_metadata_test.cc index f54fa2256e217e9aa954a10470cd461023be631d..eded2077fce965ab1c729c610764afa2228ca128 100644 --- a/tensorflow/compiler/xla/tests/hlo_metadata_test.cc +++ b/tensorflow/compiler/xla/tests/hlo_metadata_test.cc @@ -46,7 +46,7 @@ TEST_F(HloMetadataTest, MetadataPropagation) { builder.ClearOpMetadata(); Shape argument_layout = ShapeUtil::MakeShape(F32, {}); - TF_ASSIGN_OR_ASSERT_OK( + TF_ASSERT_OK_AND_ASSIGN( std::unique_ptr executable, local_client_->Compile(builder.Build().ValueOrDie(), {&argument_layout, &argument_layout}, diff --git a/tensorflow/compiler/xla/tests/hlo_test_base.cc b/tensorflow/compiler/xla/tests/hlo_test_base.cc index 62878fed5549a6720a782d01c292ff143187e9a4..8149e2b7cc72018ef8deb61305bb61ceb77200f9 100644 --- a/tensorflow/compiler/xla/tests/hlo_test_base.cc +++ b/tensorflow/compiler/xla/tests/hlo_test_base.cc @@ -23,16 +23,14 @@ limitations under the License. #include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor" #include "tensorflow/compiler/xla/layout_util.h" -#include "tensorflow/compiler/xla/legacy_flags/hlo_test_base_flags.h" +#include "tensorflow/compiler/xla/legacy_flags/debug_options_flags.h" #include "tensorflow/compiler/xla/ptr_util.h" #include "tensorflow/compiler/xla/service/backend.h" #include "tensorflow/compiler/xla/service/computation_layout.h" #include "tensorflow/compiler/xla/service/executable.h" #include "tensorflow/compiler/xla/service/hlo_computation.h" #include "tensorflow/compiler/xla/service/hlo_execution_profile.h" -#include "tensorflow/compiler/xla/service/hlo_graph_dumper.h" #include "tensorflow/compiler/xla/service/hlo_instruction.h" -#include "tensorflow/compiler/xla/service/hlo_module_config.h" #include "tensorflow/compiler/xla/service/transfer_manager.h" #include "tensorflow/compiler/xla/shape_layout.h" #include "tensorflow/compiler/xla/shape_util.h" @@ -40,6 +38,7 @@ limitations under the License. #include "tensorflow/compiler/xla/types.h" #include "tensorflow/core/common_runtime/eigen_thread_pool.h" #include "tensorflow/core/platform/logging.h" +#include "tensorflow/core/platform/test.h" #include "tensorflow/core/platform/types.h" namespace se = ::perftools::gputools; @@ -55,15 +54,6 @@ struct HloTestBase::EigenThreadPoolWrapper { HloTestBase::HloTestBase() : backend_(Backend::CreateDefaultBackend().ConsumeValueOrDie()) { - test_hlo_dumper_ = [](const HloModule& module, const string& label) { - legacy_flags::HloTestBaseFlags* flags = legacy_flags::GetHloTestBaseFlags(); - if (flags->xla_hlo_test_generate_hlo_graph) { - const bool show_addresses = true; - const bool show_layouts = true; - hlo_graph_dumper::DumpGraph(*module.entry_computation(), label, - show_addresses, show_layouts); - } - }; VLOG(1) << "executing on platform " << backend_->platform()->Name(); } @@ -74,30 +64,28 @@ HloTestBase::~HloTestBase() { } } +/* static */ +std::unique_ptr HloTestBase::CreateNewModule() { + HloModuleConfig config; + + auto debug_options = legacy_flags::GetDebugOptionsFromFlags(); + // TODO(b/38354253): Change tests to use Parameters instead of Constants. + debug_options.add_xla_disable_hlo_passes("constant_folding"); + + config.set_debug_options(debug_options); + + return MakeUnique(TestName(), VersionedComputationHandle(), + config); +} + StatusOr HloTestBase::Execute( std::unique_ptr module, tensorflow::gtl::ArraySlice arguments, Shape* result_shape) { - auto module_config = MakeUnique( - module->entry_computation()->ComputeProgramShape()); - return Execute(std::move(module), std::move(module_config), arguments, - result_shape); -} - -StatusOr HloTestBase::Execute( - std::unique_ptr hlo_module, - std::unique_ptr module_config, - tensorflow::gtl::ArraySlice arguments, - Shape* result_shape) { - VLOG(3) << "module_config layout " - << LayoutUtil::HumanString(module_config->entry_computation_layout() - .result_layout() - .layout()); TF_ASSIGN_OR_RETURN( std::unique_ptr executable, - backend_->compiler()->Compile(std::move(hlo_module), - std::move(module_config), test_hlo_dumper_, + backend_->compiler()->Compile(std::move(module), backend_->default_stream_executor())); se::Stream stream(backend_->default_stream_executor()); @@ -111,8 +99,9 @@ StatusOr HloTestBase::Execute( backend_->eigen_intra_op_thread_pool_device()); HloExecutionProfile hlo_execution_profile; - ServiceExecutableRunOptions service_run_options(run_options, - backend_->StreamBorrower()); + ServiceExecutableRunOptions service_run_options( + run_options, backend_->StreamBorrower(), + backend_->inter_op_thread_pool()); TF_ASSIGN_OR_RETURN( se::DeviceMemoryBase result, executable->ExecuteOnStream(&service_run_options, arguments, @@ -123,9 +112,7 @@ StatusOr HloTestBase::Execute( *result_shape = executable->result_shape(); - // TODO(b/36256956) Ideally tuple elements could always be distinct buffers. - if (ShapeUtil::IsTuple(*result_shape) && - backend_->transfer_manager()->TupleElementsAreDistinctBuffers()) { + if (ShapeUtil::IsTuple(*result_shape)) { // We must record element buffers of tuples as well to avoid leaks. DCHECK(!ShapeUtil::IsNestedTuple(*result_shape)); TF_ASSIGN_OR_RETURN( @@ -181,20 +168,26 @@ std::unique_ptr HloTestBase::ExecuteAndTransfer( return TransferFromDevice(result_shape, device_base); } -std::unique_ptr HloTestBase::ExecuteAndTransfer( - std::unique_ptr module, - std::unique_ptr module_config, - tensorflow::gtl::ArraySlice arguments) { - Shape result_shape; - se::DeviceMemoryBase device_base = - Execute(std::move(module), std::move(module_config), arguments, - &result_shape) - .ValueOrDie(); - return TransferFromDevice(result_shape, device_base); +/* static */ +string HloTestBase::TestName() { + return ::testing::UnitTest::GetInstance()->current_test_info()->name(); } -string HloTestBase::TestName() const { - return ::testing::UnitTest::GetInstance()->current_test_info()->name(); +int ParseDebugOptionsFlagsAndRunTests(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; + } + ::testing::InitGoogleTest(&argc, argv); + if (argc > 1) { + LOG(ERROR) << "Unknown argument " << argv[1] << "\n" << usage; + return 2; + } + return RUN_ALL_TESTS(); } } // namespace xla diff --git a/tensorflow/compiler/xla/tests/hlo_test_base.h b/tensorflow/compiler/xla/tests/hlo_test_base.h index 6119473d8158fe87b3611a3edc3490058556288a..7f3d163290aba3cfcea1b3204e6c88134e172ed7 100644 --- a/tensorflow/compiler/xla/tests/hlo_test_base.h +++ b/tensorflow/compiler/xla/tests/hlo_test_base.h @@ -24,7 +24,6 @@ limitations under the License. #include "tensorflow/compiler/xla/service/compiler.h" #include "tensorflow/compiler/xla/service/hlo_computation.h" #include "tensorflow/compiler/xla/service/hlo_module.h" -#include "tensorflow/compiler/xla/service/hlo_module_config.h" #include "tensorflow/compiler/xla/statusor.h" #include "tensorflow/compiler/xla/tests/literal_test_util.h" #include "tensorflow/compiler/xla/types.h" @@ -45,18 +44,15 @@ class HloTestBase : public ::testing::Test { ~HloTestBase() override; - // Executes the given module and returns a global data handle. - StatusOr Execute( - std::unique_ptr module, - tensorflow::gtl::ArraySlice - arguments, - Shape* result_shape); + // Creates a new HLO module for a test. The module created will have + // TestName() for its name; it will also automatically populate its debug + // options from command-line flags. It's recommended to use this method to + // create all HloModules for tests. + static std::unique_ptr CreateNewModule(); - // Variation of Execute which takes a custom module_config instead of creating - // a default one. + // Executes the given module and returns a global data handle. StatusOr Execute( std::unique_ptr module, - std::unique_ptr module_config, tensorflow::gtl::ArraySlice arguments, Shape* result_shape); @@ -65,7 +61,7 @@ class HloTestBase : public ::testing::Test { perftools::gputools::DeviceMemoryBase TransferToDevice( const Literal& literal); - // Transfers the array refered to by the given handle from the device and + // Transfers the array referred to by the given handle from the device and // returns as a Literal. std::unique_ptr TransferFromDevice( const Shape& shape, perftools::gputools::DeviceMemoryBase device_base); @@ -76,42 +72,38 @@ class HloTestBase : public ::testing::Test { tensorflow::gtl::ArraySlice arguments); - // Variation of ExecuteAndTransfer which takes a custom module_config instead - // of creating a default one. - std::unique_ptr ExecuteAndTransfer( - std::unique_ptr module, - std::unique_ptr module_config, - tensorflow::gtl::ArraySlice - arguments); + // Convenience method to force the layout of a given parameter in a module. + // The layout of parameter number 'param_no' in the 'module' is set to + // 'layout'. + void ForceParameterLayout(HloModule* module, int64 param_no, + const Layout& layout) { + ASSERT_LT(param_no, + module->mutable_entry_computation_layout()->parameter_count()); + module->mutable_entry_computation_layout() + ->mutable_parameter_layout(param_no) + ->ResetLayout(layout); + } - // Helpers for comparing ordered and unordered equality of HloInstruction - // containers. - void ExpectEqOrdered( - tensorflow::gtl::ArraySlice actual, - tensorflow::gtl::ArraySlice expected) { - std::vector expected_vec(expected.begin(), - expected.end()); - std::vector actual_vec(actual.begin(), actual.end()); - EXPECT_TRUE(testing::VectorMatcher(expected_vec)( - actual_vec)); + // Convenience method to force the layout of the computation result in a + // module. The result layout of 'module' is set to 'layout'. + void ForceResultLayout(HloModule* module, const Layout& layout) { + module->mutable_entry_computation_layout() + ->mutable_result_layout() + ->ResetLayout(layout); } - void ExpectEqUnordered( - tensorflow::gtl::ArraySlice actual, - tensorflow::gtl::ArraySlice expected) { - std::vector expected_vec(expected.begin(), - expected.end()); - std::vector actual_vec(actual.begin(), actual.end()); - EXPECT_TRUE(testing::UnorderedElementsAre( - expected_vec)(actual_vec)); + // Convenience method to clear the layout of the computation result in + // 'module'. + void ForceClearResultLayout(HloModule* module) { + module->mutable_entry_computation_layout() + ->mutable_result_layout() + ->Clear(); } - string TestName() const; + static string TestName(); std::unique_ptr backend_; - Compiler::HloDumper test_hlo_dumper_; - // This vector contains handles of all the device memory allocations performed // by the test. These are deallocated on destruction of the test object. std::vector allocations_; @@ -121,6 +113,11 @@ class HloTestBase : public ::testing::Test { std::unique_ptr thread_pool_wrapper_; }; +// Convenience function that parses XLA debug options flags from argc/argv, +// calls InitGoogleTest and then calls and returns RUN_ALL_TESTS. Intended to be +// invoked from a test main() function. +int ParseDebugOptionsFlagsAndRunTests(int argc, char** argv); + } // namespace xla #endif // TENSORFLOW_COMPILER_XLA_TESTS_HLO_TEST_BASE_H_ diff --git a/tensorflow/compiler/xla/tests/inprocess_service_test.cc b/tensorflow/compiler/xla/tests/inprocess_service_test.cc deleted file mode 100644 index ea0be07872f31b8e3357d91a164ce8727a159f63..0000000000000000000000000000000000000000 --- a/tensorflow/compiler/xla/tests/inprocess_service_test.cc +++ /dev/null @@ -1,205 +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/compiler/xla/array2d.h" -#include "tensorflow/compiler/xla/client/computation.h" -#include "tensorflow/compiler/xla/client/computation_builder.h" -#include "tensorflow/compiler/xla/client/global_data.h" -#include "tensorflow/compiler/xla/client/local_client.h" -#include "tensorflow/compiler/xla/layout_util.h" -#include "tensorflow/compiler/xla/legacy_flags/cpu_compiler_flags.h" -#include "tensorflow/compiler/xla/shape_util.h" -#include "tensorflow/compiler/xla/statusor.h" -#include "tensorflow/compiler/xla/test_helpers.h" -#include "tensorflow/compiler/xla/tests/client_library_test_base.h" -#include "tensorflow/compiler/xla/tests/literal_test_util.h" -#include "tensorflow/compiler/xla/tests/test_macros.h" -#include "tensorflow/compiler/xla/xla_data.pb.h" -#include "tensorflow/core/lib/gtl/array_slice.h" -#include "tensorflow/core/platform/test.h" -#include "tensorflow/core/platform/types.h" - -namespace xla { -namespace { - -// Tests which exercise the "InProcess" methods of xla::Client. The -// "InProcess" methods require that the client and server share the same -// process. -class InProcessServiceTest : public ClientLibraryTestBase { - protected: - std::unique_ptr ExecuteR2F32Constant( - std::initializer_list> values, - tensorflow::gtl::ArraySlice minor_to_major) { - ComputationBuilder builder(client_, TestName()); - builder.ConstantR2(values); - auto computation = builder.Build().ConsumeValueOrDie(); - CHECK_EQ(2, minor_to_major.size()); - - ExecutionOptions execution_options; - *execution_options.mutable_shape_with_output_layout() = - ShapeUtil::MakeShapeWithLayout( - F32, - /*dimensions=*/{static_cast(values.size()), - static_cast(values.begin()->size())}, - minor_to_major); - return client_->Execute(computation, {}, &execution_options) - .ConsumeValueOrDie(); - } - - ErrorSpec error_spec_{0.0001}; -}; - -XLA_TEST_F(InProcessServiceTest, TransferFromServer) { - ComputationBuilder builder(client_, TestName()); - builder.ConstantR1({1, 42, 5}); - auto computation = builder.Build().ConsumeValueOrDie(); - - auto handle = client_->Execute(computation, {}).ConsumeValueOrDie(); - - std::vector result(3, 0); - ASSERT_IS_OK(client_->TransferInProcess(*handle, result.data())); - EXPECT_MATCH(result, testing::VectorMatcher({1, 42, 5})); -} - -XLA_TEST_F(InProcessServiceTest, TransferToServer) { - std::vector input{1.0f, 2.0f, -42.0f}; - Shape shape = ShapeUtil::MakeShape(F32, {3}); - auto data_handle = client_->TransferToServerInProcess(shape, input.data()) - .ConsumeValueOrDie(); - - ComputationBuilder builder(client_, TestName()); - auto param = builder.Parameter(0, ShapeUtil::MakeShape(F32, {3}), "param"); - builder.Add(param, param); - - ComputeAndCompareR1(&builder, {2.0f, 4.0f, -84.0f}, - {data_handle.get()}, error_spec_); -} - -// TODO(b/28506710): This test case seems not to test inprocess -// methods. -TEST_F(InProcessServiceTest, GetShape) { - ComputationBuilder builder(client_, TestName()); - builder.ConstantR1({1, 42, 5}); - auto computation = builder.Build().ConsumeValueOrDie(); - - auto handle = client_->Execute(computation, {}).ConsumeValueOrDie(); - - Shape shape = client_->GetShape(*handle).ConsumeValueOrDie(); - ASSERT_EQ(S32, shape.element_type()); - ASSERT_EQ(1, ShapeUtil::Rank(shape)); - ASSERT_EQ(3, shape.dimensions(0)); -} - -XLA_TEST_F(InProcessServiceTest, GetShapeOfClientSuppliedArrayRowMajor) { - std::vector input{1.0f, 2.0f, 3.0f, 4.0f}; - Shape shape = ShapeUtil::MakeShape(F32, {2, 2}); - shape.clear_layout(); - *shape.mutable_layout() = LayoutUtil::MakeLayout({1, 0}); - auto handle = client_->TransferToServerInProcess(shape, input.data()) - .ConsumeValueOrDie(); - - Shape shape_returned = client_->GetShape(*handle).ConsumeValueOrDie(); - ASSERT_TRUE(ShapeUtil::Equal(shape, shape_returned)); -} - -XLA_TEST_F(InProcessServiceTest, GetShapeOfClientSuppliedArrayColMajor) { - std::vector input{1.0f, 2.0f, 3.0f, 4.0f}; - Shape shape = ShapeUtil::MakeShape(F32, {2, 2}); - shape.clear_layout(); - *shape.mutable_layout() = LayoutUtil::MakeLayout({0, 1}); - auto handle = client_->TransferToServerInProcess(shape, input.data()) - .ConsumeValueOrDie(); - - Shape shape_returned = client_->GetShape(*handle).ConsumeValueOrDie(); - ASSERT_TRUE(ShapeUtil::Equal(shape, shape_returned)); -} - -TEST_F(InProcessServiceTest, TransferToServerNoLayout) { - std::vector input{1.0f, 2.0f, -42.0f}; - Shape shape = ShapeUtil::MakeShape(F32, {3}); - shape.clear_layout(); - auto transfer_status = - client_->TransferToServerInProcess(shape, input.data()); - ASSERT_EQ(transfer_status.status().code(), - tensorflow::error::INVALID_ARGUMENT); -} - -XLA_TEST_F(InProcessServiceTest, ExecuteRowMajor) { - auto handle = - ExecuteR2F32Constant({{1.0, 2.0}, {3.0, 4.0}}, /*minor_to_major=*/{1, 0}); - - std::vector result(4, 0.0); - Shape shape; - ASSERT_IS_OK(client_->TransferInProcess(*handle, result.data())); - - EXPECT_MATCH(result, testing::VectorMatcher({1.0, 2.0, 3.0, 4.0})); -} - -XLA_TEST_F(InProcessServiceTest, ExecuteColumnMajor) { - auto handle = - ExecuteR2F32Constant({{1.0, 2.0}, {3.0, 4.0}}, /*minor_to_major=*/{0, 1}); - - std::vector result(4, 0); - Shape shape; - ASSERT_IS_OK(client_->TransferInProcess(*handle, result.data())); - - EXPECT_MATCH(result, testing::VectorMatcher({1.0, 3.0, 2.0, 4.0})); -} - -XLA_TEST_F(InProcessServiceTest, ExecuteAndReuseDifferentLayouts) { - // Create arrays on the server which have different layouts. Verify the - // computation still produces the correct results. - auto handle_rowmaj = - ExecuteR2F32Constant({{1.0, 2.0}, {3.0, 4.0}}, /*minor_to_major=*/{1, 0}); - - auto handle_colmaj = ExecuteR2F32Constant({{10.0, 20.0}, {30.0, 40.0}}, - /*minor_to_major=*/{0, 1}); - - ComputationBuilder builder(client_, TestName()); - auto param0 = - builder.Parameter(0, ShapeUtil::MakeShape(F32, {2, 2}), "param0"); - auto param1 = - builder.Parameter(1, ShapeUtil::MakeShape(F32, {2, 2}), "param1"); - builder.Add(param0, param1); - - Array2D expected({{11.0, 22.0}, {33.0, 44.0}}); - ComputeAndCompareR2(&builder, expected, - {handle_rowmaj.get(), handle_colmaj.get()}, - error_spec_); -} - -} // namespace -} // namespace xla - -int main(int argc, char** argv) { - std::vector flag_list; - xla::legacy_flags::AppendCpuCompilerFlags(&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; - } - testing::InitGoogleTest(&argc, argv); - if (argc > 1) { - LOG(ERROR) << "Unknown argument " << argv[1] << "\n" << usage; - return 2; - } - return RUN_ALL_TESTS(); -} diff --git a/tensorflow/compiler/xla/tests/literal_test_util.cc b/tensorflow/compiler/xla/tests/literal_test_util.cc index f7bbc0f38bb501e042542cf7f0a3d4fadb3a2a23..4d8b50fbbf715e8d491667ecb4f4f336ef2d8a68 100644 --- a/tensorflow/compiler/xla/tests/literal_test_util.cc +++ b/tensorflow/compiler/xla/tests/literal_test_util.cc @@ -24,7 +24,7 @@ limitations under the License. #include "tensorflow/compiler/xla/literal_util.h" #include "tensorflow/compiler/xla/ptr_util.h" #include "tensorflow/compiler/xla/shape_util.h" -#include "tensorflow/compiler/xla/test_helpers.h" +#include "tensorflow/compiler/xla/test.h" #include "tensorflow/compiler/xla/types.h" #include "tensorflow/core/lib/core/casts.h" #include "tensorflow/core/lib/io/path.h" @@ -41,20 +41,25 @@ namespace xla { /* static */ void LiteralTestUtil::AssertEqualShapes(const Shape& expected, const Shape& actual) { - ASSERT_EQ(ShapeUtil::Rank(expected), ShapeUtil::Rank(actual)); - ASSERT_EQ(expected.element_type(), actual.element_type()) - << PrimitiveType_Name(expected.element_type()) << " vs " - << PrimitiveType_Name(actual.element_type()); - ASSERT_EQ(expected.dimensions_size(), actual.dimensions_size()); - for (int i = 0; i < expected.dimensions_size(); ++i) { - ASSERT_EQ(expected.dimensions(i), actual.dimensions(i)) - << "mismatch in dimension #" << i - << " expected: " << ShapeUtil::HumanString(expected) - << " actual: " << ShapeUtil::HumanString(actual); - } - ASSERT_EQ(expected.tuple_shapes_size(), actual.tuple_shapes_size()); - for (int i = 0; i < expected.tuple_shapes_size(); ++i) { - AssertEqualShapes(expected.tuple_shapes(i), actual.tuple_shapes(i)); + ASSERT_EQ(ShapeUtil::IsTuple(expected), ShapeUtil::IsTuple(actual)); + if (ShapeUtil::IsTuple(expected)) { + ASSERT_EQ(ShapeUtil::TupleElementCount(expected), + ShapeUtil::TupleElementCount(actual)); + for (int i = 0; i < expected.tuple_shapes_size(); ++i) { + AssertEqualShapes(expected.tuple_shapes(i), actual.tuple_shapes(i)); + } + } else { + ASSERT_EQ(ShapeUtil::Rank(expected), ShapeUtil::Rank(actual)); + ASSERT_EQ(expected.element_type(), actual.element_type()) + << PrimitiveType_Name(expected.element_type()) << " vs " + << PrimitiveType_Name(actual.element_type()); + ASSERT_EQ(expected.dimensions_size(), actual.dimensions_size()); + for (int i = 0; i < expected.dimensions_size(); ++i) { + ASSERT_EQ(expected.dimensions(i), actual.dimensions(i)) + << "mismatch in dimension #" << i + << " expected: " << ShapeUtil::HumanString(expected) + << " actual: " << ShapeUtil::HumanString(actual); + } } } @@ -76,11 +81,11 @@ string Hostname() { // 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 -testing::AssertionResult CompareFloatsBitwiseEqual(FloatT lhs, FloatT rhs) { +::testing::AssertionResult CompareFloatsBitwiseEqual(FloatT lhs, FloatT rhs) { auto ulhs = tensorflow::bit_cast(lhs); auto urhs = tensorflow::bit_cast(rhs); if (ulhs != urhs) { - return testing::AssertionFailure() << tensorflow::strings::Printf( + return ::testing::AssertionFailure() << tensorflow::strings::Printf( "floating values are not bitwise-equal; and equality testing " "was requested: %s=%g=%a vs %s=%g=%a", tensorflow::strings::StrCat(tensorflow::strings::Hex(ulhs)) @@ -90,33 +95,33 @@ testing::AssertionResult CompareFloatsBitwiseEqual(FloatT lhs, FloatT rhs) { .c_str(), rhs, rhs); } - return testing::AssertionSuccess(); + return ::testing::AssertionSuccess(); } // 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 -testing::AssertionResult CompareEqual(NativeT lhs, NativeT rhs) { +::testing::AssertionResult CompareEqual(NativeT lhs, NativeT rhs) { if (lhs == rhs) { - return testing::AssertionSuccess(); + return ::testing::AssertionSuccess(); } ::testing::Message msg; msg << "Expected equality of these values:"; msg << "\n " << lhs; msg << "\n " << rhs; - return testing::AssertionFailure() << msg; + return ::testing::AssertionFailure() << msg; } // Specializations for floating types that do bitwise comparisons when equality // comparison is requested. template <> -testing::AssertionResult CompareEqual(float lhs, float rhs) { +::testing::AssertionResult CompareEqual(float lhs, float rhs) { return CompareFloatsBitwiseEqual(lhs, rhs); } template <> -testing::AssertionResult CompareEqual(double lhs, double rhs) { +::testing::AssertionResult CompareEqual(double lhs, double rhs) { return CompareFloatsBitwiseEqual(lhs, rhs); } @@ -128,9 +133,9 @@ bool ExpectLiteralsEqual(const Literal& expected, const Literal& actual, tensorflow::gtl::MutableArraySlice multi_index, int64 dimension) { if (dimension == expected.shape().dimensions_size()) { - NativeT expected_value = LiteralUtil::Get(expected, multi_index); - NativeT actual_value = LiteralUtil::Get(actual, multi_index); - testing::AssertionResult result = + NativeT expected_value = expected.Get(multi_index); + NativeT actual_value = actual.Get(multi_index); + ::testing::AssertionResult result = CompareEqual(expected_value, actual_value); return result; // Defines implicit coersion to bool. } @@ -147,11 +152,15 @@ bool ExpectLiteralsEqual(const Literal& expected, const Literal& actual, } // namespace /* static */ void LiteralTestUtil::ExpectEqual(const Literal& expected, - const Literal& actual) { - EXPECT_TRUE(Equal(expected, actual)) << "expected:\n" - << LiteralUtil::ToString(expected) - << "\n\tvs actual:\n" - << LiteralUtil::ToString(actual); + const Literal& actual, + const string& message) { + EXPECT_TRUE(Equal(expected, actual)) + << "expected:\n" + << expected.ToString() << "\n\tvs actual:\n" + << actual.ToString() + << (message.empty() + ? "" + : tensorflow::strings::StrCat("\nmessage: ", message)); } /* static */ void LiteralTestUtil::ExpectNotEqual(const Literal& expected, @@ -159,10 +168,12 @@ bool ExpectLiteralsEqual(const Literal& expected, const Literal& actual, EXPECT_FALSE(Equal(expected, actual)); } -/* static */ testing::AssertionResult LiteralTestUtil::Equal( +/* static */ ::testing::AssertionResult LiteralTestUtil::Equal( const Literal& expected, const Literal& actual) { - VLOG(1) << "expected: " << LiteralUtil::ToString(expected); - VLOG(1) << "actual: " << LiteralUtil::ToString(actual); + VLOG(1) << "expected:"; + XLA_VLOG_LINES(1, expected.ToString()); + VLOG(1) << "actual:"; + XLA_VLOG_LINES(1, actual.ToString()); AssertEqualShapes(expected.shape(), actual.shape()); std::vector multi_index(expected.shape().dimensions_size(), 0); @@ -207,11 +218,11 @@ bool ExpectLiteralsEqual(const Literal& expected, const Literal& actual, << "Unsupported primitive type in LiteralTestUtil::ExpectEqual: " << PrimitiveType_Name(expected.shape().element_type()); } - testing::AssertionResult result = testing::AssertionSuccess(); + ::testing::AssertionResult result = ::testing::AssertionSuccess(); if (!match) { - result = testing::AssertionFailure() - << "expected: " << LiteralUtil::ToString(expected) - << "\nactual: " << LiteralUtil::ToString(actual); + result = ::testing::AssertionFailure() + << "expected: " << expected.ToString() + << "\nactual: " << actual.ToString(); VLOG(1) << result.message(); } return result; @@ -219,8 +230,8 @@ bool ExpectLiteralsEqual(const Literal& expected, const Literal& actual, /* static */ void LiteralTestUtil::ExpectEqualTuple(const Literal& expected, const Literal& actual) { - VLOG(1) << "expected: " << LiteralUtil::ToString(expected); - VLOG(1) << "actual: " << LiteralUtil::ToString(actual); + VLOG(1) << "expected: " << expected.ToString(); + VLOG(1) << "actual: " << actual.ToString(); ASSERT_TRUE(ShapeUtil::IsTuple(expected.shape())); ASSERT_TRUE(ShapeUtil::IsTuple(actual.shape())); @@ -247,8 +258,10 @@ class NearComparator { // within the error bound. Emits useful log messages and dumps literals to // temporary files on failure. Returns true if literals match. bool ExpectNear(const Literal& expected, const Literal& actual) { - VLOG(1) << "expected: " << LiteralUtil::ToString(expected); - VLOG(1) << "actual: " << LiteralUtil::ToString(actual); + VLOG(1) << "expected:"; + XLA_VLOG_LINES(1, expected.ToString()); + VLOG(1) << "actual:"; + XLA_VLOG_LINES(1, actual.ToString()); LiteralTestUtil::AssertEqualShapes(expected.shape(), actual.shape()); @@ -262,7 +275,7 @@ class NearComparator { max_abs_err_ = 0.0; *miscompares_.mutable_shape() = ShapeUtil::ChangeElementType(actual.shape(), PRED); - miscompares_.mutable_preds()->Resize( + miscompares_.mutable_preds()->resize( ShapeUtil::ElementsIn(miscompares_.shape()), false); multi_index_.resize(expected.shape().dimensions_size(), 0); @@ -282,9 +295,9 @@ class NearComparator { if (num_miscompares_ > 0) { if (!VLOG_IS_ON(1)) { LOG(INFO) << "expected: " << ShapeUtil::HumanString(expected.shape()) - << " " << LiteralUtil::ToString(expected); + << " " << expected.ToString(); LOG(INFO) << "actual: " << ShapeUtil::HumanString(actual.shape()) - << " " << LiteralUtil::ToString(actual); + << " " << actual.ToString(); } EXPECT_TRUE(num_miscompares_ == 0) << "\nmax relative mismatch at index " @@ -314,7 +327,7 @@ class NearComparator { private: // EXPECTs that the two given scalar values are within the error bound. Keeps - // track of how many mismatches have occured to keep the size of the output + // track of how many mismatches have occurred to keep the size of the output // manageable. template bool ExpectValuesNear(NativeT expected, NativeT actual) { @@ -369,10 +382,9 @@ class NearComparator { void ExpectLiteralsNear(const Literal& expected, const Literal& actual, int64 dimension) { if (dimension == expected.shape().dimensions_size()) { - bool near = - ExpectValuesNear(LiteralUtil::Get(expected, multi_index_), - LiteralUtil::Get(actual, multi_index_)); - LiteralUtil::Set(&miscompares_, multi_index_, !near); + bool near = ExpectValuesNear(expected.Get(multi_index_), + actual.Get(multi_index_)); + miscompares_.Set(multi_index_, !near); } else { for (int64 i = 0; i < expected.shape().dimensions(dimension); ++i) { multi_index_[dimension] = i; @@ -389,7 +401,7 @@ class NearComparator { tensorflow::strings::Printf("tempfile-%s-%llx-%s", Hostname().c_str(), now_usec, name.c_str())); TF_CHECK_OK(tensorflow::WriteBinaryProto(tensorflow::Env::Default(), - filename, literal)); + filename, literal.ToProto())); LOG(ERROR) << "wrote to " << name << " file: " << filename; } @@ -421,28 +433,32 @@ class NearComparator { } // namespace -/* static */ testing::AssertionResult LiteralTestUtil::Near( +/* static */ ::testing::AssertionResult LiteralTestUtil::Near( const Literal& expected, const Literal& actual, const ErrorSpec& error) { NearComparator comparator(error); return comparator.ExpectNear(expected, actual) - ? testing::AssertionSuccess() - : testing::AssertionFailure() << "values were not near"; + ? ::testing::AssertionSuccess() + : ::testing::AssertionFailure() << "values were not near"; } /* static */ void LiteralTestUtil::ExpectNear(const Literal& expected, const Literal& actual, - const ErrorSpec& error) { - EXPECT_TRUE(Near(expected, actual, error)); + const ErrorSpec& error, + const string& message) { + EXPECT_TRUE(Near(expected, actual, error)) + << (message.empty() + ? "" + : tensorflow::strings::StrCat("\nmessage: ", message)); } -/* static */ testing::AssertionResult LiteralTestUtil::NearTuple( +/* static */ ::testing::AssertionResult LiteralTestUtil::NearTuple( const Literal& expected, const Literal& actual, const ErrorSpec& error) { - VLOG(1) << "expected: " << LiteralUtil::ToString(expected); - VLOG(1) << "actual: " << LiteralUtil::ToString(actual); + VLOG(1) << "expected: " << expected.ToString(); + VLOG(1) << "actual: " << actual.ToString(); if (!ShapeUtil::IsTuple(expected.shape()) || !ShapeUtil::IsTuple(actual.shape())) { - return testing::AssertionFailure() + return ::testing::AssertionFailure() << "tuples expected expected shape = " << expected.shape().ShortDebugString() << " actual shape = " << actual.shape().ShortDebugString(); @@ -469,7 +485,7 @@ class NearComparator { } } - return testing::AssertionSuccess(); + return ::testing::AssertionSuccess(); } /* static */ void LiteralTestUtil::ExpectNearTuple(const Literal& expected, @@ -504,8 +520,7 @@ class NearComparator { *shape_with_layout.mutable_layout() = LayoutUtil::MakeLayout(minor_to_major); // Allocate space in the new literal. - LiteralUtil::Reserve(ShapeUtil::ElementsIn(literal.shape()), - new_literal.get()); + new_literal->Reserve(ShapeUtil::ElementsIn(literal.shape())); // Copy data into new literal, element-by-element. for (int64 i = 0; i < ShapeUtil::ElementsIn(literal.shape()); ++i) { @@ -515,44 +530,36 @@ class NearComparator { IndexUtil::LinearIndexToMultidimensionalIndex(shape_with_layout, i); switch (literal.shape().element_type()) { case PRED: - LiteralUtil::Set( - new_literal.get(), to_multi_index, - LiteralUtil::Get(literal, from_multi_index)); + new_literal->Set(to_multi_index, + literal.Get(from_multi_index)); break; case U8: - LiteralUtil::Set( - new_literal.get(), to_multi_index, - LiteralUtil::Get(literal, from_multi_index)); + new_literal->Set(to_multi_index, + literal.Get(from_multi_index)); break; case U32: - LiteralUtil::Set( - new_literal.get(), to_multi_index, - LiteralUtil::Get(literal, from_multi_index)); + new_literal->Set(to_multi_index, + literal.Get(from_multi_index)); break; case S32: - LiteralUtil::Set( - new_literal.get(), to_multi_index, - LiteralUtil::Get(literal, from_multi_index)); + new_literal->Set(to_multi_index, + literal.Get(from_multi_index)); break; case U64: - LiteralUtil::Set( - new_literal.get(), to_multi_index, - LiteralUtil::Get(literal, from_multi_index)); + new_literal->Set(to_multi_index, + literal.Get(from_multi_index)); break; case S64: - LiteralUtil::Set( - new_literal.get(), to_multi_index, - LiteralUtil::Get(literal, from_multi_index)); + new_literal->Set(to_multi_index, + literal.Get(from_multi_index)); break; case F32: - LiteralUtil::Set( - new_literal.get(), to_multi_index, - LiteralUtil::Get(literal, from_multi_index)); + new_literal->Set(to_multi_index, + literal.Get(from_multi_index)); break; case F64: - LiteralUtil::Set( - new_literal.get(), to_multi_index, - LiteralUtil::Get(literal, from_multi_index)); + new_literal->Set(to_multi_index, + literal.Get(from_multi_index)); break; default: LOG(FATAL) << "Unhandled primitive element type: " diff --git a/tensorflow/compiler/xla/tests/literal_test_util.h b/tensorflow/compiler/xla/tests/literal_test_util.h index 85656a53e4400f2b0522e20a7b46922016432103..f645c4e8dcda73806a4204876716b93aa5fb7185 100644 --- a/tensorflow/compiler/xla/tests/literal_test_util.h +++ b/tensorflow/compiler/xla/tests/literal_test_util.h @@ -18,15 +18,18 @@ limitations under the License. #include #include +#include #include #include "tensorflow/compiler/xla/array2d.h" #include "tensorflow/compiler/xla/array3d.h" #include "tensorflow/compiler/xla/array4d.h" #include "tensorflow/compiler/xla/literal_util.h" +#include "tensorflow/compiler/xla/test.h" #include "tensorflow/compiler/xla/test_helpers.h" #include "tensorflow/compiler/xla/types.h" #include "tensorflow/compiler/xla/xla_data.pb.h" +#include "tensorflow/core/lib/core/errors.h" #include "tensorflow/core/lib/gtl/array_slice.h" #include "tensorflow/core/platform/macros.h" #include "tensorflow/core/platform/test.h" @@ -57,11 +60,12 @@ class LiteralTestUtil { // Asserts that the expected and actual literals are (bitwise) equal for all // elements in the literal. Also, asserts that the rank, dimensions sizes, and // primitive type are equal. - static testing::AssertionResult Equal( + static ::testing::AssertionResult Equal( const Literal& expected, const Literal& actual) TF_MUST_USE_RESULT; // Expects that expected and actual are Equal. - static void ExpectEqual(const Literal& expected, const Literal& actual); + static void ExpectEqual(const Literal& expected, const Literal& actual, + const string& message = ""); // Expects that expected and actual are Not Equal. static void ExpectNotEqual(const Literal& expected, const Literal& actual); @@ -101,13 +105,13 @@ class LiteralTestUtil { // Asserts that the expected and actual literals are within the given error // bound for all elements. Also, asserts that the rank, dimensions sizes, and // bounds are equivalent. Only supported for floating point values. - static testing::AssertionResult Near( + static ::testing::AssertionResult Near( const Literal& expected, const Literal& actual, const ErrorSpec& error) TF_MUST_USE_RESULT; // Expects expected and actual to be Near with the given error. static void ExpectNear(const Literal& expected, const Literal& actual, - const ErrorSpec& error); + const ErrorSpec& error, const string& message = ""); // Asserts the given literal are within the given error bound of the given // expected values. Only supported for floating point values. @@ -127,6 +131,12 @@ class LiteralTestUtil { std::initializer_list>> expected, const Literal& actual, const ErrorSpec& error); + template + static void ExpectR4Near( + std::initializer_list>>> + expected, + const Literal& actual, const ErrorSpec& error); // Asserts the given literal are within the given error bound to the given // array. Only supported for floating point values. @@ -147,7 +157,7 @@ class LiteralTestUtil { // tuples are within the given error bound. Tuples are matched recursively. // If the elements of the tuple are not floating-point types, the error spec // is ignored and exact equality is checked. - static testing::AssertionResult NearTuple( + static ::testing::AssertionResult NearTuple( const Literal& expected, const Literal& actual, const ErrorSpec& error) TF_MUST_USE_RESULT; @@ -170,6 +180,36 @@ class LiteralTestUtil { tensorflow::gtl::ArraySlice minor_to_major, const Literal& literal); + // Creates a literal with the supplied shape, and uses the provided value + // generator to populate the literal's values. + // Returns the new literal object, or an error Status if failed. + template < + PrimitiveType type, + typename T = typename primitive_util::PrimitiveTypeToNative::type> + static StatusOr> CreateRandomLiteral( + const Shape& shape, + const std::function)>& generator); + + // Creates a literal with the supplied shape, and initializes the literal + // values using a normal distribution with given mean and stddev standard + // deviation, and using the engine as entropy generator. + // Returns the new literal object, or an error Status if failed. + template < + PrimitiveType type, typename E, + typename T = typename primitive_util::PrimitiveTypeToNative::type> + static StatusOr> CreateRandomLiteral( + const Shape& shape, E* engine, T mean, T stddev); + + // Creates a literal with the supplied shape, and initializes the literal + // values using a normal distribution with given mean and stddev standard + // deviation. + // Returns the new literal object, or an error Status if failed. + template < + PrimitiveType type, + typename T = typename primitive_util::PrimitiveTypeToNative::type> + static StatusOr> CreateRandomLiteral( + const Shape& shape, T mean, T stddev); + private: TF_DISALLOW_COPY_AND_ASSIGN(LiteralTestUtil); }; @@ -177,20 +217,20 @@ class LiteralTestUtil { template /* static */ void LiteralTestUtil::ExpectR0Equal(NativeT expected, const Literal& actual) { - ExpectEqual(*LiteralUtil::CreateR0(expected), actual); + ExpectEqual(*Literal::CreateR0(expected), actual); } template /* static */ void LiteralTestUtil::ExpectR1Equal( tensorflow::gtl::ArraySlice expected, const Literal& actual) { - ExpectEqual(*LiteralUtil::CreateR1(expected), actual); + ExpectEqual(*Literal::CreateR1(expected), actual); } template /* static */ void LiteralTestUtil::ExpectR2Equal( std::initializer_list> expected, const Literal& actual) { - ExpectEqual(*LiteralUtil::CreateR2(expected), actual); + ExpectEqual(*Literal::CreateR2(expected), actual); } template @@ -198,46 +238,46 @@ template std::initializer_list>> expected, const Literal& actual) { - ExpectEqual(*LiteralUtil::CreateR3(expected), actual); + ExpectEqual(*Literal::CreateR3(expected), actual); } template /* static */ void LiteralTestUtil::ExpectR2EqualArray2D( const Array2D& expected, const Literal& actual) { - ExpectEqual(*LiteralUtil::CreateR2FromArray2D(expected), actual); + ExpectEqual(*Literal::CreateR2FromArray2D(expected), actual); } template /* static */ void LiteralTestUtil::ExpectR3EqualArray3D( const Array3D& expected, const Literal& actual) { - ExpectEqual(*LiteralUtil::CreateR3FromArray3D(expected), actual); + ExpectEqual(*Literal::CreateR3FromArray3D(expected), actual); } template /* static */ void LiteralTestUtil::ExpectR4EqualArray4D( const Array4D& expected, const Literal& actual) { - ExpectEqual(*LiteralUtil::CreateR4FromArray4D(expected), actual); + ExpectEqual(*Literal::CreateR4FromArray4D(expected), actual); } template /* static */ void LiteralTestUtil::ExpectR0Near(NativeT expected, const Literal& actual, const ErrorSpec& error) { - ExpectNear(*LiteralUtil::CreateR0(expected), actual, error); + ExpectNear(*Literal::CreateR0(expected), actual, error); } template /* static */ void LiteralTestUtil::ExpectR1Near( tensorflow::gtl::ArraySlice expected, const Literal& actual, const ErrorSpec& error) { - ExpectNear(*LiteralUtil::CreateR1(expected), actual, error); + ExpectNear(*Literal::CreateR1(expected), actual, error); } template /* static */ void LiteralTestUtil::ExpectR2Near( std::initializer_list> expected, const Literal& actual, const ErrorSpec& error) { - ExpectNear(*LiteralUtil::CreateR2(expected), actual, error); + ExpectNear(*Literal::CreateR2(expected), actual, error); } template @@ -245,28 +285,71 @@ template std::initializer_list>> expected, const Literal& actual, const ErrorSpec& error) { - ExpectNear(*LiteralUtil::CreateR3(expected), actual, error); + ExpectNear(*Literal::CreateR3(expected), actual, error); +} + +template +/* static */ void LiteralTestUtil::ExpectR4Near( + std::initializer_list>>> + expected, + const Literal& actual, const ErrorSpec& error) { + ExpectNear(*Literal::CreateR4(expected), actual, error); } template /* static */ void LiteralTestUtil::ExpectR2NearArray2D( const Array2D& expected, const Literal& actual, const ErrorSpec& error) { - ExpectNear(*LiteralUtil::CreateR2FromArray2D(expected), actual, error); + ExpectNear(*Literal::CreateR2FromArray2D(expected), actual, error); } template /* static */ void LiteralTestUtil::ExpectR3NearArray3D( const Array3D& expected, const Literal& actual, const ErrorSpec& error) { - ExpectNear(*LiteralUtil::CreateR3FromArray3D(expected), actual, error); + ExpectNear(*Literal::CreateR3FromArray3D(expected), actual, error); } template /* static */ void LiteralTestUtil::ExpectR4NearArray4D( const Array4D& expected, const Literal& actual, const ErrorSpec& error) { - ExpectNear(*LiteralUtil::CreateR4FromArray4D(expected), actual, error); + ExpectNear(*Literal::CreateR4FromArray4D(expected), actual, error); +} + +template +/* static */ StatusOr> +LiteralTestUtil::CreateRandomLiteral( + const Shape& shape, + const std::function)>& generator) { + using NativeT = typename primitive_util::PrimitiveTypeToNative::type; + TF_RET_CHECK(shape.element_type() == type); + std::unique_ptr literal = Literal::CreateFromShape(shape); + TF_RETURN_IF_ERROR(literal.get()->Populate( + [&](tensorflow::gtl::ArraySlice indexes) { + return generator(indexes); + })); + return std::move(literal); +} + +template +/* static */ StatusOr> +LiteralTestUtil::CreateRandomLiteral(const Shape& shape, E* engine, T mean, + T stddev) { + using NativeT = typename primitive_util::PrimitiveTypeToNative::type; + std::normal_distribution generator(mean, stddev); + return CreateRandomLiteral( + shape, [&](tensorflow::gtl::ArraySlice /*indexes*/) { + return generator(*engine); + }); +} + +template +/* static */ StatusOr> +LiteralTestUtil::CreateRandomLiteral(const Shape& shape, T mean, T stddev) { + std::minstd_rand0 engine; + return CreateRandomLiteral(shape, &engine, mean, stddev); } } // namespace xla diff --git a/tensorflow/compiler/xla/tests/literal_test_util_test.cc b/tensorflow/compiler/xla/tests/literal_test_util_test.cc index fdec11c0e98cb3c36f8e56872c0ca5e8336af7f0..2acf27ed390b0732ba40fcf505c746bd7d8b651e 100644 --- a/tensorflow/compiler/xla/tests/literal_test_util_test.cc +++ b/tensorflow/compiler/xla/tests/literal_test_util_test.cc @@ -31,9 +31,8 @@ namespace xla { namespace { TEST(LiteralTestUtilTest, ComparesEqualTuplesEqual) { - std::unique_ptr literal = LiteralUtil::MakeTuple({ - LiteralUtil::CreateR0(42).get(), - LiteralUtil::CreateR0(64).get(), + std::unique_ptr literal = Literal::MakeTuple({ + Literal::CreateR0(42).get(), Literal::CreateR0(64).get(), }); LiteralTestUtil::ExpectEqual(*literal, *literal); } @@ -43,13 +42,11 @@ TEST(LiteralTestUtilTest, ComparesUnequalTuplesUnequal) { // un-fail an assertion failure. The CHECK-failure is death, so we can make a // death assertion. auto unequal_things_are_equal = [] { - std::unique_ptr lhs = LiteralUtil::MakeTuple({ - LiteralUtil::CreateR0(42).get(), - LiteralUtil::CreateR0(64).get(), + std::unique_ptr lhs = Literal::MakeTuple({ + Literal::CreateR0(42).get(), Literal::CreateR0(64).get(), }); - std::unique_ptr rhs = LiteralUtil::MakeTuple({ - LiteralUtil::CreateR0(64).get(), - LiteralUtil::CreateR0(42).get(), + std::unique_ptr rhs = Literal::MakeTuple({ + Literal::CreateR0(64).get(), Literal::CreateR0(42).get(), }); CHECK(LiteralTestUtil::Equal(*lhs, *rhs)) << "LHS and RHS are unequal"; }; @@ -58,8 +55,8 @@ TEST(LiteralTestUtilTest, ComparesUnequalTuplesUnequal) { TEST(LiteralTestUtilTest, ExpectNearFailurePlacesResultsInTemporaryDirectory) { auto dummy_lambda = [] { - auto two = LiteralUtil::CreateR0(2); - auto four = LiteralUtil::CreateR0(4); + auto two = Literal::CreateR0(2); + auto four = Literal::CreateR0(4); ErrorSpec error(0.001); CHECK(LiteralTestUtil::Near(*two, *four, error)) << "two is not near four"; }; @@ -83,15 +80,16 @@ TEST(LiteralTestUtilTest, ExpectNearFailurePlacesResultsInTemporaryDirectory) { LOG(INFO) << "results: [" << tensorflow::str_util::Join(results, ", ") << "]"; EXPECT_EQ(3, results.size()); for (const string& result : results) { - Literal literal; + LiteralProto literal_proto; TF_CHECK_OK(tensorflow::ReadBinaryProto(tensorflow::Env::Default(), result, - &literal)); + &literal_proto)); + Literal literal(literal_proto); if (result.find("expected") != string::npos) { - EXPECT_EQ("2", LiteralUtil::ToString(literal)); + EXPECT_EQ("2", literal.ToString()); } else if (result.find("actual") != string::npos) { - EXPECT_EQ("4", LiteralUtil::ToString(literal)); + EXPECT_EQ("4", literal.ToString()); } else if (result.find("miscompares") != string::npos) { - EXPECT_EQ("true", LiteralUtil::ToString(literal)); + EXPECT_EQ("true", literal.ToString()); } else { FAIL() << "unknown file in temporary directory: " << result; } diff --git a/tensorflow/compiler/xla/tests/llvm_compiler_test.cc b/tensorflow/compiler/xla/tests/llvm_compiler_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..40e5bccb55e7ca4f924797056b9b1f92a8ff83c3 --- /dev/null +++ b/tensorflow/compiler/xla/tests/llvm_compiler_test.cc @@ -0,0 +1,63 @@ +/* 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/llvm_compiler.h" +#include "tensorflow/compiler/xla/service/hlo_instruction.h" +#include "tensorflow/compiler/xla/tests/hlo_test_base.h" +#include "tensorflow/compiler/xla/tests/test_macros.h" +#include "tensorflow/core/platform/test.h" + +namespace xla { +namespace { + +class LLVMCompilerTest : public HloTestBase {}; + +XLA_TEST_F(LLVMCompilerTest, CompilerHooks) { + int pre_opt_hook_call_count = 0; + int post_opt_hook_call_count = 0; + + auto pre_opt_hook = [&pre_opt_hook_call_count](const llvm::Module &) { + ++pre_opt_hook_call_count; + return Status::OK(); + }; + auto post_opt_hook = [&post_opt_hook_call_count](const llvm::Module &) { + ++post_opt_hook_call_count; + return Status::OK(); + }; + + // Create HLO module, and run the compiler. + auto builder = HloComputation::Builder(TestName()); + builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(42.0))); + + auto hlo_module = CreateNewModule(); + hlo_module->AddEntryComputation(builder.Build()); + + auto compiler = static_cast(backend_->compiler()); + compiler->SetPreOptimizationHook(pre_opt_hook); + compiler->SetPostOptimizationHook(post_opt_hook); + + ASSERT_TRUE( + compiler + ->Compile(std::move(hlo_module), backend_->default_stream_executor()) + .ok()); + + // Test that hooks were called. + EXPECT_EQ(1, pre_opt_hook_call_count); + EXPECT_EQ(1, post_opt_hook_call_count); +} + +} // namespace +} // namespace xla diff --git a/tensorflow/compiler/xla/tests/llvm_irgen_test_base.cc b/tensorflow/compiler/xla/tests/llvm_irgen_test_base.cc new file mode 100644 index 0000000000000000000000000000000000000000..98dd9613a78a0c8673dfed6d68b64f33da79dfe0 --- /dev/null +++ b/tensorflow/compiler/xla/tests/llvm_irgen_test_base.cc @@ -0,0 +1,82 @@ +/* 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/llvm_irgen_test_base.h" + +#include +#include + +#include "tensorflow/compiler/xla/service/llvm_ir/llvm_util.h" +#include "tensorflow/compiler/xla/tests/filecheck.h" +#include "tensorflow/core/platform/test.h" + +namespace xla { + +void LLVMIRGenTestBase::SetIrHook(bool match_optimized_ir) { + auto llvm_compiler = GetLLVMCompiler(); + using std::placeholders::_1; + + // Add the IR inspection hook to the LLVM compiler. + if (match_optimized_ir) { + llvm_compiler->SetPostOptimizationHook( + std::bind(&LLVMIRGenTestBase::IrHook, this, _1)); + } else { + llvm_compiler->SetPreOptimizationHook( + std::bind(&LLVMIRGenTestBase::IrHook, this, _1)); + } +} + +void LLVMIRGenTestBase::ResetIrHook() { + auto llvm_compiler = GetLLVMCompiler(); + + llvm_compiler->RemovePreOptimizationHook(); + llvm_compiler->RemovePostOptimizationHook(); +} + +void LLVMIRGenTestBase::CompileAndVerifyIr( + std::unique_ptr hlo_module, const string& pattern, + bool match_optimized_ir) { + SetIrHook(match_optimized_ir); + ASSERT_TRUE(CompileToExecutable(std::move(hlo_module)).ok()); + ResetIrHook(); + + StatusOr filecheck_result = RunFileCheck(ir_, pattern); + ASSERT_TRUE(filecheck_result.ok()); + EXPECT_TRUE(filecheck_result.ValueOrDie()); +} + +void LLVMIRGenTestBase::CompileAheadOfTimeAndVerifyIr( + std::unique_ptr hlo_module, const AotCompilationOptions& options, + const string& pattern, bool match_optimized_ir) { + SetIrHook(match_optimized_ir); + ASSERT_TRUE( + CompileToAotCompilationResult(std::move(hlo_module), options).ok()); + ResetIrHook(); + + StatusOr filecheck_result = RunFileCheck(ir_, pattern); + ASSERT_TRUE(filecheck_result.ok()); + EXPECT_TRUE(filecheck_result.ValueOrDie()); +} + +LLVMCompiler* LLVMIRGenTestBase::GetLLVMCompiler() const { + return static_cast(backend_->compiler()); +} + +Status LLVMIRGenTestBase::IrHook(const llvm::Module& module) { + ir_ = llvm_ir::DumpModuleToString(module); + return Status::OK(); +} + +} // namespace xla diff --git a/tensorflow/compiler/xla/tests/llvm_irgen_test_base.h b/tensorflow/compiler/xla/tests/llvm_irgen_test_base.h new file mode 100644 index 0000000000000000000000000000000000000000..f0a0df76ac03f80946507e8eefbe0548d2a9f721 --- /dev/null +++ b/tensorflow/compiler/xla/tests/llvm_irgen_test_base.h @@ -0,0 +1,67 @@ +/* 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_LLVM_IRGEN_TEST_BASE_H_ +#define TENSORFLOW_COMPILER_XLA_TESTS_LLVM_IRGEN_TEST_BASE_H_ + +#include + +#include "tensorflow/compiler/xla/service/llvm_compiler.h" +#include "tensorflow/compiler/xla/tests/codegen_test_base.h" + +namespace xla { + +// Tests that verify IR emitted by the CPU/GPU backend is as expected. +class LLVMIRGenTestBase : public CodegenTestBase { + protected: + // Compiles the given HLO module to LLVM IR and verifies the IR matches the + // given pattern. `pattern` is in the FileCheck pattern matching syntax + // (http://llvm.org/docs/CommandGuide/FileCheck.html). + // + // This function invokes the JIT compiler. + // + // If `match_optimized_ir` is true, match the version of the IR after internal + // optimizations are applied; otherwise, the IR before optimizations is + // matched. + void CompileAndVerifyIr(std::unique_ptr hlo_module, + const string& pattern, bool match_optimized_ir); + + // Compiles the given HLO module to LLVM IR and verifies the IR matches the + // given pattern. `pattern` is in the FileCheck pattern matching syntax + // (http://llvm.org/docs/CommandGuide/FileCheck.html). + // + // This function invokes the AOT compiler, with options in `options`. + // + // If `match_optimized_ir` is true, match the version of the IR after internal + // optimizations are applied; otherwise, the IR before optimizations is + // matched. + void CompileAheadOfTimeAndVerifyIr(std::unique_ptr hlo_module, + const AotCompilationOptions& options, + const string& pattern, + bool match_optimized_ir); + + private: + LLVMCompiler* GetLLVMCompiler() const; + + void SetIrHook(bool match_optimized_ir); + void ResetIrHook(); + + string ir_; + Status IrHook(const llvm::Module& module); +}; + +} // namespace xla + +#endif // TENSORFLOW_COMPILER_XLA_TESTS_LLVM_IRGEN_TEST_BASE_H_ diff --git a/tensorflow/compiler/xla/tests/local_client_aot_test_helper.cc b/tensorflow/compiler/xla/tests/local_client_aot_test_helper.cc index 7ea83a9e956ca8b5bb26ea6aaa844d2b63107328..9266760ce4cf61fd591b797abba504dbae1a1572 100644 --- a/tensorflow/compiler/xla/tests/local_client_aot_test_helper.cc +++ b/tensorflow/compiler/xla/tests/local_client_aot_test_helper.cc @@ -19,7 +19,7 @@ limitations under the License. #include #include -#include "external/llvm/include/llvm/ADT/Triple.h" +#include "llvm/ADT/Triple.h" #include "tensorflow/compiler/xla/client/client_library.h" #include "tensorflow/compiler/xla/client/computation_builder.h" #include "tensorflow/compiler/xla/service/cpu/cpu_compiler.h" @@ -42,7 +42,7 @@ xla::Computation Doubler(xla::Client* client) { int main(int argc, char** argv) { tensorflow::port::InitMain(argv[0], &argc, &argv); - auto client = xla::ClientLibrary::LocalClientOrDie(); + auto client = xla::ClientLibrary::GetOrCreateCompileOnlyClient().ValueOrDie(); xla::ComputationBuilder builder(client, "aot_test_helper"); auto opaque_shape = xla::ShapeUtil::MakeOpaqueShape(); @@ -74,7 +74,7 @@ int main(int argc, char** argv) { llvm::Triple triple(xla::llvm_ir::AsStringRef(triple_string)); xla::Computation computation = builder.Build().ConsumeValueOrDie(); - xla::LocalClient::AheadOfTimeComputationInstance instance{ + xla::CompileOnlyClient::AotComputationInstance instance{ &computation, /*argument_layouts=*/{&opaque_shape}, &r0f32}; xla::cpu::CpuAotCompilationOptions options( diff --git a/tensorflow/compiler/xla/tests/local_client_test_base.cc b/tensorflow/compiler/xla/tests/local_client_test_base.cc index 7fe4c9020f4c67ecc9888425cf0a2c358ad49e6d..49207356e3027cff52a29f962fedbd3593a4925e 100644 --- a/tensorflow/compiler/xla/tests/local_client_test_base.cc +++ b/tensorflow/compiler/xla/tests/local_client_test_base.cc @@ -17,12 +17,19 @@ limitations under the License. #include +#define EIGEN_USE_THREADS + +#include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor" #include "tensorflow/compiler/xla/client/local_client.h" #include "tensorflow/compiler/xla/map_util.h" #include "tensorflow/compiler/xla/ptr_util.h" #include "tensorflow/compiler/xla/shape_util.h" #include "tensorflow/compiler/xla/status_macros.h" #include "tensorflow/compiler/xla/test_helpers.h" +#include "tensorflow/core/common_runtime/eigen_thread_pool.h" +#include "tensorflow/core/lib/core/threadpool.h" +#include "tensorflow/core/platform/cpu_info.h" +#include "tensorflow/core/platform/env.h" #include "tensorflow/core/platform/logging.h" namespace xla { @@ -91,16 +98,34 @@ int64 TestAllocator::deallocation_count(int device_ordinal) const { return allocator_; } +// Define this in .cc file to avoid having to include eigen or forward declare +// these types in the header. +struct LocalClientTestBase::EigenThreadPoolWrapper { + explicit EigenThreadPoolWrapper() + : pool(new tensorflow::thread::ThreadPool( + tensorflow::Env::Default(), "XLAEigenTest", /*num_threads=*/2)), + wrapper(new tensorflow::EigenThreadPoolWrapper(pool.get())), + device(new Eigen::ThreadPoolDevice(wrapper.get(), + wrapper->NumThreads())) {} + + std::unique_ptr pool; + std::unique_ptr wrapper; + std::unique_ptr device; +}; + LocalClientTestBase::LocalClientTestBase( perftools::gputools::Platform* platform) : local_client_( - ClientLibrary::GetOrCreateLocalClient(platform).ValueOrDie()) { + ClientLibrary::GetOrCreateLocalClient(platform).ValueOrDie()), + thread_pool_wrapper_(new EigenThreadPoolWrapper()) { stream_executor_ = PlatformUtil::GetStreamExecutors(local_client_->platform()) .ValueOrDie()[local_client_->default_device_ordinal()]; transfer_manager_ = TransferManager::GetForPlatform(local_client_->platform()).ValueOrDie(); } +LocalClientTestBase::~LocalClientTestBase() {} + std::unique_ptr LocalClientTestBase::LiteralToScopedShapedBuffer(const Literal& literal) { return LiteralToScopedShapedBuffer(literal, @@ -166,18 +191,13 @@ LocalClientTestBase::ShapedBufferToScopedShapedBuffer( } *scoped_buffer->mutable_buffers() = shaped_buffer->buffers(); - TF_CHECK_OK( - scoped_buffer->mutable_shape_index_to_buffer_entry() - ->ForEachMutableElement( - [&shaped_buffer](const ShapeIndex& index, bool is_leaf, - size_t* buffer_entry) -> ::tensorflow::Status { - if (is_leaf) { - *buffer_entry = - shaped_buffer->shape_index_to_buffer_entry().element( - index); - } - return tensorflow::Status::OK(); - })); + scoped_buffer->mutable_shape_index_to_buffer_entry()->ForEachMutableElement( + [&shaped_buffer](const ShapeIndex& index, size_t* buffer_entry) { + if (ShapeUtil::IsLeafIndex(shaped_buffer->shape(), index)) { + *buffer_entry = + shaped_buffer->shape_index_to_buffer_entry().element(index); + } + }); return scoped_buffer; } @@ -190,8 +210,7 @@ ExecutableRunOptions LocalClientTestBase::DefaultExecutableRunOptions() const { ExecutableRunOptions run_options; run_options.set_inter_op_thread_pool( local_client_->backend().inter_op_thread_pool()); - run_options.set_intra_op_thread_pool( - local_client_->backend().eigen_intra_op_thread_pool_device()); + run_options.set_intra_op_thread_pool(thread_pool_wrapper_->device.get()); run_options.set_allocator(GetOrCreateAllocator(local_client_->platform())); return run_options; } diff --git a/tensorflow/compiler/xla/tests/local_client_test_base.h b/tensorflow/compiler/xla/tests/local_client_test_base.h index 4e7b05cea60887eec628ce9b4848321e721030e5..e3c3bb46cf26cc742b7abb39a3e457d823d829ec 100644 --- a/tensorflow/compiler/xla/tests/local_client_test_base.h +++ b/tensorflow/compiler/xla/tests/local_client_test_base.h @@ -74,8 +74,10 @@ class TestAllocator : public StreamExecutorMemoryAllocator { // A base class for tests which exercise the LocalClient interface. class LocalClientTestBase : public ::testing::Test { protected: + struct EigenThreadPoolWrapper; explicit LocalClientTestBase( perftools::gputools::Platform* platform = nullptr); + virtual ~LocalClientTestBase(); static TestAllocator* GetOrCreateAllocator( perftools::gputools::Platform* platform); @@ -142,6 +144,8 @@ class LocalClientTestBase : public ::testing::Test { TransferManager* transfer_manager_; LocalClient* local_client_; + + std::unique_ptr thread_pool_wrapper_; }; } // namespace xla diff --git a/tensorflow/compiler/xla/tests/log_test.cc b/tensorflow/compiler/xla/tests/log_test.cc index b520d89de3cacd15e3766ed6f2c468cfa367a64f..174d433a9e17312c3548668feeeb2e92712c87f8 100644 --- a/tensorflow/compiler/xla/tests/log_test.cc +++ b/tensorflow/compiler/xla/tests/log_test.cc @@ -18,7 +18,6 @@ limitations under the License. #include "tensorflow/compiler/xla/client/computation_builder.h" #include "tensorflow/compiler/xla/client/local_client.h" -#include "tensorflow/compiler/xla/legacy_flags/cpu_compiler_flags.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" @@ -47,6 +46,7 @@ TEST_F(LogTest, LogTenValues) { builder.Log(x); std::vector expected; + expected.reserve(input.size()); for (float f : input) { expected.push_back(std::log(f)); } @@ -56,20 +56,3 @@ TEST_F(LogTest, LogTenValues) { } // namespace } // namespace xla - -int main(int argc, char** argv) { - std::vector flag_list; - xla::legacy_flags::AppendCpuCompilerFlags(&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; - } - testing::InitGoogleTest(&argc, argv); - if (argc > 1) { - LOG(ERROR) << "Unknown argument " << argv[1] << "\n" << usage; - return 2; - } - return RUN_ALL_TESTS(); -} diff --git a/tensorflow/compiler/xla/tests/map_test.cc b/tensorflow/compiler/xla/tests/map_test.cc index 2433c5653a6562b9672eeff81192dfc3152dffed..01ee421baac3b17da962d9ddc7b15b8e6039200a 100644 --- a/tensorflow/compiler/xla/tests/map_test.cc +++ b/tensorflow/compiler/xla/tests/map_test.cc @@ -21,19 +21,17 @@ limitations under the License. #include "tensorflow/compiler/xla/client/global_data.h" #include "tensorflow/compiler/xla/client/lib/arithmetic.h" #include "tensorflow/compiler/xla/client/local_client.h" -#include "tensorflow/compiler/xla/legacy_flags/cpu_compiler_flags.h" #include "tensorflow/compiler/xla/literal_util.h" #include "tensorflow/compiler/xla/shape_util.h" #include "tensorflow/compiler/xla/statusor.h" +#include "tensorflow/compiler/xla/test.h" #include "tensorflow/compiler/xla/test_helpers.h" #include "tensorflow/compiler/xla/tests/client_library_test_base.h" #include "tensorflow/compiler/xla/tests/literal_test_util.h" #include "tensorflow/compiler/xla/tests/test_macros.h" #include "tensorflow/compiler/xla/tests/test_utils.h" -#include "tensorflow/compiler/xla/xla.pb.h" #include "tensorflow/compiler/xla/xla_data.pb.h" #include "tensorflow/core/platform/stream_executor_no_cuda.h" -#include "tensorflow/core/platform/test.h" #include "tensorflow/core/platform/types.h" namespace xla { @@ -42,8 +40,10 @@ namespace { class MapTest : public ClientLibraryTestBase { public: explicit MapTest(perftools::gputools::Platform* platform = nullptr) - : ClientLibraryTestBase(platform, - /*disabled_pass_names=*/{"algsimp", "inline"}) {} + : ClientLibraryTestBase(platform) { + mutable_debug_options()->add_xla_disable_hlo_passes("algsimp"); + mutable_debug_options()->add_xla_disable_hlo_passes("inline"); + } // Creates a function that adds its scalar argument with the constant 1.0. // @@ -100,8 +100,8 @@ class MapTest : public ClientLibraryTestBase { // Creates a function that adds its scalar argument with the constant 1.0 and // then multiplies by the original element. // - // /---------------\ - // / \ + // /------------------| + // / | // x {R0F32} ----> (add) ----> (mul) // / // 1.0f ---------/ @@ -147,8 +147,8 @@ class MapTest : public ClientLibraryTestBase { // Creates a function that adds three scalar arguments // - // x {R0F32} ----\ - // \ + // x {R0F32} -------| + // | // y {R0F32} ----> (add) ---> (add) // / // z {R0F32} ---------------/ @@ -168,7 +168,7 @@ class MapTest : public ClientLibraryTestBase { TEST_F(MapTest, MapEachElemPlusOneR0) { // Applies lambda (x) (+ x 1)) to an input scalar. ComputationBuilder builder(client_, TestName()); - std::unique_ptr param0_literal = LiteralUtil::CreateR0(42.0); + std::unique_ptr param0_literal = Literal::CreateR0(42.0); std::unique_ptr param0_data = client_->TransferToServer(*param0_literal).ConsumeValueOrDie(); @@ -182,7 +182,7 @@ TEST_F(MapTest, MapEachElemPlusOneR0) { XLA_TEST_F(MapTest, MapEachElemPlusOneR1S0) { // Maps (lambda (x) (+ x 1)) onto an input R1F32 vector of length 0. ComputationBuilder builder(client_, TestName()); - std::unique_ptr param0_literal = LiteralUtil::CreateR1({}); + std::unique_ptr param0_literal = Literal::CreateR1({}); std::unique_ptr param0_data = client_->TransferToServer(*param0_literal).ConsumeValueOrDie(); @@ -197,7 +197,7 @@ TEST_F(MapTest, MapEachElemPlusOneR1S4) { // Maps (lambda (x) (+ x 1)) onto an input R1F32 vector of length 4. ComputationBuilder builder(client_, TestName()); std::unique_ptr param0_literal = - LiteralUtil::CreateR1({2.2f, 3.3f, 4.4f, 5.5f}); + Literal::CreateR1({2.2f, 3.3f, 4.4f, 5.5f}); std::unique_ptr param0_data = client_->TransferToServer(*param0_literal).ConsumeValueOrDie(); @@ -211,7 +211,7 @@ TEST_F(MapTest, MapEachElemPlusOneR1S4) { TEST_F(MapTest, MapEachF32ElementToS32Constant) { ComputationBuilder builder(client_, TestName()); std::unique_ptr param0_literal = - LiteralUtil::CreateR1({2.2f, 3.3f, 4.4f, 5.5f}); + Literal::CreateR1({2.2f, 3.3f, 4.4f, 5.5f}); std::unique_ptr param0_data = client_->TransferToServer(*param0_literal).ConsumeValueOrDie(); @@ -224,7 +224,7 @@ TEST_F(MapTest, MapEachF32ElementToS32Constant) { TEST_F(MapTest, MapEachF32ElementToU32Constant) { ComputationBuilder builder(client_, TestName()); std::unique_ptr param0_literal = - LiteralUtil::CreateR1({2.2f, 3.3f, 4.4f, 5.5f}); + Literal::CreateR1({2.2f, 3.3f, 4.4f, 5.5f}); std::unique_ptr param0_data = client_->TransferToServer(*param0_literal).ConsumeValueOrDie(); @@ -238,7 +238,7 @@ TEST_F(MapTest, MapEachElemLongerChainR1) { // Maps (lambda (x) (* (+ x 1) x)) onto an input R1F32 vector. ComputationBuilder builder(client_, TestName()); std::unique_ptr param0_literal = - LiteralUtil::CreateR1({2.6f, -5.1f, 0.1f, 0.2f, 999.0f, 255.5f}); + Literal::CreateR1({2.6f, -5.1f, 0.1f, 0.2f, 999.0f, 255.5f}); std::unique_ptr param0_data = client_->TransferToServer(*param0_literal).ConsumeValueOrDie(); @@ -254,7 +254,7 @@ XLA_TEST_F(MapTest, MapMultipleMapsR1S0) { // Maps (lambda (x) (+ x 1)) onto an input R1F32 vector of length 0, and then // maps (lambda (x) (* x 2)) on the result. ComputationBuilder builder(client_, TestName()); - std::unique_ptr param0_literal = LiteralUtil::CreateR1({}); + std::unique_ptr param0_literal = Literal::CreateR1({}); std::unique_ptr param0_data = client_->TransferToServer(*param0_literal).ConsumeValueOrDie(); @@ -271,7 +271,7 @@ TEST_F(MapTest, MapMultipleMapsR1S4) { // maps (lambda (x) (* x 2)) on the result. ComputationBuilder builder(client_, TestName()); std::unique_ptr param0_literal = - LiteralUtil::CreateR1({2.2f, 3.3f, 4.4f, 5.5f}); + Literal::CreateR1({2.2f, 3.3f, 4.4f, 5.5f}); std::unique_ptr param0_data = client_->TransferToServer(*param0_literal).ConsumeValueOrDie(); @@ -286,7 +286,7 @@ TEST_F(MapTest, MapMultipleMapsR1S4) { TEST_F(MapTest, MapEachElemPlusOneR2) { // Maps (lambda (x) (+ x 1)) onto an input R2F32 vector. ComputationBuilder builder(client_, TestName()); - std::unique_ptr param0_literal = LiteralUtil::CreateR2( + std::unique_ptr param0_literal = Literal::CreateR2( {{13.25f, 14.0f}, {-7.1f, -7.2f}, {-8.8f, 8.8f}}); std::unique_ptr param0_data = client_->TransferToServer(*param0_literal).ConsumeValueOrDie(); @@ -383,11 +383,11 @@ TEST_F(MapTest, MapBinaryAdder) { // Maps (lambda (x y) (+ x y)) onto two R1F32 vectors. ComputationBuilder builder(client_, TestName()); std::unique_ptr param0_literal = - LiteralUtil::CreateR1({2.2f, 3.3f, 4.4f, 5.5f}); + Literal::CreateR1({2.2f, 3.3f, 4.4f, 5.5f}); std::unique_ptr param0_data = client_->TransferToServer(*param0_literal).ConsumeValueOrDie(); std::unique_ptr param1_literal = - LiteralUtil::CreateR1({5.1f, 4.4f, -0.1f, -5.5f}); + Literal::CreateR1({5.1f, 4.4f, -0.1f, -5.5f}); std::unique_ptr param1_data = client_->TransferToServer(*param1_literal).ConsumeValueOrDie(); @@ -432,12 +432,12 @@ XLA_TEST_F(MapTest, AddWithMixedLayouts) { XLA_TEST_F(MapTest, AddR3_3x0x2) { ComputationBuilder builder(client_, TestName()); std::unique_ptr param0_literal = - LiteralUtil::CreateR3FromArray3D(Array3D(3, 0, 2)); + Literal::CreateR3FromArray3D(Array3D(3, 0, 2)); std::unique_ptr param0_data = client_->TransferToServer(*param0_literal).ConsumeValueOrDie(); std::unique_ptr param1_literal = - LiteralUtil::CreateR3FromArray3D(Array3D(3, 0, 2)); + Literal::CreateR3FromArray3D(Array3D(3, 0, 2)); std::unique_ptr param1_data = client_->TransferToServer(*param1_literal).ConsumeValueOrDie(); @@ -454,15 +454,15 @@ TEST_F(MapTest, MapTernaryAdder) { // Maps (lambda (x y z) (+ x y z)) onto three R1F32 vectors. ComputationBuilder builder(client_, TestName()); std::unique_ptr param0_literal = - LiteralUtil::CreateR1({2.2f, 3.3f, 4.4f, 5.5f}); + Literal::CreateR1({2.2f, 3.3f, 4.4f, 5.5f}); std::unique_ptr param0_data = client_->TransferToServer(*param0_literal).ConsumeValueOrDie(); std::unique_ptr param1_literal = - LiteralUtil::CreateR1({5.1f, 4.4f, -0.1f, -5.5f}); + Literal::CreateR1({5.1f, 4.4f, -0.1f, -5.5f}); std::unique_ptr param1_data = client_->TransferToServer(*param1_literal).ConsumeValueOrDie(); std::unique_ptr param2_literal = - LiteralUtil::CreateR1({-10.0f, -100.0f, -900.0f, -400.0f}); + Literal::CreateR1({-10.0f, -100.0f, -900.0f, -400.0f}); std::unique_ptr param2_data = client_->TransferToServer(*param2_literal).ConsumeValueOrDie(); @@ -515,11 +515,11 @@ TEST_F(MapTest, MapOperantionWithBuildError) { auto error_add = sub_builder->BuildAndNoteError(); std::unique_ptr param0_literal = - LiteralUtil::CreateR1({2.2f, 3.3f, 4.4f, 5.5f}); + Literal::CreateR1({2.2f, 3.3f, 4.4f, 5.5f}); std::unique_ptr param0_data = client_->TransferToServer(*param0_literal).ConsumeValueOrDie(); std::unique_ptr param1_literal = - LiteralUtil::CreateR1({5.1f, 4.4f, -0.1f, -5.5f}); + Literal::CreateR1({5.1f, 4.4f, -0.1f, -5.5f}); std::unique_ptr param1_data = client_->TransferToServer(*param1_literal).ConsumeValueOrDie(); @@ -529,9 +529,10 @@ TEST_F(MapTest, MapOperantionWithBuildError) { StatusOr computation_status = builder.Build(); ASSERT_TRUE(!computation_status.ok()); - EXPECT_MATCH(computation_status.status().ToString(), - testing::HasSubstr("error from: ErrorAdd: binary op with " - "different element types: f32[] and u16[]")); + EXPECT_THAT( + computation_status.status().ToString(), + ::testing::HasSubstr("error from: ErrorAdd: binary op BINOP_ADD with " + "different element types: f32[] and u16[]")); } // MapTest disables inline and algsimp. MapTestWithFullOpt runs all @@ -552,8 +553,8 @@ TEST_F(MapTestWithFullOpt, MapScalarPower) { sub_builder->Pow(x, y); auto power = sub_builder->BuildAndNoteError(); - std::unique_ptr param0_literal = LiteralUtil::CreateR0(2.0f); - std::unique_ptr param1_literal = LiteralUtil::CreateR0(5.0f); + std::unique_ptr param0_literal = Literal::CreateR0(2.0f); + std::unique_ptr param1_literal = Literal::CreateR0(5.0f); std::unique_ptr param0_data = client_->TransferToServer(*param0_literal).ConsumeValueOrDie(); std::unique_ptr param1_data = @@ -579,8 +580,8 @@ TEST_F(MapTestWithFullOpt, MapSubtractOppositeOrder) { sub_builder->Sub(y, x); // note that this is y - x, not x - y auto sub_opposite = sub_builder->BuildAndNoteError(); - std::unique_ptr param0_literal = LiteralUtil::CreateR0(2.0f); - std::unique_ptr param1_literal = LiteralUtil::CreateR0(5.0f); + std::unique_ptr param0_literal = Literal::CreateR0(2.0f); + std::unique_ptr param1_literal = Literal::CreateR0(5.0f); std::unique_ptr param0_data = client_->TransferToServer(*param0_literal).ConsumeValueOrDie(); std::unique_ptr param1_data = @@ -604,7 +605,7 @@ TEST_F(MapTestWithFullOpt, MapSquare) { sub_builder->Mul(x, x); auto square = sub_builder->BuildAndNoteError(); - std::unique_ptr param0_literal = LiteralUtil::CreateR0(10.0f); + std::unique_ptr param0_literal = Literal::CreateR0(10.0f); std::unique_ptr param0_data = client_->TransferToServer(*param0_literal).ConsumeValueOrDie(); @@ -617,20 +618,3 @@ TEST_F(MapTestWithFullOpt, MapSquare) { } // namespace } // namespace xla - -int main(int argc, char** argv) { - std::vector flag_list; - xla::legacy_flags::AppendCpuCompilerFlags(&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; - } - testing::InitGoogleTest(&argc, argv); - if (argc > 1) { - LOG(ERROR) << "Unknown argument " << argv[1] << "\n" << usage; - return 2; - } - return RUN_ALL_TESTS(); -} diff --git a/tensorflow/compiler/xla/tests/matrix_ops_simple_test.cc b/tensorflow/compiler/xla/tests/matrix_ops_simple_test.cc index 8aa4029440653763de120ce10cf4145066662bf3..4c33bb2c3661f185c93f798cd4e989f0b39178c1 100644 --- a/tensorflow/compiler/xla/tests/matrix_ops_simple_test.cc +++ b/tensorflow/compiler/xla/tests/matrix_ops_simple_test.cc @@ -21,7 +21,6 @@ limitations under the License. #include "tensorflow/compiler/xla/client/computation.h" #include "tensorflow/compiler/xla/client/computation_builder.h" #include "tensorflow/compiler/xla/client/local_client.h" -#include "tensorflow/compiler/xla/legacy_flags/cpu_compiler_flags.h" #include "tensorflow/compiler/xla/literal_util.h" #include "tensorflow/compiler/xla/ptr_util.h" #include "tensorflow/compiler/xla/reference_util.h" @@ -87,8 +86,8 @@ TEST_F(MatOpsSimpleTest, ExpTwoByTwoValues) { builder.Exp(data); std::unique_ptr expected = - LiteralUtil::CreateR2({{2.71828, 1.00000}, // row 0 - {0.36788, 1.64872}}); // row 1 + Literal::CreateR2({{2.71828, 1.00000}, // row 0 + {0.36788, 1.64872}}); // row 1 ComputeAndCompareLiteral(&builder, *expected, {}, ErrorSpec(1e-5)); } @@ -115,8 +114,8 @@ TEST_F(MatOpsSimpleTest, MapTwoByTwo) { auto map = builder.Map({data}, add_half); std::unique_ptr expected = - LiteralUtil::CreateR2({{1.5, 0.5}, // row 0 - {-0.5, 1.0}}); // row 1 + Literal::CreateR2({{1.5, 0.5}, // row 0 + {-0.5, 1.0}}); // row 1 ComputeAndCompareLiteral(&builder, *expected, {}, ErrorSpec(1e-5)); } @@ -133,8 +132,8 @@ TEST_F(MatOpsSimpleTest, MaxTwoByTwoValues) { auto max = builder.Max(lhs, rhs); std::unique_ptr expected = - LiteralUtil::CreateR2({{7.0, 6.0}, // row 0 - {3.0, -4.0}}); // row 1 + Literal::CreateR2({{7.0, 6.0}, // row 0 + {3.0, -4.0}}); // row 1 ComputeAndCompareLiteral(&builder, *expected, {}, ErrorSpec(1e-6)); } @@ -158,22 +157,69 @@ TEST_F(MatOpsSimpleTest, Max32x8Linspace) { TestLinspaceMax(32, 8); } TEST_F(MatOpsSimpleTest, Max64x8Linspace) { TestLinspaceMax(64, 8); } -} // namespace -} // namespace xla - -int main(int argc, char** argv) { - std::vector flag_list; - xla::legacy_flags::AppendCpuCompilerFlags(&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; +class MatOpsDotAddTest + : public ClientLibraryTestBase, + public ::testing::WithParamInterface> {}; + +TEST_P(MatOpsDotAddTest, Dot_Add_2x2_2x2) { + bool row_major = std::get<0>(GetParam()); + bool add_lhs = std::get<1>(GetParam()); + bool transpose = std::get<2>(GetParam()); + Array2D lhs({{1.0, 2.0}, {3.0, 4.0}}); + Array2D rhs({{10.0, 11.0}, {12.0, 13.0}}); + + auto minor_to_major = [](bool row_major) -> std::vector { + return {row_major ? 1 : 0, row_major ? 0 : 1}; + }; + + auto prim_type = primitive_util::NativeToPrimitiveType(); + Shape lhs_shape = + ShapeUtil::MakeShape(prim_type, {lhs.height(), lhs.width()}); + Shape rhs_shape = + ShapeUtil::MakeShape(prim_type, {rhs.height(), rhs.width()}); + + TF_ASSERT_OK_AND_ASSIGN( + auto lhs_handle, + client_->TransferToServer(*Literal::CreateR2FromArray2DWithLayout( + lhs, LayoutUtil::MakeLayout(minor_to_major(row_major))))); + TF_ASSERT_OK_AND_ASSIGN( + auto rhs_handle, + client_->TransferToServer(*Literal::CreateR2FromArray2DWithLayout( + rhs, LayoutUtil::MakeLayout(minor_to_major(row_major))))); + + ComputationBuilder builder(client_, TestName()); + auto lhs_arg = builder.Parameter(0, lhs_shape, "lhs"); + auto lhs_mat_arg = lhs_arg; + if (transpose) { + lhs_mat_arg = builder.Transpose(lhs_mat_arg, {1, 0}); } - testing::InitGoogleTest(&argc, argv); - if (argc > 1) { - LOG(ERROR) << "Unknown argument " << argv[1] << "\n" << usage; - return 2; + auto rhs_arg = builder.Parameter(1, rhs_shape, "rhs"); + auto result = builder.Dot(lhs_mat_arg, rhs_arg); + Array2D expected; + if (add_lhs) { + result = builder.Add(result, lhs_arg); + if (transpose) { + expected = Array2D({{47, 52}, {71, 78}}); + } else { + expected = Array2D({{35, 39}, {81, 89}}); + } + } else { + result = builder.Add(result, rhs_arg); + if (transpose) { + expected = Array2D({{56, 61}, {80, 87}}); + } else { + expected = Array2D({{44, 48}, {90, 98}}); + } } - return RUN_ALL_TESTS(); + + ComputeAndCompareR2(&builder, expected, + {lhs_handle.get(), rhs_handle.get()}, + ErrorSpec(1e-6)); } + +INSTANTIATE_TEST_CASE_P(MatOpsDotAddTestInstances, MatOpsDotAddTest, + ::testing::Combine(::testing::Bool(), ::testing::Bool(), + ::testing::Bool())); + +} // namespace +} // namespace xla diff --git a/tensorflow/compiler/xla/tests/multidimensional_slice_test.cc b/tensorflow/compiler/xla/tests/multidimensional_slice_test.cc index 2cd680399b3fd4f8d8742f3142ce400b18044fbd..11c0bf7a5a5bde9edcfb7f76a5c10ac4dd77bcee 100644 --- a/tensorflow/compiler/xla/tests/multidimensional_slice_test.cc +++ b/tensorflow/compiler/xla/tests/multidimensional_slice_test.cc @@ -21,7 +21,6 @@ limitations under the License. #include "tensorflow/compiler/xla/array3d.h" #include "tensorflow/compiler/xla/client/computation_builder.h" #include "tensorflow/compiler/xla/client/local_client.h" -#include "tensorflow/compiler/xla/legacy_flags/cpu_compiler_flags.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" @@ -36,7 +35,7 @@ XLA_TEST_F(SliceTest, Slice2D) { ComputationBuilder builder(client_, "slice_2d"); auto original = builder.ConstantR2( {{1.0, 2.0, 3.0}, {4.0, 5.0, 6.0}, {7.0, 8.0, 9.0}, {10.0, 11.0, 12.0}}); - builder.Slice(original, {2, 1}, {4, 3}); + builder.Slice(original, {2, 1}, {4, 3}, {1, 1}); Array2D expected({{8.0f, 9.0f}, {11.0f, 12.0f}}); ComputeAndCompareR2(&builder, expected, {}, ErrorSpec(0.000001)); @@ -47,7 +46,7 @@ XLA_TEST_F(SliceTest, Slice3D) { Array3D array_3d( {{{1.0f, 2.0f}, {3.0f, 4.0f}}, {{5.0f, 6.0f}, {7.0f, 8.0f}}}); auto original = builder.ConstantR3FromArray3D(array_3d); - builder.Slice(original, {0, 0, 1}, {2, 1, 2}); + builder.Slice(original, {0, 0, 1}, {2, 1, 2}, {1, 1, 1}); Array3D expected_3d({{{2.0f}}, {{6.0f}}}); ComputeAndCompareR3(&builder, expected_3d, {}, ErrorSpec(0.000001)); @@ -55,20 +54,3 @@ XLA_TEST_F(SliceTest, Slice3D) { } // namespace } // namespace xla - -int main(int argc, char** argv) { - std::vector flag_list; - xla::legacy_flags::AppendCpuCompilerFlags(&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; - } - testing::InitGoogleTest(&argc, argv); - if (argc > 1) { - LOG(ERROR) << "Unknown argument " << argv[1] << "\n" << usage; - return 2; - } - return RUN_ALL_TESTS(); -} diff --git a/tensorflow/compiler/xla/tests/multioutput_fusion_test.cc b/tensorflow/compiler/xla/tests/multioutput_fusion_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..606d801c84e8c6607fa5703a4b48d77bab0b045b --- /dev/null +++ b/tensorflow/compiler/xla/tests/multioutput_fusion_test.cc @@ -0,0 +1,176 @@ +/* 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 "tensorflow/compiler/xla/client/computation_builder.h" +#include "tensorflow/compiler/xla/client/local_client.h" +#include "tensorflow/compiler/xla/literal_util.h" +#include "tensorflow/compiler/xla/primitive_util.h" +#include "tensorflow/compiler/xla/ptr_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/service/hlo_opcode.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/compiler/xla/xla_data.pb.h" +#include "tensorflow/core/lib/gtl/array_slice.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" +#include "tensorflow/core/platform/types.h" + +namespace xla { +namespace { + +using ::tensorflow::gtl::ArraySlice; + +class MultiOutputFusionTest : public HloTestBase { + protected: + MultiOutputFusionTest() { error_spec_ = ErrorSpec{0.0001, 1e-2}; } + + void RunTest2D(bool manual_fusion, int64 size) { + auto builder = HloComputation::Builder(TestName()); + auto hlo_module = CreateNewModule(); + + const Shape elem_shape0 = ShapeUtil::MakeShape(F32, {}); + const Shape elem_shape2 = ShapeUtil::MakeShape(F32, {size, size}); + + auto const0 = builder.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(8.0f))); + auto param0 = builder.AddInstruction( + HloInstruction::CreateParameter(0, elem_shape0, "0")); + + auto add1 = builder.AddInstruction(HloInstruction::CreateBinary( + elem_shape0, HloOpcode::kAdd, param0, const0)); + + HloInstruction* broadcast = builder.AddInstruction( + HloInstruction::CreateBroadcast(elem_shape2, add1, {0, 1})); + + auto param1 = builder.AddInstruction( + HloInstruction::CreateParameter(1, elem_shape2, "1")); + + HloInstruction* add2 = builder.AddInstruction(HloInstruction::CreateBinary( + elem_shape2, HloOpcode::kAdd, broadcast, param1)); + HloInstruction* sub = builder.AddInstruction(HloInstruction::CreateBinary( + elem_shape2, HloOpcode::kSubtract, param1, broadcast)); + HloInstruction* dot = builder.AddInstruction( + HloInstruction::CreateBinary(elem_shape2, HloOpcode::kDot, sub, add2)); + auto computation = hlo_module->AddEntryComputation(builder.Build(dot)); + + if (manual_fusion) { + auto tuple = computation->AddInstruction(HloInstruction::CreateTuple( + ArraySlice({sub, add2}, 0, 2))); + auto gte0 = computation->AddInstruction( + HloInstruction::CreateGetTupleElement(elem_shape2, tuple, 0)); + auto gte1 = computation->AddInstruction( + HloInstruction::CreateGetTupleElement(elem_shape2, tuple, 1)); + TF_CHECK_OK(dot->ReplaceOperandWith(0, gte0)); + TF_CHECK_OK(dot->ReplaceOperandWith(1, gte1)); + + CHECK_NE( + computation->CreateFusionInstruction( + {tuple, sub, add2, broadcast}, HloInstruction::FusionKind::kLoop), + nullptr); + } + + Literal input; + input.PopulateWithValue(2.5f, {size, size}); + auto p1 = TransferToDevice(input); + auto p0 = TransferToDevice(*Literal::CreateR0(-9.0f)); + + Literal expect; + expect.PopulateWithValue(size * 1.5f * 3.5f, {size, size}); + auto actual = ExecuteAndTransfer(std::move(hlo_module), {p0, p1}); + LiteralTestUtil::ExpectNear(expect, *actual, error_spec_); + } + + void RunTest1D(bool manual_fusion, int size) { + auto builder = HloComputation::Builder(TestName()); + auto hlo_module = CreateNewModule(); + + const Shape elem_shape_F32 = ShapeUtil::MakeShape(F32, {size}); + const Shape elem_shape_U8 = ShapeUtil::MakeShape(F64, {size}); + auto param0 = builder.AddInstruction( + HloInstruction::CreateParameter(0, elem_shape_F32, "0")); + auto param1 = builder.AddInstruction( + HloInstruction::CreateParameter(1, elem_shape_U8, "1")); + + HloInstruction* param0_U8 = builder.AddInstruction( + HloInstruction::CreateConvert(elem_shape_U8, param0)); + HloInstruction* param1_F32 = builder.AddInstruction( + HloInstruction::CreateConvert(elem_shape_F32, param1)); + HloInstruction* add = builder.AddInstruction(HloInstruction::CreateBinary( + elem_shape_F32, HloOpcode::kAdd, param0, param1_F32)); + HloInstruction* sub_U8 = + builder.AddInstruction(HloInstruction::CreateBinary( + elem_shape_U8, HloOpcode::kSubtract, param0_U8, param1)); + HloInstruction* sub = builder.AddInstruction( + HloInstruction::CreateConvert(elem_shape_F32, sub_U8)); + + HloInstruction* reshape = + builder.AddInstruction(HloInstruction::CreateReshape( + ShapeUtil::MakeShape(F32, {size, 1}), add)); + HloInstruction* dot = builder.AddInstruction(HloInstruction::CreateBinary( + ShapeUtil::MakeShape(F32, {1}), HloOpcode::kDot, sub, reshape)); + auto computation = hlo_module->AddEntryComputation(builder.Build(dot)); + + if (manual_fusion) { + auto tuple = computation->AddInstruction(HloInstruction::CreateTuple( + ArraySlice({sub_U8, add}, 0, 2))); + + auto gte0 = computation->AddInstruction( + HloInstruction::CreateGetTupleElement(elem_shape_U8, tuple, 0)); + auto gte1 = computation->AddInstruction( + HloInstruction::CreateGetTupleElement(elem_shape_F32, tuple, 1)); + TF_CHECK_OK(sub->ReplaceOperandWith(0, gte0)); + TF_CHECK_OK(reshape->ReplaceOperandWith(0, gte1)); + + CHECK_NE(computation->CreateFusionInstruction( + {tuple, sub_U8, add, param0_U8, param1_F32}, + HloInstruction::FusionKind::kLoop), + nullptr); + } + + Literal input0, input1; + input0.PopulateWithValue(2.5f, {size}); + input1.PopulateWithValue(1, {size}); + auto p0 = TransferToDevice(input0); + auto p1 = TransferToDevice(input1); + + Literal expect = *Literal::CreateR1({size * 1.5f * 3.5f}); + auto actual = ExecuteAndTransfer(std::move(hlo_module), {p0, p1}); + LiteralTestUtil::ExpectNear(expect, *actual, error_spec_); + } +}; + +XLA_TEST_F(MultiOutputFusionTest, 2DNofusion) { RunTest2D(false, 5); } +XLA_TEST_F(MultiOutputFusionTest, 2DFusion) { RunTest2D(true, 5); } +XLA_TEST_F(MultiOutputFusionTest, 2DFusionSize129) { RunTest2D(true, 129); } +XLA_TEST_F(MultiOutputFusionTest, DiffentTypesNoFusion) { RunTest1D(false, 8); } +XLA_TEST_F(MultiOutputFusionTest, DiffentTypesFusion) { RunTest1D(true, 8); } + +} // namespace +} // namespace xla diff --git a/tensorflow/compiler/xla/tests/pad_test.cc b/tensorflow/compiler/xla/tests/pad_test.cc index f044c94b8d0b780636b96e36930650655a728f0a..3fd83a4c3b104831f03366339fb7b8b5d816a3f7 100644 --- a/tensorflow/compiler/xla/tests/pad_test.cc +++ b/tensorflow/compiler/xla/tests/pad_test.cc @@ -21,7 +21,6 @@ limitations under the License. #include "tensorflow/compiler/xla/client/computation_builder.h" #include "tensorflow/compiler/xla/client/lib/arithmetic.h" #include "tensorflow/compiler/xla/client/local_client.h" -#include "tensorflow/compiler/xla/legacy_flags/cpu_compiler_flags.h" #include "tensorflow/compiler/xla/ptr_util.h" #include "tensorflow/compiler/xla/reference_util.h" #include "tensorflow/compiler/xla/tests/client_library_test_base.h" @@ -182,8 +181,8 @@ TEST_F(PadTest, Pad4DFloatArrayMinorFirstSmall) { const float pad_value = -5.123f; Array4D input_array(1, 1, 2, 3, {1, 2, 3, 4, 5, 6}); - auto input = LiteralUtil::CreateR4FromArray4D(input_array); - input = LiteralUtil::Relayout(*input, layout); + auto input = Literal::CreateR4FromArray4D(input_array); + input = input->Relayout(layout); b.Pad(b.ConstantLiteral(*input), b.ConstantR0(pad_value), padding_config); @@ -227,8 +226,8 @@ XLA_TEST_F(PadTest, Pad4DFloatArrayMinorFirstNonTrivialMinorDimensions) { input_array(0, 0, 0, 0) = 1.0f; input_array(0, 24, 6, 6) = 2.0f; input_array(0, 17, 2, 5) = 3.0f; - auto input = LiteralUtil::CreateR4FromArray4D(input_array); - input = LiteralUtil::Relayout(*input, layout); + auto input = Literal::CreateR4FromArray4D(input_array); + input = input->Relayout(layout); b.Pad(b.ConstantLiteral(*input), b.ConstantR0(pad_value), padding_config); @@ -307,7 +306,7 @@ XLA_TEST_F(PadTest, Large2DPad) { auto ones = MakeUnique>(4, 4); ones->Fill(1.0f); - auto input_literal = LiteralUtil::CreateR2FromArray2D(*ones); + auto input_literal = Literal::CreateR2FromArray2D(*ones); std::unique_ptr input_data = client_->TransferToServer(*input_literal).ConsumeValueOrDie(); @@ -333,7 +332,7 @@ XLA_TEST_F(PadTest, AllTypes2DPad) { auto operand = MakeUnique>(in_rows, in_cols); operand->FillUnique(0.0f); - auto input_literal = LiteralUtil::CreateR2FromArray2D(*operand); + auto input_literal = Literal::CreateR2FromArray2D(*operand); std::unique_ptr input_data = client_->TransferToServer(*input_literal).ConsumeValueOrDie(); @@ -364,7 +363,7 @@ XLA_TEST_F(PadTest, High2DPad) { auto operand = MakeUnique>(in_rows, in_cols); operand->FillUnique(1.0f); - auto input_literal = LiteralUtil::CreateR2FromArray2D(*operand); + auto input_literal = Literal::CreateR2FromArray2D(*operand); auto expected = ReferenceUtil::PadArray2D(*operand, padding_config, 2.718f); std::unique_ptr input_data = client_->TransferToServer(*input_literal).ConsumeValueOrDie(); @@ -396,7 +395,7 @@ XLA_TEST_F(PadTest, NegativePadding2D) { auto operand = MakeUnique>(in_rows, in_cols); operand->FillUnique(1.0f); - auto input_literal = LiteralUtil::CreateR2FromArray2D(*operand); + auto input_literal = Literal::CreateR2FromArray2D(*operand); auto expected = ReferenceUtil::PadArray2D(*operand, padding_config, 2.718f); std::unique_ptr input_data = client_->TransferToServer(*input_literal).ConsumeValueOrDie(); @@ -428,7 +427,7 @@ XLA_TEST_F(PadTest, NegativeAndInteriorPadding2D) { auto operand = MakeUnique>(in_rows, in_cols); operand->FillUnique(1.0f); - auto input_literal = LiteralUtil::CreateR2FromArray2D(*operand); + auto input_literal = Literal::CreateR2FromArray2D(*operand); auto expected = ReferenceUtil::PadArray2D(*operand, padding_config, 2.718f); std::unique_ptr input_data = client_->TransferToServer(*input_literal).ConsumeValueOrDie(); @@ -452,7 +451,7 @@ XLA_TEST_F(PadTest, ReducePad) { auto ones = MakeUnique>(2, 2, 2, 2); ones->Fill(1.0); - auto input_literal = LiteralUtil::CreateR4FromArray4D(*ones); + auto input_literal = Literal::CreateR4FromArray4D(*ones); std::unique_ptr input_data = client_->TransferToServer(*input_literal).ConsumeValueOrDie(); @@ -465,20 +464,3 @@ XLA_TEST_F(PadTest, ReducePad) { } // namespace } // namespace xla - -int main(int argc, char** argv) { - std::vector flag_list; - xla::legacy_flags::AppendCpuCompilerFlags(&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; - } - testing::InitGoogleTest(&argc, argv); - if (argc > 1) { - LOG(ERROR) << "Unknown argument " << argv[1] << "\n" << usage; - return 2; - } - return RUN_ALL_TESTS(); -} diff --git a/tensorflow/compiler/xla/tests/params_test.cc b/tensorflow/compiler/xla/tests/params_test.cc index 2f05576ceeb6c4142136235575d06aa63f22ba2b..aa84b8ff1aec9e801dfa0cc9a79b775797b98b81 100644 --- a/tensorflow/compiler/xla/tests/params_test.cc +++ b/tensorflow/compiler/xla/tests/params_test.cc @@ -24,7 +24,6 @@ limitations under the License. #include "tensorflow/compiler/xla/client/global_data.h" #include "tensorflow/compiler/xla/client/local_client.h" #include "tensorflow/compiler/xla/layout_util.h" -#include "tensorflow/compiler/xla/legacy_flags/cpu_compiler_flags.h" #include "tensorflow/compiler/xla/literal_util.h" #include "tensorflow/compiler/xla/shape_util.h" #include "tensorflow/compiler/xla/statusor.h" @@ -43,8 +42,7 @@ class ParamsTest : public ClientLibraryTestBase {}; XLA_TEST_F(ParamsTest, ConstantR0F32Param) { ComputationBuilder builder(client_, TestName()); - std::unique_ptr param0_literal = - LiteralUtil::CreateR0(3.14159f); + std::unique_ptr param0_literal = Literal::CreateR0(3.14159f); std::unique_ptr param0_data = client_->TransferToServer(*param0_literal).ConsumeValueOrDie(); @@ -56,7 +54,7 @@ XLA_TEST_F(ParamsTest, ConstantR0F32Param) { XLA_TEST_F(ParamsTest, ConstantR1S0F32Param) { ComputationBuilder builder(client_, TestName()); - std::unique_ptr param0_literal = LiteralUtil::CreateR1({}); + std::unique_ptr param0_literal = Literal::CreateR1({}); std::unique_ptr param0_data = client_->TransferToServer(*param0_literal).ConsumeValueOrDie(); @@ -69,7 +67,7 @@ XLA_TEST_F(ParamsTest, ConstantR1S0F32Param) { XLA_TEST_F(ParamsTest, ConstantR1S2F32Param) { ComputationBuilder builder(client_, TestName()); std::unique_ptr param0_literal = - LiteralUtil::CreateR1({3.14f, -100.25f}); + Literal::CreateR1({3.14f, -100.25f}); std::unique_ptr param0_data = client_->TransferToServer(*param0_literal).ConsumeValueOrDie(); @@ -82,7 +80,7 @@ XLA_TEST_F(ParamsTest, ConstantR1S2F32Param) { XLA_TEST_F(ParamsTest, ConstantR1U8Param) { ComputationBuilder builder(client_, TestName()); string str("hello world"); - std::unique_ptr param0_literal = LiteralUtil::CreateR1U8(str); + std::unique_ptr param0_literal = Literal::CreateR1U8(str); std::unique_ptr param0_data = client_->TransferToServer(*param0_literal).ConsumeValueOrDie(); @@ -95,7 +93,7 @@ XLA_TEST_F(ParamsTest, ConstantR1U8Param) { XLA_TEST_F(ParamsTest, ConstantR2_3x0_F32Param) { ComputationBuilder builder(client_, TestName()); std::unique_ptr param0_literal = - LiteralUtil::CreateR2FromArray2D(Array2D(3, 0)); + Literal::CreateR2FromArray2D(Array2D(3, 0)); std::unique_ptr param0_data = client_->TransferToServer(*param0_literal).ConsumeValueOrDie(); @@ -107,7 +105,7 @@ XLA_TEST_F(ParamsTest, ConstantR2_3x0_F32Param) { XLA_TEST_F(ParamsTest, ConstantR2F32Param) { ComputationBuilder builder(client_, TestName()); - std::unique_ptr param0_literal = LiteralUtil::CreateR2( + std::unique_ptr param0_literal = Literal::CreateR2( {{3.14f, -100.25f}, {7e8f, 7e-9f}, {30.3f, -100.0f}}); std::unique_ptr param0_data = client_->TransferToServer(*param0_literal).ConsumeValueOrDie(); @@ -123,12 +121,12 @@ XLA_TEST_F(ParamsTest, ConstantR2F32Param) { XLA_TEST_F(ParamsTest, TwoParameters) { ComputationBuilder builder(client_, TestName()); - std::unique_ptr literal0 = LiteralUtil::CreateR1({1, 2}); + std::unique_ptr literal0 = Literal::CreateR1({1, 2}); std::unique_ptr param0_data = client_->TransferToServer(*literal0).ConsumeValueOrDie(); auto param0 = builder.Parameter(0, literal0->shape(), "param0"); - std::unique_ptr literal1 = LiteralUtil::CreateR1({10, 20}); + std::unique_ptr literal1 = Literal::CreateR1({10, 20}); std::unique_ptr param1_data = client_->TransferToServer(*literal1).ConsumeValueOrDie(); auto param1 = builder.Parameter(1, literal1->shape(), "param1"); @@ -154,7 +152,7 @@ XLA_TEST_F(ParamsTest, TwoParameters) { XLA_TEST_F(ParamsTest, MissingParameter) { // Test that an error is returned when a computation with an incomplete set of // parameters (parameter numbers not contiguous from 0) is executed. - std::unique_ptr literal = LiteralUtil::CreateR0(3.14159f); + std::unique_ptr literal = Literal::CreateR0(3.14159f); std::unique_ptr data = client_->TransferToServer(*literal).ConsumeValueOrDie(); @@ -163,7 +161,7 @@ XLA_TEST_F(ParamsTest, MissingParameter) { auto computation = builder.Build().ConsumeValueOrDie(); auto execute_status = client_->Execute(computation, {data.get(), data.get()}, - /*output_layout=*/nullptr, + /*execution_options=*/nullptr, /*execution_profile=*/nullptr); ASSERT_EQ(execute_status.status().code(), tensorflow::error::FAILED_PRECONDITION); @@ -172,12 +170,12 @@ XLA_TEST_F(ParamsTest, MissingParameter) { XLA_TEST_F(ParamsTest, UnusedParameter) { ComputationBuilder builder(client_, TestName()); - std::unique_ptr literal0 = LiteralUtil::CreateR1({1, 2}); + std::unique_ptr literal0 = Literal::CreateR1({1, 2}); std::unique_ptr param0_data = client_->TransferToServer(*literal0).ConsumeValueOrDie(); auto param0 = builder.Parameter(0, literal0->shape(), "param0"); - std::unique_ptr literal1 = LiteralUtil::CreateR1({10, 20}); + std::unique_ptr literal1 = Literal::CreateR1({10, 20}); std::unique_ptr param1_data = client_->TransferToServer(*literal1).ConsumeValueOrDie(); auto param1 = builder.Parameter(1, literal1->shape(), "param1"); @@ -192,12 +190,11 @@ XLA_TEST_F(ParamsTest, UnusedParametersInUnusedExpression) { // unused expression. ComputationBuilder builder(client_, TestName()); - std::unique_ptr literal0 = LiteralUtil::CreateR1({1, 2}); + std::unique_ptr literal0 = Literal::CreateR1({1, 2}); std::unique_ptr param0_data = client_->TransferToServer(*literal0).ConsumeValueOrDie(); - std::unique_ptr literal1 = - LiteralUtil::CreateR1({10, 20, 30}); + std::unique_ptr literal1 = Literal::CreateR1({10, 20, 30}); std::unique_ptr param1_data = client_->TransferToServer(*literal1).ConsumeValueOrDie(); @@ -237,7 +234,7 @@ XLA_TEST_F(ParamsTest, HundredLargeR1Parameters) { std::vector sum_value = {{entry0, entry1}}; sum_value.resize(size); - std::unique_ptr literal = LiteralUtil::CreateR1(sum_value); + std::unique_ptr literal = Literal::CreateR1(sum_value); param_data_owner.push_back( client_->TransferToServer(*literal).ConsumeValueOrDie()); ComputationDataHandle param = @@ -246,6 +243,7 @@ XLA_TEST_F(ParamsTest, HundredLargeR1Parameters) { } std::vector param_data; + param_data.reserve(param_data_owner.size()); for (const std::unique_ptr& data : param_data_owner) { param_data.push_back(data.get()); } @@ -266,9 +264,9 @@ XLA_TEST_F(ParamsTest, std::unique_ptr data = client_ - ->TransferToServer(*LiteralUtil::MakeTuple({ - LiteralUtil::CreateR1({1, 2, 3}).get(), - LiteralUtil::CreateR1({4, 5, 6}).get(), + ->TransferToServer(*Literal::MakeTuple({ + Literal::CreateR1({1, 2, 3}).get(), + Literal::CreateR1({4, 5, 6}).get(), })) .ConsumeValueOrDie(); @@ -280,7 +278,7 @@ XLA_TEST_F(ParamsTest, // Verifies that passing a 2x2 with {0, 1} layout returns the same value back // when (transferred to the server and) passed through a parameter. XLA_TEST_F(ParamsTest, R2_2x2_Layout_01) { - std::unique_ptr literal = LiteralUtil::CreateR2({ + std::unique_ptr literal = Literal::CreateR2({ {1, 2}, {3, 4}, }); *literal->mutable_shape()->mutable_layout() = LayoutUtil::MakeLayout({0, 1}); @@ -294,7 +292,7 @@ XLA_TEST_F(ParamsTest, R2_2x2_Layout_01) { // As above, but for {1, 0} layout. XLA_TEST_F(ParamsTest, R2_2x2_Layout_10) { - std::unique_ptr literal = LiteralUtil::CreateR2({ + std::unique_ptr literal = Literal::CreateR2({ {1, 3}, {2, 4}, }); *literal->mutable_shape()->mutable_layout() = LayoutUtil::MakeLayout({1, 0}); @@ -307,7 +305,7 @@ XLA_TEST_F(ParamsTest, R2_2x2_Layout_10) { } XLA_TEST_F(ParamsTest, R2_2x2_TryToPassReverseLayoutToParameter) { - std::unique_ptr literal = LiteralUtil::CreateR2({ + std::unique_ptr literal = Literal::CreateR2({ {1, 3}, {2, 4}, }); const Shape original = literal->shape(); @@ -320,13 +318,13 @@ XLA_TEST_F(ParamsTest, R2_2x2_TryToPassReverseLayoutToParameter) { std::reverse(original_layout.begin(), original_layout.end()); *literal->mutable_shape()->mutable_layout() = LayoutUtil::MakeLayout(original_layout); - ASSERT_EQ(2, LiteralUtil::Get(*literal, {0, 1})); + ASSERT_EQ(2, literal->Get({0, 1})); } // Use the original shape in building the computation. ComputationBuilder builder(client_, TestName()); auto input = builder.Parameter(0, original, "input"); // Use the slice operator to get an off-diagonal element. - builder.Slice(input, {0, 1}, {1, 2}); + builder.Slice(input, {0, 1}, {1, 2}, {1, 1}); std::unique_ptr data = client_->TransferToServer(*literal).ConsumeValueOrDie(); @@ -338,20 +336,3 @@ XLA_TEST_F(ParamsTest, R2_2x2_TryToPassReverseLayoutToParameter) { } // namespace } // namespace xla - -int main(int argc, char** argv) { - std::vector flag_list; - xla::legacy_flags::AppendCpuCompilerFlags(&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; - } - testing::InitGoogleTest(&argc, argv); - if (argc > 1) { - LOG(ERROR) << "Unknown argument " << argv[1] << "\n" << usage; - return 2; - } - return RUN_ALL_TESTS(); -} diff --git a/tensorflow/compiler/xla/tests/plugin.bzl b/tensorflow/compiler/xla/tests/plugin.bzl new file mode 100644 index 0000000000000000000000000000000000000000..8a5d91363b619c6b214a96ad96e92742e3052541 --- /dev/null +++ b/tensorflow/compiler/xla/tests/plugin.bzl @@ -0,0 +1,36 @@ +# 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. +# ============================================================================== +"""Additional XLA devices to be included in the unit test suite.""" + +# Example: +# +# plugins = { +# "foo": { +# "deps": [ +# "//tensorflow/compiler/plugin/foo:foo_lib", +# "//tensorflow/compiler/plugin/foo:test_macros", +# ], +# "disabled_manifest": "tensorflow/compiler/plugin/foo/disabled_test_manifest.txt", +# "copts": [], +# "tags": [], +# "args": [] +# "data": [ +# "//tensorflow/compiler/plugin/foo:disabled_test_manifest.txt", +# ], +# }, +# } + +plugins = {} + diff --git a/tensorflow/compiler/xla/tests/pred_test.cc b/tensorflow/compiler/xla/tests/pred_test.cc index c9c7bf4cd6373ff8f2c2bdcf43abdcde2eabb311..3500e8dc28570fe216f53b746c3757e080aa689f 100644 --- a/tensorflow/compiler/xla/tests/pred_test.cc +++ b/tensorflow/compiler/xla/tests/pred_test.cc @@ -20,7 +20,6 @@ limitations under the License. #include "tensorflow/compiler/xla/client/computation_builder.h" #include "tensorflow/compiler/xla/client/lib/arithmetic.h" #include "tensorflow/compiler/xla/client/local_client.h" -#include "tensorflow/compiler/xla/legacy_flags/cpu_compiler_flags.h" #include "tensorflow/compiler/xla/tests/client_library_test_base.h" #include "tensorflow/core/lib/core/status_test_util.h" #include "tensorflow/core/platform/test.h" @@ -137,20 +136,3 @@ TEST_F(PredTest, AnyR2False) { } // namespace } // namespace xla - -int main(int argc, char** argv) { - std::vector flag_list; - xla::legacy_flags::AppendCpuCompilerFlags(&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; - } - testing::InitGoogleTest(&argc, argv); - if (argc > 1) { - LOG(ERROR) << "Unknown argument " << argv[1] << "\n" << usage; - return 2; - } - return RUN_ALL_TESTS(); -} diff --git a/tensorflow/compiler/xla/tests/prng_test.cc b/tensorflow/compiler/xla/tests/prng_test.cc index 0cd0f97b0621d771ae039f0be6bd6c67161b49a4..0f82291fea6559381b60a610222a869c999f64cf 100644 --- a/tensorflow/compiler/xla/tests/prng_test.cc +++ b/tensorflow/compiler/xla/tests/prng_test.cc @@ -17,10 +17,10 @@ limitations under the License. #include "tensorflow/compiler/xla/client/computation_builder.h" #include "tensorflow/compiler/xla/client/local_client.h" -#include "tensorflow/compiler/xla/legacy_flags/cpu_compiler_flags.h" #include "tensorflow/compiler/xla/literal_util.h" #include "tensorflow/compiler/xla/primitive_util.h" #include "tensorflow/compiler/xla/shape_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/util.h" @@ -55,12 +55,11 @@ void PrngTest::UniformTest(T a, T b, tensorflow::gtl::ArraySlice dims) { SetSeed(42); auto actual = ExecuteAndTransferOrDie(&builder, /*arguments=*/{}); - EXPECT_TRUE(ContainersEqual(dims, actual->shape().dimensions())); - LiteralUtil::EachCell(*actual, - [=](tensorflow::gtl::ArraySlice, T value) { - EXPECT_LE(a, value); - EXPECT_LT(value, b); - }); + EXPECT_THAT(dims, ::testing::ElementsAreArray(actual->shape().dimensions())); + actual->EachCell([=](tensorflow::gtl::ArraySlice, T value) { + EXPECT_LE(a, value); + EXPECT_LT(value, b); + }); } void PrngTest::BernoulliTest(float p, tensorflow::gtl::ArraySlice dims) { @@ -68,17 +67,16 @@ void PrngTest::BernoulliTest(float p, tensorflow::gtl::ArraySlice dims) { auto shape = ShapeUtil::MakeShape(U32, dims); builder.RngBernoulli(builder.ConstantR0(p), shape); - TF_ASSIGN_OR_ASSERT_OK(auto computation, builder.Build()); - ExecutionOptions execution_options; + TF_ASSERT_OK_AND_ASSIGN(auto computation, builder.Build()); + ExecutionOptions execution_options = execution_options_; execution_options.set_seed(42); - TF_ASSIGN_OR_ASSERT_OK( - auto actual, - client_->ExecuteAndTransfer(computation, /*arguments=*/{}, - &execution_options)); - EXPECT_TRUE(ContainersEqual(dims, actual->shape().dimensions())); + TF_ASSERT_OK_AND_ASSIGN( + auto actual, client_->ExecuteAndTransfer(computation, /*arguments=*/{}, + &execution_options)); + EXPECT_THAT(dims, ::testing::ElementsAreArray(actual->shape().dimensions())); int32 sum = 0; - LiteralUtil::EachCell( - *actual, [&sum](tensorflow::gtl::ArraySlice, uint32 value) { + actual->EachCell( + [&sum](tensorflow::gtl::ArraySlice, uint32 value) { EXPECT_TRUE(value == 0 || value == 1); sum += value; }); @@ -122,10 +120,8 @@ double PrngTest::UniformChiSquared(int32 range_size, int32 expected_count) { SetSeed(42); auto actual = ExecuteAndTransferOrDie(&builder, /*arguments=*/{}); std::vector counts(range_size, 0); - LiteralUtil::EachCell( - *actual, [&counts](tensorflow::gtl::ArraySlice, int32 value) { - ++counts[value]; - }); + actual->EachCell([&counts](tensorflow::gtl::ArraySlice, + int32 value) { ++counts[value]; }); int64 sum = 0; for (int32 i = 0; i < range_size; ++i) { sum += Square(static_cast(counts[i] - expected_count)); @@ -168,23 +164,22 @@ XLA_TEST_F(PrngTest, MapUsingRng) { ComputationBuilder builder(client_, TestName()); std::unique_ptr param0_literal = - LiteralUtil::CreateR1({2.2f, 5.3f, 4.4f, 5.5f}); - TF_ASSIGN_OR_ASSERT_OK(std::unique_ptr param0_data, - client_->TransferToServer(*param0_literal)); + Literal::CreateR1({2.2f, 5.3f, 4.4f, 5.5f}); + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr param0_data, + client_->TransferToServer(*param0_literal)); auto param0 = builder.Parameter(0, param0_literal->shape(), "param0"); auto fn = build_sum_rng(builder); builder.Map({param0}, fn); - TF_ASSIGN_OR_ASSERT_OK(auto computation, builder.Build()); + TF_ASSERT_OK_AND_ASSIGN(auto computation, builder.Build()); - ExecutionOptions execution_options; + ExecutionOptions execution_options = execution_options_; execution_options.set_seed(125); - TF_ASSIGN_OR_ASSERT_OK( - auto actual, - client_->ExecuteAndTransfer(computation, - /*arguments=*/{param0_data.get()}, - &execution_options)); + TF_ASSERT_OK_AND_ASSIGN( + auto actual, client_->ExecuteAndTransfer( + computation, + /*arguments=*/{param0_data.get()}, &execution_options)); EXPECT_EQ(actual->f32s_size(), param0_literal->f32s_size()); for (int i = 0; i < param0_literal->f32s_size(); ++i) { @@ -193,7 +188,7 @@ XLA_TEST_F(PrngTest, MapUsingRng) { } } -// This tests demonstrates the global seeding behaviour. +// This tests demonstrates the global seeding behavior. // * If a seed is passed in via Execute (ExecuteAndTransfer) then the output is // fixed (i.e., there is a single output for a given seed); // * If no seed is passed in then the output of every call can be different; @@ -207,47 +202,45 @@ XLA_TEST_F(PrngTest, PassInGlobalRngSeed) { return builder.Build(); }; - ExecutionOptions execution_options1; + ExecutionOptions execution_options1 = execution_options_; execution_options1.set_seed(42); - ExecutionOptions execution_options2; + ExecutionOptions execution_options2 = execution_options_; execution_options2.set_seed(65); std::unique_ptr result1; { - TF_ASSIGN_OR_ASSERT_OK(auto computation, build_computation()); - TF_ASSIGN_OR_ASSERT_OK( - result1, - client_->ExecuteAndTransfer(computation, /*arguments=*/{}, - &execution_options1)); + TF_ASSERT_OK_AND_ASSIGN(auto computation, build_computation()); + TF_ASSERT_OK_AND_ASSIGN( + result1, client_->ExecuteAndTransfer(computation, /*arguments=*/{}, + &execution_options1)); } std::unique_ptr result2; std::unique_ptr result3; { - TF_ASSIGN_OR_ASSERT_OK(auto computation, build_computation()); - TF_ASSIGN_OR_ASSERT_OK( - result2, - client_->ExecuteAndTransfer(computation, /*arguments=*/{}, - &execution_options1)); - TF_ASSIGN_OR_ASSERT_OK( - result3, - client_->ExecuteAndTransfer(computation, /*arguments=*/{}, - &execution_options1)); + TF_ASSERT_OK_AND_ASSIGN(auto computation, build_computation()); + TF_ASSERT_OK_AND_ASSIGN( + result2, client_->ExecuteAndTransfer(computation, /*arguments=*/{}, + &execution_options1)); + TF_ASSERT_OK_AND_ASSIGN( + result3, client_->ExecuteAndTransfer(computation, /*arguments=*/{}, + &execution_options1)); } std::unique_ptr result4; std::unique_ptr result5; std::unique_ptr result6; { - TF_ASSIGN_OR_ASSERT_OK(auto computation, build_computation()); - TF_ASSIGN_OR_ASSERT_OK( - result4, - client_->ExecuteAndTransfer(computation, /*arguments=*/{}, - &execution_options2)); - TF_ASSIGN_OR_ASSERT_OK( - result5, client_->ExecuteAndTransfer(computation, /*arguments=*/{})); - TF_ASSIGN_OR_ASSERT_OK( - result6, client_->ExecuteAndTransfer(computation, /*arguments=*/{})); + TF_ASSERT_OK_AND_ASSIGN(auto computation, build_computation()); + TF_ASSERT_OK_AND_ASSIGN( + result4, client_->ExecuteAndTransfer(computation, /*arguments=*/{}, + &execution_options2)); + TF_ASSERT_OK_AND_ASSIGN( + result5, client_->ExecuteAndTransfer(computation, /*arguments=*/{}, + &execution_options_)); + TF_ASSERT_OK_AND_ASSIGN( + result6, client_->ExecuteAndTransfer(computation, /*arguments=*/{}, + &execution_options_)); } LiteralTestUtil::ExpectEqual(*result1, *result2); @@ -271,22 +264,16 @@ XLA_TEST_F(PrngTest, TenValuesN01) { // TODO(b/25995601): Test that resultant values are reasonable } -} // namespace -} // namespace xla +XLA_TEST_F(PrngTest, RngUniformCrash) { + ComputationBuilder builder(client_, TestName()); -int main(int argc, char** argv) { - std::vector flag_list; - xla::legacy_flags::AppendCpuCompilerFlags(&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; - } - testing::InitGoogleTest(&argc, argv); - if (argc > 1) { - LOG(ERROR) << "Unknown argument " << argv[1] << "\n" << usage; - return 2; - } - return RUN_ALL_TESTS(); + // This used to crash XLA during LLVM IR generation for CPUs. + auto rng_uniform = builder.RngUniform(builder.ConstantR0(0), + builder.ConstantR0(1000 * 1000), + ShapeUtil::MakeShape(S32, {})); + SetSeed(0); + ExecuteAndTransferOrDie(&builder, /*arguments=*/{}); } + +} // namespace +} // namespace xla diff --git a/tensorflow/compiler/xla/tests/query_inferred_shape_test.cc b/tensorflow/compiler/xla/tests/query_inferred_shape_test.cc index eb7e63705b27395263c8f1db16bfb3910b7f44d6..212512207cfdc4d2ebdc4e7fd8f5794852cc6a79 100644 --- a/tensorflow/compiler/xla/tests/query_inferred_shape_test.cc +++ b/tensorflow/compiler/xla/tests/query_inferred_shape_test.cc @@ -17,7 +17,6 @@ limitations under the License. #include "tensorflow/compiler/xla/client/computation_builder.h" #include "tensorflow/compiler/xla/client/local_client.h" -#include "tensorflow/compiler/xla/legacy_flags/cpu_compiler_flags.h" #include "tensorflow/compiler/xla/shape_util.h" #include "tensorflow/compiler/xla/statusor.h" #include "tensorflow/compiler/xla/test_helpers.h" @@ -42,20 +41,3 @@ TEST_F(QueryInferredShapeTest, OnePlusOneShape) { } // namespace } // namespace xla - -int main(int argc, char** argv) { - std::vector flag_list; - xla::legacy_flags::AppendCpuCompilerFlags(&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; - } - testing::InitGoogleTest(&argc, argv); - if (argc > 1) { - LOG(ERROR) << "Unknown argument " << argv[1] << "\n" << usage; - return 2; - } - return RUN_ALL_TESTS(); -} diff --git a/tensorflow/compiler/xla/tests/reduce_precision_test.cc b/tensorflow/compiler/xla/tests/reduce_precision_test.cc new file mode 100644 index 0000000000000000000000000000000000000000..4756ba096896806ece8fe35d18c4eaef041b8830 --- /dev/null +++ b/tensorflow/compiler/xla/tests/reduce_precision_test.cc @@ -0,0 +1,372 @@ +/* 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 "tensorflow/compiler/xla/array2d.h" +#include "tensorflow/compiler/xla/client/computation_builder.h" +#include "tensorflow/compiler/xla/client/global_data.h" +#include "tensorflow/compiler/xla/client/local_client.h" +#include "tensorflow/compiler/xla/layout_util.h" +#include "tensorflow/compiler/xla/literal_util.h" +#include "tensorflow/compiler/xla/service/reduce_precision_insertion.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/compiler/xla/xla_data.pb.h" +#include "tensorflow/core/lib/core/casts.h" +#include "tensorflow/core/platform/types.h" + +namespace xla { +namespace { + +// Tests to confirm that the ReducePrecision operation produces the expected +// numerical values. +class ReducePrecisionAccuracyTest : public ClientLibraryTestBase, + public ::testing::WithParamInterface { +}; + +// For reduction to IEEE-f16, we want to test the following cases, in both +// positive and negative variants. (Note: IEEE-f16 is 5 exponent bits and 10 +// mantissa bits.) +// +// Vectors of exponent and mantissa sizes to test. We want to test IEEE-f32 (a +// no-op), IEEE-f16, and exponent-reduction-only and mantissa-reduction-only +// variants of IEEE-f16. +static const int exponent_sizes[] = {8, 5, 5, 8}; +static const int mantissa_sizes[] = {23, 10, 23, 10}; + +string TestDataToString(const ::testing::TestParamInfo data) { + int i = data.param; + return tensorflow::strings::StrCat(exponent_sizes[i], "_exponent_bits_", + mantissa_sizes[i], "_mantissa_bits"); +} + +// The FPVAL macro allows us to write out the binary representation of the +// input and expected values in a more readable manner. The mantissa bits +// are separated into the "high" bits (retained with reduction to IEEE-f16) +// and the "low" bits (truncated with reduction to IEEE-f16). +#define FPVAL(EXPONENT, HIGH_MANTISSA, LOW_MANTISSA) \ + ((0b##EXPONENT << 23) + (0b##HIGH_MANTISSA << 13) + (0b##LOW_MANTISSA)) + +// Each element in the test-value array consists of four numbers. The first is +// the input value and the following are the expected output values for the +// various precision-reduction cases. +static const uint32_t test_values[][4] = { + // True zero. + { + FPVAL(00000000, 0000000000, 0000000000000), // 0.0 + FPVAL(00000000, 0000000000, 0000000000000), // 0.0 + FPVAL(00000000, 0000000000, 0000000000000), // 0.0 + FPVAL(00000000, 0000000000, 0000000000000) // 0.0 + }, + // Largest exponent that underflows to zero. + { + FPVAL(01110000, 0000000000, 0000000000000), // 3.05176e-05 + FPVAL(00000000, 0000000000, 0000000000000), // 0.0 + FPVAL(00000000, 0000000000, 0000000000000), // 0.0 + FPVAL(01110000, 0000000000, 0000000000000) // 3.05176e-05 + }, + // Largest value that rounds to a denormal and thus clamps to zero. + { + FPVAL(01110000, 1111111111, 0111111111111), // 6.10203e-05 + FPVAL(00000000, 0000000000, 0000000000000), // 0.0 + FPVAL(00000000, 0000000000, 0000000000000), // 0.0 + FPVAL(01110000, 1111111111, 0000000000000) // 6.10054e-05 + }, + // Smallest value that doesn't underflow to zero, due to mantissa rounding + // up and incrementing the exponent out of the denormal range. + { + FPVAL(01110000, 1111111111, 1000000000000), // 6.10203e-05 + FPVAL(01110001, 0000000000, 0000000000000), // 6.10352e-05 + FPVAL(00000000, 0000000000, 0000000000000), // 0.0 + FPVAL(01110001, 0000000000, 0000000000000) // 6.10352e-05 + }, + // Smallest value that doesn't underflow to zero even without mantissa + // rounding. + { + FPVAL(01110001, 0000000000, 0000000000000), // 6.10352e-05 + FPVAL(01110001, 0000000000, 0000000000000), // 6.10352e-05 + FPVAL(01110001, 0000000000, 0000000000000), // 6.10352e-05 + FPVAL(01110001, 0000000000, 0000000000000) // 6.10352e-05 + }, + // One (to make sure bias-handling is done correctly. + { + FPVAL(01111111, 0000000000, 0000000000000), // 1.0 + FPVAL(01111111, 0000000000, 0000000000000), // 1.0 + FPVAL(01111111, 0000000000, 0000000000000), // 1.0 + FPVAL(01111111, 0000000000, 0000000000000) // 1.0 + }, + // Values in a space where ties round down due to ties-to-even: + // Value with highest mantissa that rounds down. + { + FPVAL(01111111, 0000000000, 1000000000000), // 1.00049 + FPVAL(01111111, 0000000000, 0000000000000), // 1.0 + FPVAL(01111111, 0000000000, 1000000000000), // 1.00049 + FPVAL(01111111, 0000000000, 0000000000000) // 1.0 + }, + // Value with lowest mantissa that rounds up. + { + FPVAL(01111111, 0000000000, 1000000000001), // 1.00049 + FPVAL(01111111, 0000000001, 0000000000000), // 1.00098 + FPVAL(01111111, 0000000000, 1000000000001), // 1.00049 + FPVAL(01111111, 0000000001, 0000000000000) // 1.00098 + }, + // Values in a space where ties round up due to ties-to-even: + // Value with highest mantissa that rounds down. + { + FPVAL(01111111, 0000000001, 0111111111111), // 1.00146 + FPVAL(01111111, 0000000001, 0000000000000), // 1.00098 + FPVAL(01111111, 0000000001, 0111111111111), // 1.00146 + FPVAL(01111111, 0000000001, 0000000000000) // 1.00098 + }, + // Value with a mantissa that rounds up. + { + FPVAL(01111111, 0000000001, 1000000000000), // 1.00146 + FPVAL(01111111, 0000000010, 0000000000000), // 1.00195 + FPVAL(01111111, 0000000001, 1000000000000), // 1.00146 + FPVAL(01111111, 0000000010, 0000000000000) // 1.00195 + }, + // Largest value that does not overflow to infinity. + { + FPVAL(10001110, 1111111111, 0111111111111), // 65520.0 + FPVAL(10001110, 1111111111, 0000000000000), // 65504.0 + FPVAL(10001110, 1111111111, 0111111111111), // 65520.0 + FPVAL(10001110, 1111111111, 0000000000000) // 65504.0 + }, + // Smallest value that overflows to infinity due to mantissa rounding up. + { + FPVAL(10001110, 1111111111, 1000000000000), // 65520.0 + FPVAL(11111111, 0000000000, 0000000000000), // Inf + FPVAL(10001110, 1111111111, 1000000000000), // 65520.0 + FPVAL(10001111, 0000000000, 0000000000000) // 65536.0 + }, + // Smallest value that overflows to infinity, without mantissa rounding. + { + FPVAL(10001111, 0000000000, 0000000000000), // 65536.0 + FPVAL(11111111, 0000000000, 0000000000000), // Inf + FPVAL(11111111, 0000000000, 0000000000000), // Inf + FPVAL(10001111, 0000000000, 0000000000000) // 65536.0 + }, + // Smallest value that overflows to infinity due to mantissa rounding up, + // even when exponent bits aren't reduced. + { + FPVAL(11111110, 1111111111, 1000000000000), // 3.40199e+38 + FPVAL(11111111, 0000000000, 0000000000000), // Inf + FPVAL(11111111, 0000000000, 0000000000000), // Inf + FPVAL(11111111, 0000000000, 0000000000000) // Inf + }, + // True infinity. + { + FPVAL(11111111, 0000000000, 0000000000000), // Inf + FPVAL(11111111, 0000000000, 0000000000000), // Inf + FPVAL(11111111, 0000000000, 0000000000000), // Inf + FPVAL(11111111, 0000000000, 0000000000000) // Inf + }, + // NAN with a 1 in the preserved bits. + { + FPVAL(11111111, 1000000000, 0000000000000), // NaN + FPVAL(11111111, 1000000000, 0000000000000), // NaN + FPVAL(11111111, 1000000000, 0000000000000), // NaN + FPVAL(11111111, 1000000000, 0000000000000) // NaN + }, + // NAN with a 1 in the truncated bits. + { + FPVAL(11111111, 0000000000, 0000000000001), // NaN + FPVAL(11111111, 0000000000, 0000000000001), // NaN + FPVAL(11111111, 0000000000, 0000000000001), // NaN + FPVAL(11111111, 0000000000, 0000000000001) // NaN + }, + // NAN with all ones, causing rounding overflow. + { + FPVAL(11111111, 1111111111, 1111111111111), // NaN + FPVAL(11111111, 1111111111, 1111111111111), // NaN + FPVAL(11111111, 1111111111, 1111111111111), // NaN + FPVAL(11111111, 1111111111, 1111111111111) // NaN + }}; + +XLA_TEST_P(ReducePrecisionAccuracyTest, ReducePrecisionF32) { + int index = GetParam(); + int exponent_bits = exponent_sizes[index]; + int mantissa_bits = mantissa_sizes[index]; + + std::vector input_values; + std::vector expected_values; + + 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])); + // 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)); + expected_values.push_back( + tensorflow::bit_cast(test_value[index] ^ sign_bit)); + } + + // This is required for proper handling of NaN values. + SetFastMathDisabled(true); + + ComputationBuilder builder(client_, TestName()); + + std::unique_ptr a_literal = Literal::CreateR1({input_values}); + std::unique_ptr a_data = + client_->TransferToServer(*a_literal).ConsumeValueOrDie(); + auto a = builder.Parameter(0, a_literal->shape(), "a"); + + auto reduce_precision = + builder.ReducePrecision(a, exponent_bits, mantissa_bits); + + ComputeAndCompareR1(&builder, expected_values, {a_data.get()}); +} + +INSTANTIATE_TEST_CASE_P(ReducePrecisionAccuracyTest, + ReducePrecisionAccuracyTest, + ::testing::Values(0, 1, 2, 3), TestDataToString); + +// Tests to confirm that the compiler optimization functions add the expected +// ReducePrecisionInsertion passes. +class ReducePrecisionInsertionTest : public ClientLibraryTestBase {}; + +XLA_TEST_F(ReducePrecisionInsertionTest, ReducePrecisionBeforeFusion) { + ComputationBuilder builder(client_, TestName()); + + std::unique_ptr a_literal = Literal::CreateR1({1.00001}); + std::unique_ptr a_data = + client_->TransferToServer(*a_literal).ConsumeValueOrDie(); + auto a = builder.Parameter(0, a_literal->shape(), "a"); + + // Abs doesn't affect resolution. + auto abs = builder.Abs(a); + + // Near 1.0, Log(x) approximates x - 1; this lets us confirm that the + // reduce-precision operation showed up in the correct place in the + // graph. + auto log = builder.Log(abs); + + // Insert precision-reduction after the Abs(x) operation, rounding that + // result to exactly 1.0f. + auto reduce_precision_pass = execution_options_.mutable_debug_options() + ->add_hlo_reduce_precision_options(); + *reduce_precision_pass = ReducePrecisionInsertion::make_options_proto( + HloReducePrecisionOptions::OP_OUTPUTS, 5, 10, + [](const HloOpcode opcode) { return opcode == HloOpcode::kAbs; }); + + ComputeAndCompareR1(&builder, {0.0f}, {a_data.get()}); +} + +XLA_TEST_F(ReducePrecisionInsertionTest, ReducePrecisionSkippedAfterFusion) { + ComputationBuilder builder(client_, TestName()); + + std::unique_ptr a_literal = Literal::CreateR1({1.00001}); + std::unique_ptr a_data = + client_->TransferToServer(*a_literal).ConsumeValueOrDie(); + auto a = builder.Parameter(0, a_literal->shape(), "a"); + + // These two operations should be fused by any reasonable backend. + auto abs = builder.Abs(a); + auto neg = builder.Neg(abs); + + // Add a pass after operation fusion, suffixing kAbs operations. This + // should not see into the fusion nodes and thus should not affect the + // result. + auto reduce_precision_pass = execution_options_.mutable_debug_options() + ->add_hlo_reduce_precision_options(); + *reduce_precision_pass = ReducePrecisionInsertion::make_options_proto( + HloReducePrecisionOptions::UNFUSED_OP_OUTPUTS, 5, 10, + [](const HloOpcode opcode) { return opcode == HloOpcode::kAbs; }); + + ComputeAndCompareR1(&builder, {-1.00001f}, {a_data.get()}); +} + +XLA_TEST_F(ReducePrecisionInsertionTest, ReducePrecisionAddedAfterFusion) { + ComputationBuilder builder(client_, TestName()); + + std::unique_ptr a_literal = Literal::CreateR1({1.00001}); + std::unique_ptr a_data = + client_->TransferToServer(*a_literal).ConsumeValueOrDie(); + auto a = builder.Parameter(0, a_literal->shape(), "a"); + + // These two operations should be fused by any reasonable backend. + auto abs = builder.Abs(a); + auto neg = builder.Neg(abs); + + // Add a pass after operation fusion, suffixing kFusion operations. + auto reduce_precision_pass = execution_options_.mutable_debug_options() + ->add_hlo_reduce_precision_options(); + *reduce_precision_pass = ReducePrecisionInsertion::make_options_proto( + HloReducePrecisionOptions::UNFUSED_OP_OUTPUTS, 5, 10, + [](const HloOpcode opcode) { return opcode == HloOpcode::kFusion; }); + + ComputeAndCompareR1(&builder, {-1.0f}, {a_data.get()}); +} + +XLA_TEST_F(ReducePrecisionInsertionTest, ReducePrecisionSkippedFusionContains) { + ComputationBuilder builder(client_, TestName()); + + std::unique_ptr a_literal = Literal::CreateR1({1.00001}); + std::unique_ptr a_data = + client_->TransferToServer(*a_literal).ConsumeValueOrDie(); + auto a = builder.Parameter(0, a_literal->shape(), "a"); + + // These two operations should be fused by any reasonable backend. + auto abs = builder.Abs(a); + auto neg = builder.Neg(abs); + + // Add a pass suffixing fusion nodes containing kCos operations. This + // should have no effect. + auto reduce_precision_pass = execution_options_.mutable_debug_options() + ->add_hlo_reduce_precision_options(); + *reduce_precision_pass = ReducePrecisionInsertion::make_options_proto( + HloReducePrecisionOptions::FUSION_OUTPUTS_BY_CONTENT, 5, 10, + [](const HloOpcode opcode) { return opcode == HloOpcode::kCos; }); + + ComputeAndCompareR1(&builder, {-1.00001f}, {a_data.get()}); +} + +XLA_TEST_F(ReducePrecisionInsertionTest, ReducePrecisionAddedFusionContains) { + ComputationBuilder builder(client_, TestName()); + + std::unique_ptr a_literal = Literal::CreateR1({1.00001}); + std::unique_ptr a_data = + client_->TransferToServer(*a_literal).ConsumeValueOrDie(); + auto a = builder.Parameter(0, a_literal->shape(), "a"); + + // These two operations should be fused by any reasonable backend. + auto abs = builder.Abs(a); + auto neg = builder.Neg(abs); + + // Add a pass suffixing fusion nodes containing kAbs operations. This + // should see the kAbs operation within the above fusion node. + auto reduce_precision_pass = execution_options_.mutable_debug_options() + ->add_hlo_reduce_precision_options(); + *reduce_precision_pass = ReducePrecisionInsertion::make_options_proto( + HloReducePrecisionOptions::FUSION_OUTPUTS_BY_CONTENT, 5, 10, + [](const HloOpcode opcode) { return opcode == HloOpcode::kAbs; }); + + ComputeAndCompareR1(&builder, {-1.0f}, {a_data.get()}); +} + +} // namespace +} // namespace xla diff --git a/tensorflow/compiler/xla/tests/reduce_test.cc b/tensorflow/compiler/xla/tests/reduce_test.cc index 34fce21758b98c52831ac4ddb168d3e1538e9f1d..2271f32c5946f3d3e7e6b43b089e68ab3101b61b 100644 --- a/tensorflow/compiler/xla/tests/reduce_test.cc +++ b/tensorflow/compiler/xla/tests/reduce_test.cc @@ -40,7 +40,6 @@ limitations under the License. #include "tensorflow/compiler/xla/client/lib/arithmetic.h" #include "tensorflow/compiler/xla/client/local_client.h" #include "tensorflow/compiler/xla/layout_util.h" -#include "tensorflow/compiler/xla/legacy_flags/cpu_compiler_flags.h" #include "tensorflow/compiler/xla/literal_util.h" #include "tensorflow/compiler/xla/reference_util.h" #include "tensorflow/compiler/xla/shape_util.h" @@ -61,14 +60,14 @@ namespace { class ReduceTest : public ClientLibraryTestBase { protected: ReduceTest() { - // Implementation note: layed out z >> y >> x by default. + // Implementation note: laid out z >> y >> x by default. // clang-format off - literal_2d_ = LiteralUtil::CreateR2({ + literal_2d_ = Literal::CreateR2({ // x0 x1 x2 { 1.f, 2.f, 3.f}, // y0 { 4.f, 5.f, 6.f}, // y1 }); - literal_3d_ = LiteralUtil::CreateR3Projected({ + literal_3d_ = Literal::CreateR3Projected({ // x0 x1 x2 { 1.f, 2.f, 3.f}, // y0 { 4.f, 5.f, 6.f}, // y1 @@ -97,7 +96,7 @@ class ReduceTest : public ClientLibraryTestBase { } } std::unique_ptr input_literal = - LiteralUtil::CreateR1(AsSlice(input_data)); + Literal::CreateR1(AsSlice(input_data)); std::unique_ptr input_global_data = client_->TransferToServer(*input_literal).ConsumeValueOrDie(); @@ -129,7 +128,7 @@ class ReduceTest : public ClientLibraryTestBase { builder.Reduce(pred_values, init_value, reduce, /*dimensions_to_reduce=*/{0}); - std::unique_ptr input_literal = LiteralUtil::CreateR1(input_data); + std::unique_ptr input_literal = Literal::CreateR1(input_data); std::unique_ptr input_global_data = client_->TransferToServer(*input_literal).ConsumeValueOrDie(); @@ -156,9 +155,9 @@ class ReduceTest : public ClientLibraryTestBase { Array2D input_data(rows, cols); input_data.FillRandom(3.14f, 0.04); std::unique_ptr input_literal = - LiteralUtil::CreateR2FromArray2D(input_data); - input_literal = LiteralUtil::Relayout( - *input_literal, LayoutUtil::MakeLayout({minor, major})); + Literal::CreateR2FromArray2D(input_data); + input_literal = + input_literal->Relayout(LayoutUtil::MakeLayout({minor, major})); std::unique_ptr input_global_data = client_->TransferToServer(*input_literal).ConsumeValueOrDie(); @@ -184,9 +183,9 @@ class ReduceTest : public ClientLibraryTestBase { Array2D input_data(rows, cols); input_data.FillRandom(3.14f, 0.04); std::unique_ptr input_literal = - LiteralUtil::CreateR2FromArray2D(input_data); - input_literal = LiteralUtil::Relayout( - *input_literal, LayoutUtil::MakeLayout({minor, major})); + Literal::CreateR2FromArray2D(input_data); + input_literal = + input_literal->Relayout(LayoutUtil::MakeLayout({minor, major})); std::unique_ptr input_global_data = client_->TransferToServer(*input_literal).ConsumeValueOrDie(); @@ -202,6 +201,102 @@ class ReduceTest : public ClientLibraryTestBase { ErrorSpec(0.01, 1e-4)); } + template + void ComputeAndCompareGeneric( + typename std::enable_if::value, + ComputationBuilder>::type* builder, + tensorflow::gtl::ArraySlice expected, + tensorflow::gtl::ArraySlice arguments) { + ComputeAndCompareR1(builder, expected, arguments, + ErrorSpec(0.01, 1e-4)); + } + + template + void ComputeAndCompareGeneric( + typename std::enable_if::value, + ComputationBuilder>::type* builder, + tensorflow::gtl::ArraySlice expected, + tensorflow::gtl::ArraySlice arguments) { + ComputeAndCompareR1(builder, expected, arguments); + } + + template + void RunVectorizedReduceTestForType( + const std::function& + reduction_function_generator, + const std::function& + reference_reduction_function, + const NativeT& initial_value) { + const int rows = 64, cols = 128; + const int minor = 1, major = 0; + ComputationBuilder builder(client_, TestName()); + Computation reduction_function = reduction_function_generator(&builder); + const Shape input_shape = ShapeUtil::MakeShape( + xla::primitive_util::NativeToPrimitiveType(), {rows, cols}); + auto input = builder.Parameter(0, input_shape, "input"); + auto zero = builder.ConstantR0(initial_value); + builder.Reduce(input, zero, reduction_function, + /*dimensions_to_reduce=*/{0}); + + Array2D input_data(rows, cols); + input_data.FillUnique(initial_value); + std::unique_ptr input_literal = + Literal::CreateR2FromArray2D(input_data); + input_literal = + input_literal->Relayout(LayoutUtil::MakeLayout({minor, major})); + std::unique_ptr input_global_data = + client_->TransferToServer(*input_literal).ConsumeValueOrDie(); + + // NativeT can be bool, and std::vector does not convert to + // ArraySlice. + std::unique_ptr expected(new NativeT[cols]); + for (int64 colno = 0; colno < cols; ++colno) { + NativeT column_result = initial_value; + for (int64 rowno = 0; rowno < rows; ++rowno) { + column_result = reference_reduction_function(column_result, + input_data(rowno, colno)); + } + expected[colno] = column_result; + } + + ComputeAndCompareGeneric( + &builder, tensorflow::gtl::ArraySlice(expected.get(), cols), + {input_global_data.get()}); + } + + void RunVectorizedReduceTest( + const std::function& + reduction_function_generator_for_type, + const std::function& + reference_reduction_function_for_floats, + const std::function& + reference_reduction_function_for_ints, + const std::function& + reference_reduction_function_for_uints, + float floating_point_identity, int32 signed_int_identity, + uint32 unsigned_int_identity) { + // Float version + RunVectorizedReduceTestForType( + [&](ComputationBuilder* builder) { + return reduction_function_generator_for_type(F32, builder); + }, + reference_reduction_function_for_floats, floating_point_identity); + + // Signed int version + RunVectorizedReduceTestForType( + [&](ComputationBuilder* builder) { + return reduction_function_generator_for_type(S32, builder); + }, + reference_reduction_function_for_ints, signed_int_identity); + + // Unsigned int version + RunVectorizedReduceTestForType( + [&](ComputationBuilder* builder) { + return reduction_function_generator_for_type(U32, builder); + }, + reference_reduction_function_for_uints, unsigned_int_identity); + } + std::unique_ptr literal_2d_; std::unique_ptr literal_3d_; uint32 seed_ = 0xdeadbeef; @@ -211,9 +306,9 @@ XLA_TEST_F(ReduceTest, ReduceR1_0_F32_To_R0) { RunR1ToR0Test(0); } XLA_TEST_F(ReduceTest, ReduceR1_1_F32_To_R0) { RunR1ToR0Test(1); } XLA_TEST_F(ReduceTest, ReduceR1_2_F32_To_R0) { RunR1ToR0Test(2); } XLA_TEST_F(ReduceTest, ReduceR1_16_F32_To_R0) { RunR1ToR0Test(16); } -XLA_TEST_F(ReduceTest, ReduceR1_240_F32_To_R0) { RunR1ToR0Test(240); } XLA_TEST_F(ReduceTest, ReduceR1_128_F32_To_R0) { RunR1ToR0Test(128); } XLA_TEST_F(ReduceTest, ReduceR1_129_F32_To_R0) { RunR1ToR0Test(129); } +XLA_TEST_F(ReduceTest, ReduceR1_240_F32_To_R0) { RunR1ToR0Test(240); } XLA_TEST_F(ReduceTest, ReduceR1_256_F32_To_R0) { RunR1ToR0Test(256); } XLA_TEST_F(ReduceTest, ReduceR1_1024_F32_To_R0) { RunR1ToR0Test(1024); } XLA_TEST_F(ReduceTest, ReduceR1_2048_F32_To_R0) { RunR1ToR0Test(2048); } @@ -221,6 +316,9 @@ XLA_TEST_F(ReduceTest, ReduceR1_16K_F32_To_R0) { RunR1ToR0Test(16 * 1024); } XLA_TEST_F(ReduceTest, ReduceR1_16KP1_F32_To_R0) { RunR1ToR0Test(16 * 1024 + 1); } +XLA_TEST_F(ReduceTest, ReduceR1_64K_F32_To_R0) { RunR1ToR0Test(64 * 1024); } +XLA_TEST_F(ReduceTest, ReduceR1_1M_F32_To_R0) { RunR1ToR0Test(1024 * 1024); } +XLA_TEST_F(ReduceTest, ReduceR1_16M_F32_To_R0) { RunR1ToR0Test(4096 * 4096); } XLA_TEST_F(ReduceTest, ReduceR2_0x0_To_R0) { RunR2ToR0Test(0, 0); } XLA_TEST_F(ReduceTest, ReduceR2_0x2_To_R0) { RunR2ToR0Test(0, 2); } @@ -302,9 +400,8 @@ XLA_TEST_F(ReduceTest, ReduceElementwiseR2_111x50_To_R1) { Array2D input_data(rows, cols); input_data.FillRandom(3.14f, 0.04); std::unique_ptr input_literal = - LiteralUtil::CreateR2FromArray2D(input_data); - input_literal = - LiteralUtil::Relayout(*input_literal, LayoutUtil::MakeLayout({0, 1})); + Literal::CreateR2FromArray2D(input_data); + input_literal = input_literal->Relayout(LayoutUtil::MakeLayout({0, 1})); std::unique_ptr input_global_data = client_->TransferToServer(*input_literal).ConsumeValueOrDie(); @@ -335,9 +432,8 @@ XLA_TEST_F(ReduceTest, TransposeAndReduceElementwiseR2_111x50_To_R1) { Array2D input_data(rows, cols); input_data.FillRandom(3.14f, 0.04); std::unique_ptr input_literal = - LiteralUtil::CreateR2FromArray2D(input_data); - input_literal = - LiteralUtil::Relayout(*input_literal, LayoutUtil::MakeLayout({0, 1})); + Literal::CreateR2FromArray2D(input_data); + input_literal = input_literal->Relayout(LayoutUtil::MakeLayout({0, 1})); std::unique_ptr input_global_data = client_->TransferToServer(*input_literal).ConsumeValueOrDie(); @@ -368,7 +464,7 @@ XLA_TEST_F(ReduceTest, Reshape_111x2x25Reduce_111x50_To_R1) { Array3D input_data(rows, 2, cols / 2); input_data.FillRandom(3.14f, 0.04); std::unique_ptr input_literal = - LiteralUtil::CreateR3FromArray3D(input_data); + Literal::CreateR3FromArray3D(input_data); std::unique_ptr input_global_data = client_->TransferToServer(*input_literal).ConsumeValueOrDie(); @@ -431,7 +527,7 @@ XLA_TEST_F(ReduceTest, MaxReduce2DToR0) { auto max = CreateScalarMaxComputation(F32, &builder); Array2D input(300, 250); input.FillRandom(214.0f); - auto input_literal = LiteralUtil::CreateR2FromArray2D(input); + auto input_literal = Literal::CreateR2FromArray2D(input); builder.Reduce(builder.ConstantLiteral(*input_literal), builder.ConstantR0(FLT_MIN), max, {0, 1}); auto input_max = FLT_MIN; @@ -446,7 +542,7 @@ XLA_TEST_F(ReduceTest, MinReduce2DToR0) { auto min = CreateScalarMinComputation(F32, &builder); Array2D input(150, 130); input.FillRandom(214.0f); - auto input_literal = LiteralUtil::CreateR2FromArray2D(input); + auto input_literal = Literal::CreateR2FromArray2D(input); builder.Reduce(builder.ConstantLiteral(*input_literal), builder.ConstantR0(FLT_MAX), min, {0, 1}); @@ -456,6 +552,32 @@ XLA_TEST_F(ReduceTest, MinReduce2DToR0) { ComputeAndCompareR0(&builder, input_min, {}, ErrorSpec(0.0001)); } +XLA_TEST_F(ReduceTest, UnsignedInt_MinReduce) { + ComputationBuilder builder(client_, TestName()); + Array2D input({{1}, {2}}); + auto min = CreateScalarMinComputation(U32, &builder); + auto input_literal = Literal::CreateR2FromArray2D(input); + auto initial_value = + builder.ConstantR0(std::numeric_limits::max()); + + builder.Reduce(builder.ConstantLiteral(*input_literal), initial_value, min, + {0, 1}); + ComputeAndCompareR0(&builder, 1, {}); +} + +XLA_TEST_F(ReduceTest, UnsignedInt_MaxReduce) { + ComputationBuilder builder(client_, TestName()); + Array2D input({{1}, {2}}); + auto max = CreateScalarMaxComputation(U32, &builder); + auto input_literal = Literal::CreateR2FromArray2D(input); + auto initial_value = + builder.ConstantR0(std::numeric_limits::min()); + + builder.Reduce(builder.ConstantLiteral(*input_literal), initial_value, max, + {0, 1}); + ComputeAndCompareR0(&builder, 2, {}); +} + // Reduces a matrix among dimension 1. XLA_TEST_F(ReduceTest, Reduce2DAmong1) { ComputationBuilder builder(client_, TestName()); @@ -567,6 +689,58 @@ XLA_TEST_F(ReduceTest, ReduceR3AmongDim2) { ComputeAndCompareR2(&builder, expected, {}, ErrorSpec(0.0001)); } +XLA_TEST_F(ReduceTest, VectorizedReduce_Add) { + RunVectorizedReduceTest(CreateScalarAddComputation, + [](float a, float b) { return a + b; }, + [](int32 a, int32 b) { + return static_cast(static_cast(a) + + static_cast(b)); + }, + [](uint32 a, uint32 b) { return a + b; }, 0.0, 0, 0); +} + +XLA_TEST_F(ReduceTest, VectorizedReduce_Multiply) { + RunVectorizedReduceTest(CreateScalarMultiplyComputation, + [](float a, float b) { return a * b; }, + [](int32 a, int32 b) { + return static_cast(static_cast(a) * + static_cast(b)); + }, + [](uint32 a, uint32 b) { return a * b; }, 1.0, 1, 1); +} + +XLA_TEST_F(ReduceTest, VectorizedReduce_Max) { + RunVectorizedReduceTest(CreateScalarMaxComputation, + [](float a, float b) { return std::max(a, b); }, + [](int32 a, int32 b) { return std::max(a, b); }, + [](uint32 a, uint32 b) { return std::max(a, b); }, + std::numeric_limits::min(), + std::numeric_limits::min(), + std::numeric_limits::min()); +} + +XLA_TEST_F(ReduceTest, VectorizedReduce_Min) { + RunVectorizedReduceTest(CreateScalarMinComputation, + [](float a, float b) { return std::min(a, b); }, + [](int32 a, int32 b) { return std::min(a, b); }, + [](uint32 a, uint32 b) { return std::min(a, b); }, + std::numeric_limits::max(), + std::numeric_limits::max(), + std::numeric_limits::max()); +} + +XLA_TEST_F(ReduceTest, VectorizedReduce_LogicalAnd) { + RunVectorizedReduceTestForType(CreateScalarLogicalAndComputation, + [](bool a, bool b) { return a && b; }, + true); +} + +XLA_TEST_F(ReduceTest, VectorizedReduce_LogicalOr) { + RunVectorizedReduceTestForType(CreateScalarLogicalOrComputation, + [](bool a, bool b) { return a || b; }, + false); +} + class ReduceR3ToR2Test : public ReduceTest, public ::testing::WithParamInterface {}; @@ -574,11 +748,12 @@ XLA_TEST_P(ReduceR3ToR2Test, ReduceR3ToR2) { ComputationBuilder builder(client_, TestName()); const auto& bounds = GetParam().bounds; Array3D input_array(bounds[0], bounds[1], bounds[2]); - input_array.FillRandom(3.14f, 0.05); + // input_array.FillRandom(3.14f, 0.05); + input_array.Fill(1.0f); - auto input_literal = LiteralUtil::CreateR3FromArray3D(input_array); - input_literal = LiteralUtil::Relayout( - *input_literal, LayoutUtil::MakeLayout(GetParam().layout)); + auto input_literal = Literal::CreateR3FromArray3D(input_array); + input_literal = + input_literal->Relayout(LayoutUtil::MakeLayout(GetParam().layout)); std::unique_ptr input_data = client_->TransferToServer(*input_literal).ConsumeValueOrDie(); @@ -620,22 +795,24 @@ INSTANTIATE_TEST_CASE_P( BoundsLayout{{2, 300, 784}, {2, 1, 0}, {1}}, BoundsLayout{{2, 300, 784}, {2, 1, 0}, {0}})); -} // namespace -} // namespace xla +// TODO(b/64093391) Disabled on GPU due to an assertion failure when running +// IrEmitterUnnested::EmitInitializer() for the Reduce operator. Failed on +// 2017-07-26. +XLA_TEST_F(ReduceTest, DISABLED_ON_GPU(OperationOnConstantAsInitValue)) { + ComputationBuilder builder(client_, TestName()); + Computation max_f32 = CreateScalarMaxComputation(F32, &builder); -int main(int argc, char** argv) { - std::vector flag_list; - xla::legacy_flags::AppendCpuCompilerFlags(&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; - } - testing::InitGoogleTest(&argc, argv); - if (argc > 1) { - LOG(ERROR) << "Unknown argument " << argv[1] << "\n" << usage; - return 2; - } - return RUN_ALL_TESTS(); + auto a = builder.ConstantR0(2.0f); + auto a2 = builder.Abs(a); + + std::unique_ptr b_literal = Literal::CreateR1({1.0f, 4.0f}); + std::unique_ptr b_data = + client_->TransferToServer(*b_literal).ConsumeValueOrDie(); + auto b = builder.Parameter(0, b_literal->shape(), "b"); + auto max = builder.Reduce(b, a2, max_f32, {0}); + + ComputeAndCompareR0(&builder, 4.0f, {b_data.get()}); } + +} // namespace +} // namespace xla diff --git a/tensorflow/compiler/xla/tests/reduce_window_test.cc b/tensorflow/compiler/xla/tests/reduce_window_test.cc index 56501e43b5c5d965ea4305f2ca88909b253ed273..6ef5c4a8c8b0f103cdfdbf4b12344d34b964cac2 100644 --- a/tensorflow/compiler/xla/tests/reduce_window_test.cc +++ b/tensorflow/compiler/xla/tests/reduce_window_test.cc @@ -25,10 +25,10 @@ limitations under the License. #include "tensorflow/compiler/xla/client/lib/arithmetic.h" #include "tensorflow/compiler/xla/client/local_client.h" #include "tensorflow/compiler/xla/client/padding.h" -#include "tensorflow/compiler/xla/legacy_flags/cpu_compiler_flags.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/xla_data.pb.h" @@ -43,7 +43,7 @@ class ReduceWindowTest : public ClientLibraryTestBase { public: ReduceWindowTest() : builder_(client_, TestName()) {} - void ReduceWindowAdd(ComputationDataHandle input, + void ReduceWindowAdd(const ComputationDataHandle& input, tensorflow::gtl::ArraySlice window_dimensions, tensorflow::gtl::ArraySlice window_strides, Padding padding) { @@ -52,21 +52,21 @@ class ReduceWindowTest : public ClientLibraryTestBase { window_dimensions, window_strides, padding); } - void ReduceWindowMax(ComputationDataHandle input, + void ReduceWindowMax(const ComputationDataHandle& input, tensorflow::gtl::ArraySlice window_dimensions, tensorflow::gtl::ArraySlice window_strides, Padding padding) { builder_.ReduceWindow( - input, builder_.ConstantLiteral(LiteralUtil::MinValue(F32)), + input, builder_.ConstantLiteral(Literal::MinValue(F32)), CreateScalarMax(), window_dimensions, window_strides, padding); } - void ReduceWindowMin(ComputationDataHandle input, + void ReduceWindowMin(const ComputationDataHandle& input, tensorflow::gtl::ArraySlice window_dimensions, tensorflow::gtl::ArraySlice window_strides, Padding padding) { builder_.ReduceWindow(input, - builder_.ConstantLiteral(LiteralUtil::MaxValue(F32)), + builder_.ConstantLiteral(Literal::MaxValue(F32)), CreateScalarMinComputation(F32, &builder_), window_dimensions, window_strides, padding); } @@ -74,6 +74,12 @@ class ReduceWindowTest : public ClientLibraryTestBase { ComputationBuilder builder_; }; +TEST_F(ReduceWindowTest, Min3In5Stride2) { + const auto input = builder_.ConstantR1({10000, 1000, 100, 10, 1}); + ReduceWindowMin(input, {3}, {2}, Padding::kValid); + ComputeAndCompareR1(&builder_, {100, 1}, {}, ErrorSpec(0.0001)); +} + XLA_TEST_F(ReduceWindowTest, ZeroElementSmall) { Array4D input_array(1, 0, 2, 1); @@ -131,6 +137,26 @@ TEST_F(ReduceWindowTest, Along2ndMinorDim) { ComputeAndCompareR4(&builder_, *res, {}, ErrorSpec(1e-3, 1e-3)); } +TEST_F(ReduceWindowTest, AmongMajor2Dims) { + Array4D input_array(4, 4, 6, 8); + input_array.FillWithMinorDimNum(); + + int win_len = 3; + int win_stride = 1; + + Padding padding = Padding::kSame; + const auto input_data_handle = + builder_.ConstantR4FromArray4D(input_array); + // Reduce only along the x and y dimensions, according to the win_len. + ReduceWindowAdd(input_data_handle, {win_len, win_len, 1, 1}, + {win_stride, win_stride, 1, 1}, padding); + + auto result = ReferenceUtil::ReduceWindow4DAdd( + input_array, 0.0f, {win_len, win_len, 1, 1}, + {win_stride, win_stride, 1, 1}, padding); + ComputeAndCompareR4(&builder_, *result, {}, ErrorSpec(1e-3, 1e-3)); +} + TEST_F(ReduceWindowTest, AmongMajor2DimsMediumSize) { Array4D input_array(9, 12, 4, 89); input_array.FillRandom(2.0f); @@ -182,191 +208,6 @@ TEST_F(ReduceWindowTest, DISABLED_AmongMajor2DimsMediumSizeLargePadding) { ComputeAndCompareR4(&builder_, *result, {}, ErrorSpec(1e-3, 1e-3)); } -// TODO(b/31809540): Implement minor dim reduction to reduce num of reshapes. -TEST_F(ReduceWindowTest, ReduceR4AmongXYMinorSmall) { - Array4D input_array(2, 2, 4, 16); - - Array2D yx({{0.f, 1.f, 2.f, 3.f, 4.f, 5.f, 6.f, 7.f, 8.f, 9.f, 10.f, - 11.f, 12.f, 13.f, 14.f, 15.f}, - {16.f, 17.f, 18.f, 19.f, 20.f, 21.f, 22.f, 23.f, 24.f, - 25.f, 26.f, 27.f, 28.f, 29.f, 30.f, 31.f}, - {32.f, 33.f, 34.f, 35.f, 36.f, 37.f, 38.f, 39.f, 40.f, - 41.f, 42.f, 43.f, 44.f, 45.f, 46.f, 47.f}, - {48.f, 49.f, 50.f, 51.f, 52.f, 53.f, 54.f, 55.f, 56.f, - 57.f, 58.f, 59.f, 60.f, 61.f, 62.f, 63.f}}); - input_array.FillWithYX(yx); - - int win_len = 2; - int win_stride = 2; - const auto input = builder_.ConstantR4FromArray4D(input_array); - Padding padding = Padding::kValid; - ReduceWindowAdd(input, {1, 1, win_len, win_len}, - {1, 1, win_stride, win_stride}, padding); - - auto res = ReferenceUtil::ReduceWindow4DAdd( - input_array, 0.0f, {1, 1, win_len, win_len}, - {1, 1, win_stride, win_stride}, padding); - ComputeAndCompareR4(&builder_, *res, {}, ErrorSpec(1e-3, 1e-3)); -} - -// TODO(b/31809540): Implement minor dim reduction to reduce num of reshapes. -TEST_F(ReduceWindowTest, ReduceR4AmongXYMinorSmallOverlapped) { - constexpr int64 p = 2; - constexpr int64 z = 2; - constexpr int64 y = 4; - constexpr int64 x = 16; - Array4D input_array(p, z, y, x); - - Array2D yx({{0.f, 1.f, 2.f, 3.f, 4.f, 5.f, 6.f, 7.f, 8.f, 9.f, 10.f, - 11.f, 12.f, 13.f, 14.f, 15.f}, - {16.f, 17.f, 18.f, 19.f, 20.f, 21.f, 22.f, 23.f, 24.f, - 25.f, 26.f, 27.f, 28.f, 29.f, 30.f, 31.f}, - {32.f, 33.f, 34.f, 35.f, 36.f, 37.f, 38.f, 39.f, 40.f, - 41.f, 42.f, 43.f, 44.f, 45.f, 46.f, 47.f}, - {48.f, 49.f, 50.f, 51.f, 52.f, 53.f, 54.f, 55.f, 56.f, - 57.f, 58.f, 59.f, 60.f, 61.f, 62.f, 63.f}}); - input_array.FillWithYX(yx); - - int win_len = 4; - int win_stride = 2; - const auto input = builder_.ConstantR4FromArray4D(input_array); - ReduceWindowAdd(input, {1, 1, win_len, win_len}, - {1, 1, win_stride, win_stride}, Padding::kValid); - - // Expected result - Array2D yx_result({{408.f, 440.f, 472.f, 504.f, 536.f, 568.f, 600.f}}); - Array4D expected(p, z, 1, 7); - expected.FillWithYX(yx_result); - ComputeAndCompareR4(&builder_, expected, {}, ErrorSpec(1e-3, 1e-3)); -} - -TEST_F(ReduceWindowTest, MaxTrivial) { - const auto input = builder_.ConstantR1({42}); - ReduceWindowMax(input, {1}, {1}, Padding::kValid); - ComputeAndCompareR1(&builder_, {42}, {}, ErrorSpec(0.0001)); -} - -TEST_F(ReduceWindowTest, Add3In3) { - const auto input = builder_.ConstantR1({20, 100, 3}); - ReduceWindowAdd(input, {3}, {1}, Padding::kValid); - ComputeAndCompareR1(&builder_, {123}, {}, ErrorSpec(0.0001)); -} - -TEST_F(ReduceWindowTest, Add4In16Stride4) { - const auto input = builder_.ConstantR1( - {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16}); - ReduceWindowAdd(input, {4}, {4}, Padding::kValid); - ComputeAndCompareR1(&builder_, {10, 26, 42, 58}, {}, - ErrorSpec(0.0001)); -} - -TEST_F(ReduceWindowTest, DISABLED_ON_CPU(DISABLED_ON_GPU(Min3In5Stride2))) { - const auto input = builder_.ConstantR1({10000, 1000, 100, 10, 1}); - ReduceWindowMin(input, {3}, {2}, Padding::kValid); - ComputeAndCompareR1(&builder_, {100, 1}, {}, ErrorSpec(0.0001)); -} - -TEST_F(ReduceWindowTest, Max3In3) { - const auto input = builder_.ConstantR1({20, 100, 3}); - ReduceWindowMax(input, {3}, {1}, Padding::kValid); - ComputeAndCompareR1(&builder_, {100}, {}, ErrorSpec(0.0001)); -} - -TEST_F(ReduceWindowTest, Add2In3) { - const auto input = builder_.ConstantR1({100, 10, 1}); - ReduceWindowAdd(input, {2}, {1}, Padding::kValid); - ComputeAndCompareR1(&builder_, {110, 11}, {}, ErrorSpec(0.0001)); -} - -TEST_F(ReduceWindowTest, Add3In5Stride2) { - const auto input = builder_.ConstantR1({10000, 1000, 100, 10, 1}); - ReduceWindowAdd(input, {3}, {2}, Padding::kValid); - ComputeAndCompareR1(&builder_, {11100, 111}, {}, ErrorSpec(0.0001)); -} - -TEST_F(ReduceWindowTest, Max4In16Stride4) { - const auto input = builder_.ConstantR1( - {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16}); - ReduceWindowMax(input, {4}, {4}, Padding::kValid); - ComputeAndCompareR1(&builder_, {4, 8, 12, 16}, {}, ErrorSpec(0.0001)); -} - -TEST_F(ReduceWindowTest, Max4In16Stride3) { - const auto input = builder_.ConstantR1( - {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16}); - ReduceWindowMax(input, {4}, {3}, Padding::kValid); - ComputeAndCompareR1(&builder_, {4, 7, 10, 13, 16}, {}, - ErrorSpec(0.0001)); -} - -TEST_F(ReduceWindowTest, Max4In16Stride8) { - const auto input = builder_.ConstantR1( - {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16}); - ReduceWindowMax(input, {4}, {8}, Padding::kValid); - ComputeAndCompareR1(&builder_, {4, 12}, {}, ErrorSpec(0.0001)); -} - -TEST_F(ReduceWindowTest, Max3In5Stride2) { - const auto input = builder_.ConstantR1({10000, 1000, 100, 10, 1}); - ReduceWindowMax(input, {3}, {2}, Padding::kValid); - ComputeAndCompareR1(&builder_, {10000, 100}, {}, ErrorSpec(0.0001)); -} - -TEST_F(ReduceWindowTest, Max3In5Stride1) { - const auto input = builder_.ConstantR1({10000, 1000, 100, 10, 101}); - ReduceWindowMax(input, {3}, {1}, Padding::kValid); - ComputeAndCompareR1(&builder_, {10000, 1000, 101}, {}, - ErrorSpec(0.0001)); -} - -TEST_F(ReduceWindowTest, Add3In4Stride2) { - const auto input = builder_.ConstantR1({1000, 100, 10, 1}); - ReduceWindowAdd(input, {3}, {2}, Padding::kValid); - ComputeAndCompareR1(&builder_, {1110}, {}, ErrorSpec(0.0001)); -} - -XLA_TEST_F(ReduceWindowTest, Add2In3SamePad) { - const auto input = builder_.ConstantR1({100, 10, 1}); - ReduceWindowAdd(input, {2}, {1}, Padding::kSame); - ComputeAndCompareR1(&builder_, {110, 11, 1}, {}, ErrorSpec(0.0001)); -} - -XLA_TEST_F(ReduceWindowTest, Add3In3SamePad) { - const auto input = builder_.ConstantR1({100, 10, 1}); - ReduceWindowAdd(input, {3}, {1}, Padding::kSame); - ComputeAndCompareR1(&builder_, {110, 111, 11}, {}, ErrorSpec(0.0001)); -} - -XLA_TEST_F(ReduceWindowTest, Add3In3Stride3SamePad) { - const auto input = builder_.ConstantR1({100, 10, 1}); - ReduceWindowAdd(input, {3}, {2}, Padding::kSame); - ComputeAndCompareR1(&builder_, {110, 11}, {}, ErrorSpec(0.0001)); -} - -TEST_F(ReduceWindowTest, Add2x2In2x2Overlapped) { - Array2D input_array({{1.2f, -2.5f, 0.9f, 1.0f}, - {3.7f, 0.2f, -1.0f, -0.2f}, - {-0.4f, 2.7f, 1.1f, 2.2f}, - {0.6f, 1.7f, 1.4f, -0.2f}}); - auto input = builder_.ConstantR2FromArray2D(input_array); - ReduceWindowAdd(input, {2, 2}, {1, 1}, Padding::kValid); - Array2D expected( - {{2.6f, -2.4f, 0.7f}, {6.2f, 3.0f, 2.1f}, {4.6f, 6.9f, 4.5f}}); - ComputeAndCompareR2(&builder_, expected, {}, ErrorSpec(0.0001)); -} - -TEST_F(ReduceWindowTest, Add2x2In2x2Disjoint) { - Array2D input_array({{1.2f, -2.5f, 0.9f, 1.0f}, - {3.7f, 0.2f, -1.0f, -0.2f}, - {-0.4f, 2.7f, 1.1f, 2.2f}, - {0.6f, 1.7f, 1.4f, -0.2f}}); - auto input = builder_.ConstantR2FromArray2D(input_array); - ReduceWindowAdd(input, {2, 2}, {2, 2}, Padding::kValid); - Array2D expected({ - {2.6f, 0.7f}, {4.6f, 4.5f}, - }); - ComputeAndCompareR2(&builder_, expected, {}, ErrorSpec(0.0001)); -} XLA_TEST_F(ReduceWindowTest, Add1x1x2In2x1x2) { Array3D input_array(2, 1, 2); @@ -458,22 +299,676 @@ XLA_TEST_F(ReduceWindowTest, NonstandardReduceFunction) { ComputeAndCompareR4(&builder_, *expected, {}, ErrorSpec(1e-3, 1e-3)); } -} // namespace -} // namespace xla +TEST_F(ReduceWindowTest, R4UnitWindow) { + Array4D input_array(13, 12, 8, 15); + input_array.Fill(1.0f); + std::unique_ptr input_literal = + Literal::CreateR4FromArray4DWithLayout( + input_array, LayoutUtil::MakeLayout({0, 3, 2, 1})); + ComputationDataHandle input = + builder_.Parameter(0, input_literal->shape(), "operand"); -int main(int argc, char** argv) { - std::vector flag_list; - xla::legacy_flags::AppendCpuCompilerFlags(&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; - } - testing::InitGoogleTest(&argc, argv); - if (argc > 1) { - LOG(ERROR) << "Unknown argument " << argv[1] << "\n" << usage; - return 2; + Padding padding = Padding::kSame; + ReduceWindowAdd(input, {1, 1, 7, 1}, {1, 4, 1, 1}, padding); + + auto res = ReferenceUtil::ReduceWindow4DAdd(input_array, 0.0f, {1, 1, 7, 1}, + {1, 4, 1, 1}, padding); + + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr input_data, + client_->TransferToServer(*input_literal)); + ComputeAndCompareR4(&builder_, *res, {input_data.get()}, + ErrorSpec(1e-3, 1e-3)); +} + +XLA_TEST_F(HloTestBase, R6Add) { + auto b = HloComputation::Builder(TestName()); + + std::vector input_dims(6, 8); + std::unique_ptr arg_literal = + Literal::CreateFullWithMonotonicDim0MajorLayout(input_dims, 1.0f); + auto input = + b.AddInstruction(HloInstruction::CreateConstant(std::move(arg_literal))); + + auto init_value = b.AddInstruction( + HloInstruction::CreateConstant(Literal::CreateR0(0.f))); + + HloComputation::Builder add_computation("add"); + Shape scalar_shape = ShapeUtil::MakeShape(F32, {}); + auto param_lhs = add_computation.AddInstruction( + HloInstruction::CreateParameter(0, scalar_shape, "lhs")); + auto param_rhs = add_computation.AddInstruction( + HloInstruction::CreateParameter(1, scalar_shape, "rhs")); + add_computation.AddInstruction(HloInstruction::CreateBinary( + scalar_shape, HloOpcode::kAdd, param_lhs, param_rhs)); + + auto module = CreateNewModule(); + auto add_func = module->AddEmbeddedComputation(add_computation.Build()); + + WindowDimension trivial_dim; + trivial_dim.set_size(1); + trivial_dim.set_stride(1); + trivial_dim.set_padding_low(0); + trivial_dim.set_padding_high(0); + trivial_dim.set_window_dilation(1); + trivial_dim.set_base_dilation(1); + + WindowDimension active_dim; + active_dim.set_size(3); + active_dim.set_stride(1); + active_dim.set_padding_low(0); + active_dim.set_padding_high(0); + active_dim.set_window_dilation(1); + active_dim.set_base_dilation(1); + + Window window; + *window.add_dimensions() = trivial_dim; + *window.add_dimensions() = trivial_dim; + *window.add_dimensions() = active_dim; + *window.add_dimensions() = active_dim; + *window.add_dimensions() = trivial_dim; + *window.add_dimensions() = trivial_dim; + + Shape shape = ShapeUtil::MakeShape(F32, {8, 8, 6, 6, 8, 8}); + b.AddInstruction(HloInstruction::CreateReduceWindow(shape, input, init_value, + window, add_func)); + + std::vector output_dims = {8, 8, 6, 6, 8, 8}; + std::unique_ptr expected = + Literal::CreateFullWithMonotonicDim0MajorLayout(output_dims, 9.0f); + + module->AddEntryComputation(b.Build()); + auto actual = ExecuteAndTransfer(std::move(module), {}); + + LiteralTestUtil::ExpectNear(*actual, *expected, ErrorSpec(1e-3, 1e-3)); +} + +XLA_TEST_F(ReduceWindowTest, R4SecondMinorStride) { + Array4D input_array(2, 1, 27, 119); + input_array.FillRandom(2.0f); + std::unique_ptr input_literal = + Literal::CreateR4FromArray4DWithLayout( + input_array, LayoutUtil::MakeLayout({3, 2, 1, 0})); + ComputationDataHandle input = + builder_.Parameter(0, input_literal->shape(), "operand"); + + int win_len = 1; + int stride = 8; + Padding padding = Padding::kSame; + ReduceWindowAdd(input, {1, 1, win_len, 1}, {1, 1, stride, 1}, padding); + + auto res = ReferenceUtil::ReduceWindow4DAdd( + input_array, 0.0f, {1, 1, win_len, 1}, {1, 1, stride, 1}, padding); + + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr input_data, + client_->TransferToServer(*input_literal)); + ComputeAndCompareR4(&builder_, *res, {input_data.get()}, + ErrorSpec(1e-3, 1e-3)); +} + +XLA_TEST_F(ReduceWindowTest, R4SecondMinorUnitStride) { + Array4D input_array(3, 2, 4, 64); + input_array.FillRandom(2.0f); + std::unique_ptr input_literal = + Literal::CreateR4FromArray4DWithLayout( + input_array, LayoutUtil::MakeLayout({3, 2, 1, 0})); + ComputationDataHandle input = + builder_.Parameter(0, input_literal->shape(), "operand"); + + int win_len = 3; + int stride = 1; + Padding padding = Padding::kSame; + ReduceWindowAdd(input, {1, 1, win_len, 1}, {1, 1, stride, 1}, padding); + + auto res = ReferenceUtil::ReduceWindow4DAdd( + input_array, 0.0f, {1, 1, win_len, 1}, {1, 1, stride, 1}, padding); + + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr input_data, + client_->TransferToServer(*input_literal)); + ComputeAndCompareR4(&builder_, *res, {input_data.get()}, + ErrorSpec(1e-3, 1e-3)); +} + +XLA_TEST_F(ReduceWindowTest, R4SecondMinorWin) { + Array4D input_array(1, 3, 12, 200); + input_array.FillRandom(2.0f); + std::unique_ptr input_literal = + Literal::CreateR4FromArray4DWithLayout( + input_array, LayoutUtil::MakeLayout({3, 2, 1, 0})); + ComputationDataHandle input = + builder_.Parameter(0, input_literal->shape(), "operand"); + + int win_len = 8; + int stride = 5; + Padding padding = Padding::kSame; + ReduceWindowAdd(input, {1, 1, win_len, 1}, {1, 1, stride, 1}, padding); + + auto res = ReferenceUtil::ReduceWindow4DAdd( + input_array, 0.0f, {1, 1, win_len, 1}, {1, 1, stride, 1}, padding); + + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr input_data, + client_->TransferToServer(*input_literal)); + ComputeAndCompareR4(&builder_, *res, {input_data.get()}, + ErrorSpec(1e-3, 1e-3)); +} + +TEST_F(ReduceWindowTest, AmongMajor2DimsMultipleMinor) { + Array4D input_array(6, 4, 10, 130); + input_array.FillRandom(2.0f); + + int win_len = 3; + int win_stride = 2; + + Padding padding = Padding::kSame; + const auto input_data_handle = + builder_.ConstantR4FromArray4D(input_array); + // Reduce only along the x and y dimensions, according to the win_len. + ReduceWindowAdd(input_data_handle, {win_len, win_len, 1, 1}, + {win_stride, win_stride, 1, 1}, padding); + + auto result = ReferenceUtil::ReduceWindow4DAdd( + input_array, 0.0f, {win_len, win_len, 1, 1}, + {win_stride, win_stride, 1, 1}, padding); + ComputeAndCompareR4(&builder_, *result, {}, ErrorSpec(1e-3, 1e-3)); +} + +XLA_TEST_F(ReduceWindowTest, Add24In1152_NoOverlap) { + std::vector input_vector(128 * 9, 1); + const auto input = builder_.ConstantR1(input_vector); + ReduceWindowAdd(input, {32}, {128}, Padding::kValid); + ComputeAndCompareR1(&builder_, {32, 32, 32, 32, 32, 32, 32, 32, 32}, + {}, ErrorSpec(0.0001)); +} + +XLA_TEST_F(ReduceWindowTest, Add128In128Stride128) { + const auto input = builder_.ConstantR1( + {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, + 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, + 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, + 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}); + ReduceWindowAdd(input, {128}, {128}, Padding::kValid); + ComputeAndCompareR1(&builder_, {1088}, {}, ErrorSpec(0.0001)); +} + +// Regression test for a bug that appeared in Inception (b/34784899). +TEST_F(ReduceWindowTest, R2ReduceWindowInceptionFromBroadcast) { + Array2D input_array(14, 14, 1.0f); + ComputationDataHandle input = + builder_.Broadcast(builder_.ConstantLiteral(Literal::One(F32)), {14, 14}); + + int win_len = 3; + int stride = 1; + Padding padding = Padding::kSame; + ReduceWindowAdd(input, {win_len, win_len}, {stride, stride}, padding); + + auto res = ReferenceUtil::ReduceWindow2DAdd( + input_array, 0.0f, {win_len, win_len}, {stride, stride}, padding); + + ComputeAndCompareR2(&builder_, *res, {}, ErrorSpec(1e-3, 1e-3)); +} + +TEST_F(ReduceWindowTest, R2ReduceWindowNonOverlappingFromBroadcast) { + Array2D input_array(6, 4, 1.0f); + ComputationDataHandle input = + builder_.Broadcast(builder_.ConstantLiteral(Literal::One(F32)), {6, 4}); + + Padding padding = Padding::kSame; + ReduceWindowAdd(input, {4, 2}, {3, 3}, padding); + + auto res = ReferenceUtil::ReduceWindow2DAdd(input_array, 0.0f, {4, 2}, {3, 3}, + padding); + + ComputeAndCompareR2(&builder_, *res, {}, ErrorSpec(1e-3, 1e-3)); +} + +enum Reducer { kAdd, kMax }; + +struct R4ReduceWindowTestData { + int64 base_bounds[4]; + int64 window_bounds[4]; + int64 strides[4]; + int64 pad_low[4]; + int64 pad_high[4]; + + Reducer reducer; +}; + +string R4ReduceWindowTestDataToString( + const ::testing::TestParamInfo& data) { + string str = tensorflow::strings::StrCat( + "base_bounds_", + tensorflow::str_util::Join(data.param.base_bounds, "x"), // + "__window_bounds_", + tensorflow::str_util::Join(data.param.window_bounds, "x"), // + "__strides_", tensorflow::str_util::Join(data.param.strides, "x"), // + "__pad_low_", tensorflow::str_util::Join(data.param.pad_low, "x"), // + "__pad_high_", tensorflow::str_util::Join(data.param.pad_high, "x"), // + (data.param.reducer == kAdd) ? "add" : "max"); + CHECK(data.param.reducer == kAdd || data.param.reducer == kMax); + + // Test names are not allowed to contain the '-' character. + std::replace(str.begin(), str.end(), '-', 'n'); + return str; +} + +class R4ReduceWindowTest + : public ClientLibraryTestBase, + public ::testing::WithParamInterface { + protected: + void DoIt() { + ComputationBuilder b(client_, TestName()); + const auto& param = GetParam(); + + const float kInitValue = 0.0f; + + Array4D input(param.base_bounds[0], param.base_bounds[1], + param.base_bounds[2], param.base_bounds[3]); + input.FillIota(1); + std::unique_ptr input_literal = + Literal::CreateR4FromArray4D(input); + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr input_arg, + client_->TransferToServer(*input_literal)); + + std::vector> padding(4); + for (int i = 0; i < 4; ++i) { + padding[i] = {param.pad_low[i], param.pad_high[i]}; + } + + auto parameter = b.Parameter(0, input_literal->shape(), "p0"); + auto pad_value = b.ConstantR0(kInitValue); + CHECK(param.reducer == kAdd || param.reducer == kMax); + auto computation = param.reducer == kAdd + ? CreateScalarAddComputation(F32, &b) + : CreateScalarMaxComputation(F32, &b); + b.ReduceWindowWithGeneralPadding( + /*operand=*/parameter, + /*init_value=*/pad_value, + /*computation=*/computation, + /*window_dimensions=*/param.window_bounds, + /*window_strides=*/param.strides, + /*padding=*/padding); + + CHECK(param.reducer == kAdd || param.reducer == kMax); + auto reduce_func = param.reducer == kAdd + ? +[](float a, float b) { return a + b; } + : +[](float a, float b) { return std::max(a, b); }; + std::unique_ptr> expected = + ReferenceUtil::ReduceWindow4DGeneric( + /*operand=*/input, + /*init=*/kInitValue, + /*reduce_func=*/reduce_func, + /*window=*/param.window_bounds, + /*stride=*/param.strides, + /*padding=*/padding); + ComputeAndCompareR4(&b, *expected, {input_arg.get()}, + ErrorSpec(1e-3, 1e-3)); } - return RUN_ALL_TESTS(); +}; + +TEST_P(R4ReduceWindowTest, DoIt) { DoIt(); } + +// base_bounds, window_bounds, strides, pad_low, pad_high +const R4ReduceWindowTestData kR4ReduceWindowTestValues[] = { + // Minimal edge case. + R4ReduceWindowTestData{/*base_bounds=*/{1, 1, 1, 1}, + /*window_bounds=*/{1, 1, 1, 1}, + /*strides=*/{1, 1, 1, 1}, + /*pad_low=*/{0, 0, 0, 0}, + /*pad_high=*/{0, 0, 0, 0}, + /*reducer=*/kAdd}, + + // Zero base bound edge case. + R4ReduceWindowTestData{/*base_bounds=*/{1, 0, 1, 1}, + /*window_bounds=*/{1, 1, 1, 1}, + /*strides=*/{1, 1, 1, 1}, + /*pad_low=*/{0, 0, 0, 0}, + /*pad_high=*/{0, 0, 0, 0}, + /*reducer=*/kAdd}, + + // With non-1x1 window. + R4ReduceWindowTestData{/*base_bounds=*/{4, 6, 17, 140}, + /*window_bounds=*/{2, 3, 1, 1}, + /*strides=*/{1, 1, 1, 1}, + /*pad_low=*/{0, 0, 0, 0}, + /*pad_high=*/{0, 0, 0, 0}, + /*reducer=*/kAdd}, + + // With max instead of add. + R4ReduceWindowTestData{/*base_bounds=*/{4, 6, 17, 140}, + /*window_bounds=*/{2, 3, 1, 1}, + /*strides=*/{1, 1, 1, 1}, + /*pad_low=*/{0, 0, 0, 0}, + /*pad_high=*/{0, 0, 0, 0}, + /*reducer=*/kMax}, + + // With stride. + R4ReduceWindowTestData{/*base_bounds=*/{4, 10, 17, 140}, + /*window_bounds=*/{3, 2, 1, 1}, + /*strides=*/{2, 4, 1, 1}, + /*pad_low=*/{0, 0, 0, 0}, + /*pad_high=*/{0, 0, 0, 0}, + /*reducer=*/kAdd}, + + // With low padding. + R4ReduceWindowTestData{/*base_bounds=*/{4, 6, 17, 140}, + /*window_bounds=*/{3, 2, 1, 1}, + /*strides=*/{2, 2, 1, 1}, + /*pad_low=*/{3, 2, 0, 0}, + /*pad_high=*/{0, 0, 0, 0}, + /*reducer=*/kAdd}, + + // With high padding. + R4ReduceWindowTestData{/*base_bounds=*/{4, 6, 17, 140}, + /*window_bounds=*/{3, 2, 1, 1}, + /*strides=*/{2, 2, 1, 1}, + /*pad_low=*/{0, 0, 0, 0}, + /*pad_high=*/{2, 3, 0, 0}, + /*reducer=*/kAdd}, + + // Window touches both sides of the padding simultaneously. + R4ReduceWindowTestData{/*base_bounds=*/{1, 1, 17, 140}, + /*window_bounds=*/{3, 3, 1, 1}, + /*strides=*/{1, 1, 1, 1}, + /*pad_low=*/{1, 1, 0, 0}, + /*pad_high=*/{1, 1, 0, 0}, + /*reducer=*/kAdd}, + + // Window is entirely in the padding for some positions. + R4ReduceWindowTestData{/*base_bounds=*/{1, 1, 17, 140}, + /*window_bounds=*/{3, 3, 1, 1}, + /*strides=*/{1, 1, 1, 1}, + /*pad_low=*/{4, 4, 0, 0}, + /*pad_high=*/{4, 4, 0, 0}, + /*reducer=*/kAdd}, + + // Zero base bound with padding edge case. + R4ReduceWindowTestData{/*base_bounds=*/{2, 0, 3, 4}, + /*window_bounds=*/{1, 1, 1, 1}, + /*strides=*/{1, 1, 1, 1}, + /*pad_low=*/{0, 1, 0, 0}, + /*pad_high=*/{0, 0, 0, 0}, + /*reducer=*/kAdd}, + + // With stride, low padding and high padding. + R4ReduceWindowTestData{/*base_bounds=*/{4, 3, 17, 140}, + /*window_bounds=*/{3, 4, 1, 1}, + /*strides=*/{3, 1, 1, 1}, + /*pad_low=*/{10, 1, 0, 0}, + /*pad_high=*/{2, 3, 0, 0}, + /*reducer=*/kAdd}, + + // With second minor dimension == 9. + R4ReduceWindowTestData{/*base_bounds=*/{2, 3, 9, 127}, + /*window_bounds=*/{1, 1, 1, 1}, + /*strides=*/{1, 1, 1, 1}, + /*pad_low=*/{0, 0, 0, 0}, + /*pad_high=*/{0, 0, 0, 0}, + /*reducer=*/kAdd}, + + // With minor dimension == 129. + R4ReduceWindowTestData{/*base_bounds=*/{3, 2, 7, 129}, + /*window_bounds=*/{1, 1, 1, 1}, + /*strides=*/{1, 1, 1, 1}, + /*pad_low=*/{0, 0, 0, 0}, + /*pad_high=*/{0, 0, 0, 0}, + /*reducer=*/kAdd}, + + // With minor dims reduction and non-overlapped stride. + R4ReduceWindowTestData{/*base_bounds=*/{2, 2, 4, 16}, + /*window_bounds=*/{1, 1, 2, 2}, + /*strides=*/{1, 1, 2, 2}, + /*pad_low=*/{0, 0, 0, 0}, + /*pad_high=*/{0, 0, 0, 0}, + /*reducer=*/kAdd}, + + // With minor dims reduction and overlapped stride. + R4ReduceWindowTestData{/*base_bounds=*/{2, 2, 4, 16}, + /*window_bounds=*/{1, 1, 4, 4}, + /*strides=*/{1, 1, 2, 2}, + /*pad_low=*/{0, 0, 0, 0}, + /*pad_high=*/{0, 0, 0, 0}, + /*reducer=*/kAdd}, +}; + +INSTANTIATE_TEST_CASE_P(R4ReduceWindowTestInstantiation, R4ReduceWindowTest, + ::testing::ValuesIn(kR4ReduceWindowTestValues), + R4ReduceWindowTestDataToString); + +class R4ReduceWindowLargeTest : public R4ReduceWindowTest {}; + +XLA_TEST_P(R4ReduceWindowLargeTest, DoIt) { DoIt(); } + +// Test cases that are large/slow/failed. +const R4ReduceWindowTestData kR4ReduceWindowLargeTestValues[] = { + R4ReduceWindowTestData{/*base_bounds=*/{28, 28, 256, 128}, + /*window_bounds=*/{3, 3, 1, 1}, + /*strides=*/{1, 1, 1, 1}, + /*pad_low=*/{1, 1, 0, 0}, + /*pad_high=*/{1, 1, 0, 0}, + /*reducer=*/kMax}, + + R4ReduceWindowTestData{/*base_bounds=*/{112, 112, 64, 128}, + /*window_bounds=*/{3, 3, 1, 1}, + /*strides=*/{2, 2, 1, 1}, + /*pad_low=*/{0, 0, 0, 0}, + /*pad_high=*/{1, 1, 0, 0}, + /*reducer=*/kAdd}, +}; + +INSTANTIATE_TEST_CASE_P(R4ReduceWindowLargeTestInstantiation, + R4ReduceWindowLargeTest, + ::testing::ValuesIn(kR4ReduceWindowLargeTestValues), + R4ReduceWindowTestDataToString); + +struct R2ReduceWindowTestData { + int64 base_bounds[2]; + int64 window_bounds[2]; + int64 strides[2]; + int64 layout[2]; + Padding padding; + Reducer reducer; +} kR2TestCases[] = { + {/*base_bounds=*/{4, 18}, /*window_bounds=*/{2, 4}, + /*strides=*/{1, 2}, /*layout=*/{0, 1}, + /*padding=*/Padding::kSame, /*reducer=*/Reducer::kAdd}, + {/*base_bounds=*/{2, 5}, /*window_bounds=*/{2, 4}, + /*strides=*/{1, 1}, /*layout=*/{0, 1}, + /*padding=*/Padding::kSame, /*reducer=*/Reducer::kAdd}, + {/*base_bounds=*/{1, 3}, /*window_bounds=*/{2, 3}, + /*strides=*/{1, 1}, /*layout=*/{0, 1}, + /*padding=*/Padding::kSame, /*reducer=*/Reducer::kAdd}, + {/*base_bounds=*/{3, 129}, /*window_bounds=*/{1, 100}, + /*strides=*/{2, 99}, /*layout=*/{0, 1}, + /*padding=*/Padding::kSame, /*reducer=*/Reducer::kAdd}, + {/*base_bounds=*/{6, 152}, /*window_bounds=*/{2, 25}, + /*strides=*/{5, 4}, /*layout=*/{0, 1}, + /*padding=*/Padding::kSame, /*reducer=*/Reducer::kAdd}, + {/*base_bounds=*/{6, 4}, /*window_bounds=*/{4, 2}, + /*strides=*/{3, 3}, /*layout=*/{0, 1}, + /*padding=*/Padding::kSame, /*reducer=*/Reducer::kAdd}, + {/*base_bounds=*/{5, 147}, /*window_bounds=*/{1, 36}, + /*strides=*/{4, 5}, /*layout=*/{1, 0}, + /*padding=*/Padding::kSame, /*reducer=*/Reducer::kAdd}, + {/*base_bounds=*/{4, 153}, /*window_bounds=*/{2, 93}, + /*strides=*/{1, 1}, /*layout=*/{1, 0}, + /*padding=*/Padding::kSame, /*reducer=*/Reducer::kAdd}, + // Regression test for a bug that appeared in Inception (b/34784899). + {/*base_bounds=*/{28, 28}, /*window_bounds=*/{3, 3}, + /*strides=*/{1, 1}, /*layout=*/{1, 0}, + /*padding=*/Padding::kSame, /*reducer=*/Reducer::kAdd}, + // Regression test for a bug that appeared in Inception (b/34784899). + {/*base_bounds=*/{4, 32}, /*window_bounds=*/{2, 2}, + /*strides=*/{2, 2}, /*layout=*/{1, 0}, + /*padding=*/Padding::kValid, /*reducer=*/Reducer::kAdd}, + {/*base_bounds=*/{4, 4}, /*window_bounds=*/{2, 2}, + /*strides=*/{1, 1}, /*layout=*/{1, 0}, + /*padding=*/Padding::kValid, /*reducer=*/Reducer::kAdd}, +}; + +string R2ReduceWindowTestDataToString( + const ::testing::TestParamInfo& data) { + string str = tensorflow::strings::StrCat( + "base_bounds_", + tensorflow::str_util::Join(data.param.base_bounds, "x"), // + "__window_bounds_", + tensorflow::str_util::Join(data.param.window_bounds, "x"), // + "__strides_", tensorflow::str_util::Join(data.param.strides, "x"), // + "__padding_", data.param.padding == Padding::kSame ? "same" : "valid", // + "__layout_", data.param.layout[0], "_", data.param.layout[1], // + "__reducer_", data.param.reducer == kAdd ? "add" : "max"); + return str; } + +class R2ReduceWindowTest + : public ClientLibraryTestBase, + public ::testing::WithParamInterface {}; + +TEST_P(R2ReduceWindowTest, Add) { + ComputationBuilder b(client_, TestName()); + const auto& param = GetParam(); + CHECK(param.reducer == kAdd); + + const float kInitValue = 0.0f; + Array2D input(param.base_bounds[0], param.base_bounds[1], 1.0f); + std::unique_ptr input_literal = + Literal::CreateR2FromArray2DWithLayout( + input, LayoutUtil::MakeLayout(param.layout)); + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr input_arg, + client_->TransferToServer(*input_literal)); + b.ReduceWindow(/*operand=*/ + b.Parameter(0, input_literal->shape(), "p0"), + /*init_value=*/b.ConstantR0(kInitValue), + /*computation=*/CreateScalarAddComputation(F32, &b), + /*window_dimensions=*/param.window_bounds, + /*window_strides=*/param.strides, /*padding=*/param.padding); + + auto expected = ReferenceUtil::ReduceWindow2DAdd( + /*operand=*/input, /*init=*/kInitValue, /*window=*/param.window_bounds, + /*stride=*/param.strides, /*padding=*/param.padding); + + ComputeAndCompareR2(&b, *expected, {input_arg.get()}, + ErrorSpec(1e-3, 1e-3)); +} + +INSTANTIATE_TEST_CASE_P(R2ReduceWindowTestInstantiation, R2ReduceWindowTest, + ::testing::ValuesIn(kR2TestCases), + R2ReduceWindowTestDataToString); + +struct R1ReduceWindowTestData { + int64 base_bounds[1]; + int64 window_bounds[1]; + int64 strides[1]; + Padding padding; + Reducer reducer; +} kR1TestCases[] = { + {/*base_bounds=*/{1}, /*window_bounds=*/{1}, + /*strides=*/{1}, + /*padding=*/Padding::kValid, /*reducer=*/Reducer::kAdd}, + + {/*base_bounds=*/{3}, /*window_bounds=*/{3}, + /*strides=*/{1}, + /*padding=*/Padding::kValid, /*reducer=*/Reducer::kAdd}, + + {/*base_bounds=*/{3}, /*window_bounds=*/{2}, + /*strides=*/{1}, + /*padding=*/Padding::kValid, /*reducer=*/Reducer::kAdd}, + + {/*base_bounds=*/{5}, /*window_bounds=*/{1}, + /*strides=*/{1}, + /*padding=*/Padding::kValid, /*reducer=*/Reducer::kMax}, + + {/*base_bounds=*/{16}, /*window_bounds=*/{4}, + /*strides=*/{4}, + /*padding=*/Padding::kValid, /*reducer=*/Reducer::kMax}, + + {/*base_bounds=*/{16}, /*window_bounds=*/{4}, + /*strides=*/{3}, + /*padding=*/Padding::kValid, /*reducer=*/Reducer::kAdd}, + + {/*base_bounds=*/{128 * 2}, /*window_bounds=*/{30}, + /*strides=*/{27}, + /*padding=*/Padding::kValid, /*reducer=*/Reducer::kAdd}, + + {/*base_bounds=*/{128 * 17}, /*window_bounds=*/{7}, + /*strides=*/{64}, + /*padding=*/Padding::kValid, /*reducer=*/Reducer::kAdd}, + + {/*base_bounds=*/{128 * 2}, /*window_bounds=*/{32}, + /*strides=*/{56}, + /*padding=*/Padding::kValid, /*reducer=*/Reducer::kAdd}, + + {/*base_bounds=*/{3}, /*window_bounds=*/{2}, + /*strides=*/{1}, + /*padding=*/Padding::kSame, /*reducer=*/Reducer::kAdd}, + + {/*base_bounds=*/{5}, /*window_bounds=*/{3}, + /*strides=*/{2}, + /*padding=*/Padding::kSame, /*reducer=*/Reducer::kAdd}, + + {/*base_bounds=*/{16}, /*window_bounds=*/{4}, + /*strides=*/{3}, + /*padding=*/Padding::kSame, /*reducer=*/Reducer::kAdd}, +}; + +string R1ReduceWindowTestDataToString( + const ::testing::TestParamInfo& data) { + string str = tensorflow::strings::StrCat( + "base_bounds_", + tensorflow::str_util::Join(data.param.base_bounds, "x"), // + "__window_bounds_", + tensorflow::str_util::Join(data.param.window_bounds, "x"), // + "__strides_", tensorflow::str_util::Join(data.param.strides, "x"), // + "__padding_", data.param.padding == Padding::kSame ? "same" : "valid", // + "__reducer_", data.param.reducer == kAdd ? "add" : "max"); + return str; +} + +class R1ReduceWindowTest + : public ClientLibraryTestBase, + public ::testing::WithParamInterface {}; + +TEST_P(R1ReduceWindowTest, DoIt) { + ComputationBuilder b(client_, TestName()); + const auto& param = GetParam(); + CHECK(param.reducer == kAdd || param.reducer == kMax); + + const float kInitValue = 0.0f; + std::vector input_vector(param.base_bounds[0]); + std::iota(std::begin(input_vector), std::end(input_vector), 0); + std::unique_ptr input_literal = + Literal::CreateR1(tensorflow::gtl::ArraySlice(input_vector)); + TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr input_arg, + client_->TransferToServer(*input_literal)); + + auto computation = param.reducer == kAdd + ? CreateScalarAddComputation(F32, &b) + : CreateScalarMaxComputation(F32, &b); + b.ReduceWindow(/*operand=*/ + b.Parameter(0, input_literal->shape(), "p0"), + /*init_value=*/b.ConstantR0(kInitValue), + /*computation=*/computation, + /*window_dimensions=*/param.window_bounds, + /*window_strides=*/param.strides, /*padding=*/param.padding); + + auto reduce_func = param.reducer == kAdd + ? +[](float a, float b) { return a + b; } + : +[](float a, float b) { return std::max(a, b); }; + auto expected = ReferenceUtil::ReduceWindow1DGeneric( + /*operand=*/tensorflow::gtl::ArraySlice(input_vector), + /*init=*/kInitValue, + /*reduce_func=*/reduce_func, + /*window=*/param.window_bounds, + /*stride=*/param.strides, /*padding=*/param.padding); + + ComputeAndCompareR1(&b, tensorflow::gtl::ArraySlice(*expected), + {input_arg.get()}, ErrorSpec(1e-3, 1e-3)); +} + +INSTANTIATE_TEST_CASE_P(R1ReduceWindowTestInstantiation, R1ReduceWindowTest, + ::testing::ValuesIn(kR1TestCases), + R1ReduceWindowTestDataToString); +} // namespace +} // namespace xla diff --git a/tensorflow/compiler/xla/tests/replay_test.cc b/tensorflow/compiler/xla/tests/replay_test.cc index 802087b5086883183930736cea53976098265e87..92efd2947d6384d4ffaf6dc0134ddaf313ddedf7 100644 --- a/tensorflow/compiler/xla/tests/replay_test.cc +++ b/tensorflow/compiler/xla/tests/replay_test.cc @@ -19,7 +19,6 @@ limitations under the License. #include "tensorflow/compiler/xla/client/computation_builder.h" #include "tensorflow/compiler/xla/client/global_data.h" #include "tensorflow/compiler/xla/client/local_client.h" -#include "tensorflow/compiler/xla/legacy_flags/cpu_compiler_flags.h" #include "tensorflow/compiler/xla/literal_util.h" #include "tensorflow/compiler/xla/protobuf_util.h" #include "tensorflow/compiler/xla/service/session.pb.h" @@ -60,7 +59,8 @@ TEST_F(ReplayTest, TwoPlusTwoReplay) { // Run it. std::unique_ptr literal = - client_->ExecuteAndTransfer(replayed, /*arguments=*/{}) + client_ + ->ExecuteAndTransfer(replayed, /*arguments=*/{}, &execution_options_) .ConsumeValueOrDie(); // Expect 4. @@ -91,15 +91,16 @@ XLA_TEST_F(ReplayTest, XPlusYReplayWithParameters) { // Run it. std::unique_ptr x_data = - client_->TransferToServer(*LiteralUtil::CreateR0(2)) + client_->TransferToServer(*Literal::CreateR0(2)) .ConsumeValueOrDie(); std::unique_ptr y_data = - client_->TransferToServer(*LiteralUtil::CreateR0(3)) + client_->TransferToServer(*Literal::CreateR0(3)) .ConsumeValueOrDie(); std::unique_ptr literal = client_ ->ExecuteAndTransfer(replayed, - /*arguments=*/{x_data.get(), y_data.get()}) + /*arguments=*/{x_data.get(), y_data.get()}, + &execution_options_) .ConsumeValueOrDie(); // Expect 5. @@ -140,7 +141,8 @@ TEST_F(ReplayTest, MapPlusTwoOverR1) { // Run it. std::unique_ptr literal = - client_->ExecuteAndTransfer(replayed, /*arguments=*/{}) + client_ + ->ExecuteAndTransfer(replayed, /*arguments=*/{}, &execution_options_) .ConsumeValueOrDie(); // Expect result. @@ -149,20 +151,3 @@ TEST_F(ReplayTest, MapPlusTwoOverR1) { } // namespace } // namespace xla - -int main(int argc, char** argv) { - std::vector flag_list; - xla::legacy_flags::AppendCpuCompilerFlags(&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; - } - testing::InitGoogleTest(&argc, argv); - if (argc > 1) { - LOG(ERROR) << "Unknown argument " << argv[1] << "\n" << usage; - return 2; - } - return RUN_ALL_TESTS(); -} diff --git a/tensorflow/compiler/xla/tests/reshape_motion_test.cc b/tensorflow/compiler/xla/tests/reshape_motion_test.cc index ce309eb7439e8295b1300c19e24b502755d42a4f..e045e164e2e2db7d3480e7c2d1e20f461820ae67 100644 --- a/tensorflow/compiler/xla/tests/reshape_motion_test.cc +++ b/tensorflow/compiler/xla/tests/reshape_motion_test.cc @@ -25,7 +25,6 @@ limitations under the License. #include "tensorflow/compiler/xla/client/global_data.h" #include "tensorflow/compiler/xla/client/local_client.h" #include "tensorflow/compiler/xla/layout_util.h" -#include "tensorflow/compiler/xla/legacy_flags/cpu_compiler_flags.h" #include "tensorflow/compiler/xla/literal_util.h" #include "tensorflow/compiler/xla/reference_util.h" #include "tensorflow/compiler/xla/shape_util.h" @@ -58,20 +57,3 @@ TEST_F(ReshapeMotionTest, ElementwiseOfReshapesWithNonSameInputShapes) { } // namespace } // namespace xla - -int main(int argc, char** argv) { - std::vector flag_list; - xla::legacy_flags::AppendCpuCompilerFlags(&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; - } - testing::InitGoogleTest(&argc, argv); - if (argc > 1) { - LOG(ERROR) << "Unknown argument " << argv[1] << "\n" << usage; - return 2; - } - return RUN_ALL_TESTS(); -} diff --git a/tensorflow/compiler/xla/tests/reshape_test.cc b/tensorflow/compiler/xla/tests/reshape_test.cc index 18e6e2d3f1d6aedb68f83b8058517398760c39ba..bb7160e3a03053a4f3d8da712c1424e50f37dfeb 100644 --- a/tensorflow/compiler/xla/tests/reshape_test.cc +++ b/tensorflow/compiler/xla/tests/reshape_test.cc @@ -25,19 +25,17 @@ limitations under the License. #include "tensorflow/compiler/xla/client/global_data.h" #include "tensorflow/compiler/xla/client/local_client.h" #include "tensorflow/compiler/xla/layout_util.h" -#include "tensorflow/compiler/xla/legacy_flags/cpu_compiler_flags.h" #include "tensorflow/compiler/xla/literal_util.h" #include "tensorflow/compiler/xla/reference_util.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/test_helpers.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/gtl/array_slice.h" -#include "tensorflow/core/platform/test.h" #include "tensorflow/core/platform/types.h" namespace xla { @@ -68,6 +66,22 @@ XLA_TEST_F(ReshapeTest, SingleElementArrayToScalar) { ComputeAndCompareR0(&builder, 1.0f, {}, zero_error_spec_); } +XLA_TEST_F(ReshapeTest, ScalarToSingleElementArray) { + ComputationBuilder builder(client_, TestName()); + + std::unique_ptr param0_literal = Literal::CreateR0(1.0f); + std::unique_ptr param0_data = + client_->TransferToServer(*param0_literal).ConsumeValueOrDie(); + + auto a = builder.Parameter(0, ShapeUtil::MakeShape(F32, {}), "param0"); + a = builder.Neg(a); + auto reshape = + builder.Reshape(/*operand=*/a, /*dimensions=*/{}, /*new_sizes=*/{1}); + + ComputeAndCompareR1(&builder, {-1.0f}, {param0_data.get()}, + zero_error_spec_); +} + XLA_TEST_F(ReshapeTest, Trivial0x3) { ComputationBuilder builder(client_, TestName()); auto a = builder.ConstantR2FromArray2D(Array2D(0, 3)); @@ -76,6 +90,24 @@ XLA_TEST_F(ReshapeTest, Trivial0x3) { ComputeAndCompareR1(&builder, {}, {}, zero_error_spec_); } +// TODO(b/29185393): Make this work with the GPU backend. The GPU backend +// does not handle zero-sized shapes correctly. Failed last on 2017-05-15 +// with an incorrect result rank. +XLA_TEST_F(ReshapeTest, DISABLED_ON_GPU(Trivial0x3WithParameter)) { + ComputationBuilder builder(client_, TestName()); + + std::unique_ptr param0_literal = + Literal::CreateR2FromArray2D(Array2D(0, 3)); + std::unique_ptr param0_data = + client_->TransferToServer(*param0_literal).ConsumeValueOrDie(); + + auto a = builder.Parameter(0, ShapeUtil::MakeShape(F32, {0, 3}), "param0"); + auto result = builder.Collapse(/*operand=*/a, /*dimensions=*/{0, 1}); + + ComputeAndCompareR1(&builder, {}, {param0_data.get()}, + zero_error_spec_); +} + XLA_TEST_F(ReshapeTest, Trivial3x0) { ComputationBuilder builder(client_, TestName()); auto a = builder.ConstantR2FromArray2D(Array2D(3, 0)); @@ -369,7 +401,7 @@ XLA_TEST_F(ReshapeTest, FullyConnectedCollapseDesugared) { XLA_TEST_F(ReshapeTest, ToScalar) { for (int rank = 0; rank < 8; ++rank) { ComputationBuilder b(client_, TestName()); - auto input = LiteralUtil::CreateR1({83.0f}); + auto input = Literal::CreateR1({83.0f}); std::vector ones(rank, 1); // this is {1, ..., 1}. std::vector dimensions(rank); std::iota(dimensions.begin(), dimensions.end(), 0); @@ -383,15 +415,15 @@ XLA_TEST_F(ReshapeTest, ToScalar) { XLA_TEST_F(ReshapeTest, BadDimensions) { ComputationBuilder b(client_, TestName()); b.Reshape(b.ConstantR1({1}), {}, {}); - EXPECT_MATCH(ExecuteToString(&b, {}), - testing::HasSubstr("dimensions not a permutation")); + EXPECT_THAT(ExecuteToString(&b, {}), + ::testing::HasSubstr("dimensions not a permutation")); } XLA_TEST_F(ReshapeTest, BadNewSizes) { ComputationBuilder b(client_, TestName()); b.Reshape(b.ConstantR1({1, 2}), {1}, {}); - EXPECT_MATCH(ExecuteToString(&b, {}), - testing::HasSubstr("mismatched element counts")); + EXPECT_THAT(ExecuteToString(&b, {}), + ::testing::HasSubstr("mismatched element counts")); } XLA_TEST_F(ReshapeTest, R4Dim0MinorLayoutToR2Dim0MajorLayout) { @@ -401,7 +433,7 @@ XLA_TEST_F(ReshapeTest, R4Dim0MinorLayoutToR2Dim0MajorLayout) { builder.Reshape(a, /*dimensions=*/{0, 1, 2, 3}, /*new_sizes=*/{2, 8}); // clang-format off - auto literal = LiteralUtil::CreateR4FromArray4DWithLayout(Array4D{ + auto literal = Literal::CreateR4FromArray4DWithLayout(Array4D{ { { {0, 1}, @@ -433,7 +465,7 @@ XLA_TEST_F(ReshapeTest, R4Dim0MinorLayoutToR2Dim0MajorLayout) { }); Computation computation = builder.Build().ConsumeValueOrDie(); - ExecutionOptions execution_options; + ExecutionOptions execution_options = execution_options_; *execution_options.mutable_shape_with_output_layout() = ShapeUtil::MakeShapeWithLayout(F32, {2, 8}, {1, 0}); std::unique_ptr actual = @@ -441,12 +473,12 @@ XLA_TEST_F(ReshapeTest, R4Dim0MinorLayoutToR2Dim0MajorLayout) { ->ExecuteAndTransfer(computation, {input.get()}, &execution_options) .ConsumeValueOrDie(); std::unique_ptr expected = - LiteralUtil::CreateR2FromArray2D(expected_array); + Literal::CreateR2FromArray2D(expected_array); LiteralTestUtil::ExpectEqual(*expected, *actual); } XLA_TEST_F(ReshapeTest, R2ToR4_3x8_To_3x2x1x4) { - std::unique_ptr input = LiteralUtil::CreateR2({ + std::unique_ptr input = Literal::CreateR2({ {0, 1, 2, 3, 4, 5, 6, 7}, {100, 101, 102, 103, 104, 105, 106, 107}, {200, 201, 202, 203, 204, 205, 206, 207}, @@ -474,7 +506,7 @@ XLA_TEST_F(ReshapeTest, R2ToR4_3x8_To_3x2x1x4) { // Tests R2->R4 reshape with the reshape dimensions {1, 0}. XLA_TEST_F(ReshapeTest, R2ToR4_3x8_To_3x2x1x4_Dimensions_10) { - std::unique_ptr input = LiteralUtil::CreateR2({ + std::unique_ptr input = Literal::CreateR2({ {0, 1, 2, 3, 4, 5, 6, 7}, {100, 101, 102, 103, 104, 105, 106, 107}, {200, 201, 202, 203, 204, 205, 206, 207}, @@ -508,7 +540,7 @@ XLA_TEST_F(ReshapeTest, R4ToR2_2x1x1x1_To_2x1) { [&rng, &distribution](tensorflow::gtl::ArraySlice /* indices */, float* cell) { *cell = distribution(rng); }); std::unique_ptr input_literal = - LiteralUtil::CreateR4FromArray4DWithLayout( + Literal::CreateR4FromArray4DWithLayout( input, LayoutUtil::MakeLayout({3, 2, 1, 0})); std::unique_ptr input_data = client_->TransferToServer(*input_literal).ConsumeValueOrDie(); @@ -531,7 +563,7 @@ XLA_TEST_F(ReshapeTest, R4ToR2_2x1x4x1_To_4x2) { [&rng, &distribution](tensorflow::gtl::ArraySlice /* indices */, float* cell) { *cell = distribution(rng); }); std::unique_ptr input_literal = - LiteralUtil::CreateR4FromArray4DWithLayout( + Literal::CreateR4FromArray4DWithLayout( input, LayoutUtil::MakeLayout({3, 2, 1, 0})); std::unique_ptr input_data = client_->TransferToServer(*input_literal).ConsumeValueOrDie(); @@ -555,7 +587,7 @@ XLA_TEST_F(ReshapeTest, R4ToR2_5x10x2x3_To_5x60_Dimensions_0213) { [&rng, &distribution](tensorflow::gtl::ArraySlice /* indices */, float* cell) { *cell = distribution(rng); }); std::unique_ptr input_literal = - LiteralUtil::CreateR4FromArray4DWithLayout( + Literal::CreateR4FromArray4DWithLayout( input, LayoutUtil::MakeLayout({3, 2, 1, 0})); std::unique_ptr input_data = client_->TransferToServer(*input_literal).ConsumeValueOrDie(); @@ -569,7 +601,7 @@ XLA_TEST_F(ReshapeTest, R4ToR2_5x10x2x3_To_5x60_Dimensions_0213) { expected_array(indices[0], indices[2] * 30 + indices[1] * 3 + indices[3]) = *cell; }); - auto expected = LiteralUtil::CreateR2FromArray2D(expected_array); + auto expected = Literal::CreateR2FromArray2D(expected_array); ComputeAndCompareLiteral(&builder, *expected, {input_data.get()}); } @@ -581,7 +613,7 @@ XLA_TEST_F(ReshapeTest, NoopReshape) { [&rng, &distribution](tensorflow::gtl::ArraySlice /* indices */, float* cell) { *cell = distribution(rng); }); std::unique_ptr input_literal = - LiteralUtil::CreateR4FromArray4DWithLayout( + Literal::CreateR4FromArray4DWithLayout( input_array, LayoutUtil::MakeLayout({1, 2, 3, 0})); std::unique_ptr input_data = client_->TransferToServer(*input_literal).ConsumeValueOrDie(); @@ -592,7 +624,7 @@ XLA_TEST_F(ReshapeTest, NoopReshape) { /*new_sizes=*/{7, 2, 3, 5}); Computation computation = builder.Build().ConsumeValueOrDie(); - ExecutionOptions execution_options; + ExecutionOptions execution_options = execution_options_; *execution_options.mutable_shape_with_output_layout() = ShapeUtil::MakeShapeWithLayout(F32, {7, 2, 3, 5}, {2, 3, 0, 1}); std::unique_ptr output_literal = @@ -608,7 +640,7 @@ XLA_TEST_F(ReshapeTest, NoopReshape) { } XLA_TEST_F(ReshapeTest, R4ToR4Reshape_Trivial) { - auto literal_1x2x3x4 = LiteralUtil::CreateR4( + auto literal_1x2x3x4 = Literal::CreateR4( {{{{1, 2, 3, 4}, {5, 6, 7, 8}, {9, 10, 11, 12}}, {{13, 14, 15, 16}, {17, 18, 19, 20}, {21, 22, 23, 24}}}}); @@ -621,7 +653,7 @@ XLA_TEST_F(ReshapeTest, R4ToR4Reshape_Trivial) { } XLA_TEST_F(ReshapeTest, R4ToR4Reshape) { - auto literal_1x2x3x4 = LiteralUtil::CreateR4( + auto literal_1x2x3x4 = Literal::CreateR4( {{{{1, 2, 3, 4}, {5, 6, 7, 8}, {9, 10, 11, 12}}, {{13, 14, 15, 16}, {17, 18, 19, 20}, {21, 22, 23, 24}}}}); @@ -631,7 +663,7 @@ XLA_TEST_F(ReshapeTest, R4ToR4Reshape) { /*new_sizes=*/{2, 4, 3, 1}); // clang-format off - auto expected_2x4x3x1 = LiteralUtil::CreateR4( + auto expected_2x4x3x1 = Literal::CreateR4( {{{{1}, {5}, {9}}, {{2}, {6}, {10}}, {{3}, {7}, {11}}, @@ -655,7 +687,7 @@ XLA_TEST_F(ReshapeTest, R4TwoMinorTransposeSimple) { [&rng, &distribution](tensorflow::gtl::ArraySlice /* indices */, float* cell) { *cell = distribution(rng); }); std::unique_ptr input_literal = - LiteralUtil::CreateR4FromArray4DWithLayout( + Literal::CreateR4FromArray4DWithLayout( input, LayoutUtil::MakeLayout({3, 2, 1, 0})); std::unique_ptr input_data = client_->TransferToServer(*input_literal).ConsumeValueOrDie(); @@ -664,9 +696,9 @@ XLA_TEST_F(ReshapeTest, R4TwoMinorTransposeSimple) { auto a = builder.Parameter(0, input_literal->shape(), "a"); builder.Reshape(a, /*dimensions=*/{0, 1, 3, 2}, /*new_sizes=*/new_bounds); - std::unique_ptr expected = LiteralUtil::Relayout( - *LiteralTestUtil::Reshape(new_bounds, {2, 3, 1, 0}, *input_literal), - LayoutUtil::MakeLayout({3, 2, 1, 0})); + std::unique_ptr expected = + LiteralTestUtil::Reshape(new_bounds, {2, 3, 1, 0}, *input_literal) + ->Relayout(LayoutUtil::MakeLayout({3, 2, 1, 0})); // Specify the requested output shape explicitly to ensure that this reshape // actually corresponds to a two minor transpose. @@ -684,7 +716,7 @@ XLA_TEST_F(ReshapeTest, R4TwoMinorTransposeMajorFirstEffectiveR2) { [&rng, &distribution](tensorflow::gtl::ArraySlice /* indices */, float* cell) { *cell = distribution(rng); }); std::unique_ptr input_literal = - LiteralUtil::CreateR4FromArray4DWithLayout( + Literal::CreateR4FromArray4DWithLayout( input, LayoutUtil::MakeLayout({3, 2, 1, 0})); std::unique_ptr input_data = client_->TransferToServer(*input_literal).ConsumeValueOrDie(); @@ -693,9 +725,9 @@ XLA_TEST_F(ReshapeTest, R4TwoMinorTransposeMajorFirstEffectiveR2) { auto a = builder.Parameter(0, input_literal->shape(), "a"); builder.Reshape(a, /*dimensions=*/{0, 1, 3, 2}, /*new_sizes=*/new_bounds); - std::unique_ptr expected = LiteralUtil::Relayout( - *LiteralTestUtil::Reshape(new_bounds, {2, 3, 1, 0}, *input_literal), - LayoutUtil::MakeLayout({3, 2, 1, 0})); + std::unique_ptr expected = + LiteralTestUtil::Reshape(new_bounds, {2, 3, 1, 0}, *input_literal) + ->Relayout(LayoutUtil::MakeLayout({3, 2, 1, 0})); // Specify the requested output shape explicitly to ensure that this reshape // actually corresponds to a two minor transpose. @@ -713,7 +745,7 @@ XLA_TEST_F(ReshapeTest, R4TwoMinorTransposeMajorFirstMinorEffectiveR1) { [&rng, &distribution](tensorflow::gtl::ArraySlice /* indices */, float* cell) { *cell = distribution(rng); }); std::unique_ptr input_literal = - LiteralUtil::CreateR4FromArray4DWithLayout( + Literal::CreateR4FromArray4DWithLayout( input, LayoutUtil::MakeLayout({3, 2, 1, 0})); std::unique_ptr input_data = client_->TransferToServer(*input_literal).ConsumeValueOrDie(); @@ -722,9 +754,9 @@ XLA_TEST_F(ReshapeTest, R4TwoMinorTransposeMajorFirstMinorEffectiveR1) { auto a = builder.Parameter(0, input_literal->shape(), "a"); builder.Reshape(a, /*dimensions=*/{0, 1, 3, 2}, /*new_sizes=*/new_bounds); - std::unique_ptr expected = LiteralUtil::Relayout( - *LiteralTestUtil::Reshape(new_bounds, {2, 3, 1, 0}, *input_literal), - LayoutUtil::MakeLayout({3, 2, 1, 0})); + std::unique_ptr expected = + LiteralTestUtil::Reshape(new_bounds, {2, 3, 1, 0}, *input_literal) + ->Relayout(LayoutUtil::MakeLayout({3, 2, 1, 0})); // Specify the requested output shape explicitly to ensure that this reshape // actually corresponds to a two minor transpose. @@ -743,7 +775,7 @@ XLA_TEST_F(ReshapeTest, R4TwoMinorTransposeMajorFirstMinorEffectiveR1InR2) { [&rng, &distribution](tensorflow::gtl::ArraySlice /* indices */, float* cell) { *cell = distribution(rng); }); std::unique_ptr input_literal = - LiteralUtil::CreateR4FromArray4DWithLayout( + Literal::CreateR4FromArray4DWithLayout( input, LayoutUtil::MakeLayout({3, 2, 1, 0})); std::unique_ptr input_data = client_->TransferToServer(*input_literal).ConsumeValueOrDie(); @@ -752,9 +784,9 @@ XLA_TEST_F(ReshapeTest, R4TwoMinorTransposeMajorFirstMinorEffectiveR1InR2) { auto a = builder.Parameter(0, input_literal->shape(), "a"); builder.Reshape(a, /*dimensions=*/{0, 1, 3, 2}, /*new_sizes=*/new_bounds); - std::unique_ptr expected = LiteralUtil::Relayout( - *LiteralTestUtil::Reshape(new_bounds, {2, 3, 1, 0}, *input_literal), - LayoutUtil::MakeLayout({3, 2, 1, 0})); + std::unique_ptr expected = + LiteralTestUtil::Reshape(new_bounds, {2, 3, 1, 0}, *input_literal) + ->Relayout(LayoutUtil::MakeLayout({3, 2, 1, 0})); // Specify the requested output shape explicitly to ensure that this reshape // actually corresponds to a two minor transpose. @@ -772,7 +804,7 @@ XLA_TEST_F(ReshapeTest, R4TwoMinorTransposeTrivialR2) { [&rng, &distribution](tensorflow::gtl::ArraySlice /* indices */, float* cell) { *cell = distribution(rng); }); std::unique_ptr input_literal = - LiteralUtil::CreateR4FromArray4DWithLayout( + Literal::CreateR4FromArray4DWithLayout( input, LayoutUtil::MakeLayout({0, 1, 2, 3})); std::unique_ptr input_data = client_->TransferToServer(*input_literal).ConsumeValueOrDie(); @@ -781,9 +813,9 @@ XLA_TEST_F(ReshapeTest, R4TwoMinorTransposeTrivialR2) { auto a = builder.Parameter(0, input_literal->shape(), "a"); builder.Reshape(a, /*dimensions=*/{1, 0, 2, 3}, /*new_sizes=*/new_bounds); - std::unique_ptr expected = LiteralUtil::Relayout( - *LiteralTestUtil::Reshape(new_bounds, {1, 0, 2, 3}, *input_literal), - input_literal->shape().layout()); + std::unique_ptr expected = + LiteralTestUtil::Reshape(new_bounds, {1, 0, 2, 3}, *input_literal) + ->Relayout(input_literal->shape().layout()); // Specify the requested output shape explicitly to ensure that this reshape // actually corresponds to a two minor transpose. @@ -793,20 +825,3 @@ XLA_TEST_F(ReshapeTest, R4TwoMinorTransposeTrivialR2) { } // namespace } // namespace xla - -int main(int argc, char** argv) { - std::vector flag_list; - xla::legacy_flags::AppendCpuCompilerFlags(&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; - } - testing::InitGoogleTest(&argc, argv); - if (argc > 1) { - LOG(ERROR) << "Unknown argument " << argv[1] << "\n" << usage; - return 2; - } - return RUN_ALL_TESTS(); -} diff --git a/tensorflow/compiler/xla/tests/reverse_test.cc b/tensorflow/compiler/xla/tests/reverse_test.cc index 63dd4421fade73089d41f9f0d2aa217b7e785eb2..1f6cfc85ccd25bb22db51411f7376489c14c3603 100644 --- a/tensorflow/compiler/xla/tests/reverse_test.cc +++ b/tensorflow/compiler/xla/tests/reverse_test.cc @@ -19,7 +19,6 @@ limitations under the License. #include "tensorflow/compiler/xla/array4d.h" #include "tensorflow/compiler/xla/client/computation_builder.h" #include "tensorflow/compiler/xla/client/local_client.h" -#include "tensorflow/compiler/xla/legacy_flags/cpu_compiler_flags.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" @@ -154,20 +153,3 @@ TEST_F(ReverseTest, Reverse4DFloatArrayOnDim01) { } // namespace } // namespace xla - -int main(int argc, char** argv) { - std::vector flag_list; - xla::legacy_flags::AppendCpuCompilerFlags(&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; - } - testing::InitGoogleTest(&argc, argv); - if (argc > 1) { - LOG(ERROR) << "Unknown argument " << argv[1] << "\n" << usage; - return 2; - } - return RUN_ALL_TESTS(); -} 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 5b734c0f400d1c1d937d09c0e5d5e796f84fb8b4..8cbfcc6f5c4272706a0f9fd809041516bf32432b 100644 --- a/tensorflow/compiler/xla/tests/round_trip_packed_literal_test.cc +++ b/tensorflow/compiler/xla/tests/round_trip_packed_literal_test.cc @@ -18,7 +18,6 @@ limitations under the License. #include "tensorflow/compiler/xla/client/global_data.h" #include "tensorflow/compiler/xla/client/local_client.h" #include "tensorflow/compiler/xla/layout_util.h" -#include "tensorflow/compiler/xla/legacy_flags/cpu_compiler_flags.h" #include "tensorflow/compiler/xla/literal_util.h" #include "tensorflow/compiler/xla/packed_literal_reader.h" #include "tensorflow/compiler/xla/shape_util.h" @@ -65,8 +64,8 @@ TEST_F(RoundTripPackedLiteralTest, RoundTripsR1F32Length2) { reader.Read(ShapeUtil::MakeShape(F32, {2})).ConsumeValueOrDie(); EXPECT_TRUE(reader.IsExhausted()); - EXPECT_EQ(42.0, LiteralUtil::Get(*actual, {0})); - EXPECT_EQ(24.0, LiteralUtil::Get(*actual, {1})); + EXPECT_EQ(42.0, actual->Get({0})); + EXPECT_EQ(24.0, actual->Get({1})); } TEST_F(RoundTripPackedLiteralTest, RoundTripsR2F32Size2x2Dim0Minor) { @@ -95,10 +94,10 @@ TEST_F(RoundTripPackedLiteralTest, RoundTripsR2F32Size2x2Dim0Minor) { .ConsumeValueOrDie(); EXPECT_TRUE(reader.IsExhausted()); - EXPECT_EQ(42.0f, LiteralUtil::Get(*actual, {0, 0})); - EXPECT_EQ(24.0f, LiteralUtil::Get(*actual, {0, 1})); - EXPECT_EQ(64.0f, LiteralUtil::Get(*actual, {1, 0})); - EXPECT_EQ(46.0f, LiteralUtil::Get(*actual, {1, 1})); + EXPECT_EQ(42.0f, actual->Get({0, 0})); + EXPECT_EQ(24.0f, actual->Get({0, 1})); + EXPECT_EQ(64.0f, actual->Get({1, 0})); + EXPECT_EQ(46.0f, actual->Get({1, 1})); std::unique_ptr round_tripped = RoundTripToServer(*actual); LiteralTestUtil::ExpectEqual(*round_tripped, *actual); @@ -130,10 +129,10 @@ TEST_F(RoundTripPackedLiteralTest, RoundTripsR2F32Size2x2Dim1Minor) { .ConsumeValueOrDie(); EXPECT_TRUE(reader.IsExhausted()); - EXPECT_EQ(42.0f, LiteralUtil::Get(*actual, {0, 0})); - EXPECT_EQ(24.0f, LiteralUtil::Get(*actual, {1, 0})); - EXPECT_EQ(64.0f, LiteralUtil::Get(*actual, {0, 1})); - EXPECT_EQ(46.0f, LiteralUtil::Get(*actual, {1, 1})); + EXPECT_EQ(42.0f, actual->Get({0, 0})); + EXPECT_EQ(24.0f, actual->Get({1, 0})); + EXPECT_EQ(64.0f, actual->Get({0, 1})); + EXPECT_EQ(46.0f, actual->Get({1, 1})); std::unique_ptr round_tripped = RoundTripToServer(*actual); LiteralTestUtil::ExpectEqual(*round_tripped, *actual); @@ -141,20 +140,3 @@ TEST_F(RoundTripPackedLiteralTest, RoundTripsR2F32Size2x2Dim1Minor) { } // namespace } // namespace xla - -int main(int argc, char** argv) { - std::vector flag_list; - xla::legacy_flags::AppendCpuCompilerFlags(&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; - } - testing::InitGoogleTest(&argc, argv); - if (argc > 1) { - LOG(ERROR) << "Unknown argument " << argv[1] << "\n" << usage; - return 2; - } - return RUN_ALL_TESTS(); -} diff --git a/tensorflow/compiler/xla/tests/round_trip_transfer_test.cc b/tensorflow/compiler/xla/tests/round_trip_transfer_test.cc index 04a8bab0eb8218c07f7ad77ef25c35c48c459e35..32db45f8a66266712ba4091c2aa6368f0b822bd2 100644 --- a/tensorflow/compiler/xla/tests/round_trip_transfer_test.cc +++ b/tensorflow/compiler/xla/tests/round_trip_transfer_test.cc @@ -23,13 +23,11 @@ limitations under the License. #include "tensorflow/compiler/xla/array4d.h" #include "tensorflow/compiler/xla/client/global_data.h" #include "tensorflow/compiler/xla/client/local_client.h" -#include "tensorflow/compiler/xla/legacy_flags/cpu_compiler_flags.h" #include "tensorflow/compiler/xla/literal_util.h" #include "tensorflow/compiler/xla/statusor.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/platform/test.h" #include "tensorflow/core/platform/types.h" @@ -48,62 +46,61 @@ class RoundTripTransferTest : public ClientLibraryTestBase { }; TEST_F(RoundTripTransferTest, R0S32) { - RoundTripTest(*LiteralUtil::CreateR0(42)); + RoundTripTest(*Literal::CreateR0(42)); } TEST_F(RoundTripTransferTest, R0F32) { - RoundTripTest(*LiteralUtil::CreateR0(42.0)); + RoundTripTest(*Literal::CreateR0(42.0)); } TEST_F(RoundTripTransferTest, R1F32_Len0) { - RoundTripTest(*LiteralUtil::CreateR1({})); + RoundTripTest(*Literal::CreateR1({})); } TEST_F(RoundTripTransferTest, R1F32_Len2) { - RoundTripTest(*LiteralUtil::CreateR1({42.0, 64.0})); + RoundTripTest(*Literal::CreateR1({42.0, 64.0})); } TEST_F(RoundTripTransferTest, R1F32_Len256) { std::vector values(256); std::iota(values.begin(), values.end(), 1.0); - RoundTripTest(*LiteralUtil::CreateR1(values)); + RoundTripTest(*Literal::CreateR1(values)); } TEST_F(RoundTripTransferTest, R1F32_Len1024) { std::vector values(1024); std::iota(values.begin(), values.end(), 1.0); - RoundTripTest(*LiteralUtil::CreateR1(values)); + RoundTripTest(*Literal::CreateR1(values)); } TEST_F(RoundTripTransferTest, R1F32_Len1025) { std::vector values(1025); std::iota(values.begin(), values.end(), 1.0); - RoundTripTest(*LiteralUtil::CreateR1(values)); + RoundTripTest(*Literal::CreateR1(values)); } TEST_F(RoundTripTransferTest, R1F32_Len4096) { std::vector values(4096); std::iota(values.begin(), values.end(), 1.0); - RoundTripTest(*LiteralUtil::CreateR1(values)); + RoundTripTest(*Literal::CreateR1(values)); } TEST_F(RoundTripTransferTest, R2F32_Len10x0) { - RoundTripTest( - *LiteralUtil::CreateR2FromArray2D(Array2D(10, 0))); + RoundTripTest(*Literal::CreateR2FromArray2D(Array2D(10, 0))); } TEST_F(RoundTripTransferTest, R2F32_Len2x2) { - RoundTripTest(*LiteralUtil::CreateR2({{42.0, 64.0}, {77.0, 88.0}})); + RoundTripTest(*Literal::CreateR2({{42.0, 64.0}, {77.0, 88.0}})); } TEST_F(RoundTripTransferTest, R3F32) { RoundTripTest( - *LiteralUtil::CreateR3({{{1.0, 2.0}, {1.0, 2.0}, {1.0, 2.0}}, - {{3.0, 4.0}, {3.0, 4.0}, {3.0, 4.0}}})); + *Literal::CreateR3({{{1.0, 2.0}, {1.0, 2.0}, {1.0, 2.0}}, + {{3.0, 4.0}, {3.0, 4.0}, {3.0, 4.0}}})); } TEST_F(RoundTripTransferTest, R4F32) { - RoundTripTest(*LiteralUtil::CreateR4({{ + RoundTripTest(*Literal::CreateR4({{ {{10, 11, 12, 13}, {14, 15, 16, 17}}, {{18, 19, 20, 21}, {22, 23, 24, 25}}, {{26, 27, 28, 29}, {30, 31, 32, 33}}, @@ -111,54 +108,34 @@ TEST_F(RoundTripTransferTest, R4F32) { } TEST_F(RoundTripTransferTest, EmptyTuple) { - RoundTripTest(*LiteralUtil::MakeTuple({})); + RoundTripTest(*Literal::MakeTuple({})); } TEST_F(RoundTripTransferTest, TupleOfR1F32) { - RoundTripTest( - *LiteralUtil::MakeTuple({LiteralUtil::CreateR1({1, 2}).get(), - LiteralUtil::CreateR1({3, 4}).get()})); + RoundTripTest(*Literal::MakeTuple({Literal::CreateR1({1, 2}).get(), + Literal::CreateR1({3, 4}).get()})); } TEST_F(RoundTripTransferTest, TupleOfR1F32_Len0_Len2) { - RoundTripTest( - *LiteralUtil::MakeTuple({LiteralUtil::CreateR1({}).get(), - LiteralUtil::CreateR1({3, 4}).get()})); + RoundTripTest(*Literal::MakeTuple({Literal::CreateR1({}).get(), + Literal::CreateR1({3, 4}).get()})); } TEST_F(RoundTripTransferTest, TupleOfR0F32AndR1S32) { - RoundTripTest( - *LiteralUtil::MakeTuple({LiteralUtil::CreateR0(1.0).get(), - LiteralUtil::CreateR1({2, 3}).get()})); + RoundTripTest(*Literal::MakeTuple({Literal::CreateR0(1.0).get(), + Literal::CreateR1({2, 3}).get()})); } // Below two tests are added to identify the cost of large data transfers. TEST_F(RoundTripTransferTest, R2F32_Large) { - RoundTripTest(*LiteralUtil::CreateR2F32Linspace(-1.0f, 1.0f, 512, 512)); + RoundTripTest(*Literal::CreateR2F32Linspace(-1.0f, 1.0f, 512, 512)); } TEST_F(RoundTripTransferTest, R4F32_Large) { Array4D array4d(2, 2, 256, 256); array4d.FillWithMultiples(1.0f); - RoundTripTest(*LiteralUtil::CreateR4FromArray4D(array4d)); + RoundTripTest(*Literal::CreateR4FromArray4D(array4d)); } } // namespace } // namespace xla - -int main(int argc, char** argv) { - std::vector flag_list; - xla::legacy_flags::AppendCpuCompilerFlags(&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; - } - testing::InitGoogleTest(&argc, argv); - if (argc > 1) { - LOG(ERROR) << "Unknown argument " << argv[1] << "\n" << usage; - return 2; - } - return RUN_ALL_TESTS(); -} diff --git a/tensorflow/compiler/xla/tests/scalar_computations_test.cc b/tensorflow/compiler/xla/tests/scalar_computations_test.cc index 0005c0c9e23fff52ee3f92d07cefeed92c6c5399..f3cbc0132389c0e2629c57f1cba1f5124c7eaa8d 100644 --- a/tensorflow/compiler/xla/tests/scalar_computations_test.cc +++ b/tensorflow/compiler/xla/tests/scalar_computations_test.cc @@ -20,7 +20,6 @@ limitations under the License. #include "tensorflow/compiler/xla/client/computation_builder.h" #include "tensorflow/compiler/xla/client/global_data.h" #include "tensorflow/compiler/xla/client/local_client.h" -#include "tensorflow/compiler/xla/legacy_flags/cpu_compiler_flags.h" #include "tensorflow/compiler/xla/literal_util.h" #include "tensorflow/compiler/xla/status_macros.h" #include "tensorflow/compiler/xla/statusor.h" @@ -70,35 +69,35 @@ class ScalarComputationsTest : public ClientLibraryTestBase { } }; -TEST_F(ScalarComputationsTest, NegateScalarF32) { +XLA_TEST_F(ScalarComputationsTest, NegateScalarF32) { ComputationBuilder builder(client_, TestName()); builder.Neg(builder.ConstantR0(2.1f)); ComputeAndCompareR0(&builder, -2.1f, {}, error_spec_); } -TEST_F(ScalarComputationsTest, NegateScalarS32) { +XLA_TEST_F(ScalarComputationsTest, NegateScalarS32) { ComputationBuilder builder(client_, TestName()); builder.Neg(builder.ConstantR0(2)); ComputeAndCompareR0(&builder, -2, {}); } -TEST_F(ScalarComputationsTest, AddTwoScalarsF32) { +XLA_TEST_F(ScalarComputationsTest, AddTwoScalarsF32) { ComputationBuilder builder(client_, TestName()); builder.Add(builder.ConstantR0(2.1f), builder.ConstantR0(5.5f)); ComputeAndCompareR0(&builder, 7.6f, {}, error_spec_); } -TEST_F(ScalarComputationsTest, AddTwoScalarsS32) { +XLA_TEST_F(ScalarComputationsTest, AddTwoScalarsS32) { ComputationBuilder builder(client_, TestName()); builder.Add(builder.ConstantR0(2), builder.ConstantR0(5)); ComputeAndCompareR0(&builder, 7, {}); } -TEST_F(ScalarComputationsTest, AddTwoScalarsU32) { +XLA_TEST_F(ScalarComputationsTest, AddTwoScalarsU32) { ComputationBuilder builder(client_, TestName()); builder.Add(builder.ConstantR0(35), builder.ConstantR0(57)); @@ -124,7 +123,7 @@ XLA_TEST_F(ScalarComputationsTest, AddTwoScalarsU64) { XLA_TEST_F(ScalarComputationsTest, AddTwoScalarsS64) { ComputationBuilder builder(client_, TestName()); const int64 a = static_cast(1) << 62; - const int64 b = a + 1; + const int64 b = a - 1; builder.Add(builder.ConstantR0(a), builder.ConstantR0(b)); ComputeAndCompareR0(&builder, a + b, {}); @@ -138,21 +137,21 @@ XLA_TEST_F(ScalarComputationsTest, AddTwoScalarsF64) { ComputeAndCompareR0(&builder, 3.75, {}); } -TEST_F(ScalarComputationsTest, SubtractTwoScalarsF32) { +XLA_TEST_F(ScalarComputationsTest, SubtractTwoScalarsF32) { ComputationBuilder builder(client_, TestName()); builder.Sub(builder.ConstantR0(2.1f), builder.ConstantR0(5.5f)); ComputeAndCompareR0(&builder, -3.4f, {}, error_spec_); } -TEST_F(ScalarComputationsTest, SubtractTwoScalarsS32) { +XLA_TEST_F(ScalarComputationsTest, SubtractTwoScalarsS32) { ComputationBuilder builder(client_, TestName()); builder.Sub(builder.ConstantR0(2), builder.ConstantR0(5)); ComputeAndCompareR0(&builder, -3, {}); } -TEST_F(ScalarComputationsTest, MulThreeScalarsF32) { +XLA_TEST_F(ScalarComputationsTest, MulThreeScalarsF32) { ComputationBuilder builder(client_, TestName()); builder.Mul(builder.Mul(builder.ConstantR0(2.1f), builder.ConstantR0(5.5f)), @@ -161,7 +160,7 @@ TEST_F(ScalarComputationsTest, MulThreeScalarsF32) { ComputeAndCompareR0(&builder, 5.775f, {}, error_spec_); } -TEST_F(ScalarComputationsTest, MulTwoScalarsS32) { +XLA_TEST_F(ScalarComputationsTest, MulTwoScalarsS32) { std::vector data = {0, 1, -1, @@ -185,7 +184,7 @@ TEST_F(ScalarComputationsTest, MulTwoScalarsS32) { } } -TEST_F(ScalarComputationsTest, MulTwoScalarsU32) { +XLA_TEST_F(ScalarComputationsTest, MulTwoScalarsU32) { std::vector data = {0, 1, 0xDEADBEEF, 1234, 0x1a243514, 0xFFFFFFFF, 0x80808080}; @@ -200,7 +199,7 @@ TEST_F(ScalarComputationsTest, MulTwoScalarsU32) { } } -TEST_F(ScalarComputationsTest, MulThreeScalarsS32) { +XLA_TEST_F(ScalarComputationsTest, MulThreeScalarsS32) { ComputationBuilder builder(client_, TestName()); builder.Mul( builder.Mul(builder.ConstantR0(2), builder.ConstantR0(5)), @@ -209,11 +208,11 @@ TEST_F(ScalarComputationsTest, MulThreeScalarsS32) { ComputeAndCompareR0(&builder, 10, {}); } -TEST_F(ScalarComputationsTest, MulThreeScalarsF32Params) { +XLA_TEST_F(ScalarComputationsTest, MulThreeScalarsF32Params) { ComputationBuilder builder(client_, TestName()); - std::unique_ptr a_literal = LiteralUtil::CreateR0(2.1f); - std::unique_ptr b_literal = LiteralUtil::CreateR0(5.5f); - std::unique_ptr c_literal = LiteralUtil::CreateR0(0.5f); + std::unique_ptr a_literal = Literal::CreateR0(2.1f); + std::unique_ptr b_literal = Literal::CreateR0(5.5f); + std::unique_ptr c_literal = Literal::CreateR0(0.5f); std::unique_ptr a_data = client_->TransferToServer(*a_literal).ConsumeValueOrDie(); @@ -232,7 +231,7 @@ TEST_F(ScalarComputationsTest, MulThreeScalarsF32Params) { error_spec_); } -TEST_F(ScalarComputationsTest, DivideTwoScalarsF32) { +XLA_TEST_F(ScalarComputationsTest, DivideTwoScalarsF32) { ComputationBuilder builder(client_, TestName()); builder.Div(builder.ConstantR0(5.0f), builder.ConstantR0(2.5f)); @@ -279,6 +278,36 @@ XLA_TEST_P(DivS32Test, RemainderTwoScalarsS32) { ComputeAndCompareR0(&builder, p.remainder, {}); } +XLA_TEST_P(DivS32Test, DivideTwoScalarsNonConstS32) { + DivS32Params p = GetParam(); + ComputationBuilder builder(client_, TestName()); + ComputationDataHandle dividend; + ComputationDataHandle divisor; + auto dividendd = + CreateR0Parameter(p.dividend, 0, "dividend", &builder, ÷nd); + auto divisord = + CreateR0Parameter(p.divisor, 1, "divisor", &builder, &divisor); + builder.Div(dividend, divisor); + + ComputeAndCompareR0(&builder, p.quotient, + {dividendd.get(), divisord.get()}); +} + +XLA_TEST_P(DivS32Test, RemainderTwoScalarsNonConstDivisorS32) { + DivS32Params p = GetParam(); + ComputationBuilder builder(client_, TestName()); + ComputationDataHandle dividend; + ComputationDataHandle divisor; + auto dividendd = + CreateR0Parameter(p.dividend, 0, "dividend", &builder, ÷nd); + auto divisord = + CreateR0Parameter(p.divisor, 1, "divisor", &builder, &divisor); + builder.Rem(dividend, divisor); + + ComputeAndCompareR0(&builder, p.remainder, + {dividendd.get(), divisord.get()}); +} + INSTANTIATE_TEST_CASE_P( DivS32Test_Instantiation, DivS32Test, ::testing::Values( @@ -308,13 +337,95 @@ INSTANTIATE_TEST_CASE_P( DivS32Params{INT32_MIN, -0x40000000, 2, 0}, // DivS32Params{INT32_MIN + 1, -0x40000000, 1, -0x3fffffff})); -TEST_F(ScalarComputationsTest, RemainderTwoScalarsNonConstDividendS32) { +XLA_TEST_F(ScalarComputationsTest, DivU32s) { + // clang-format off + // Some interesting values to test. + std::vector vals = { + 0, 1, 2, 17, 101, 3333, 0x7FFFFFFF, 0x80000000, UINT32_MAX - 1, UINT32_MAX}; + // clang-format on + + Computation div_computation; + { + ComputationBuilder builder(client_, TestName()); + + ComputationDataHandle dividend = + builder.Parameter(0, ShapeUtil::MakeShape(U32, {}), "dividend"); + ComputationDataHandle divisor = + builder.Parameter(1, ShapeUtil::MakeShape(U32, {}), "divisor"); + builder.Div(dividend, divisor); + TF_ASSERT_OK_AND_ASSIGN(div_computation, builder.Build()); + } + + for (uint32 divisor : vals) { + if (divisor != 0) { + for (uint32 dividend : vals) { + auto dividend_literal = Literal::CreateR0(dividend); + auto divisor_literal = Literal::CreateR0(divisor); + TF_ASSERT_OK_AND_ASSIGN(auto dividend_data, + client_->TransferToServer(*dividend_literal)); + TF_ASSERT_OK_AND_ASSIGN(auto divisor_data, + client_->TransferToServer(*divisor_literal)); + auto actual_literal = + client_ + ->ExecuteAndTransfer(div_computation, + {dividend_data.get(), divisor_data.get()}, + &execution_options_) + .ConsumeValueOrDie(); + auto expected_literal = Literal::CreateR0(dividend / divisor); + LiteralTestUtil::ExpectEqual(*expected_literal, *actual_literal); + } + } + } +} + +XLA_TEST_F(ScalarComputationsTest, RemU32s) { + // clang-format off + // Some interesting values to test. + std::vector vals = { + 0, 1, 2, 17, 101, 3333, 0x7FFFFFFF, 0x80000000, UINT32_MAX - 1, UINT32_MAX}; + // clang-format on + + Computation rem_computation; + { + ComputationBuilder builder(client_, TestName()); + + ComputationDataHandle dividend = + builder.Parameter(0, ShapeUtil::MakeShape(U32, {}), "dividend"); + ComputationDataHandle divisor = + builder.Parameter(1, ShapeUtil::MakeShape(U32, {}), "divisor"); + builder.Rem(dividend, divisor); + TF_ASSERT_OK_AND_ASSIGN(rem_computation, builder.Build()); + } + + for (uint32 divisor : vals) { + if (divisor != 0) { + for (uint32 dividend : vals) { + auto dividend_literal = Literal::CreateR0(dividend); + auto divisor_literal = Literal::CreateR0(divisor); + TF_ASSERT_OK_AND_ASSIGN(auto dividend_data, + client_->TransferToServer(*dividend_literal)); + TF_ASSERT_OK_AND_ASSIGN(auto divisor_data, + client_->TransferToServer(*divisor_literal)); + auto actual_literal = + client_ + ->ExecuteAndTransfer(rem_computation, + {dividend_data.get(), divisor_data.get()}, + &execution_options_) + .ConsumeValueOrDie(); + auto expected_literal = Literal::CreateR0(dividend % divisor); + LiteralTestUtil::ExpectEqual(*expected_literal, *actual_literal); + } + } + } +} + +XLA_TEST_F(ScalarComputationsTest, RemainderTwoScalarsNonConstDividendS32) { ComputationBuilder builder(client_, TestName()); auto x = builder.Parameter(0, ShapeUtil::MakeShape(S32, {}), "x"); builder.Rem(x, builder.ConstantR0(80000)); - std::unique_ptr literal = LiteralUtil::CreateR0(87919); - TF_ASSIGN_OR_ASSERT_OK(auto input_data, client_->TransferToServer(*literal)); + std::unique_ptr literal = Literal::CreateR0(87919); + TF_ASSERT_OK_AND_ASSIGN(auto input_data, client_->TransferToServer(*literal)); ComputeAndCompareR0(&builder, 7919, {input_data.get()}); } @@ -335,7 +446,7 @@ XLA_TEST_F(ScalarComputationsTest, RemTwoScalarsU32) { ComputeAndCompareR0(&builder, 2, {}); } -TEST_F(ScalarComputationsTest, LogicalAnd) { +XLA_TEST_F(ScalarComputationsTest, LogicalAnd) { for (bool x : {false, true}) { for (bool y : {false, true}) { ComputationBuilder builder(client_, TestName()); @@ -347,7 +458,7 @@ TEST_F(ScalarComputationsTest, LogicalAnd) { } } -TEST_F(ScalarComputationsTest, LogicalOr) { +XLA_TEST_F(ScalarComputationsTest, LogicalOr) { for (bool x : {false, true}) { for (bool y : {false, true}) { ComputationBuilder builder(client_, TestName()); @@ -359,7 +470,7 @@ TEST_F(ScalarComputationsTest, LogicalOr) { } } -TEST_F(ScalarComputationsTest, LogicalNot) { +XLA_TEST_F(ScalarComputationsTest, LogicalNot) { for (bool x : {false, true}) { ComputationBuilder builder(client_, TestName()); builder.LogicalNot(builder.ConstantR0(x)); @@ -368,7 +479,7 @@ TEST_F(ScalarComputationsTest, LogicalNot) { } } -TEST_F(ScalarComputationsTest, SelectScalarTrue) { +XLA_TEST_F(ScalarComputationsTest, SelectScalarTrue) { ComputationBuilder builder(client_, TestName()); builder.Select(builder.ConstantR0(true), // The predicate. builder.ConstantR0(123.0f), // The value on true. @@ -377,7 +488,7 @@ TEST_F(ScalarComputationsTest, SelectScalarTrue) { ComputeAndCompareR0(&builder, 123.0f, {}, error_spec_); } -TEST_F(ScalarComputationsTest, SelectScalarFalse) { +XLA_TEST_F(ScalarComputationsTest, SelectScalarFalse) { ComputationBuilder builder(client_, TestName()); builder.Select(builder.ConstantR0(false), // The predicate. builder.ConstantR0(123.0f), // The value on true. @@ -388,7 +499,7 @@ TEST_F(ScalarComputationsTest, SelectScalarFalse) { // This test is an explicit version of what is happening in the following // templatized comparison tests. -TEST_F(ScalarComputationsTest, CompareGtScalar) { +XLA_TEST_F(ScalarComputationsTest, CompareGtScalar) { ComputationBuilder builder(client_, TestName()); builder.Gt(builder.ConstantR0(2.0f), builder.ConstantR0(1.0f)); @@ -396,30 +507,30 @@ TEST_F(ScalarComputationsTest, CompareGtScalar) { } // S32 comparisons. -TEST_F(ScalarComputationsTest, CompareEqS32Greater) { +XLA_TEST_F(ScalarComputationsTest, CompareEqS32Greater) { TestCompare(2, 1, false, &ComputationBuilder::Eq); } -TEST_F(ScalarComputationsTest, CompareEqS32Equal) { +XLA_TEST_F(ScalarComputationsTest, CompareEqS32Equal) { TestCompare(3, 3, true, &ComputationBuilder::Eq); } -TEST_F(ScalarComputationsTest, CompareNeS32) { +XLA_TEST_F(ScalarComputationsTest, CompareNeS32) { TestCompare(2, 1, true, &ComputationBuilder::Ne); } -TEST_F(ScalarComputationsTest, CompareGeS32) { +XLA_TEST_F(ScalarComputationsTest, CompareGeS32) { TestCompare(2, 1, true, &ComputationBuilder::Ge); } -TEST_F(ScalarComputationsTest, CompareGtS32) { +XLA_TEST_F(ScalarComputationsTest, CompareGtS32) { TestCompare(1, 5, false, &ComputationBuilder::Gt); } -TEST_F(ScalarComputationsTest, CompareLeS32) { +XLA_TEST_F(ScalarComputationsTest, CompareLeS32) { TestCompare(2, 1, false, &ComputationBuilder::Le); } -TEST_F(ScalarComputationsTest, CompareLtS32) { +XLA_TEST_F(ScalarComputationsTest, CompareLtS32) { TestCompare(9, 7, false, &ComputationBuilder::Lt); TestCompare(std::numeric_limits::min(), std::numeric_limits::max(), true, @@ -427,105 +538,105 @@ TEST_F(ScalarComputationsTest, CompareLtS32) { } // U32 comparisons. -TEST_F(ScalarComputationsTest, CompareEqU32False) { +XLA_TEST_F(ScalarComputationsTest, CompareEqU32False) { TestCompare(2, 1, false, &ComputationBuilder::Eq); } -TEST_F(ScalarComputationsTest, CompareNeU32) { +XLA_TEST_F(ScalarComputationsTest, CompareNeU32) { TestCompare(2, 1, true, &ComputationBuilder::Ne); } -TEST_F(ScalarComputationsTest, CompareGeU32Greater) { +XLA_TEST_F(ScalarComputationsTest, CompareGeU32Greater) { TestCompare(2, 1, true, &ComputationBuilder::Ge); } -TEST_F(ScalarComputationsTest, CompareGeU32Equal) { +XLA_TEST_F(ScalarComputationsTest, CompareGeU32Equal) { TestCompare(3, 3, true, &ComputationBuilder::Ge); } -TEST_F(ScalarComputationsTest, CompareGtU32) { +XLA_TEST_F(ScalarComputationsTest, CompareGtU32) { TestCompare(1, 5, false, &ComputationBuilder::Gt); TestCompare(5, 5, false, &ComputationBuilder::Gt); TestCompare(5, 1, true, &ComputationBuilder::Gt); } -TEST_F(ScalarComputationsTest, CompareLeU32) { +XLA_TEST_F(ScalarComputationsTest, CompareLeU32) { TestCompare(2, 1, false, &ComputationBuilder::Le); } -TEST_F(ScalarComputationsTest, CompareLtU32) { +XLA_TEST_F(ScalarComputationsTest, CompareLtU32) { TestCompare(9, 7, false, &ComputationBuilder::Lt); TestCompare(0, std::numeric_limits::max(), true, &ComputationBuilder::Lt); } // F32 comparisons. -TEST_F(ScalarComputationsTest, CompareEqF32False) { +XLA_TEST_F(ScalarComputationsTest, CompareEqF32False) { TestCompare(2.0, 1.3, false, &ComputationBuilder::Eq); } -TEST_F(ScalarComputationsTest, CompareNeF32) { +XLA_TEST_F(ScalarComputationsTest, CompareNeF32) { TestCompare(2.0, 1.3, true, &ComputationBuilder::Ne); } -TEST_F(ScalarComputationsTest, CompareGeF32Greater) { +XLA_TEST_F(ScalarComputationsTest, CompareGeF32Greater) { TestCompare(2.0, 1.9, true, &ComputationBuilder::Ge); } -TEST_F(ScalarComputationsTest, CompareGeF32Equal) { +XLA_TEST_F(ScalarComputationsTest, CompareGeF32Equal) { TestCompare(3.5, 3.5, true, &ComputationBuilder::Ge); } -TEST_F(ScalarComputationsTest, CompareGtF32) { +XLA_TEST_F(ScalarComputationsTest, CompareGtF32) { TestCompare(1.0, 5.2, false, &ComputationBuilder::Gt); } -TEST_F(ScalarComputationsTest, CompareLeF32) { +XLA_TEST_F(ScalarComputationsTest, CompareLeF32) { TestCompare(2.0, 1.2, false, &ComputationBuilder::Le); } -TEST_F(ScalarComputationsTest, CompareLtF32) { +XLA_TEST_F(ScalarComputationsTest, CompareLtF32) { TestCompare(9.0, 7.2, false, &ComputationBuilder::Lt); } // F32 comparisons with exceptional values. The test names encode the // left/right operands at the end, and use Minf and Mzero for -inf and -0.0. -TEST_F(ScalarComputationsTest, CompareLtF32MinfMzero) { +XLA_TEST_F(ScalarComputationsTest, CompareLtF32MinfMzero) { TestCompare(-INFINITY, -0.0, true, &ComputationBuilder::Lt); } -TEST_F(ScalarComputationsTest, CompareLtF32MzeroZero) { +XLA_TEST_F(ScalarComputationsTest, CompareLtF32MzeroZero) { // Comparisons of 0.0 to -0.0 consider them equal in IEEE 754. TestCompare(-0.0, 0.0, false, &ComputationBuilder::Lt); } -TEST_F(ScalarComputationsTest, CompareLtF32ZeroInf) { +XLA_TEST_F(ScalarComputationsTest, CompareLtF32ZeroInf) { TestCompare(0.0, INFINITY, true, &ComputationBuilder::Lt); } -TEST_F(ScalarComputationsTest, CompareGeF32MinfMzero) { +XLA_TEST_F(ScalarComputationsTest, CompareGeF32MinfMzero) { TestCompare(-INFINITY, -0.0, false, &ComputationBuilder::Ge); } -TEST_F(ScalarComputationsTest, CompareGeF32MzeroZero) { +XLA_TEST_F(ScalarComputationsTest, CompareGeF32MzeroZero) { // Comparisons of 0.0 to -0.0 consider them equal in IEEE 754. TestCompare(-0.0, 0.0, true, &ComputationBuilder::Ge); } -TEST_F(ScalarComputationsTest, CompareGeF32ZeroInf) { +XLA_TEST_F(ScalarComputationsTest, CompareGeF32ZeroInf) { TestCompare(0.0, INFINITY, false, &ComputationBuilder::Ge); } -TEST_F(ScalarComputationsTest, ExpScalar) { +XLA_TEST_F(ScalarComputationsTest, ExpScalar) { ComputationBuilder builder(client_, TestName()); builder.Exp(builder.ConstantR0(2.0f)); ComputeAndCompareR0(&builder, 7.3890562, {}, error_spec_); } -TEST_F(ScalarComputationsTest, LogScalar) { +XLA_TEST_F(ScalarComputationsTest, LogScalar) { ComputationBuilder builder(client_, "log"); builder.Log(builder.ConstantR0(2.0f)); ComputeAndCompareR0(&builder, 0.6931471, {}, error_spec_); } -TEST_F(ScalarComputationsTest, TanhScalar) { +XLA_TEST_F(ScalarComputationsTest, TanhScalar) { ComputationBuilder builder(client_, TestName()); builder.Tanh(builder.ConstantR0(2.0f)); @@ -539,14 +650,14 @@ XLA_TEST_F(ScalarComputationsTest, TanhDoubleScalar) { ComputeAndCompareR0(&builder, 0.96402758, {}, error_spec_); } -TEST_F(ScalarComputationsTest, PowScalar) { +XLA_TEST_F(ScalarComputationsTest, PowScalar) { ComputationBuilder builder(client_, TestName()); builder.Pow(builder.ConstantR0(2.0f), builder.ConstantR0(3.0f)); ComputeAndCompareR0(&builder, 8.0, {}, error_spec_); } -TEST_F(ScalarComputationsTest, ClampScalarHigh) { +XLA_TEST_F(ScalarComputationsTest, ClampScalarHigh) { ComputationBuilder builder(client_, TestName()); builder.Clamp(builder.ConstantR0(2.0f), // The lower bound. builder.ConstantR0(5.0f), // The operand to be clamped. @@ -555,7 +666,7 @@ TEST_F(ScalarComputationsTest, ClampScalarHigh) { ComputeAndCompareR0(&builder, 3.0, {}, error_spec_); } -TEST_F(ScalarComputationsTest, ClampScalarMiddle) { +XLA_TEST_F(ScalarComputationsTest, ClampScalarMiddle) { ComputationBuilder builder(client_, TestName()); builder.Clamp(builder.ConstantR0(2.0f), // The lower bound. builder.ConstantR0(2.5f), // The operand to be clamped. @@ -564,7 +675,7 @@ TEST_F(ScalarComputationsTest, ClampScalarMiddle) { ComputeAndCompareR0(&builder, 2.5, {}, error_spec_); } -TEST_F(ScalarComputationsTest, ClampScalarLow) { +XLA_TEST_F(ScalarComputationsTest, ClampScalarLow) { ComputationBuilder builder(client_, TestName()); builder.Clamp(builder.ConstantR0(2.0f), // The lower bound. builder.ConstantR0(-5.0f), // The operand to be clamped. @@ -573,57 +684,57 @@ TEST_F(ScalarComputationsTest, ClampScalarLow) { ComputeAndCompareR0(&builder, 2.0, {}, error_spec_); } -TEST_F(ScalarComputationsTest, MinS32Above) { +XLA_TEST_F(ScalarComputationsTest, MinS32Above) { TestMinMax(10, 3, 3, &ComputationBuilder::Min); } -TEST_F(ScalarComputationsTest, MinS32Below) { +XLA_TEST_F(ScalarComputationsTest, MinS32Below) { TestMinMax(-100, 3, -100, &ComputationBuilder::Min); } -TEST_F(ScalarComputationsTest, MaxS32Above) { +XLA_TEST_F(ScalarComputationsTest, MaxS32Above) { TestMinMax(10, 3, 10, &ComputationBuilder::Max); } -TEST_F(ScalarComputationsTest, MaxS32Below) { +XLA_TEST_F(ScalarComputationsTest, MaxS32Below) { TestMinMax(-100, 3, 3, &ComputationBuilder::Max); } -TEST_F(ScalarComputationsTest, MinU32Above) { +XLA_TEST_F(ScalarComputationsTest, MinU32Above) { const uint32 large = std::numeric_limits::max(); TestMinMax(large, 3, 3, &ComputationBuilder::Min); } -TEST_F(ScalarComputationsTest, MinU32Below) { +XLA_TEST_F(ScalarComputationsTest, MinU32Below) { TestMinMax(0, 5, 0, &ComputationBuilder::Min); } -TEST_F(ScalarComputationsTest, MaxU32Above) { +XLA_TEST_F(ScalarComputationsTest, MaxU32Above) { const uint32 large = std::numeric_limits::max(); TestMinMax(large, 3, large, &ComputationBuilder::Max); } -TEST_F(ScalarComputationsTest, MaxU32Below) { +XLA_TEST_F(ScalarComputationsTest, MaxU32Below) { TestMinMax(0, 5, 5, &ComputationBuilder::Max); } -TEST_F(ScalarComputationsTest, MinF32Above) { +XLA_TEST_F(ScalarComputationsTest, MinF32Above) { TestMinMax(10.1f, 3.1f, 3.1f, &ComputationBuilder::Min); } -TEST_F(ScalarComputationsTest, MinF32Below) { +XLA_TEST_F(ScalarComputationsTest, MinF32Below) { TestMinMax(-100.1f, 3.1f, -100.1f, &ComputationBuilder::Min); } -TEST_F(ScalarComputationsTest, MaxF32Above) { +XLA_TEST_F(ScalarComputationsTest, MaxF32Above) { TestMinMax(10.1f, 3.1f, 10.1f, &ComputationBuilder::Max); } -TEST_F(ScalarComputationsTest, MaxF32Below) { +XLA_TEST_F(ScalarComputationsTest, MaxF32Below) { TestMinMax(-100.1f, 3.1f, 3.1f, &ComputationBuilder::Max); } -TEST_F(ScalarComputationsTest, ComplicatedArithmeticExpressionF32) { +XLA_TEST_F(ScalarComputationsTest, ComplicatedArithmeticExpressionF32) { // Compute the expression (1 * (3 - 1) * (7 + 0) - 4) / 20. ComputationBuilder b(client_, TestName()); b.Div( @@ -636,7 +747,7 @@ TEST_F(ScalarComputationsTest, ComplicatedArithmeticExpressionF32) { ComputeAndCompareR0(&b, 0.5, {}, error_spec_); } -TEST_F(ScalarComputationsTest, ComplicatedArithmeticExpressionS32) { +XLA_TEST_F(ScalarComputationsTest, ComplicatedArithmeticExpressionS32) { // Compute the expression 1 * (3 - 1) * (7 + 0) - 4. ComputationBuilder b(client_, TestName()); b.Sub(b.Mul(b.ConstantR0(1), @@ -647,9 +758,9 @@ TEST_F(ScalarComputationsTest, ComplicatedArithmeticExpressionS32) { ComputeAndCompareR0(&b, 10, {}); } -TEST_F(ScalarComputationsTest, SqrtF320) { +XLA_TEST_F(ScalarComputationsTest, SqrtF320) { ComputationBuilder builder(client_, TestName()); - Literal zero_literal = LiteralUtil::Zero(PrimitiveType::F32); + Literal zero_literal = Literal::Zero(PrimitiveType::F32); std::unique_ptr zero_data = client_->TransferToServer(zero_literal).ConsumeValueOrDie(); @@ -663,20 +774,3 @@ TEST_F(ScalarComputationsTest, SqrtF320) { } // namespace } // namespace xla - -int main(int argc, char** argv) { - std::vector flag_list; - xla::legacy_flags::AppendCpuCompilerFlags(&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; - } - testing::InitGoogleTest(&argc, argv); - if (argc > 1) { - LOG(ERROR) << "Unknown argument " << argv[1] << "\n" << usage; - return 2; - } - return RUN_ALL_TESTS(); -} diff --git a/tensorflow/compiler/xla/tests/select_and_scatter_test.cc b/tensorflow/compiler/xla/tests/select_and_scatter_test.cc index fb1effc8c46b85c691458fbb2e7dc647e12ad09f..dbef97146c708b3dff0b937d9e20a204dabf7596 100644 --- a/tensorflow/compiler/xla/tests/select_and_scatter_test.cc +++ b/tensorflow/compiler/xla/tests/select_and_scatter_test.cc @@ -25,7 +25,6 @@ limitations under the License. #include "tensorflow/compiler/xla/client/local_client.h" #include "tensorflow/compiler/xla/client/padding.h" #include "tensorflow/compiler/xla/layout_util.h" -#include "tensorflow/compiler/xla/legacy_flags/cpu_compiler_flags.h" #include "tensorflow/compiler/xla/literal_util.h" #include "tensorflow/compiler/xla/reference_util.h" #include "tensorflow/compiler/xla/status_macros.h" @@ -376,20 +375,3 @@ XLA_TEST_F(SelectAndScatterTest, R1F32OverlappingWindowMinScatter) { } // namespace } // namespace xla - -int main(int argc, char** argv) { - std::vector flag_list; - xla::legacy_flags::AppendCpuCompilerFlags(&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; - } - testing::InitGoogleTest(&argc, argv); - if (argc > 1) { - LOG(ERROR) << "Unknown argument " << argv[1] << "\n" << usage; - return 2; - } - return RUN_ALL_TESTS(); -} diff --git a/tensorflow/compiler/xla/tests/select_test.cc b/tensorflow/compiler/xla/tests/select_test.cc index 5ec9ac95faeed1c5ac9e1ca579f7f6ee619a28cb..009e7d24c5cbface4da910e2366db1ff749d5d68 100644 --- a/tensorflow/compiler/xla/tests/select_test.cc +++ b/tensorflow/compiler/xla/tests/select_test.cc @@ -19,7 +19,6 @@ limitations under the License. #include "tensorflow/compiler/xla/client/computation_builder.h" #include "tensorflow/compiler/xla/client/global_data.h" #include "tensorflow/compiler/xla/client/local_client.h" -#include "tensorflow/compiler/xla/legacy_flags/cpu_compiler_flags.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" @@ -257,20 +256,3 @@ TEST_F(SelectTest, SelectR1F32WithScalarPredicateFalse) { } } // namespace } // namespace xla - -int main(int argc, char** argv) { - std::vector flag_list; - xla::legacy_flags::AppendCpuCompilerFlags(&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; - } - testing::InitGoogleTest(&argc, argv); - if (argc > 1) { - LOG(ERROR) << "Unknown argument " << argv[1] << "\n" << usage; - return 2; - } - return RUN_ALL_TESTS(); -} diff --git a/tensorflow/compiler/xla/tests/set_return_value_test.cc b/tensorflow/compiler/xla/tests/set_return_value_test.cc index e15d744d9534459b1ef9b515cf9f114233a85781..29f79ec28a1ae6fcd5299846e85eec992ad2e46f 100644 --- a/tensorflow/compiler/xla/tests/set_return_value_test.cc +++ b/tensorflow/compiler/xla/tests/set_return_value_test.cc @@ -17,7 +17,6 @@ limitations under the License. #include "tensorflow/compiler/xla/client/computation_builder.h" #include "tensorflow/compiler/xla/client/local_client.h" -#include "tensorflow/compiler/xla/legacy_flags/cpu_compiler_flags.h" #include "tensorflow/compiler/xla/tests/client_library_test_base.h" #include "tensorflow/compiler/xla/tests/literal_test_util.h" #include "tensorflow/core/lib/core/status.h" @@ -97,20 +96,3 @@ TEST_F(SetReturnValueTest, SetValueMultipleTimesAndModify) { } // namespace } // namespace xla - -int main(int argc, char** argv) { - std::vector flag_list; - xla::legacy_flags::AppendCpuCompilerFlags(&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; - } - testing::InitGoogleTest(&argc, argv); - if (argc > 1) { - LOG(ERROR) << "Unknown argument " << argv[1] << "\n" << usage; - return 2; - } - return RUN_ALL_TESTS(); -} diff --git a/tensorflow/compiler/xla/tests/slice_test.cc b/tensorflow/compiler/xla/tests/slice_test.cc index d63582fb98a7f4d085246a1172f22cb2b021cf1e..5da6104cfa755b4043abd718d04b2dd203bc809b 100644 --- a/tensorflow/compiler/xla/tests/slice_test.cc +++ b/tensorflow/compiler/xla/tests/slice_test.cc @@ -21,7 +21,6 @@ limitations under the License. #include "tensorflow/compiler/xla/array2d.h" #include "tensorflow/compiler/xla/client/computation_builder.h" #include "tensorflow/compiler/xla/client/local_client.h" -#include "tensorflow/compiler/xla/legacy_flags/cpu_compiler_flags.h" #include "tensorflow/compiler/xla/reference_util.h" #include "tensorflow/compiler/xla/tests/client_library_test_base.h" #include "tensorflow/compiler/xla/tests/literal_test_util.h" @@ -32,96 +31,51 @@ limitations under the License. namespace xla { namespace { -class SliceTest : public ClientLibraryTestBase { - protected: - template - void RunSliceTenToTwo() { - std::vector constant; - for (int i = 0; i < 10; ++i) { - constant.push_back(static_cast(i)); - } - - ComputationBuilder builder(client_, TestName()); - auto original = builder.ConstantR1(constant); - builder.Slice(original, {2}, {4}); - - const std::vector expected = {static_cast(2), - static_cast(3)}; - ComputeAndCompareR1(&builder, expected, {}); - } -}; - -XLA_TEST_F(SliceTest, SliceZeroToZeroF32) { - ComputationBuilder builder(client_, TestName()); - auto original = builder.ConstantR1({}); - builder.Slice(original, {0}, {0}); - - ComputeAndCompareR1(&builder, {}, {}); -} - -XLA_TEST_F(SliceTest, SliceTenToZeroF32) { - ComputationBuilder builder(client_, TestName()); - std::vector constant(10, 0.3); - auto original = builder.ConstantR1(constant); - builder.Slice(original, {7}, {7}); +class SliceTest : public ClientLibraryTestBase {}; - ComputeAndCompareR1(&builder, {}, {}); -} - -TEST_F(SliceTest, SliceTenToTwoF32) { RunSliceTenToTwo(); } - -XLA_TEST_F(SliceTest, SliceTenToTwoF64) { RunSliceTenToTwo(); } - -TEST_F(SliceTest, SliceTenToTwoU32) { RunSliceTenToTwo(); } - -TEST_F(SliceTest, SliceTenToTwoS32) { RunSliceTenToTwo(); } - -XLA_TEST_F(SliceTest, SliceTenToTwoU64) { RunSliceTenToTwo(); } - -XLA_TEST_F(SliceTest, SliceTenToTwoS64) { RunSliceTenToTwo(); } - -TEST_F(SliceTest, SliceTenToTen) { - const std::vector values = {0.0, 1.0, 2.0, 3.0, 4.0, - 5.0, 6.0, 7.0, 8.0, 9.0}; +TEST_F(SliceTest, Slice3x3x3_To_3x3x1_F32) { + Array3D values(3, 3, 3); + values.FillIota(0); ComputationBuilder builder(client_, TestName()); - auto original = builder.ConstantR1(values); - builder.Slice(original, {0}, {10}); + auto original = builder.ConstantR3FromArray3D(values); + builder.Slice(original, {0, 0, 0}, {3, 3, 1}, {1, 1, 1}); - ComputeAndCompareR1(&builder, values, {}, ErrorSpec(0.000001)); + Array3D expected{ + {{0.0}, {3.0}, {6.0}}, {{9.0}, {12.0}, {15.0}}, {{18.0}, {21.0}, {24.0}}}; + ComputeAndCompareR3(&builder, expected, {}, ErrorSpec(0.000001)); } -TEST_F(SliceTest, SliceLastFourOf1024) { - std::vector values(1024); - std::iota(values.begin(), values.end(), 0.0); +TEST_F(SliceTest, Slice3x3x3_To_3x1x3_F32) { + Array3D values(3, 3, 3); + values.FillIota(0); ComputationBuilder builder(client_, TestName()); - auto original = builder.ConstantR1(values); - builder.Slice(original, {1024 - 4}, {1024}); + auto original = builder.ConstantR3FromArray3D(values); + builder.Slice(original, {0, 0, 0}, {3, 1, 3}, {1, 1, 1}); - const std::vector expected = {1020, 1021, 1022, 1023}; - ComputeAndCompareR1(&builder, expected, {}, ErrorSpec(0.000001)); + Array3D expected{ + {{0.0, 1.0, 2.0}}, {{9.0, 10.0, 11.0}}, {{18.0, 19.0, 20.0}}}; + ComputeAndCompareR3(&builder, expected, {}, ErrorSpec(0.000001)); } -// TODO(b/28491443): Fix wrong result on CPU and GPU. Failed on -// 2016-05-01. Also b/28508652 -TEST_F(SliceTest, DISABLED_SliceUnaligned1024In4096Values) { - std::vector values(4096); - std::iota(values.begin(), values.end(), 0.0); +TEST_F(SliceTest, Slice3x3x3_To_1x3x3_F32) { + Array3D values(3, 3, 3); + values.FillIota(0); ComputationBuilder builder(client_, TestName()); - auto original = builder.ConstantR1(values); - builder.Slice(original, {7}, {7 + 1024}); + auto original = builder.ConstantR3FromArray3D(values); + builder.Slice(original, {0, 0, 0}, {1, 3, 3}, {1, 1, 1}); - std::vector expected(1024); - std::iota(values.begin(), values.end(), 7.0); - ComputeAndCompareR1(&builder, expected, {}, ErrorSpec(0.000001)); + Array3D expected{ + {{{0.0, 1.0, 2.0}, {3.0, 4.0, 5.0}, {6.0, 7.0, 8.0}}}}; + ComputeAndCompareR3(&builder, expected, {}, ErrorSpec(0.000001)); } XLA_TEST_F(SliceTest, Slice0x0to0x0F32) { ComputationBuilder builder(client_, TestName()); auto original = builder.ConstantR2FromArray2D(Array2D(0, 0)); - builder.Slice(original, {0, 0}, {0, 0}); + builder.Slice(original, {0, 0}, {0, 0}, {1, 1}); ComputeAndCompareR2(&builder, Array2D(0, 0), {}); } @@ -129,7 +83,7 @@ XLA_TEST_F(SliceTest, Slice0x0to0x0F32) { XLA_TEST_F(SliceTest, Slice0x20to0x5F32) { ComputationBuilder builder(client_, TestName()); auto original = builder.ConstantR2FromArray2D(Array2D(0, 20)); - builder.Slice(original, {0, 15}, {0, 20}); + builder.Slice(original, {0, 15}, {0, 20}, {1, 1}); ComputeAndCompareR2(&builder, Array2D(0, 5), {}); } @@ -137,7 +91,7 @@ XLA_TEST_F(SliceTest, Slice0x20to0x5F32) { XLA_TEST_F(SliceTest, Slice3x0to2x0F32) { ComputationBuilder builder(client_, TestName()); auto original = builder.ConstantR2FromArray2D(Array2D(3, 0)); - builder.Slice(original, {1, 0}, {3, 0}); + builder.Slice(original, {1, 0}, {3, 0}, {1, 1}); ComputeAndCompareR2(&builder, Array2D(2, 0), {}); } @@ -152,7 +106,7 @@ XLA_TEST_F(SliceTest, SliceQuadrantOf256x256) { ComputationBuilder builder(client_, TestName()); auto original = builder.ConstantR2FromArray2D(values); - builder.Slice(original, {128, 128}, {256, 256}); + builder.Slice(original, {128, 128}, {256, 256}, {1, 1}); Array2D expected(128, 128); for (int row = 0; row < 128; ++row) { @@ -170,7 +124,7 @@ TEST_F(SliceTest, Slice_1x4096_To_1x1024) { ComputationBuilder builder(client_, TestName()); auto original = builder.ConstantR2FromArray2D(values); - builder.Slice(original, {0, 3072}, {1, 4096}); + builder.Slice(original, {0, 3072}, {1, 4096}, {1, 1}); Array2D expected(1, 1024); std::iota(expected.data(), expected.data() + 1024, 3072.0); @@ -191,7 +145,7 @@ TEST_F(SliceTest, Slice_16x4_To_16x2) { } ComputationBuilder builder(client_, TestName()); auto original = builder.ConstantR2FromArray2D(values); - builder.Slice(original, {0, 0}, {16, 2}); + builder.Slice(original, {0, 0}, {16, 2}, {1, 1}); ComputeAndCompareR2(&builder, expected, {}, ErrorSpec(0.000001)); } @@ -199,19 +153,99 @@ TEST_F(SliceTest, Slice_16x4_To_16x2) { TEST_F(SliceTest, SliceR4ThreeDimsMiddleMinor) { Array4D values(2, 2, 24, 256); values.FillRandom(3.14f); - auto expected = - ReferenceUtil::Slice4D(values, {{1, 0, 8, 0}}, {{2, 2, 16, 128}}); + auto expected = ReferenceUtil::Slice4D( + values, {{1, 0, 8, 0}}, {{2, 2, 16, 128}}, /*strides=*/{{1, 1, 1, 1}}); ComputationBuilder builder(client_, TestName()); auto original = builder.ConstantR4FromArray4D(values); - builder.Slice(original, {1, 0, 8, 0}, {2, 2, 16, 128}); + builder.Slice(original, {1, 0, 8, 0}, {2, 2, 16, 128}, {1, 1, 1, 1}); ComputeAndCompareR4(&builder, *expected, {}, ErrorSpec(0.000001)); } +struct R1Spec { + int64 input_dim0; + int64 slice_start; + int64 slice_limit; + int64 slice_stride; +}; + +// Parameterized test that generates R1 values, slices them according +// to the R1Spec, and compares the result with a computed version. +class SliceR1Test : public ClientLibraryTestBase, + public ::testing::WithParamInterface { + protected: + template + void Run(const R1Spec& spec) { + std::vector input(spec.input_dim0); + std::iota(input.begin(), input.end(), NativeT()); + + ComputationBuilder builder(client_, TestName()); + auto original = builder.ConstantR1(input); + builder.Slice(original, {spec.slice_start}, {spec.slice_limit}, + {spec.slice_stride}); + + std::vector expected; + for (int i = spec.slice_start; i < spec.slice_limit; + i += spec.slice_stride) { + expected.push_back(i); + } + + ComputeAndCompareR1(&builder, expected, {}); + } +}; + +XLA_TEST_P(SliceR1Test, DoIt_F32) { + Run(GetParam()); +} + +XLA_TEST_P(SliceR1Test, DoIt_F64) { + Run(GetParam()); +} + +XLA_TEST_P(SliceR1Test, DoIt_U32) { + Run(GetParam()); +} + +XLA_TEST_P(SliceR1Test, DoIt_S32) { + Run(GetParam()); +} + +XLA_TEST_P(SliceR1Test, DoIt_U64) { + Run(GetParam()); +} + +XLA_TEST_P(SliceR1Test, DoIt_S64) { + Run(GetParam()); +} + +INSTANTIATE_TEST_CASE_P( // + SliceR1TestInstantiation, // + SliceR1Test, // + ::testing::Values( // + R1Spec{10, 0, 0, 1}, // + R1Spec{10, 7, 7, 1}, // + R1Spec{10, 2, 4, 1}, // + R1Spec{10, 2, 4, 1}, // + R1Spec{10, 2, 4, 1}, // + R1Spec{10, 2, 4, 1}, // + R1Spec{10, 2, 4, 1}, // + R1Spec{10, 2, 4, 1}, // + R1Spec{10, 0, 10, 1}, // + R1Spec{1024, 1024 - 4, 1024, 1}, // + R1Spec{4096, 7, 7 + 1024, 1}, // + R1Spec{10, 0, 10, 2}, // + R1Spec{10, 0, 10, 3}, // + R1Spec{10, 0, 10, 4}, // + R1Spec{10, 0, 10, 5}, // + R1Spec{10, 0, 10, 10} // + ) // +); + struct R2Spec { int64 input_dim0; int64 input_dim1; std::array slice_starts; std::array slice_limits; + std::array slice_strides; Layout layout; }; @@ -220,17 +254,17 @@ struct R2Spec { class SliceR2Test : public ClientLibraryTestBase, public ::testing::WithParamInterface {}; -TEST_P(SliceR2Test, DoIt) { +XLA_TEST_P(SliceR2Test, DoIt) { const R2Spec& spec = GetParam(); Array2D input(spec.input_dim0, spec.input_dim1); input.FillUnique(); ComputationBuilder builder(client_, TestName()); - auto a = builder.ConstantR2FromArray2D(input); - builder.Slice(a, spec.slice_starts, spec.slice_limits); + auto a = builder.ConstantR2FromArray2DWithLayout(input, spec.layout); + builder.Slice(a, spec.slice_starts, spec.slice_limits, spec.slice_strides); - std::unique_ptr> expected = - ReferenceUtil::Slice2D(input, spec.slice_starts, spec.slice_limits); + std::unique_ptr> expected = ReferenceUtil::Slice2D( + input, spec.slice_starts, spec.slice_limits, spec.slice_strides); ComputeAndCompareR2(&builder, *expected, {}); } @@ -238,19 +272,35 @@ TEST_P(SliceR2Test, DoIt) { INSTANTIATE_TEST_CASE_P( SliceR2TestInstantiation, SliceR2Test, ::testing::Values( - R2Spec {4, 12, {{0, 3}}, {{4, 6}}, LayoutUtil::MakeLayout({0, 1})}, - R2Spec {4, 12, {{0, 3}}, {{4, 6}}, LayoutUtil::MakeLayout({1, 0})}, - R2Spec {16, 4, {{0, 2}}, {{16, 4}}, LayoutUtil::MakeLayout({0, 1})}, - R2Spec {16, 4, {{0, 2}}, {{16, 4}}, LayoutUtil::MakeLayout({1, 0})}, - R2Spec {256, 400, {{0, 300}}, {{256, 400}}, + R2Spec {4, 12, {{0, 3}}, {{4, 6}}, {{1, 1}}, + LayoutUtil::MakeLayout({0, 1})}, + R2Spec {4, 12, {{0, 3}}, {{4, 6}}, {{1, 1}}, + LayoutUtil::MakeLayout({1, 0})}, + R2Spec {16, 4, {{0, 2}}, {{16, 4}}, {{1, 1}}, + LayoutUtil::MakeLayout({0, 1})}, + R2Spec {16, 4, {{0, 2}}, {{16, 4}}, {{1, 1}}, + LayoutUtil::MakeLayout({1, 0})}, + R2Spec {256, 400, {{0, 300}}, {{256, 400}}, {{1, 1}}, LayoutUtil::MakeLayout({1, 0})}, - R2Spec {500, 400, {{111, 123}}, {{300, 257}}, + R2Spec {500, 400, {{111, 123}}, {{300, 257}}, {{1, 1}}, LayoutUtil::MakeLayout({1, 0})}, - R2Spec {500, 400, {{111, 123}}, {{300, 400}}, + R2Spec {500, 400, {{111, 123}}, {{300, 400}}, {{1, 1}}, LayoutUtil::MakeLayout({1, 0})}, - R2Spec {384, 512, {{128, 256}}, {{256, 384}}, + R2Spec {384, 512, {{128, 256}}, {{256, 384}}, {{1, 1}}, LayoutUtil::MakeLayout({1, 0})}, - R2Spec {357, 512, {{111, 256}}, {{301, 384}}, + R2Spec {357, 512, {{111, 256}}, {{301, 384}}, {{1, 1}}, + LayoutUtil::MakeLayout({1, 0})}, + R2Spec {10, 10, {{0, 0}}, {{10, 10}}, {{1, 2}}, + LayoutUtil::MakeLayout({0, 1})}, + R2Spec {10, 10, {{0, 0}}, {{10, 10}}, {{1, 2}}, + LayoutUtil::MakeLayout({1, 0})}, + R2Spec {10, 10, {{0, 0}}, {{10, 10}}, {{2, 1}}, + LayoutUtil::MakeLayout({0, 1})}, + R2Spec {10, 10, {{0, 0}}, {{10, 10}}, {{2, 1}}, + LayoutUtil::MakeLayout({1, 0})}, + R2Spec {10, 10, {{0, 0}}, {{10, 10}}, {{2, 2}}, + LayoutUtil::MakeLayout({0, 1})}, + R2Spec {10, 10, {{0, 0}}, {{10, 10}}, {{2, 2}}, LayoutUtil::MakeLayout({1, 0})} ) ); @@ -258,20 +308,3 @@ INSTANTIATE_TEST_CASE_P( } // namespace } // namespace xla - -int main(int argc, char** argv) { - std::vector flag_list; - xla::legacy_flags::AppendCpuCompilerFlags(&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; - } - testing::InitGoogleTest(&argc, argv); - if (argc > 1) { - LOG(ERROR) << "Unknown argument " << argv[1] << "\n" << usage; - return 2; - } - return RUN_ALL_TESTS(); -} diff --git a/tensorflow/compiler/xla/tests/test_macros.cc b/tensorflow/compiler/xla/tests/test_macros.cc new file mode 100644 index 0000000000000000000000000000000000000000..173fb1b0008c9e6edaa1902a5eb3ca5f054a2a67 --- /dev/null +++ b/tensorflow/compiler/xla/tests/test_macros.cc @@ -0,0 +1,97 @@ +/* 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/test_macros.h" + +#include +#include +#include +#include + +#include "tensorflow/core/lib/strings/str_util.h" + +namespace xla { +namespace { + +// Mapping from test name; i.e. MyTest.MyTestCase to platforms on which it is +// disabled. +using ManifestT = std::unordered_map>; + +ManifestT ReadManifest() { + ManifestT manifest; + + string path = XLA_DISABLED_MANIFEST; + if (path.empty()) { + return manifest; + } + + std::ifstream file_stream(path); + // Note: parens are required to disambiguate vs function decl. + string contents((std::istreambuf_iterator(file_stream)), + std::istreambuf_iterator()); + + std::vector lines = tensorflow::str_util::Split(contents, '\n'); + for (string& line : lines) { + auto comment = line.find("//"); + if (comment != string::npos) { + line = line.substr(0, comment); + } + if (line.empty()) { + continue; + } + tensorflow::str_util::StripTrailingWhitespace(&line); + std::vector pieces = tensorflow::str_util::Split(line, ' '); + CHECK_GE(pieces.size(), 1); + auto& platforms = manifest[pieces[0]]; + for (int64 i = 1; i < pieces.size(); ++i) { + platforms.push_back(pieces[i]); + } + } + return manifest; +} + +} // namespace + +string PrependDisabledIfIndicated(const string& test_case_name, + const string& test_name) { + // TODO(leary): this code reads the manifest for every test case instantiated + // in every file. Consider switching to a singleton or using a compile-time + // genrule instead. + ManifestT manifest = ReadManifest(); + + // First try full match: test_case_name.test_name + // If that fails, try to find just the test_case_name; this would disable all + // tests in the test case. + auto it = manifest.find( + tensorflow::strings::StrCat(test_case_name, ".", test_name)); + if (it == manifest.end()) { + it = manifest.find(test_case_name); + if (it == manifest.end()) { + return test_name; + } + } + + const std::vector& disabled_platforms = it->second; + string platform_string = XLA_PLATFORM; + if (std::find(disabled_platforms.begin(), disabled_platforms.end(), + platform_string) != disabled_platforms.end()) { + return "DISABLED_" + test_name; + } + + // We didn't hit in the disabled manifest entries, so don't disable it. + return test_name; +} + +} // namespace xla diff --git a/tensorflow/compiler/xla/tests/test_macros.h b/tensorflow/compiler/xla/tests/test_macros.h index 7f987a21cac04ed2bbb0631deb22bb2cd3e67f35..3878ac1013ef1459cbe3c92a48fc6149b6a4948e 100644 --- a/tensorflow/compiler/xla/tests/test_macros.h +++ b/tensorflow/compiler/xla/tests/test_macros.h @@ -33,13 +33,6 @@ limitations under the License. #include "tensorflow/compiler/xla/types.h" #include "tensorflow/core/platform/test.h" -// Use this macro instead of directly using TEST_P for parameterized tests, -// otherwise DISABLED_ON_* macros nested in TEST_P will not get expanded since -// TEST_P stringifies its argument. That makes the test disabled for all targets -// when any one of the DISABLED_ON_* macro is used, and the test will just pass. -// TODO(b/29122096): Remove this once TEST_P fixes this problem. -#define XLA_TEST_P(test_case_name, test_name) TEST_P(test_case_name, test_name) - #define DISABLED_ON_CPU(X) X #define DISABLED_ON_CPU_PARALLEL(X) X #define DISABLED_ON_GPU(X) X @@ -71,6 +64,91 @@ limitations under the License. // clang-format on -#define XLA_TEST_F(test_fixture, test_name) TEST_F(test_fixture, test_name) - +namespace xla { + +// Reads a disabled manifest file (and retains it as a singleton) to resolve +// whether test cases should be disabled on a particular platform. +string PrependDisabledIfIndicated(const string& test_case_name, + const string& test_name); + +} // namespace xla + +// This is the internal "gtest" class instantiation -- it is identical to the +// GTEST_TEST_ macro, except that we intercept the test name for potential +// modification by PrependDisabledIfIndicated. That file can use an arbitrary +// 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, \ + 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); \ + void GTEST_TEST_CLASS_NAME_(test_case_name, test_name)::TestBody() + +// This is identical to the TEST_F macro from "gtest", but it potentially +// disables the test based on an external manifest file, DISABLED_MANIFEST. +// +// 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()) + +// Likewise, this is identical to the TEST_P macro from "gtest", but +// potentially disables the test based on the DISABLED_MANIFEST file. +// +// We have to wrap this in an outer layer so that any DISABLED_ON_* macros will +// be properly expanded before the stringification occurs. +#define XLA_TEST_P_IMPL_(test_case_name, test_name) \ + class GTEST_TEST_CLASS_NAME_(test_case_name, test_name) \ + : public test_case_name { \ + public: \ + GTEST_TEST_CLASS_NAME_(test_case_name, test_name)() {} \ + virtual void TestBody(); \ + \ + private: \ + static int AddToRegistry() { \ + ::testing::UnitTest::GetInstance() \ + ->parameterized_test_registry() \ + .GetTestCasePatternHolder( \ + #test_case_name, \ + ::testing::internal::CodeLocation(__FILE__, __LINE__)) \ + ->AddTestPattern( \ + #test_case_name, \ + PrependDisabledIfIndicated(#test_case_name, #test_name).c_str(), \ + new ::testing::internal::TestMetaFactory()); \ + return 0; \ + } \ + static int gtest_registering_dummy_ GTEST_ATTRIBUTE_UNUSED_; \ + GTEST_DISALLOW_COPY_AND_ASSIGN_(GTEST_TEST_CLASS_NAME_(test_case_name, \ + test_name)); \ + }; \ + int GTEST_TEST_CLASS_NAME_(test_case_name, \ + test_name)::gtest_registering_dummy_ = \ + GTEST_TEST_CLASS_NAME_(test_case_name, test_name)::AddToRegistry(); \ + void GTEST_TEST_CLASS_NAME_(test_case_name, test_name)::TestBody() + +#define XLA_TEST_P(test_case_name, test_name) \ + XLA_TEST_P_IMPL_(test_case_name, test_name) #endif // TENSORFLOW_COMPILER_XLA_TESTS_TEST_MACROS_H_ diff --git a/tensorflow/compiler/xla/tests/test_utils.h b/tensorflow/compiler/xla/tests/test_utils.h index 6a23df4d3c35a17a56b4ce816f79eaa642831f90..f3a522b05ebae4f1f86d6d7ddbac6e1749d3e286 100644 --- a/tensorflow/compiler/xla/tests/test_utils.h +++ b/tensorflow/compiler/xla/tests/test_utils.h @@ -61,7 +61,7 @@ std::unique_ptr CreateR2LiteralWithLayout( auto literal = MakeUnique(); const int64 d0 = values.size(); const int64 d1 = values.begin()->size(); - LiteralUtil::PopulateWithValue(0, {d0, d1}, literal.get()); + literal.get()->PopulateWithValue(0, {d0, d1}); *literal->mutable_shape()->mutable_layout() = LayoutUtil::MakeLayout(minor_to_major); TF_CHECK_OK(ShapeUtil::ValidateShape(literal->shape())); @@ -70,7 +70,7 @@ std::unique_ptr CreateR2LiteralWithLayout( for (auto inner_list : values) { int64 dim1 = 0; for (auto value : inner_list) { - LiteralUtil::Set(literal.get(), {dim0, dim1}, value); + literal.get()->Set({dim0, dim1}, value); ++dim1; } ++dim0; @@ -88,7 +88,7 @@ std::unique_ptr CreateR3LiteralWithLayout( const int64 d0 = values.size(); const int64 d1 = values.begin()->size(); const int64 d2 = values.begin()->begin()->size(); - LiteralUtil::PopulateWithValue(0, {d0, d1, d2}, literal.get()); + literal.get()->PopulateWithValue(0, {d0, d1, d2}); *literal->mutable_shape()->mutable_layout() = LayoutUtil::MakeLayout(minor_to_major); TF_CHECK_OK(ShapeUtil::ValidateShape(literal->shape())); @@ -99,7 +99,7 @@ std::unique_ptr CreateR3LiteralWithLayout( for (auto inner_inner_list : inner_list) { int64 dim2 = 0; for (auto value : inner_inner_list) { - LiteralUtil::Set(literal.get(), {dim0, dim1, dim2}, value); + literal.get()->Set({dim0, dim1, dim2}, value); ++dim2; } ++dim1; diff --git a/tensorflow/compiler/xla/tests/transpose_test.cc b/tensorflow/compiler/xla/tests/transpose_test.cc index 79f251bbc486d093d8865792406dc44077fac42a..fe5a1778a2cecff0121cee4d8b406c5b23a13e40 100644 --- a/tensorflow/compiler/xla/tests/transpose_test.cc +++ b/tensorflow/compiler/xla/tests/transpose_test.cc @@ -18,7 +18,6 @@ limitations under the License. #include "tensorflow/compiler/xla/array2d.h" #include "tensorflow/compiler/xla/client/computation_builder.h" #include "tensorflow/compiler/xla/client/local_client.h" -#include "tensorflow/compiler/xla/legacy_flags/cpu_compiler_flags.h" #include "tensorflow/compiler/xla/reference_util.h" #include "tensorflow/compiler/xla/tests/client_library_test_base.h" #include "tensorflow/compiler/xla/tests/hlo_test_base.h" @@ -184,20 +183,3 @@ TEST_F(TransposeTest, TransposeConstant021_MultipleTilesPerLayer) { } // namespace } // namespace xla - -int main(int argc, char** argv) { - std::vector flag_list; - xla::legacy_flags::AppendCpuCompilerFlags(&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; - } - testing::InitGoogleTest(&argc, argv); - if (argc > 1) { - LOG(ERROR) << "Unknown argument " << argv[1] << "\n" << usage; - return 2; - } - return RUN_ALL_TESTS(); -} diff --git a/tensorflow/compiler/xla/tests/tuple_test.cc b/tensorflow/compiler/xla/tests/tuple_test.cc index cea9316a6d684bb8512ceac62bd2e7e666fb934e..8269a5d588f7dfbdcfc2fd2b33d9dd0882a95b75 100644 --- a/tensorflow/compiler/xla/tests/tuple_test.cc +++ b/tensorflow/compiler/xla/tests/tuple_test.cc @@ -20,7 +20,6 @@ limitations under the License. #include "tensorflow/compiler/xla/client/computation.h" #include "tensorflow/compiler/xla/client/computation_builder.h" #include "tensorflow/compiler/xla/client/local_client.h" -#include "tensorflow/compiler/xla/legacy_flags/cpu_compiler_flags.h" #include "tensorflow/compiler/xla/literal_util.h" #include "tensorflow/compiler/xla/shape_util.h" #include "tensorflow/compiler/xla/statusor.h" @@ -39,6 +38,25 @@ class TupleTest : public ClientLibraryTestBase { ErrorSpec error_spec_{0.0001}; }; +// Tests a tuple-shaped constant. +XLA_TEST_F(TupleTest, TupleConstant) { + ComputationBuilder builder(client_, TestName()); + + const float constant_scalar = 7.3f; + std::initializer_list constant_vector = {1.1f, 2.0f, 3.3f}; + std::initializer_list> constant_matrix = { + {1.1f, 2.2f, 3.5f}, // row 0 + {4.8f, 5.0f, 6.7f}, // row 1 + }; + auto value = + Literal::MakeTuple({Literal::CreateR0(constant_scalar).get(), + Literal::CreateR1(constant_vector).get(), + Literal::CreateR2(constant_matrix).get()}); + + auto result = builder.ConstantLiteral(*value); + ComputeAndCompareTuple(&builder, *value, {}, error_spec_); +} + // Tests the creation of tuple data. XLA_TEST_F(TupleTest, TupleCreate) { ComputationBuilder builder(client_, TestName()); @@ -53,10 +71,10 @@ XLA_TEST_F(TupleTest, TupleCreate) { builder.ConstantR1(constant_vector), builder.ConstantR2(constant_matrix)}); - auto expected = LiteralUtil::MakeTuple( - {LiteralUtil::CreateR0(constant_scalar).get(), - LiteralUtil::CreateR1(constant_vector).get(), - LiteralUtil::CreateR2(constant_matrix).get()}); + auto expected = + Literal::MakeTuple({Literal::CreateR0(constant_scalar).get(), + Literal::CreateR1(constant_vector).get(), + Literal::CreateR2(constant_matrix).get()}); ComputeAndCompareTuple(&builder, *expected, {}, error_spec_); } @@ -67,9 +85,8 @@ XLA_TEST_F(TupleTest, TupleCreateWithZeroElementEntry) { auto result = builder.Tuple( {builder.ConstantR0(7.0), builder.ConstantR1({})}); - auto expected = - LiteralUtil::MakeTuple({LiteralUtil::CreateR0(7.0).get(), - LiteralUtil::CreateR1({}).get()}); + auto expected = Literal::MakeTuple({Literal::CreateR0(7.0).get(), + Literal::CreateR1({}).get()}); ComputeAndCompareTuple(&builder, *expected, {}, error_spec_); } @@ -77,7 +94,7 @@ XLA_TEST_F(TupleTest, TupleCreateWithZeroElementEntry) { XLA_TEST_F(TupleTest, EmptyTupleCreate) { ComputationBuilder builder(client_, TestName()); auto result = builder.Tuple({}); - auto expected = LiteralUtil::MakeTuple({}); + auto expected = Literal::MakeTuple({}); ComputeAndCompareTuple(&builder, *expected, {}, error_spec_); } @@ -146,12 +163,37 @@ XLA_TEST_F(TupleTest, TupleGTEToTuple) { builder.ConstantR2(constant_matrix)}); auto new_tuple = builder.Tuple({builder.GetTupleElement(tuple_data, 1), builder.GetTupleElement(tuple_data, 0)}); - auto expected = LiteralUtil::MakeTuple( - {LiteralUtil::CreateR2(constant_matrix).get(), - LiteralUtil::CreateR1(constant_vector).get()}); + auto expected = + Literal::MakeTuple({Literal::CreateR2(constant_matrix).get(), + Literal::CreateR1(constant_vector).get()}); ComputeAndCompareTuple(&builder, *expected, {}, error_spec_); } +XLA_TEST_F(TupleTest, SelectBetweenPredTuples) { + ComputationBuilder b(client_, TestName()); + ComputationDataHandle v1, v2; + + for (bool direction : {false, true}) { + std::unique_ptr v1_data = + CreateR0Parameter(0.0f, /*parameter_number=*/0, /*name=*/"v1", + /*builder=*/&b, /*data_handle=*/&v1); + std::unique_ptr v2_data = + CreateR0Parameter(1.0f, /*parameter_number=*/1, /*name=*/"v2", + /*builder=*/&b, /*data_handle=*/&v2); + auto v1_gt = b.Gt(v1, v2); // false + auto v2_gt = b.Gt(v2, v1); // true + auto v1_v2 = b.Tuple({v1_gt, v2_gt}); // {false, true} + auto v2_v1 = b.Tuple({v2_gt, v1_gt}); // {true, false} + auto select = b.Select(direction ? v1_gt : v2_gt, v1_v2, v2_v1); + auto expected = + Literal::MakeTuple({Literal::CreateR0(direction).get(), + Literal::CreateR0(!direction).get()}); + + ComputeAndCompareTuple(&b, *expected, {v1_data.get(), v2_data.get()}, + error_spec_); + } +} + // Builds two new tuples from an existing tuple (by means of GetTupleElement), // then adds up the components of the new tuples. XLA_TEST_F(TupleTest, TupleGTEToTupleToGTEAdd) { @@ -212,9 +254,8 @@ XLA_TEST_F(TupleTest, DISABLED_ON_CPU_PARALLEL(SelectBetweenTuplesOnFalse)) { auto select = builder.Select(builder.ConstantR0(false), tuple12, tuple21); - auto expected = - LiteralUtil::MakeTuple({LiteralUtil::CreateR1(vec2).get(), - LiteralUtil::CreateR1(vec1).get()}); + auto expected = Literal::MakeTuple({Literal::CreateR1(vec2).get(), + Literal::CreateR1(vec1).get()}); ComputeAndCompareTuple(&builder, *expected, {}, error_spec_); } @@ -258,9 +299,8 @@ XLA_TEST_F(TupleTest, DISABLED_ON_CPU_PARALLEL(SelectBetweenTuplesOnTrue)) { auto select = builder.Select(builder.ConstantR0(true), tuple12, tuple21); - auto expected = - LiteralUtil::MakeTuple({LiteralUtil::CreateR1(vec1).get(), - LiteralUtil::CreateR1(vec2).get()}); + auto expected = Literal::MakeTuple({Literal::CreateR1(vec1).get(), + Literal::CreateR1(vec2).get()}); ComputeAndCompareTuple(&builder, *expected, {}, error_spec_); } @@ -339,9 +379,8 @@ XLA_TEST_F(TupleTest, auto select = builder.Select(builder.ConstantR0(false), tuple12, tuple21); - auto expected = - LiteralUtil::MakeTuple({LiteralUtil::CreateR1(vec2).get(), - LiteralUtil::CreateR1(vec1).get()}); + auto expected = Literal::MakeTuple({Literal::CreateR1(vec2).get(), + Literal::CreateR1(vec1).get()}); ComputeAndCompareTuple(&builder, *expected, {}, error_spec_); } @@ -352,13 +391,13 @@ XLA_TEST_F(TupleTest, NestedTuples) { auto outer_tuple = builder.Tuple({inner_tuple, builder.ConstantR1({22.0, 44.0})}); - auto expected_v1 = LiteralUtil::CreateR1({1.0, 2.0}); - auto expected_s = LiteralUtil::CreateR0(42.0); + auto expected_v1 = Literal::CreateR1({1.0, 2.0}); + auto expected_s = Literal::CreateR0(42.0); auto expected_inner_tuple = - LiteralUtil::MakeTuple({expected_v1.get(), expected_s.get()}); - auto expected_v2 = LiteralUtil::CreateR1({22.0, 44.0}); + Literal::MakeTuple({expected_v1.get(), expected_s.get()}); + auto expected_v2 = Literal::CreateR1({22.0, 44.0}); auto expected = - LiteralUtil::MakeTuple({expected_inner_tuple.get(), expected_v2.get()}); + Literal::MakeTuple({expected_inner_tuple.get(), expected_v2.get()}); ComputeAndCompareTuple(&builder, *expected, {}, error_spec_); } @@ -378,14 +417,14 @@ XLA_TEST_F(TupleTest, GetTupleElementOfNestedTuple) { std::unique_ptr data = client_ - ->TransferToServer(*LiteralUtil::MakeTuple({ - LiteralUtil::MakeTuple( + ->TransferToServer(*Literal::MakeTuple({ + Literal::MakeTuple( { - LiteralUtil::CreateR1({1.0, 2.0, 3.0}).get(), - LiteralUtil::CreateR1({4.0, 5.0, 6.0}).get(), + Literal::CreateR1({1.0, 2.0, 3.0}).get(), + Literal::CreateR1({4.0, 5.0, 6.0}).get(), }) .get(), - LiteralUtil::CreateR1({7.0, 8.0, 9.0}).get(), + Literal::CreateR1({7.0, 8.0, 9.0}).get(), })) .ConsumeValueOrDie(); @@ -396,20 +435,3 @@ XLA_TEST_F(TupleTest, GetTupleElementOfNestedTuple) { } // namespace } // namespace xla - -int main(int argc, char** argv) { - std::vector flag_list; - xla::legacy_flags::AppendCpuCompilerFlags(&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; - } - testing::InitGoogleTest(&argc, argv); - if (argc > 1) { - LOG(ERROR) << "Unknown argument " << argv[1] << "\n" << usage; - return 2; - } - return RUN_ALL_TESTS(); -} diff --git a/tensorflow/compiler/xla/tests/unary_op_test.cc b/tensorflow/compiler/xla/tests/unary_op_test.cc index fdbaa0d1786f8b575e5434d3e2b0c010821be8f8..efae13a43a058b03a45174c8260bce2ed70cb75c 100644 --- a/tensorflow/compiler/xla/tests/unary_op_test.cc +++ b/tensorflow/compiler/xla/tests/unary_op_test.cc @@ -19,7 +19,6 @@ limitations under the License. #include "tensorflow/compiler/xla/client/computation_builder.h" #include "tensorflow/compiler/xla/client/global_data.h" #include "tensorflow/compiler/xla/client/local_client.h" -#include "tensorflow/compiler/xla/legacy_flags/cpu_compiler_flags.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" @@ -86,12 +85,12 @@ XLA_TEST_F(UnaryOpTest, AbsTestR1Size0) { AbsSize0TestHelper(); } -TEST_F(UnaryOpTest, AbsTestR1) { +XLA_TEST_F(UnaryOpTest, AbsTestR1) { AbsTestHelper(); AbsTestHelper(); } -TEST_F(UnaryOpTest, AbsTestR0) { +XLA_TEST_F(UnaryOpTest, AbsTestR0) { ComputationBuilder builder(client_, TestName()); auto argi = builder.ConstantR0(-5); auto absi = builder.Abs(argi); @@ -105,7 +104,7 @@ TEST_F(UnaryOpTest, AbsTestR0) { ComputeAndCompareR0(&builder, 8.0f, {}); } -TEST_F(UnaryOpTest, SignTestR0) { +XLA_TEST_F(UnaryOpTest, SignTestR0) { ComputationBuilder builder(client_, TestName()); auto argi = builder.ConstantR0(-5); auto absi = builder.Sign(argi); @@ -119,17 +118,17 @@ TEST_F(UnaryOpTest, SignTestR0) { ComputeAndCompareR0(&builder, -2.0f, {}); } -TEST_F(UnaryOpTest, SignTestR1) { +XLA_TEST_F(UnaryOpTest, SignTestR1) { SignTestHelper(); SignTestHelper(); } -TEST_F(UnaryOpTest, SignAbsTestR1) { +XLA_TEST_F(UnaryOpTest, SignAbsTestR1) { SignAbsTestHelper(); SignAbsTestHelper(); } -TEST_F(UnaryOpTest, UnsignedAbsTestR1) { +XLA_TEST_F(UnaryOpTest, UnsignedAbsTestR1) { ComputationBuilder builder(client_, TestName()); auto arg = builder.ConstantR1( {2, 25, 0, 123, std::numeric_limits::max()}); @@ -139,7 +138,7 @@ TEST_F(UnaryOpTest, UnsignedAbsTestR1) { &builder, {2, 25, 0, 123, std::numeric_limits::max()}, {}); } -TEST_F(UnaryOpTest, UnsignedSignTestR1) { +XLA_TEST_F(UnaryOpTest, UnsignedSignTestR1) { ComputationBuilder builder(client_, TestName()); auto arg = builder.ConstantR1( {2, 25, 0, 123, std::numeric_limits::max()}); @@ -148,7 +147,7 @@ TEST_F(UnaryOpTest, UnsignedSignTestR1) { ComputeAndCompareR1(&builder, {1, 1, 0, 1, 1}, {}); } -TEST_F(UnaryOpTest, SignAbsTestR2) { +XLA_TEST_F(UnaryOpTest, SignAbsTestR2) { ComputationBuilder builder(client_, TestName()); auto arg = builder.ConstantR2({{1.0, -2.0}, {-3.0, 4.0}}); auto sign = builder.Sign(arg); @@ -160,20 +159,3 @@ TEST_F(UnaryOpTest, SignAbsTestR2) { } // namespace } // namespace xla - -int main(int argc, char** argv) { - std::vector flag_list; - xla::legacy_flags::AppendCpuCompilerFlags(&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; - } - testing::InitGoogleTest(&argc, argv); - if (argc > 1) { - LOG(ERROR) << "Unknown argument " << argv[1] << "\n" << usage; - return 2; - } - return RUN_ALL_TESTS(); -} diff --git a/tensorflow/compiler/xla/tests/vector_ops_reduce_test.cc b/tensorflow/compiler/xla/tests/vector_ops_reduce_test.cc index 7f3d7d9cb4c9514890c07c7016354ca686d9e0c5..32ba067a10df6c15348344da813e6a960f05491c 100644 --- a/tensorflow/compiler/xla/tests/vector_ops_reduce_test.cc +++ b/tensorflow/compiler/xla/tests/vector_ops_reduce_test.cc @@ -22,7 +22,6 @@ limitations under the License. #include "tensorflow/compiler/xla/client/computation_builder.h" #include "tensorflow/compiler/xla/client/lib/arithmetic.h" #include "tensorflow/compiler/xla/client/local_client.h" -#include "tensorflow/compiler/xla/legacy_flags/cpu_compiler_flags.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" @@ -216,20 +215,3 @@ TEST_F(VecOpsReduceTest, AddReduceR3F32AllDims) { } // namespace } // namespace xla - -int main(int argc, char** argv) { - std::vector flag_list; - xla::legacy_flags::AppendCpuCompilerFlags(&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; - } - testing::InitGoogleTest(&argc, argv); - if (argc > 1) { - LOG(ERROR) << "Unknown argument " << argv[1] << "\n" << usage; - return 2; - } - return RUN_ALL_TESTS(); -} diff --git a/tensorflow/compiler/xla/tests/vector_ops_simple_test.cc b/tensorflow/compiler/xla/tests/vector_ops_simple_test.cc index 4ab4c84aa567153e2e2444c2ffaed21784f2fe5a..48a85f16a22cd7536222b8c03c4ebad2bb77d240 100644 --- a/tensorflow/compiler/xla/tests/vector_ops_simple_test.cc +++ b/tensorflow/compiler/xla/tests/vector_ops_simple_test.cc @@ -23,7 +23,6 @@ limitations under the License. #include "tensorflow/compiler/xla/client/global_data.h" #include "tensorflow/compiler/xla/client/lib/arithmetic.h" #include "tensorflow/compiler/xla/client/local_client.h" -#include "tensorflow/compiler/xla/legacy_flags/cpu_compiler_flags.h" #include "tensorflow/compiler/xla/shape_util.h" #include "tensorflow/compiler/xla/statusor.h" #include "tensorflow/compiler/xla/test_helpers.h" @@ -41,13 +40,15 @@ namespace { class VecOpsSimpleTest : public ClientLibraryTestBase { public: explicit VecOpsSimpleTest(perftools::gputools::Platform* platform = nullptr) - : ClientLibraryTestBase(platform, - /*disabled_pass_names=*/{"algsimp", "inline"}) {} + : ClientLibraryTestBase(platform) { + mutable_debug_options()->add_xla_disable_hlo_passes("algsimp"); + mutable_debug_options()->add_xla_disable_hlo_passes("inline"); + } ErrorSpec error_spec_{0.0001}; }; -TEST_F(VecOpsSimpleTest, ExpTenValues) { +XLA_TEST_F(VecOpsSimpleTest, ExpTenValues) { ComputationBuilder builder(client_, TestName()); auto x = builder.ConstantR1( {2.1, -2.6, 2.6, -4.0, 2.1, 2.3, -5.0, -0.9, -2.4, 1.6}); @@ -60,10 +61,11 @@ TEST_F(VecOpsSimpleTest, ExpTenValues) { ComputeAndCompareR1(&builder, expected, {}, error_spec_); } -TEST_F(VecOpsSimpleTest, ExpManyValues) { +XLA_TEST_F(VecOpsSimpleTest, ExpManyValues) { for (int count : {63, 64, 65, 127, 128, 129, 17 * 4096}) { ComputationBuilder builder(client_, TestName()); std::vector exponents; + exponents.reserve(count); for (int i = 0; i < count; ++i) { exponents.push_back(i / static_cast(count)); } @@ -71,6 +73,7 @@ TEST_F(VecOpsSimpleTest, ExpManyValues) { auto exp = builder.Exp(x); std::vector expected; + expected.reserve(exponents.size()); for (float exponent : exponents) { expected.push_back(std::exp(exponent)); } @@ -80,7 +83,7 @@ TEST_F(VecOpsSimpleTest, ExpManyValues) { } } -TEST_F(VecOpsSimpleTest, ExpIn4D) { +XLA_TEST_F(VecOpsSimpleTest, ExpIn4D) { ComputationBuilder builder(client_, TestName()); Array4D exponents(2, 2, 2, 2); @@ -102,7 +105,7 @@ TEST_F(VecOpsSimpleTest, ExpIn4D) { ErrorSpec(/*aabs=*/1e-2, /*arel=*/1e-3)); } -TEST_F(VecOpsSimpleTest, NegateTenFloatValues) { +XLA_TEST_F(VecOpsSimpleTest, NegateTenFloatValues) { ComputationBuilder builder(client_, TestName()); auto x = builder.ConstantR1( {2.1, -2.6, 2.6, -4.0, 2.1, 2.3, -5.0, -0.9, -2.4, 1.6}); @@ -113,7 +116,7 @@ TEST_F(VecOpsSimpleTest, NegateTenFloatValues) { ComputeAndCompareR1(&builder, expected, {}, error_spec_); } -TEST_F(VecOpsSimpleTest, NegateTenInt32Values) { +XLA_TEST_F(VecOpsSimpleTest, NegateTenInt32Values) { ComputationBuilder builder(client_, TestName()); auto x = builder.ConstantR1({2, -2, 12, -4, 5, 20, -15, 0, -2, 1}); builder.Neg(x); @@ -122,7 +125,7 @@ TEST_F(VecOpsSimpleTest, NegateTenInt32Values) { ComputeAndCompareR1(&builder, expected, {}); } -TEST_F(VecOpsSimpleTest, NegateUint32Values) { +XLA_TEST_F(VecOpsSimpleTest, NegateUint32Values) { ComputationBuilder builder(client_, TestName()); auto x = builder.ConstantR1( {0, 1, 42, static_cast(-1), static_cast(-12)}); @@ -132,7 +135,7 @@ TEST_F(VecOpsSimpleTest, NegateUint32Values) { ComputeAndCompareR1(&builder, expected, {}); } -TEST_F(VecOpsSimpleTest, SquareTenValues) { +XLA_TEST_F(VecOpsSimpleTest, SquareTenValues) { ComputationBuilder builder(client_, TestName()); auto x = builder.ConstantR1( {2.1, -2.6, 2.6, -4.0, 2.1, 2.3, -5.0, -0.9, -2.4, 1.6}); @@ -143,7 +146,7 @@ TEST_F(VecOpsSimpleTest, SquareTenValues) { ComputeAndCompareR1(&builder, expected, {}, error_spec_); } -TEST_F(VecOpsSimpleTest, ReciprocalTenValues) { +XLA_TEST_F(VecOpsSimpleTest, ReciprocalTenValues) { ComputationBuilder builder(client_, TestName()); auto x = builder.ConstantR1( {2.1, -2.6, 2.6, -4.0, 2.1, 2.3, -5.0, -0.9, -2.4, 1.6}); @@ -184,7 +187,7 @@ XLA_TEST_F(VecOpsSimpleTest, InvSqrtSevenValues) { ComputeAndCompareR1(&builder, expected, {}, error_spec_); } -TEST_F(VecOpsSimpleTest, AddTenValuesViaMap) { +XLA_TEST_F(VecOpsSimpleTest, AddTenValuesViaMap) { ComputationBuilder builder(client_, TestName()); auto add = CreateScalarAddComputation(F32, &builder); @@ -199,7 +202,7 @@ TEST_F(VecOpsSimpleTest, AddTenValuesViaMap) { ComputeAndCompareR1(&builder, expected, {}, error_spec_); } -TEST_F(VecOpsSimpleTest, MaxTenValues) { +XLA_TEST_F(VecOpsSimpleTest, MaxTenValues) { ComputationBuilder builder(client_, TestName()); auto x = builder.ConstantR1( {2.1, -2.6, 2.6, -4.0, 2.1, 2.3, -5.0, -0.9, -2.4, 1.6}); @@ -212,7 +215,7 @@ TEST_F(VecOpsSimpleTest, MaxTenValues) { ComputeAndCompareR1(&builder, expected, {}); } -TEST_F(VecOpsSimpleTest, MaxTenValuesFromParams) { +XLA_TEST_F(VecOpsSimpleTest, MaxTenValuesFromParams) { // Similar to MaxTenValues, except that the inputs come from params rather // than constants. ComputationBuilder builder(client_, TestName()); @@ -230,7 +233,7 @@ TEST_F(VecOpsSimpleTest, MaxTenValuesFromParams) { error_spec_); } -TEST_F(VecOpsSimpleTest, Max15000ValuesFromParams) { +XLA_TEST_F(VecOpsSimpleTest, Max15000ValuesFromParams) { // Similar to MaxTenValuesFromParams, except that the data size passed in and // out is large. ComputationBuilder builder(client_, TestName()); @@ -270,7 +273,7 @@ TEST_F(VecOpsSimpleTest, Max15000ValuesFromParams) { error_spec_); } -TEST_F(VecOpsSimpleTest, MaxTenValuesWithScalar) { +XLA_TEST_F(VecOpsSimpleTest, MaxTenValuesWithScalar) { ComputationBuilder builder(client_, TestName()); auto x = builder.ConstantR1( {2.1, -2.6, 2.6, -4.0, 2.1, 2.3, -5.0, -0.9, -2.4, 1.6}); @@ -282,7 +285,7 @@ TEST_F(VecOpsSimpleTest, MaxTenValuesWithScalar) { ComputeAndCompareR1(&builder, expected, {}); } -TEST_F(VecOpsSimpleTest, MinTenValues) { +XLA_TEST_F(VecOpsSimpleTest, MinTenValues) { ComputationBuilder builder(client_, TestName()); auto x = builder.ConstantR1( {2.1, -2.6, 2.6, -4.0, 2.1, 2.3, -5.0, -0.9, -2.4, 1.6}); @@ -295,7 +298,7 @@ TEST_F(VecOpsSimpleTest, MinTenValues) { ComputeAndCompareR1(&builder, expected, {}); } -TEST_F(VecOpsSimpleTest, MinMaxTenValues) { +XLA_TEST_F(VecOpsSimpleTest, MinMaxTenValues) { ComputationBuilder builder(client_, TestName()); auto zero = builder.ConstantR0(0); auto one = builder.ConstantR0(1); @@ -308,7 +311,7 @@ TEST_F(VecOpsSimpleTest, MinMaxTenValues) { ComputeAndCompareR1(&builder, expected, {}); } -TEST_F(VecOpsSimpleTest, ClampTenValuesConstant) { +XLA_TEST_F(VecOpsSimpleTest, ClampTenValuesConstant) { ComputationBuilder builder(client_, TestName()); auto zero = builder.ConstantR0(0); auto one = builder.ConstantR0(1); @@ -321,7 +324,7 @@ TEST_F(VecOpsSimpleTest, ClampTenValuesConstant) { ComputeAndCompareR1(&builder, expected, {}); } -TEST_F(VecOpsSimpleTest, ClampTwoValuesConstant) { +XLA_TEST_F(VecOpsSimpleTest, ClampTwoValuesConstant) { ComputationBuilder builder(client_, TestName()); auto zero = builder.ConstantR1({0.0f, 0.0f}); auto one = builder.ConstantR1({1.0f, 1.0f}); @@ -332,7 +335,7 @@ TEST_F(VecOpsSimpleTest, ClampTwoValuesConstant) { ComputeAndCompareR1(&builder, expected, {}); } -TEST_F(VecOpsSimpleTest, ClampTenValuesConstantNonzeroLower) { +XLA_TEST_F(VecOpsSimpleTest, ClampTenValuesConstantNonzeroLower) { ComputationBuilder builder(client_, TestName()); auto one = builder.ConstantR0(1); auto two = builder.ConstantR0(2); @@ -345,7 +348,7 @@ TEST_F(VecOpsSimpleTest, ClampTenValuesConstantNonzeroLower) { ComputeAndCompareR1(&builder, expected, {}); } -TEST_F(VecOpsSimpleTest, MapTenValues) { +XLA_TEST_F(VecOpsSimpleTest, MapTenValues) { Computation add_half; { // add_half(x) = x + 0.5 @@ -433,20 +436,3 @@ XLA_TEST_F(VecOpsSimpleTest, VectorPredicateNotEqual) { } // namespace } // namespace xla - -int main(int argc, char** argv) { - std::vector flag_list; - xla::legacy_flags::AppendCpuCompilerFlags(&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; - } - testing::InitGoogleTest(&argc, argv); - if (argc > 1) { - LOG(ERROR) << "Unknown argument " << argv[1] << "\n" << usage; - return 2; - } - return RUN_ALL_TESTS(); -} diff --git a/tensorflow/compiler/xla/tests/while_test.cc b/tensorflow/compiler/xla/tests/while_test.cc index 4cff1990865bcf1214a403a6241accbf82f06d00..cafaf5bcc64a1b8170b91c8ecc89c4166dac68ba 100644 --- a/tensorflow/compiler/xla/tests/while_test.cc +++ b/tensorflow/compiler/xla/tests/while_test.cc @@ -22,10 +22,10 @@ limitations under the License. #include "tensorflow/compiler/xla/client/computation_builder.h" #include "tensorflow/compiler/xla/client/lib/arithmetic.h" #include "tensorflow/compiler/xla/client/local_client.h" -#include "tensorflow/compiler/xla/legacy_flags/cpu_compiler_flags.h" #include "tensorflow/compiler/xla/literal_util.h" #include "tensorflow/compiler/xla/service/platform_util.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/tests/client_library_test_base.h" #include "tensorflow/compiler/xla/tests/literal_test_util.h" @@ -80,6 +80,70 @@ TEST_F(WhileTest, WhileWithScalarResult) { ComputeAndCompareR0(&builder, 5, {}); } +TEST_F(WhileTest, WhileWithScalarResultNonConstInit) { + auto result_shape = ShapeUtil::MakeShape(S32, {}); + auto orig_shape = ShapeUtil::MakeShape(S32, {2}); + + // Create a computation for the condition: repeat for 5 iterations. + Computation condition; + { + ComputationBuilder builder(client_, "condition"); + auto prev = builder.Parameter(0, result_shape, "prev"); + builder.Gt(builder.ConstantR0(5), prev); + condition = builder.Build().ConsumeValueOrDie(); + } + + // Create a computation for the body: add 1 to the result variable. + Computation body; + { + ComputationBuilder builder(client_, "body"); + auto prev = builder.Parameter(0, result_shape, "prev"); + auto input = builder.ConstantR0(1); + auto result = builder.Add(input, prev); + body = builder.Build().ConsumeValueOrDie(); + } + + // Create a While node with computations for the condition and the body. + ComputationBuilder builder(client_, TestName()); + auto init = builder.Reduce(builder.ConstantR1(2, 1), + builder.ConstantR0(0), + CreateScalarAddComputation(S32, &builder), {0}); + auto result = builder.While(condition, body, init); + auto shape = builder.GetShape(result).ConsumeValueOrDie(); + + ComputeAndCompareR0(&builder, 5, {}); +} + +TEST_F(WhileTest, WhileWithPredicateResult) { + auto result_shape = ShapeUtil::MakeShape(PRED, {}); + + // Create a computation for the condition: run until condition is true. + Computation condition; + { + ComputationBuilder builder(client_, "condition"); + auto prev = builder.Parameter(0, result_shape, "prev"); + builder.Ne(builder.ConstantR0(true), prev); + condition = builder.Build().ConsumeValueOrDie(); + } + + // Create a computation for the body: or condition with true. + Computation body; + { + ComputationBuilder builder(client_, "body"); + auto prev = builder.Parameter(0, result_shape, "prev"); + auto result = builder.LogicalOr(prev, builder.ConstantR0(true)); + body = builder.Build().ConsumeValueOrDie(); + } + + // Create a While node with computations for the condition and the body. + ComputationBuilder builder(client_, TestName()); + auto init = builder.Ne(builder.ConstantR0(false), + builder.ConstantR0(true)); + auto result = builder.While(condition, body, init); + + ComputeAndCompareR0(&builder, true, {}); +} + // Tests a while node when the result type T is a vector. // // All constants are chosen to produce exact results. @@ -191,6 +255,70 @@ TEST_F(WhileTest, WhileWithVectorResult) { ComputeAndCompareR1(&builder, expected, {}, ErrorSpec(0.0001)); } +// Tests a while node when the result type is a vector which is part +// of the result tuple. +// +// All constants are chosen to produce exact results. +// vector result(8, 0.0f); +// while (result.sum() < 15.5f) { +// result = result + vector(8, 0.125f); +// } +// tuple = tuple { while } +TEST_F(WhileTest, WhileWithVectorResultIntoTuple) { + Shape result_shape = ShapeUtil::MakeShape(F32, {8}); + + // Create a computation for the reduction. + Computation add; + { + ComputationBuilder builder(client_, "add"); + auto x = builder.Parameter(0, ShapeUtil::MakeShape(F32, {}), "x"); + auto y = builder.Parameter(1, ShapeUtil::MakeShape(F32, {}), "y"); + builder.Add(x, y); + add = builder.Build().ConsumeValueOrDie(); + } + + // Create a computation for the condition. + // Repeat until the sum of the result vector is less than 5.5f. + Computation condition; + { + ComputationBuilder builder(client_, "condition"); + auto prev = builder.Parameter(0, result_shape, "prev"); + auto sum = builder.Reduce(prev, builder.ConstantR0(0.0f), add, + /*dimensions_to_reduce=*/{0}); + auto test = builder.Gt(builder.ConstantR0(15.5f), sum); + condition = builder.Build().ConsumeValueOrDie(); + } + + // Create a computation for the body. + // Add a constant vector of 1.f to the result vector. + Computation body; + { + ComputationBuilder builder(client_, "body"); + auto prev = builder.Parameter(0, result_shape, "prev"); + auto input = builder.ConstantR1(8, 0.125f); + auto result = builder.Add(input, prev); + body = builder.Build().ConsumeValueOrDie(); + } + + // Create a While node with computations for the condition and the body. + ComputationBuilder builder(client_, "while"); + auto init = builder.ConstantR1(8, 0.f); + auto result = builder.While(condition, body, init); + VLOG(2) << "while = " + << ShapeUtil::HumanString( + *builder.GetShape(result).ConsumeValueOrDie()); + builder.Tuple({result}); + + // Individual elements with increase by 1/8 each time through the loop, so + // the sum will increase by 1.0. It will first be >15.5 when the elements + // have all reached 2.0. + auto expected_data = + Literal::CreateR1({2.f, 2.f, 2.f, 2.f, 2.f, 2.f, 2.f, 2.f}); + auto expected = Literal::MakeTuple({expected_data.get()}); + VLOG(2) << "expected = " << ShapeUtil::HumanString(expected->shape()); + ComputeAndCompareTuple(&builder, *expected, {}, ErrorSpec(0.0001)); +} + // Tests a while node when the result type T is a Tuple. // // tuple> result(0, vector(10, 0.0f)); @@ -238,15 +366,286 @@ TEST_F(WhileTest, WhileWithTupleResult) { VLOG(2) << "while = " << ShapeUtil::HumanString( *builder.GetShape(result).ConsumeValueOrDie()); - auto expected_counter = LiteralUtil::CreateR0(5); - auto expected_data = LiteralUtil::CreateR1( + auto expected_counter = Literal::CreateR0(5); + auto expected_data = Literal::CreateR1( {5.0f, 5.0f, 5.0f, 5.0f, 5.0f, 5.0f, 5.0f, 5.0f, 5.0f, 5.0f}); auto expected = - LiteralUtil::MakeTuple({expected_counter.get(), expected_data.get()}); + Literal::MakeTuple({expected_counter.get(), expected_data.get()}); VLOG(2) << "expected = " << ShapeUtil::HumanString(expected->shape()); ComputeAndCompareTuple(&builder, *expected, {}, ErrorSpec(0.0001)); } +TEST_F(WhileTest, WhileWithPredicateTupleResult) { + std::vector shape_elements = {ShapeUtil::MakeShape(S32, {}), + ShapeUtil::MakeShape(PRED, {})}; + Shape result_shape = ShapeUtil::MakeTupleShape(shape_elements); + + // Create a computation for the condition. + // Repeat for 5 iterations. + Computation condition; + { + ComputationBuilder builder(client_, "condition"); + auto prev = builder.Parameter(0, result_shape, "prev"); + auto iteration = builder.GetTupleElement(prev, 0); + builder.Gt(builder.ConstantR0(5), iteration); + condition = builder.Build().ConsumeValueOrDie(); + } + + // Create a computation for the body. + // Add 1 to the iteration variable and or the predicate with true + Computation body; + { + ComputationBuilder builder(client_, "body"); + auto prev = builder.Parameter(0, result_shape, "prev"); + auto iteration = builder.GetTupleElement(prev, 0); + auto pred = builder.GetTupleElement(prev, 1); + auto new_pred = builder.LogicalOr(pred, builder.ConstantR0(true)); + auto result = builder.Tuple( + {builder.Add(iteration, builder.ConstantR0(1)), new_pred}); + body = builder.Build().ConsumeValueOrDie(); + } + + // Create a While node with computations for the condition and the body. + ComputationBuilder builder(client_, "while"); + auto init = builder.Tuple({builder.ConstantR0(0), + builder.Ne(builder.ConstantR0(false), + builder.ConstantR0(true))}); + auto result = builder.While(condition, body, init); + VLOG(2) << "while = " + << ShapeUtil::HumanString( + *builder.GetShape(result).ConsumeValueOrDie()); + + auto expected_counter = Literal::CreateR0(5); + auto expected_predicate = Literal::CreateR0(true); + auto expected = + Literal::MakeTuple({expected_counter.get(), expected_predicate.get()}); + ComputeAndCompareTuple(&builder, *expected, {}, ErrorSpec(0)); +} + +// Tests two while nodes when the result type T is a Tuple and the second +// while node uses the result of the first while node which is used in two +// nodes. +// tuple> w0(0, vector(10, 0.0f)); +// w0 = while (get<0>(w0) < c1) { +// get<0>(w0) = get<0>(w0) + 1; +// get<1>(w0) = get<1>(w0) + vector(10, 1.0f); +// } +// tuple> w1(get<0>(w0), get<1>(w0)); +// w1 = while (get<0>(w1) < c2) { +// get<0>(w1) = get<0>(w1) + 1; +// get<1>(w1) = get<1>(w1) + vector(10, 1.0f); +// } +// result = get<1>(w0) + get<1>(w1) +TEST_F(WhileTest, TwoWhileWithTupleResult) { + std::vector shape_elements = {ShapeUtil::MakeShape(S32, {}), + ShapeUtil::MakeShape(F32, {10})}; + Shape result_shape = ShapeUtil::MakeTupleShape(shape_elements); + + // Create a computation for the condition. + // Repeat for 5 iterations. + Computation condition; + const int c1 = 5; + { + ComputationBuilder builder(client_, "condition"); + auto prev = builder.Parameter(0, result_shape, "prev"); + auto iteration = builder.GetTupleElement(prev, 0); + builder.Lt(iteration, builder.ConstantR0(c1)); + TF_ASSERT_OK_AND_ASSIGN(condition, builder.Build()); + } + + Computation condition2; + const int c2 = 7; + { + ComputationBuilder builder(client_, "condition2"); + auto prev = builder.Parameter(0, result_shape, "prev"); + auto iteration = builder.GetTupleElement(prev, 0); + builder.Lt(iteration, builder.ConstantR0(c2)); + TF_ASSERT_OK_AND_ASSIGN(condition2, builder.Build()); + } + + // Create a computation for the body. + // Add 1 to the iteration variable and add a constant vector of 1.0f to + // the weight variable, both of which are tuple elements. + Computation body; + { + ComputationBuilder builder(client_, "body"); + auto prev = builder.Parameter(0, result_shape, "prev"); + auto iteration = builder.GetTupleElement(prev, 0); + auto weights = builder.GetTupleElement(prev, 1); + auto input = builder.ConstantR1(10, 1.f); + auto new_weights = builder.Add(weights, input); + auto result = builder.Tuple( + {builder.Add(iteration, builder.ConstantR0(1)), new_weights}); + TF_ASSERT_OK_AND_ASSIGN(body, builder.Build()); + } + + Computation body2; + { + ComputationBuilder builder(client_, "body"); + auto prev = builder.Parameter(0, result_shape, "prev"); + auto iteration = builder.GetTupleElement(prev, 0); + auto weights = builder.GetTupleElement(prev, 1); + auto input = builder.ConstantR1(10, 1.f); + auto new_weights = builder.Add(weights, input); + auto result = builder.Tuple( + {builder.Add(iteration, builder.ConstantR0(1)), new_weights}); + TF_ASSERT_OK_AND_ASSIGN(body2, builder.Build()); + } + + // Create a While node with computations for the condition and the body. + ComputationBuilder builder(client_, "while"); + auto init = builder.Tuple( + {builder.ConstantR0(0), builder.ConstantR1(10, 0.f)}); + auto while1 = builder.While(condition, body, init); + + auto while2 = builder.While(condition2, body2, while1); + + auto while_result1 = builder.GetTupleElement(while1, 1); + auto while_result2 = builder.GetTupleElement(while2, 1); + VLOG(2) << "while_result2 = " + << ShapeUtil::HumanString( + *builder.GetShape(while_result2).ConsumeValueOrDie()); + auto result = builder.Add(while_result1, while_result2); + VLOG(2) << "result = " + << ShapeUtil::HumanString( + *builder.GetShape(result).ConsumeValueOrDie()); + const float sum = c1 + c2; + std::vector expected(10, sum); + ComputeAndCompareR1(&builder, expected, {}, ErrorSpec(0.0001)); +} + +// Test while nodes that share the while body computation. +TEST_F(WhileTest, TwoWhileLoopsAndSharedBody) { + std::vector shape_elements = {ShapeUtil::MakeShape(S32, {}), + ShapeUtil::MakeShape(F32, {10})}; + Shape result_shape = ShapeUtil::MakeTupleShape(shape_elements); + + // Create a computation for the condition. + // Repeat for 5 iterations. + Computation condition; + const int c1 = 5; + { + ComputationBuilder builder(client_, "condition"); + auto prev = builder.Parameter(0, result_shape, "prev"); + auto iteration = builder.GetTupleElement(prev, 0); + builder.Lt(iteration, builder.ConstantR0(c1)); + TF_ASSERT_OK_AND_ASSIGN(condition, builder.Build()); + } + + Computation condition2; + const int c2 = 7; + { + ComputationBuilder builder(client_, "condition2"); + auto prev = builder.Parameter(0, result_shape, "prev"); + auto iteration = builder.GetTupleElement(prev, 0); + builder.Lt(iteration, builder.ConstantR0(c2)); + TF_ASSERT_OK_AND_ASSIGN(condition2, builder.Build()); + } + + // Create a computation for the body. + // Add 1 to the iteration variable and add a constant vector of 1.0f to + // the weight variable, both of which are tuple elements. + Computation body; + { + ComputationBuilder builder(client_, "body"); + auto prev = builder.Parameter(0, result_shape, "prev"); + auto iteration = builder.GetTupleElement(prev, 0); + auto weights = builder.GetTupleElement(prev, 1); + auto input = builder.ConstantR1(10, 1.f); + auto new_weights = builder.Add(weights, input); + auto result = builder.Tuple( + {builder.Add(iteration, builder.ConstantR0(1)), new_weights}); + TF_ASSERT_OK_AND_ASSIGN(body, builder.Build()); + } + + // Create a While node with computations for the condition and the body. + ComputationBuilder builder(client_, "while"); + auto init = builder.Tuple( + {builder.ConstantR0(0), builder.ConstantR1(10, 0.f)}); + auto while1 = builder.While(condition, body, init); + + auto while2 = builder.While(condition2, body, while1); + + auto while_result1 = builder.GetTupleElement(while1, 1); + auto while_result2 = builder.GetTupleElement(while2, 1); + VLOG(2) << "while_result2 = " + << ShapeUtil::HumanString( + *builder.GetShape(while_result2).ConsumeValueOrDie()); + auto result = builder.Add(while_result1, while_result2); + VLOG(2) << "result = " + << ShapeUtil::HumanString( + *builder.GetShape(result).ConsumeValueOrDie()); + const float sum = c1 + c2; + std::vector expected(10, sum); + ComputeAndCompareR1(&builder, expected, {}, ErrorSpec(0.0001)); +} + +// Test while nodes that share the while body computation. +// TODO(b/37245345): Fails on GPU backend. +TEST_F(WhileTest, DISABLED_ON_GPU(WhileLoopsWithSharedBodyAndInit)) { + std::vector shape_elements = {ShapeUtil::MakeShape(S32, {}), + ShapeUtil::MakeShape(F32, {10})}; + Shape result_shape = ShapeUtil::MakeTupleShape(shape_elements); + + // Create a computation for the condition. + // Repeat for 5 iterations. + Computation condition; + const int c1 = 5; + { + ComputationBuilder builder(client_, "condition"); + auto prev = builder.Parameter(0, result_shape, "prev"); + auto iteration = builder.GetTupleElement(prev, 0); + builder.Lt(iteration, builder.ConstantR0(c1)); + TF_ASSERT_OK_AND_ASSIGN(condition, builder.Build()); + } + + Computation condition2; + const int c2 = 7; + { + ComputationBuilder builder(client_, "condition2"); + auto prev = builder.Parameter(0, result_shape, "prev"); + auto iteration = builder.GetTupleElement(prev, 0); + builder.Lt(iteration, builder.ConstantR0(c2)); + TF_ASSERT_OK_AND_ASSIGN(condition2, builder.Build()); + } + + // Create a computation for the body. + // Add 1 to the iteration variable and add a constant vector of 1.0f to + // the weight variable, both of which are tuple elements. + Computation body; + { + ComputationBuilder builder(client_, "body"); + auto prev = builder.Parameter(0, result_shape, "prev"); + auto iteration = builder.GetTupleElement(prev, 0); + auto weights = builder.GetTupleElement(prev, 1); + auto input = builder.ConstantR1(10, 1.f); + auto new_weights = builder.Add(weights, input); + auto result = builder.Tuple( + {builder.Add(iteration, builder.ConstantR0(1)), new_weights}); + TF_ASSERT_OK_AND_ASSIGN(body, builder.Build()); + } + + // Create a While node with computations for the condition and the body. + ComputationBuilder builder(client_, "while"); + auto init = builder.Tuple( + {builder.ConstantR0(0), builder.ConstantR1(10, 0.f)}); + auto while1 = builder.While(condition, body, init); + auto while2 = builder.While(condition2, body, init); + + auto while_result1 = builder.GetTupleElement(while1, 1); + auto while_result2 = builder.GetTupleElement(while2, 1); + VLOG(2) << "while_result2 = " + << ShapeUtil::HumanString( + *builder.GetShape(while_result2).ConsumeValueOrDie()); + auto result = builder.Add(while_result1, while_result2); + VLOG(2) << "result = " + << ShapeUtil::HumanString( + *builder.GetShape(result).ConsumeValueOrDie()); + const float sum = c1 + c2; + std::vector expected(10, sum); + ComputeAndCompareR1(&builder, expected, {}, ErrorSpec(0.0001)); +} + // WhileTest that uses DynamicUpdateSlice instruction in body computation. // Loop state tuple element 1 has as its single user operand(0) of // DynamicUpdateSlice, which will trigger in-place dynamic slice update on GPU. @@ -299,11 +698,11 @@ XLA_TEST_F(WhileTest, WhileWithDynamicUpdateSlice) { << ShapeUtil::HumanString( *builder.GetShape(result).ConsumeValueOrDie()); - auto expected_counter = LiteralUtil::CreateR0(5); - auto expected_data = LiteralUtil::CreateR1( + auto expected_counter = Literal::CreateR0(5); + auto expected_data = Literal::CreateR1( {1.0f, 1.0f, 2.0f, 2.0f, 3.0f, 3.0f, 4.0f, 4.0f, 5.0f, 5.0f}); auto expected = - LiteralUtil::MakeTuple({expected_counter.get(), expected_data.get()}); + Literal::MakeTuple({expected_counter.get(), expected_data.get()}); VLOG(2) << "expected = " << ShapeUtil::HumanString(expected->shape()); ComputeAndCompareTuple(&builder, *expected, {}, ErrorSpec(0.0001)); } @@ -315,7 +714,8 @@ XLA_TEST_F(WhileTest, WhileWithDynamicUpdateSlice) { // result += (1, U[0, 100], U[0, 100], U[0, 100], U[0, 100], U[0, 100]); // } // -// This test misuses a vector to represent a pair: +// This test misuses a vector WhileTest.WhileLoopsWithSharedBodyto represent a +// pair: // ((iteration, (random vector))). // // Note: this test currently only tests generating random values within a loop. @@ -329,7 +729,8 @@ TEST_F(WhileTest, WhileWithPrngScalarResult) { auto build_condition = [this, v6s32](int count) { ComputationBuilder builder(client_, TestName()); auto prev = builder.Reshape( - builder.Slice(builder.Parameter(0, v6s32, "prev"), {0}, {1}), {0}, {}); + builder.Slice(builder.Parameter(0, v6s32, "prev"), {0}, {1}, {1}), {0}, + {}); builder.Gt(builder.ConstantR0(count), prev); return builder.Build().ConsumeValueOrDie(); }; @@ -359,11 +760,11 @@ TEST_F(WhileTest, WhileWithPrngScalarResult) { }; for (int i = 1; i < 4; ++i) { - TF_ASSIGN_OR_ASSERT_OK(auto computation, while_loop(i)); + TF_ASSERT_OK_AND_ASSIGN(auto computation, while_loop(i)); - ExecutionOptions execution_options; + ExecutionOptions execution_options = execution_options_; execution_options.set_seed(65); - TF_ASSIGN_OR_ASSERT_OK( + TF_ASSERT_OK_AND_ASSIGN( auto result, client_->ExecuteAndTransfer(computation, {}, &execution_options)); } @@ -447,8 +848,10 @@ void BM_WhileLoop(int num_iters) { LocalClient* client = ClientLibrary::GetOrCreateLocalClient(platform).ValueOrDie(); + const int64 seq_len = 100; Shape loop_state_shape = ShapeUtil::MakeTupleShape( - {ShapeUtil::MakeShape(S32, {}), ShapeUtil::MakeShape(F32, {10})}); + {ShapeUtil::MakeShape(S32, {}), + ShapeUtil::MakeShape(F32, {seq_len, 1024, 1024})}); // Create while condition computation with 'loop_limit'. const int32 loop_limit = 100; @@ -466,20 +869,27 @@ void BM_WhileLoop(int num_iters) { { ComputationBuilder builder(client, "body"); auto prev = builder.Parameter(0, loop_state_shape, "prev"); + // TupleElement 0 auto iteration = builder.GetTupleElement(prev, 0); - auto weights = builder.GetTupleElement(prev, 1); - auto one = builder.ConstantR0(1); - auto next_iteration = builder.Add(iteration, one); - auto one_vec = builder.ConstantR1(10, 1.f); - auto new_weights = builder.Add(weights, one_vec); - auto result = builder.Tuple({next_iteration, new_weights}); + auto out0 = builder.Add(iteration, builder.ConstantR0(1)); + // TupleElement 1 + auto input = builder.GetTupleElement(prev, 1); + // Update. + auto one = builder.ConstantR0(1.0); + auto update = builder.Broadcast(one, {1, 1024, 1024}); + // Starts = iteration * 2; + auto starts = builder.ConstantR1({0, 0, 0}); + // UpdateSlice. + auto out1 = builder.DynamicUpdateSlice(input, update, starts); + auto result = builder.Tuple({out0, out1}); body = builder.Build().ConsumeValueOrDie(); } // Create a While instruction. ComputationBuilder builder(client, "while"); - auto init = builder.Tuple( - {builder.ConstantR0(0), builder.ConstantR1(10, 0.f)}); + auto zero = builder.ConstantR0(0.0); + auto input = builder.Broadcast(zero, {seq_len, 1024, 1024}); + auto init = builder.Tuple({builder.ConstantR0(0), input}); builder.While(condition, body, init); auto computation = builder.Build().ConsumeValueOrDie(); @@ -511,21 +921,3 @@ BENCHMARK(BM_WhileLoop); } // namespace } // namespace xla - -int main(int argc, char** argv) { - std::vector flag_list; - xla::legacy_flags::AppendCpuCompilerFlags(&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; - } - testing::InitGoogleTest(&argc, argv); - if (argc > 1) { - LOG(ERROR) << "Unknown argument " << argv[1] << "\n" << usage; - return 2; - } - tensorflow::testing::RunBenchmarks(); - return RUN_ALL_TESTS(); -} diff --git a/tensorflow/compiler/xla/tests/xla_internal_test_main.cc b/tensorflow/compiler/xla/tests/xla_internal_test_main.cc new file mode 100644 index 0000000000000000000000000000000000000000..92b2b1ee778f8b0f8104e7d7ff27a5c11db59768 --- /dev/null +++ b/tensorflow/compiler/xla/tests/xla_internal_test_main.cc @@ -0,0 +1,34 @@ +/* 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/legacy_flags/debug_options_flags.h" +#include "tensorflow/core/platform/logging.h" +#include "tensorflow/core/platform/test.h" + +GTEST_API_ int main(int argc, char** argv) { + std::vector flag_list; + xla::legacy_flags::AppendDebugOptionsFlags(&flag_list); + auto usage = tensorflow::Flags::Usage(argv[0], flag_list); + if (!tensorflow::Flags::Parse(&argc, argv, flag_list)) { + LOG(ERROR) << "\n" << usage; + return 2; + } + + testing::InitGoogleTest(&argc, argv); + if (argc > 1) { + LOG(ERROR) << "Unknown argument " << argv[1] << "\n" << usage; + return 2; + } + return RUN_ALL_TESTS(); +} diff --git a/tensorflow/compiler/xla/text_literal_reader.cc b/tensorflow/compiler/xla/text_literal_reader.cc index 7876272467890b56c2cca71f64e66303eb8ac632..4d060895d357493327ec50b38016478c65fef94d 100644 --- a/tensorflow/compiler/xla/text_literal_reader.cc +++ b/tensorflow/compiler/xla/text_literal_reader.cc @@ -104,8 +104,7 @@ StatusOr> TextLiteralReader::ReadAllLines() { auto result = MakeUnique(); const float fill = std::numeric_limits::quiet_NaN(); - LiteralUtil::PopulateWithValue(fill, AsInt64Slice(shape.dimensions()), - result.get()); + result->PopulateWithValue(fill, AsInt64Slice(shape.dimensions())); std::vector pieces; std::vector coordinates; std::vector coordinate_values; @@ -147,7 +146,7 @@ StatusOr> TextLiteralReader::ReadAllLines() { "\"%s\"", shape.dimensions_size(), coordinate_values.size(), line.c_str()); } - LiteralUtil::Set(result.get(), coordinate_values, value); + result->Set(coordinate_values, value); } return std::move(result); } diff --git a/tensorflow/compiler/xla/text_literal_reader.h b/tensorflow/compiler/xla/text_literal_reader.h index 3cfbb2c7fbf83d62e6797255cd693b26b522f061..e45e5291c9b10803f5e5008b72c7dd0116a0dea0 100644 --- a/tensorflow/compiler/xla/text_literal_reader.h +++ b/tensorflow/compiler/xla/text_literal_reader.h @@ -18,6 +18,7 @@ limitations under the License. #include +#include "tensorflow/compiler/xla/literal_util.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/text_literal_reader_test.cc b/tensorflow/compiler/xla/text_literal_reader_test.cc index 94d0f2646b15930f78c44fbb3d2b49fd6033a545..23070b663870a2b78b38663e09a32fcb28d9c2dc 100644 --- a/tensorflow/compiler/xla/text_literal_reader_test.cc +++ b/tensorflow/compiler/xla/text_literal_reader_test.cc @@ -19,10 +19,10 @@ limitations under the License. #include "tensorflow/compiler/xla/literal_util.h" #include "tensorflow/compiler/xla/shape_util.h" +#include "tensorflow/compiler/xla/test.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/test.h" namespace xla { namespace { @@ -46,12 +46,12 @@ TEST(TextLiteralReaderTest, ReadsR3File) { TextLiteralReader::ReadPath(fname).ConsumeValueOrDie(); EXPECT_TRUE( ShapeUtil::Equal(ShapeUtil::MakeShape(F32, {1, 2, 3}), literal->shape())); - EXPECT_EQ(42.5, LiteralUtil::Get(*literal, {0, 0, 0})); - EXPECT_EQ(43.5, LiteralUtil::Get(*literal, {0, 0, 1})); - EXPECT_EQ(44.5, LiteralUtil::Get(*literal, {0, 0, 2})); - EXPECT_EQ(45.5, LiteralUtil::Get(*literal, {0, 1, 0})); - EXPECT_EQ(46.5, LiteralUtil::Get(*literal, {0, 1, 1})); - EXPECT_EQ(47.5, LiteralUtil::Get(*literal, {0, 1, 2})); + EXPECT_EQ(42.5, literal->Get({0, 0, 0})); + EXPECT_EQ(43.5, literal->Get({0, 0, 1})); + EXPECT_EQ(44.5, literal->Get({0, 0, 2})); + EXPECT_EQ(45.5, literal->Get({0, 1, 0})); + EXPECT_EQ(46.5, literal->Get({0, 1, 1})); + EXPECT_EQ(47.5, literal->Get({0, 1, 2})); } } // namespace diff --git a/tensorflow/compiler/xla/text_literal_writer.cc b/tensorflow/compiler/xla/text_literal_writer.cc index a5097e41cb3cb3fe1c10e3c21c00c2242087deba..3fee467594d8423c707abf07a0622a738437830a 100644 --- a/tensorflow/compiler/xla/text_literal_writer.cc +++ b/tensorflow/compiler/xla/text_literal_writer.cc @@ -45,9 +45,9 @@ namespace xla { tensorflow::Status status; tensorflow::WritableFile* f_ptr = f.get(); - LiteralUtil::EachCellAsString( - literal, [f_ptr, &status](tensorflow::gtl::ArraySlice indices, - const string& value) { + literal.EachCellAsString( + [f_ptr, &status](tensorflow::gtl::ArraySlice indices, + const string& value) { if (!status.ok()) { return; } diff --git a/tensorflow/compiler/xla/text_literal_writer.h b/tensorflow/compiler/xla/text_literal_writer.h index 545bd22da918068244b23d35215fd95de89eaa56..7375493f4309c9bf75fc9d724626267dff7ce5ed 100644 --- a/tensorflow/compiler/xla/text_literal_writer.h +++ b/tensorflow/compiler/xla/text_literal_writer.h @@ -16,6 +16,7 @@ limitations under the License. #ifndef TENSORFLOW_COMPILER_XLA_TEXT_LITERAL_WRITER_H_ #define TENSORFLOW_COMPILER_XLA_TEXT_LITERAL_WRITER_H_ +#include "tensorflow/compiler/xla/literal_util.h" #include "tensorflow/compiler/xla/types.h" #include "tensorflow/compiler/xla/xla_data.pb.h" #include "tensorflow/core/lib/core/status.h" diff --git a/tensorflow/compiler/xla/text_literal_writer_test.cc b/tensorflow/compiler/xla/text_literal_writer_test.cc index 9dce4d13bb0e21d399795c5310e30b7ab64ea4ea..70cf2fb1b8a1b4f2ecfdaeaef3a00ddc974e2652 100644 --- a/tensorflow/compiler/xla/text_literal_writer_test.cc +++ b/tensorflow/compiler/xla/text_literal_writer_test.cc @@ -19,18 +19,18 @@ limitations under the License. #include #include "tensorflow/compiler/xla/literal_util.h" +#include "tensorflow/compiler/xla/test.h" #include "tensorflow/compiler/xla/test_helpers.h" #include "tensorflow/compiler/xla/types.h" #include "tensorflow/core/lib/io/path.h" #include "tensorflow/core/platform/env.h" #include "tensorflow/core/platform/logging.h" -#include "tensorflow/core/platform/test.h" namespace xla { namespace { TEST(TextLiteralWriterTest, WritesFloatLiteral) { - auto literal = LiteralUtil::CreateR2({ + auto literal = Literal::CreateR2({ {3.14, 2.17}, {1.23, 4.56}, }); string path = diff --git a/tensorflow/compiler/xla/tools/BUILD b/tensorflow/compiler/xla/tools/BUILD index 46eab7f02bb12ca39e5713e7b0f96bfa178e9102..a946d335ca6c8be583528caf9bcc97baf6245ae8 100644 --- a/tensorflow/compiler/xla/tools/BUILD +++ b/tensorflow/compiler/xla/tools/BUILD @@ -36,7 +36,7 @@ cc_library( "//tensorflow/compiler/xla/client:client_library", "//tensorflow/compiler/xla/client:computation", "//tensorflow/compiler/xla/client:local_client", - "//tensorflow/compiler/xla/legacy_flags:service_flags", + "//tensorflow/compiler/xla/legacy_flags:debug_options_flags", "//tensorflow/compiler/xla/service", "//tensorflow/compiler/xla/service:session_proto", "//tensorflow/core:lib", @@ -106,6 +106,14 @@ cc_binary( ], ) +cc_binary( + name = "replay_computation_hlo_evaluator", + deps = [ + ":replay_computation_library", + "//tensorflow/compiler/plugin/executor:plugin_lib", + ], +) + cc_binary( name = "show_literal", srcs = ["show_literal.cc"], @@ -153,6 +161,7 @@ cc_binary( "//tensorflow/compiler/xla/client:computation", "//tensorflow/compiler/xla/client:local_client", "//tensorflow/compiler/xla/service", + "//tensorflow/compiler/xla/service:computation_tracker", "//tensorflow/compiler/xla/service:session_proto", "//tensorflow/core:lib", ], @@ -176,6 +185,24 @@ cc_binary( ], ) +cc_binary( + name = "dumped_computation_to_tf_graphdef", + srcs = ["dumped_computation_to_tf_graphdef.cc"], + deps = [ + "//tensorflow/compiler/xla:statusor", + "//tensorflow/compiler/xla:types", + "//tensorflow/compiler/xla/client", + "//tensorflow/compiler/xla/client:client_library", + "//tensorflow/compiler/xla/client:computation", + "//tensorflow/compiler/xla/client:local_client", + "//tensorflow/compiler/xla/legacy_flags:debug_options_flags", + "//tensorflow/compiler/xla/service", + "//tensorflow/compiler/xla/service:hlo_graph_dumper", + "//tensorflow/compiler/xla/service:session_proto", + "//tensorflow/core:lib", + ], +) + # ----------------------------------------------------------------------------- filegroup( diff --git a/tensorflow/compiler/xla/tools/dumped_computation_to_graphviz.cc b/tensorflow/compiler/xla/tools/dumped_computation_to_graphviz.cc index 10efa9f3e8d856493b2db23195188da6fba65244..21ae8583d7cd3343230dcaff7dc17456e9e3e702 100644 --- a/tensorflow/compiler/xla/tools/dumped_computation_to_graphviz.cc +++ b/tensorflow/compiler/xla/tools/dumped_computation_to_graphviz.cc @@ -32,7 +32,7 @@ limitations under the License. #include "tensorflow/compiler/xla/client/client_library.h" #include "tensorflow/compiler/xla/client/computation.h" #include "tensorflow/compiler/xla/client/local_client.h" -#include "tensorflow/compiler/xla/legacy_flags/service_flags.h" +#include "tensorflow/compiler/xla/legacy_flags/debug_options_flags.h" #include "tensorflow/compiler/xla/service/service.h" #include "tensorflow/compiler/xla/service/session.pb.h" #include "tensorflow/compiler/xla/statusor.h" @@ -53,8 +53,11 @@ void RealMain(tensorflow::gtl::ArraySlice args) { TF_CHECK_OK( tensorflow::ReadBinaryProto(tensorflow::Env::Default(), arg, &module)); Computation computation = client->LoadSnapshot(module).ConsumeValueOrDie(); + DebugOptions debug_options = legacy_flags::GetDebugOptionsFromFlags(); + debug_options.set_xla_generate_hlo_graph(".*"); ComputationStats stats = - client->GetComputationStats(computation).ConsumeValueOrDie(); + client->GetComputationStats(computation, debug_options) + .ConsumeValueOrDie(); fprintf(stdout, ">>> %s :: %s\n", arg, stats.DebugString().c_str()); } } @@ -63,12 +66,16 @@ void RealMain(tensorflow::gtl::ArraySlice args) { } // 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); - xla::legacy_flags::ServiceFlags* flags = xla::legacy_flags::GetServiceFlags(); - flags->xla_generate_hlo_graph = ".*"; - flags->xla_hlo_graph_layout = true; - tensorflow::gtl::ArraySlice args(argv, argc); args.pop_front(); // Pop off the binary name, argv[0] xla::tools::RealMain(args); diff --git a/tensorflow/compiler/xla/tools/dumped_computation_to_operation_list.cc b/tensorflow/compiler/xla/tools/dumped_computation_to_operation_list.cc index 4c242abc9b7ce12c97fc3a5337163a6c641d5fdf..8d7f7fd1237f92053ee3b88e76ebae52b4a3a879 100644 --- a/tensorflow/compiler/xla/tools/dumped_computation_to_operation_list.cc +++ b/tensorflow/compiler/xla/tools/dumped_computation_to_operation_list.cc @@ -81,6 +81,7 @@ void RealMain(tensorflow::gtl::ArraySlice args) { client->GetComputationShape(computation).ConsumeValueOrDie(); std::vector layouts; + layouts.reserve(program_shape->parameters_size()); for (int i = 0; i < program_shape->parameters_size(); ++i) { layouts.push_back(&program_shape->parameters(i)); } diff --git a/tensorflow/compiler/xla/tools/dumped_computation_to_text.cc b/tensorflow/compiler/xla/tools/dumped_computation_to_text.cc index 8b96e13489774539b50022808975db56c5ddc6f7..2a3a8803283c62d12d8e2d213aa1730e8bd33244 100644 --- a/tensorflow/compiler/xla/tools/dumped_computation_to_text.cc +++ b/tensorflow/compiler/xla/tools/dumped_computation_to_text.cc @@ -21,6 +21,7 @@ limitations under the License. #include "tensorflow/compiler/xla/client/client_library.h" #include "tensorflow/compiler/xla/client/computation.h" #include "tensorflow/compiler/xla/client/local_client.h" +#include "tensorflow/compiler/xla/service/computation_tracker.h" #include "tensorflow/compiler/xla/service/service.h" #include "tensorflow/compiler/xla/service/session.pb.h" #include "tensorflow/compiler/xla/statusor.h" @@ -34,7 +35,7 @@ limitations under the License. namespace xla { namespace tools { -void RealMain(tensorflow::gtl::ArraySlice args) { +void RealMain(tensorflow::gtl::ArraySlice args, bool compile) { LocalClient* client = ClientLibrary::LocalClientOrDie(); LocalService* local_service = ClientLibrary::GetXlaService(client->platform()); @@ -50,23 +51,37 @@ void RealMain(tensorflow::gtl::ArraySlice args) { } Computation computation = computation_status.ConsumeValueOrDie(); - std::unique_ptr program_shape = - client->GetComputationShape(computation).ConsumeValueOrDie(); + if (compile) { + std::unique_ptr program_shape = + client->GetComputationShape(computation).ConsumeValueOrDie(); - std::vector layouts; - for (int i = 0; i < program_shape->parameters_size(); ++i) { - layouts.push_back(&program_shape->parameters(i)); - } - StatusOr> executable = - local_service->CompileExecutable( - computation.handle(), layouts, &program_shape->result(), - /*device_ordinal=*/0, /*has_hybrid_result=*/true); + std::vector layouts; + layouts.reserve(program_shape->parameters_size()); + for (int i = 0; i < program_shape->parameters_size(); ++i) { + layouts.push_back(&program_shape->parameters(i)); + } + StatusOr> executable = + local_service->CompileExecutable( + computation.handle(), layouts, &program_shape->result(), + /*device_ordinal=*/0, /*has_hybrid_result=*/true); + + const HloModule& module = executable.ValueOrDie()->module(); - const HloModule& module = executable.ValueOrDie()->module(); + fprintf(stdout, "HLO compiled for %s backend:\n%s\n", + local_service->backend().platform()->Name().c_str(), + module.ToString().c_str()); + } else { + const ComputationTracker& tracker = local_service->computation_tracker(); + UserComputation* user_computation = + tracker.Resolve(computation.handle()).ConsumeValueOrDie(); + VersionedComputationHandle versioned_handle = + user_computation->GetVersionedHandle(); + std::unique_ptr module = + tracker.BuildHloModule(versioned_handle, HloModuleConfig()) + .ConsumeValueOrDie(); - fprintf(stdout, "HLO for %s backend:\n%s\n", - local_service->backend().platform()->Name().c_str(), - module.ToString().c_str()); + fprintf(stdout, "%s\n", module->ToString().c_str()); + } } } @@ -74,10 +89,21 @@ void RealMain(tensorflow::gtl::ArraySlice args) { } // namespace xla int main(int argc, char** argv) { - tensorflow::port::InitMain(argv[0], &argc, &argv); + bool compile = false; + std::vector flag_list = { + {"compile", &compile, + "If true, compile the computation using the default client before " + "dumping the HLO. Otherwise dump the raw (uncompiled) HLO."}, + }; + const xla::string usage = tensorflow::Flags::Usage(argv[0], flag_list); + bool parsed_flags_ok = tensorflow::Flags::Parse(&argc, argv, flag_list); + QCHECK(parsed_flags_ok) << "\n" << usage; + + tensorflow::port::InitMain(usage.c_str(), &argc, &argv); + QCHECK(argc > 1) << "\nERROR: must specify at least one module\n" << usage; tensorflow::gtl::ArraySlice args(argv, argc); args.pop_front(); // Pop off the binary name, argv[0] - xla::tools::RealMain(args); + xla::tools::RealMain(args, compile); 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 new file mode 100644 index 0000000000000000000000000000000000000000..51f90b07c66f7d839f587350726333b9dbe6a9f0 --- /dev/null +++ b/tensorflow/compiler/xla/tools/dumped_computation_to_tf_graphdef.cc @@ -0,0 +1,84 @@ +/* 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_tf_graph some_binary_snapshot_proto* +// +// Dumps a tensorflow GraphDef in text format for a snapshot computation. The +// dumped graph is an HLO computation with HLO instructions as nodes and can be +// visualized on Tensorboard. Upload the dumped files on Tensorboard. +// +// some_binary_snapshot_proto is obtained by serializing the SessionModule from +// ServiceInterface::SnapshotComputation to disk. + +#include +#include +#include + +#include "tensorflow/compiler/xla/client/client.h" +#include "tensorflow/compiler/xla/client/client_library.h" +#include "tensorflow/compiler/xla/client/computation.h" +#include "tensorflow/compiler/xla/client/local_client.h" +#include "tensorflow/compiler/xla/legacy_flags/debug_options_flags.h" +#include "tensorflow/compiler/xla/service/service.h" +#include "tensorflow/compiler/xla/service/session.pb.h" +#include "tensorflow/compiler/xla/statusor.h" +#include "tensorflow/compiler/xla/types.h" +#include "tensorflow/core/lib/gtl/array_slice.h" +#include "tensorflow/core/platform/env.h" +#include "tensorflow/core/platform/init_main.h" +#include "tensorflow/core/platform/logging.h" + +using tensorflow::Env; + +namespace xla { +namespace tools { + +void RealMain(tensorflow::gtl::ArraySlice args) { + Client* client = ClientLibrary::LocalClientOrDie(); + for (char* arg : args) { + SessionModule module; + TF_CHECK_OK( + tensorflow::ReadBinaryProto(tensorflow::Env::Default(), arg, &module)); + Computation computation = client->LoadSnapshot(module).ConsumeValueOrDie(); + DebugOptions debug_options = legacy_flags::GetDebugOptionsFromFlags(); + debug_options.set_xla_generate_hlo_graph(".*"); + debug_options.set_xla_hlo_dump_as_graphdef(true); + 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); + + tensorflow::gtl::ArraySlice args(argv, argc); + args.pop_front(); // Pop off the binary name, argv[0] + xla::tools::RealMain(args); + return 0; +} diff --git a/tensorflow/compiler/xla/tools/replay_computation.cc b/tensorflow/compiler/xla/tools/replay_computation.cc index ffb2d5aefba7e3388ed6534e273a98ce6d648303..bd93e114b73aeb38c04c2c6a5169b9dc82d27e51 100644 --- a/tensorflow/compiler/xla/tools/replay_computation.cc +++ b/tensorflow/compiler/xla/tools/replay_computation.cc @@ -47,6 +47,7 @@ limitations under the License. #include "tensorflow/compiler/xla/statusor.h" #include "tensorflow/compiler/xla/types.h" #include "tensorflow/compiler/xla/xla_data.pb.h" +#include "tensorflow/core/lib/core/threadpool.h" #include "tensorflow/core/lib/gtl/array_slice.h" #include "tensorflow/core/platform/env.h" #include "tensorflow/core/platform/init_main.h" @@ -55,65 +56,101 @@ limitations under the License. namespace xla { namespace tools { +namespace { // Invokes the given computation passing arbitrary data for every (unbound) // parameter if use_fake_data, Otherwise use recorded data if available. +// +// Similarly, infeeds fake data of shape fake_infeed_shape if it is provided; +// otherwise, no infeed is performed. StatusOr> ReplayComputation( - const SessionModule& module, bool use_fake_data, Client* client) { + const SessionModule& module, tensorflow::StringPiece fake_infeed_shape, + bool use_fake_data, Client* client) { TF_ASSIGN_OR_RETURN(Computation computation, client->LoadSnapshot(module)); std::vector> arguments; if (use_fake_data) { arguments = MakeFakeArgumentsOrDie(computation, client); } else { // use recorded data if available - for (const Literal& literal : module.arguments()) { + for (const auto& proto : module.arguments()) { + Literal literal(proto); TF_ASSIGN_OR_RETURN(std::unique_ptr data, client->TransferToServer(literal)); arguments.push_back(std::move(data)); } } + // We only instantiate the thread pool if the user has requested that a + // concurrent infeed occur via the fake_infeed_shape. + tensorflow::gtl::optional pool; + + if (!fake_infeed_shape.empty()) { + pool.emplace(tensorflow::Env::Default(), "infeed", + /*num_threads=*/1); + pool->Schedule([fake_infeed_shape, client]() { + StatusOr shape_status = + ShapeUtil::ParseShapeString(fake_infeed_shape); + TF_CHECK_OK(shape_status.status()); + Shape shape = std::move(shape_status).ValueOrDie(); + StatusOr> data_status = MakeFakeLiteral(shape); + TF_CHECK_OK(data_status.status()); + std::unique_ptr data = std::move(data_status).ValueOrDie(); + while (true) { + TF_CHECK_OK(client->TransferToInfeed(*data)); + } + }); + } + std::vector execute_arguments; + execute_arguments.reserve(arguments.size()); for (auto& argument : arguments) { execute_arguments.push_back(argument.get()); } return client->ExecuteAndTransfer(computation, execute_arguments); } -void RealMain(tensorflow::gtl::ArraySlice args, bool use_fake_data) { +int RealMain(tensorflow::gtl::ArraySlice args, + tensorflow::StringPiece fake_infeed_shape, bool use_fake_data) { Client* client = ClientLibrary::LocalClientOrDie(); tensorflow::Env* env = tensorflow::Env::Default(); + int exit_status = EXIT_SUCCESS; for (char* arg : args) { SessionModule module; TF_CHECK_OK(tensorflow::ReadBinaryProto(env, arg, &module)); StatusOr> result_status = - ReplayComputation(module, use_fake_data, client); + ReplayComputation(module, fake_infeed_shape, use_fake_data, client); if (!result_status.ok()) { fprintf(stderr, "%s: error: %s\n", arg, result_status.status().ToString().c_str()); + exit_status = EXIT_FAILURE; continue; } std::unique_ptr result = result_status.ConsumeValueOrDie(); fprintf(stdout, "%s: %s :: %s:%s\n", arg, module.entry().name().c_str(), ShapeUtil::HumanString(result->shape()).c_str(), - LiteralUtil::ToString(*result).c_str()); + result->ToString().c_str()); if (module.has_result()) { fprintf(stdout, "was %s:%s\n", ShapeUtil::HumanString(module.result().shape()).c_str(), - LiteralUtil::ToString(module.result()).c_str()); + Literal(module.result()).ToString().c_str()); } } + return exit_status; } +} // namespace } // namespace tools } // namespace xla int main(int argc, char** argv) { // Flags + string fake_infeed_shape; bool use_fake_data = false; const std::vector flag_list = { tensorflow::Flag("use_fake_data", &use_fake_data, "Replay computation using fake data"), + tensorflow::Flag("fake_infeed_shape", &fake_infeed_shape, + "Shape of fake data to construct for (infinite) infeed"), }; xla::string usage = tensorflow::Flags::Usage(argv[0], flag_list); bool parse_ok = tensorflow::Flags::Parse(&argc, argv, flag_list); @@ -124,6 +161,5 @@ int main(int argc, char** argv) { tensorflow::gtl::ArraySlice args(argv, argc); args.pop_front(); // Pop off the binary name, argv[0] - xla::tools::RealMain(args, use_fake_data); - return 0; + return xla::tools::RealMain(args, fake_infeed_shape, use_fake_data); } diff --git a/tensorflow/compiler/xla/tools/show_literal.cc b/tensorflow/compiler/xla/tools/show_literal.cc index cf363913b159b6b04e66c8709073e4b130818946..b50cb5e28eac14ed99af566939f8bd64e393ff64 100644 --- a/tensorflow/compiler/xla/tools/show_literal.cc +++ b/tensorflow/compiler/xla/tools/show_literal.cc @@ -37,9 +37,10 @@ int main(int argc, char **argv) { << " "; } - xla::Literal literal; + xla::LiteralProto literal_proto; TF_CHECK_OK(tensorflow::ReadBinaryProto(tensorflow::Env::Default(), argv[1], - &literal)); - LOG(INFO) << "literal: " << literal.ShortDebugString(); - fprintf(stderr, "%s\n", xla::LiteralUtil::ToString(literal).c_str()); + &literal_proto)); + xla::Literal literal(literal_proto); + LOG(INFO) << "literal: " << literal_proto.ShortDebugString(); + fprintf(stderr, "%s\n", literal.ToString().c_str()); } diff --git a/tensorflow/compiler/xla/tools/show_text_literal.cc b/tensorflow/compiler/xla/tools/show_text_literal.cc index 2d983b407c64ab5547d722abcc2c564a7963f730..bbe9902aa17a585c4bad5b732330305dfdd45302 100644 --- a/tensorflow/compiler/xla/tools/show_text_literal.cc +++ b/tensorflow/compiler/xla/tools/show_text_literal.cc @@ -40,7 +40,7 @@ int main(int argc, char **argv) { xla::TextLiteralReader::ReadPath(argv[1]).ConsumeValueOrDie(); LOG(INFO) << "literal: " << literal->ShortDebugString(); - fprintf(stderr, "%s\n", xla::LiteralUtil::ToString(*literal).c_str()); + fprintf(stderr, "%s\n", literal->ToString().c_str()); if (literal->shape().element_type() == xla::F32) { float min = *std::min_element(literal->f32s().begin(), literal->f32s().end()); diff --git a/tensorflow/compiler/xla/types.h b/tensorflow/compiler/xla/types.h index 8258031a2c5119d085a483a0826f7284897dcee3..ea8b4b7b989b72034f33920a7d8c1a75e15a7dd1 100644 --- a/tensorflow/compiler/xla/types.h +++ b/tensorflow/compiler/xla/types.h @@ -16,8 +16,11 @@ limitations under the License. #ifndef TENSORFLOW_COMPILER_XLA_TYPES_H_ #define TENSORFLOW_COMPILER_XLA_TYPES_H_ +#include "third_party/eigen3/Eigen/Core" #include "tensorflow/core/platform/types.h" +#include + namespace xla { using ::tensorflow::string; @@ -32,6 +35,8 @@ using ::tensorflow::uint16; using ::tensorflow::uint32; using ::tensorflow::uint64; +using ::Eigen::half; + } // namespace xla #endif // TENSORFLOW_COMPILER_XLA_TYPES_H_ diff --git a/tensorflow/compiler/xla/util.cc b/tensorflow/compiler/xla/util.cc index a711b5035d842cd26945b2dac1159392813d56ab..f0e635bda4cac14ebe48935acb13c0aa4cb918e1 100644 --- a/tensorflow/compiler/xla/util.cc +++ b/tensorflow/compiler/xla/util.cc @@ -15,9 +15,10 @@ limitations under the License. #include "tensorflow/compiler/xla/util.h" +#include #include +#include -#include "tensorflow/compiler/xla/legacy_flags/util_flags.h" #include "tensorflow/compiler/xla/types.h" #include "tensorflow/core/lib/core/errors.h" #include "tensorflow/core/lib/strings/numbers.h" @@ -30,18 +31,12 @@ limitations under the License. namespace xla { namespace { -// Adds a backtrace to the provided status iff the xla_status_add_backtrace flag -// is set. This is useful for quickly tracing status errors observed coming out -// of the service. -Status MaybeAddBacktrace(Status prior) { - DCHECK(!prior.ok()); - if (legacy_flags::GetUtilFlags()->xla_status_add_backtrace) { - return Status{prior.code(), - tensorflow::strings::StrCat(prior.error_message(), " :: ", - tensorflow::CurrentStackTrace())}; - } else { - return prior; - } +// Logs the provided status message with a backtrace. +Status WithLogBacktrace(const Status& status) { + CHECK(!status.ok()); + VLOG(1) << status.ToString(); + VLOG(1) << tensorflow::CurrentStackTrace(); + return status; } } // namespace @@ -84,7 +79,7 @@ Status InvalidArgument(const char* format, ...) { va_start(args, format); tensorflow::strings::Appendv(&message, format, args); va_end(args); - return MaybeAddBacktrace(tensorflow::errors::InvalidArgument(message)); + return WithLogBacktrace(tensorflow::errors::InvalidArgument(message)); } Status Unimplemented(const char* format, ...) { @@ -93,7 +88,7 @@ Status Unimplemented(const char* format, ...) { va_start(args, format); tensorflow::strings::Appendv(&message, format, args); va_end(args); - return MaybeAddBacktrace(tensorflow::errors::Unimplemented(message)); + return WithLogBacktrace(tensorflow::errors::Unimplemented(message)); } Status InternalError(const char* format, ...) { @@ -102,7 +97,7 @@ Status InternalError(const char* format, ...) { va_start(args, format); tensorflow::strings::Appendv(&message, format, args); va_end(args); - return MaybeAddBacktrace(tensorflow::errors::Internal(message)); + return WithLogBacktrace(tensorflow::errors::Internal(message)); } Status FailedPrecondition(const char* format, ...) { @@ -111,7 +106,16 @@ Status FailedPrecondition(const char* format, ...) { va_start(args, format); tensorflow::strings::Appendv(&message, format, args); va_end(args); - return MaybeAddBacktrace(tensorflow::errors::FailedPrecondition(message)); + return WithLogBacktrace(tensorflow::errors::FailedPrecondition(message)); +} + +Status Cancelled(const char* format, ...) { + string message; + va_list args; + va_start(args, format); + tensorflow::strings::Appendv(&message, format, args); + va_end(args); + return WithLogBacktrace(tensorflow::errors::Cancelled(message)); } Status ResourceExhausted(const char* format, ...) { @@ -120,7 +124,7 @@ Status ResourceExhausted(const char* format, ...) { va_start(args, format); tensorflow::strings::Appendv(&message, format, args); va_end(args); - return MaybeAddBacktrace(tensorflow::errors::ResourceExhausted(message)); + return WithLogBacktrace(tensorflow::errors::ResourceExhausted(message)); } Status NotFound(const char* format, ...) { @@ -129,7 +133,7 @@ Status NotFound(const char* format, ...) { va_start(args, format); tensorflow::strings::Appendv(&message, format, args); va_end(args); - return MaybeAddBacktrace(tensorflow::errors::NotFound(message)); + return WithLogBacktrace(tensorflow::errors::NotFound(message)); } Status Unavailable(const char* format, ...) { @@ -138,7 +142,7 @@ Status Unavailable(const char* format, ...) { va_start(args, format); tensorflow::strings::Appendv(&message, format, args); va_end(args); - return MaybeAddBacktrace(tensorflow::errors::Unavailable(message)); + return WithLogBacktrace(tensorflow::errors::Unavailable(message)); } string Reindent(tensorflow::StringPiece original, @@ -153,16 +157,26 @@ string Reindent(tensorflow::StringPiece original, }); } +bool IsPermutation(tensorflow::gtl::ArraySlice permutation, int64 rank) { + if (rank != permutation.size()) { + return false; + } + std::vector output(permutation.size(), -1); + for (auto index : permutation) { + CHECK_GE(index, 0); + CHECK_LT(index, rank); + output[index] = 0; + } + return std::find(output.begin(), output.end(), -1) == output.end(); +} + std::vector InversePermutation( tensorflow::gtl::ArraySlice input_permutation) { + DCHECK(IsPermutation(input_permutation, input_permutation.size())); std::vector output_permutation(input_permutation.size(), -1); for (size_t i = 0; i < input_permutation.size(); ++i) { output_permutation[input_permutation[i]] = i; } - DCHECK_EQ( - 0, std::count(output_permutation.begin(), output_permutation.end(), -1)); - DCHECK(std::is_permutation(input_permutation.begin(), input_permutation.end(), - output_permutation.begin())); return output_permutation; } @@ -196,9 +210,20 @@ PaddingConfig MakeNoPaddingConfig(int64 rank) { return padding_config; } -string HumanReadableNumFlops(double flops, double nanoseconds) { +bool HasInteriorPadding(const PaddingConfig& config) { + for (const auto& dim : config.dimensions()) { + if (dim.interior_padding() != 0) { + return true; + } + } + return false; +} + +namespace { +string HumanReadableNumOps(double flops, double nanoseconds, + tensorflow::StringPiece op_prefix) { if (nanoseconds == 0) { - return "NaN FLOP/s"; + return tensorflow::strings::StrCat("NaN ", op_prefix, "OP/s"); } double nano_flops = flops / nanoseconds; string throughput = tensorflow::strings::HumanReadableNum( @@ -209,9 +234,18 @@ string HumanReadableNumFlops(double flops, double nanoseconds) { sp.ends_with("b")) { *throughput.rbegin() = 'G'; } - throughput += "FLOP/s"; + throughput += tensorflow::strings::StrCat(op_prefix, "OP/s"); return throughput; } +} // namespace + +string HumanReadableNumFlops(double flops, double nanoseconds) { + return HumanReadableNumOps(flops, nanoseconds, "FL"); +} + +string HumanReadableNumTranscendentalOps(double trops, double nanoseconds) { + return HumanReadableNumOps(trops, nanoseconds, "TR"); +} void LogLines(int sev, tensorflow::StringPiece text, const char* fname, int lineno) { diff --git a/tensorflow/compiler/xla/util.h b/tensorflow/compiler/xla/util.h index 55a66a7499571b4979ff375a8199cb329a799ef7..872715566adbf8de12c27080d6c9858a5169df70 100644 --- a/tensorflow/compiler/xla/util.h +++ b/tensorflow/compiler/xla/util.h @@ -31,6 +31,7 @@ limitations under the License. #include "tensorflow/core/lib/gtl/array_slice.h" #include "tensorflow/core/lib/math/math_util.h" #include "tensorflow/core/lib/strings/numbers.h" +#include "tensorflow/core/lib/strings/strcat.h" #include "tensorflow/core/platform/logging.h" #include "tensorflow/core/platform/macros.h" #include "tensorflow/core/platform/protobuf.h" @@ -38,6 +39,13 @@ limitations under the License. namespace xla { +// Ranks greater than 8 are very rare, so use InlinedVector to store +// the bounds and indices. And for the rare cases of ranks greater than 8, +// the InlinedVector will just behave like an std::vector<> and allocate the +// memory to store its values. +static constexpr int kInlineRank = 8; +using DimensionVector = tensorflow::gtl::InlinedVector; + // RAII timer that logs with a given label the wall clock time duration in human // readable form. This differs from base's ElapsedTimer primarily in that it // spits out the human-readable duration form. @@ -120,6 +128,14 @@ bool ContainersEqual(const Container1T& c1, const Container2T& c2) { std::equal(std::begin(c1), std::end(c1), std::begin(c2))); } +template +bool ContainersEqual(const Container1T& c1, + std::initializer_list il) { + tensorflow::gtl::ArraySlice c2{il}; + return ContainersEqual(c1, c2); +} + // Compares two containers for equality. Returns true iff the two containers // have the same size and all their elements compare equal using the predicate // p. Like std::equal, but forces size equality. @@ -130,6 +146,18 @@ bool ContainersEqual(const Container1T& c1, const Container2T& c2, std::equal(std::begin(c1), std::end(c1), std::begin(c2), p)); } +// Performs a copy of count values from src to dest, using different strides for +// source and destination. The source starting index is src_base, while the +// destination one is dest_base. +template +void StridedCopy(tensorflow::gtl::MutableArraySlice dest, int64 dest_base, + int64 dest_stride, tensorflow::gtl::ArraySlice src, + int64 src_base, int64 src_stride, int64 count) { + for (; count > 0; --count, dest_base += dest_stride, src_base += src_stride) { + dest[dest_base] = static_cast(src[src_base]); + } +} + // Adds some context information to the error message in a // Status. This is useful as Statuses are // propagated upwards. @@ -143,6 +171,7 @@ Status InvalidArgument(const char* format, ...) TF_PRINTF_ATTRIBUTE(1, 2); Status Unimplemented(const char* format, ...) TF_PRINTF_ATTRIBUTE(1, 2); Status InternalError(const char* format, ...) TF_PRINTF_ATTRIBUTE(1, 2); Status FailedPrecondition(const char* format, ...) TF_PRINTF_ATTRIBUTE(1, 2); +Status Cancelled(const char* format, ...) TF_PRINTF_ATTRIBUTE(1, 2); Status ResourceExhausted(const char* format, ...) TF_PRINTF_ATTRIBUTE(1, 2); Status NotFound(const char* format, ...) TF_PRINTF_ATTRIBUTE(1, 2); Status Unavailable(const char* format, ...) TF_PRINTF_ATTRIBUTE(1, 2); @@ -156,6 +185,9 @@ Status Unavailable(const char* format, ...) TF_PRINTF_ATTRIBUTE(1, 2); string Reindent(tensorflow::StringPiece original, tensorflow::StringPiece indentation); +// Checks whether permutation is a permutation of the [0, rank) integer range. +bool IsPermutation(tensorflow::gtl::ArraySlice permutation, int64 rank); + // Applies `permutation` on `input` and returns the permuted array. // For each i, output[permutation[i]] = input[i]. // @@ -164,17 +196,24 @@ string Reindent(tensorflow::StringPiece original, // 2. permutation.size() == input.size(). template